From 69864e1864eb3d246ab4ebf9d1173bfb4266620c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 9 Jul 2026 16:55:04 +0100 Subject: [PATCH] chore: regenerate notebooks + catalogue from scripts The workspace's generated artifacts had drifted globally from the scripts/ source of truth. Regenerated every notebook plus llms-full.txt and workspace_index.json via PyAutoBuild's generate.py; no scripts/ source changes. Generation verified deterministic (two full runs byte-identical) and all notebooks parse as valid JSON. Follow-up from PyAutoLabs/PyAutoLens#592. Issue #255. Co-Authored-By: Claude Opus 4.8 --- llms-full.txt | 2 +- notebooks/cluster/csv_api.ipynb | 1264 +++--- notebooks/cluster/lenstool/README.md | 67 + notebooks/cluster/lenstool/data.ipynb | 520 +++ notebooks/cluster/lenstool/modeling.ipynb | 679 +++ notebooks/cluster/likelihood_function.ipynb | 2258 ++++----- notebooks/cluster/modeling.ipynb | 1770 ++++---- notebooks/cluster/simulator.ipynb | 1868 ++++---- .../group/data_preparation/start_here.ipynb | 925 ++-- .../double_einstein_ring/chaining.ipynb | 921 ++-- .../advanced/double_einstein_ring/fit.ipynb | 841 ++-- .../likelihood_function.ipynb | 721 +-- .../double_einstein_ring/modeling.ipynb | 967 ++-- .../double_einstein_ring/simulator.ipynb | 861 ++-- .../advanced/double_einstein_ring/slam.ipynb | 1787 ++++---- .../advanced/mass_stellar_dark/chaining.ipynb | 779 ++-- .../advanced/mass_stellar_dark/fit.ipynb | 1087 ++--- .../likelihood_function.ipynb | 901 ++-- .../advanced/mass_stellar_dark/modeling.ipynb | 1051 ++--- .../mass_stellar_dark/simulator.ipynb | 859 ++-- .../advanced/mass_stellar_dark/slam.ipynb | 1515 ++++--- .../operated_light_profile/modeling.ipynb | 1107 ++--- .../operated_light_profile/simulator.ipynb | 929 ++-- .../features/advanced/shapelets/fit.ipynb | 857 ++-- .../advanced/shapelets/modeling.ipynb | 993 ++-- .../advanced/sky_background/fit.ipynb | 651 +-- .../advanced/sky_background/modeling.ipynb | 949 ++-- .../advanced/sky_background/simulator.ipynb | 907 ++-- .../advanced/subhalo/detect/start_here.ipynb | 2335 +++++----- .../features/advanced/subhalo/simulator.ipynb | 997 ++-- .../features/linear_light_profiles/fit.ipynb | 757 ++-- .../likelihood_function.ipynb | 1067 ++--- .../linear_light_profiles/modeling.ipynb | 1023 +++-- .../features/linear_light_profiles/slam.ipynb | 2321 +++++----- .../multi_gaussian_expansion/fit.ipynb | 973 ++-- .../likelihood_function.ipynb | 1175 ++--- .../multi_gaussian_expansion/modeling.ipynb | 1077 ++--- .../multi_gaussian_expansion/simulator.ipynb | 661 +-- .../multi_gaussian_expansion/slam.ipynb | 1311 +++--- .../source_science.ipynb | 1019 +++-- .../features/no_lens_light/modeling.ipynb | 979 ++-- .../features/no_lens_light/simulator.ipynb | 1071 ++--- .../group/features/no_lens_light/slam.ipynb | 1573 +++---- .../features/pixelization/adaptive.ipynb | 1165 ++--- .../pixelization/cpu_fast_modeling.ipynb | 961 ++-- .../features/pixelization/delaunay.ipynb | 1207 ++--- .../group/features/pixelization/fit.ipynb | 1023 +++-- .../pixelization/likelihood_function.ipynb | 1217 ++--- .../features/pixelization/modeling.ipynb | 1053 ++--- .../group/features/pixelization/slam.ipynb | 2075 ++++----- .../pixelization/source_science.ipynb | 1279 +++--- .../group/features/scaling_relation/fit.ipynb | 1094 ++--- .../likelihood_function.ipynb | 874 ++-- .../features/scaling_relation/modeling.ipynb | 1187 ++--- .../modeling_for_luminosities.ipynb | 955 ++-- .../features/scaling_relation/simulator.ipynb | 1043 +++-- notebooks/group/fit.ipynb | 1589 +++---- notebooks/group/likelihood_function.ipynb | 1609 +++---- notebooks/group/modeling.ipynb | 2133 ++++----- notebooks/group/simulator.ipynb | 1185 ++--- notebooks/group/slam.ipynb | 2457 +++++----- notebooks/group/source_science.ipynb | 1037 +++-- notebooks/guides/advanced/add_a_profile.ipynb | 1683 +++---- .../guides/advanced/custom_analysis.ipynb | 1663 +++---- notebooks/guides/advanced/multi_plane.ipynb | 1595 +++---- notebooks/guides/advanced/over_sampling.ipynb | 1590 ++++--- .../advanced/over_sampling_chaining.ipynb | 1027 +++-- notebooks/guides/data_structures.ipynb | 1709 +++---- notebooks/guides/galaxies.ipynb | 875 ++-- notebooks/guides/hpc/example_cpu.ipynb | 1253 ++--- notebooks/guides/lens_calc.ipynb | 1819 ++++---- .../advanced/expectation_propagation.ipynb | 1183 ++--- .../guides/modeling/advanced/graphical.ipynb | 893 ++-- .../modeling/advanced/hierarchical.ipynb | 1047 ++--- notebooks/guides/modeling/bug_fix.ipynb | 497 +- notebooks/guides/modeling/chaining.ipynb | 1187 ++--- notebooks/guides/modeling/cookbook.ipynb | 1273 +++--- notebooks/guides/modeling/customize.ipynb | 911 ++-- notebooks/guides/modeling/searches.ipynb | 853 ++-- .../guides/modeling/slam_start_here.ipynb | 1637 +++---- .../plotters_double_einstein_ring.ipynb | 777 ++-- .../plot/advanced/plotters_pixelization.ipynb | 923 ++-- notebooks/guides/plot/examples/mat_plot.ipynb | 633 +-- notebooks/guides/plot/examples/plotters.ipynb | 1357 +++--- notebooks/guides/plot/examples/searches.ipynb | 2415 +++++----- notebooks/guides/plot/examples/visuals.ipynb | 903 ++-- notebooks/guides/plot/start_here.ipynb | 779 ++-- notebooks/guides/point_source_pairing.ipynb | 313 ++ notebooks/guides/profiles/light.ipynb | 1797 ++++---- .../profiles/light_and_mass_profiles.ipynb | 1963 ++++---- notebooks/guides/profiles/mass.ipynb | 1721 +++---- notebooks/guides/results/_quick_fit.ipynb | 411 +- .../results/aggregator/data_fitting.ipynb | 733 +-- .../results/aggregator/galaxies_fits.ipynb | 823 ++-- .../results/aggregator/interferometer.ipynb | 209 +- .../guides/results/aggregator/models.ipynb | 729 +-- .../guides/results/aggregator/queries.ipynb | 591 +-- .../guides/results/aggregator/samples.ipynb | 1979 ++++---- .../aggregator/samples_via_aggregator.ipynb | 1841 ++++---- .../guides/results/database/start_here.ipynb | 1047 ++--- .../guides/results/latent_variables.ipynb | 869 ++-- notebooks/guides/results/start_here.ipynb | 1761 +++---- .../guides/results/workflow/csv_make.ipynb | 979 ++-- .../guides/results/workflow/fits_make.ipynb | 951 ++-- .../guides/results/workflow/png_make.ipynb | 843 ++-- notebooks/guides/tracer.ipynb | 1977 ++++---- notebooks/guides/units/cosmology.ipynb | 969 ++-- notebooks/guides/units/flux.ipynb | 967 ++-- .../units/mass_to_light_ratio_units.ipynb | 367 +- .../data_preparation/examples/data.ipynb | 731 +-- .../data_preparation/examples/noise_map.ipynb | 433 +- .../optional/extra_galaxies_centres.ipynb | 543 +-- .../examples/optional/info.ipynb | 407 +- .../examples/optional/lens_light_centre.ipynb | 473 +- .../examples/optional/mask.ipynb | 625 +-- .../optional/mask_extra_galaxies.ipynb | 577 +-- .../examples/optional/positions.ipynb | 481 +- .../data_preparation/examples/psf.ipynb | 523 ++- .../gui/extra_galaxies_centres.ipynb | 551 +-- .../gui/lens_light_centre.ipynb | 565 +-- .../imaging/data_preparation/gui/mask.ipynb | 415 +- .../gui/mask_extra_galaxies.ipynb | 551 +-- .../data_preparation/gui/positions.ipynb | 581 +-- .../manual/mask_irregular.ipynb | 429 +- .../imaging/data_preparation/start_here.ipynb | 817 ++-- .../double_einstein_ring/chaining.ipynb | 2289 +++++----- .../advanced/double_einstein_ring/fit.ipynb | 977 ++-- .../likelihood_function.ipynb | 771 ++-- .../double_einstein_ring/modeling.ipynb | 1175 ++--- .../double_einstein_ring/simulator.ipynb | 777 ++-- .../advanced/double_einstein_ring/slam.ipynb | 1535 ++++--- .../advanced/los_halos/simulator.ipynb | 1191 ++--- .../advanced/los_halos/simulator_jax.ipynb | 1223 ++--- .../advanced/mass_stellar_dark/chaining.ipynb | 879 ++-- .../advanced/mass_stellar_dark/fit.ipynb | 1047 ++--- .../likelihood_function.ipynb | 793 ++-- .../advanced/mass_stellar_dark/modeling.ipynb | 1027 +++-- .../mass_stellar_dark/simulator.ipynb | 737 +-- .../advanced/mass_stellar_dark/slam.ipynb | 1463 +++--- .../operated_light_profile/modeling.ipynb | 985 ++-- .../operated_light_profile/simulator.ipynb | 785 ++-- .../features/advanced/shapelets/fit.ipynb | 1399 +++--- .../advanced/shapelets/modeling.ipynb | 1707 +++---- .../advanced/sky_background/fit.ipynb | 633 +-- .../advanced/sky_background/modeling.ipynb | 899 ++-- .../advanced/sky_background/simulator.ipynb | 751 +-- .../advanced/subhalo/detect/database.ipynb | 799 ++-- .../advanced/subhalo/detect/start_here.ipynb | 2119 ++++----- .../sensitivity/slam_source_parametric.ipynb | 2479 +++++----- .../sensitivity/slam_source_pixelized.ipynb | 2859 ++++++------ .../subhalo/sensitivity/start_here.ipynb | 2031 +++++---- .../features/advanced/subhalo/simulator.ipynb | 997 ++-- .../features/extra_galaxies/modeling.ipynb | 1523 ++++--- .../features/extra_galaxies/simulator.ipynb | 1119 ++--- .../features/extra_galaxies/slam.ipynb | 1527 ++++--- .../features/linear_light_profiles/fit.ipynb | 829 ++-- .../likelihood_function.ipynb | 2205 ++++----- .../linear_light_profiles/modeling.ipynb | 1375 +++--- .../features/linear_light_profiles/slam.ipynb | 1357 +++--- .../multi_gaussian_expansion/fit.ipynb | 1103 ++--- .../likelihood_function.ipynb | 2585 +++++------ .../multi_gaussian_expansion/modeling.ipynb | 1491 +++--- .../multi_gaussian_expansion/simulator.ipynb | 995 ++-- .../multi_gaussian_expansion/slam.ipynb | 971 ++-- .../source_science.ipynb | 1109 ++--- .../features/no_lens_light/modeling.ipynb | 969 ++-- .../features/no_lens_light/simulator.ipynb | 991 ++-- .../imaging/features/no_lens_light/slam.ipynb | 1151 ++--- .../features/pixelization/adaptive.ipynb | 1417 +++--- .../pixelization/cpu_fast_modeling.ipynb | 1615 +++---- .../features/pixelization/delaunay.ipynb | 4033 +++++++++-------- .../imaging/features/pixelization/fit.ipynb | 1767 ++++---- .../pixelization/likelihood_function.ipynb | 3233 ++++++------- .../features/pixelization/modeling.ipynb | 1551 +++---- .../imaging/features/pixelization/slam.ipynb | 1365 +++--- .../pixelization/source_science.ipynb | 1571 +++---- .../features/scaling_relation/fit.ipynb | 1085 ++--- .../likelihood_function.ipynb | 895 ++-- .../features/scaling_relation/modeling.ipynb | 1225 ++--- .../features/scaling_relation/simulator.ipynb | 949 ++-- ...mulator_manual_signal_to_noise_ratio.ipynb | 837 ++-- notebooks/imaging/fit.ipynb | 1554 ++++--- notebooks/imaging/likelihood_function.ipynb | 1910 ++++---- notebooks/imaging/modeling.ipynb | 1998 ++++---- notebooks/imaging/simulator.ipynb | 1372 +++--- notebooks/imaging/simulator_sample.ipynb | 723 +-- notebooks/imaging/source_science.ipynb | 905 ++-- notebooks/imaging/start_here.ipynb | 9 + notebooks/interferometer/casa_reduction.ipynb | 1219 ++--- .../interferometer/data_preparation.ipynb | 787 ++-- .../features/advanced/shapelets/fit.ipynb | 799 ++-- .../advanced/shapelets/modeling.ipynb | 1069 ++--- .../features/datacube/data_preparation.ipynb | 833 ++-- .../features/datacube/delaunay.ipynb | 983 ++-- .../datacube/likelihood_function.ipynb | 2019 +++++---- .../features/datacube/modeling.ipynb | 995 ++-- .../datacube/modeling_parametric.ipynb | 863 ++-- .../features/datacube/simulator.ipynb | 1223 ++--- .../features/datacube/start_here.ipynb | 1065 ++--- .../features/extra_galaxies/modeling.ipynb | 977 ++-- .../features/extra_galaxies/simulator.ipynb | 943 ++-- .../features/extra_galaxies/slam.ipynb | 1623 +++---- .../features/linear_light_profiles/fit.ipynb | 783 ++-- .../likelihood_function.ipynb | 1711 +++---- .../linear_light_profiles/modeling.ipynb | 1345 +++--- .../features/linear_light_profiles/slam.ipynb | 1509 +++--- .../multi_gaussian_expansion/fit.ipynb | 803 ++-- .../likelihood_function.ipynb | 1565 +++---- .../multi_gaussian_expansion/modeling.ipynb | 1073 ++--- .../multi_gaussian_expansion/slam.ipynb | 1511 +++--- .../features/pixelization/delaunay.ipynb | 3527 +++++++------- .../features/pixelization/fit.ipynb | 1681 +++---- .../pixelization/likelihood_function.ipynb | 3199 ++++++------- .../many_visibilities_preparation.ipynb | 699 +-- .../features/pixelization/modeling.ipynb | 1517 ++++--- .../features/pixelization/slam.ipynb | 1497 +++--- .../pixelization/source_science.ipynb | 1545 ++++--- .../features/subhalo/detect/start_here.ipynb | 1897 ++++---- .../subhalo/sensitivity/start_here.ipynb | 711 +-- .../features/subhalo/simulator.ipynb | 775 ++-- notebooks/interferometer/fit.ipynb | 1359 +++--- .../interferometer/likelihood_function.ipynb | 1733 +++---- notebooks/interferometer/modeling.ipynb | 1577 +++---- notebooks/interferometer/simulator.ipynb | 1225 ++--- notebooks/interferometer/source_science.ipynb | 973 ++-- .../features/dataset_offsets/modeling.ipynb | 1005 ++-- .../features/dataset_offsets/simulator.ipynb | 1031 +++-- .../imaging_and_interferometer/modeling.ipynb | 919 ++-- .../simulator.ipynb | 787 ++-- .../multi/features/one_by_one/modeling.ipynb | 1059 ++--- .../features/pixelization/modeling.ipynb | 1001 ++-- .../features/pixelization/simulator.ipynb | 935 ++-- .../features/same_wavelength/modeling.ipynb | 879 ++-- .../features/same_wavelength/simulator.ipynb | 829 ++-- .../multi/features/slam/independent.ipynb | 1965 ++++---- .../multi/features/slam/simultaneous.ipynb | 1791 ++++---- .../wavelength_dependence/modeling.ipynb | 1017 +++-- .../wavelength_dependence/simulator.ipynb | 1069 ++--- notebooks/multi/modeling.ipynb | 1353 +++--- notebooks/multi/plot.ipynb | 539 ++- notebooks/multi/simulator.ipynb | 1089 ++--- .../features/deblending/modeling.ipynb | 1121 ++--- .../features/deblending/simulator.ipynb | 1313 +++--- notebooks/point_source/features/fluxes.ipynb | 845 ++-- .../features/multiple_sources/modeling.ipynb | 1079 ++--- .../features/multiple_sources/simulator.ipynb | 1023 +++-- .../point_source/features/time_delays.ipynb | 1065 ++--- notebooks/point_source/fit.ipynb | 1915 ++++---- notebooks/point_source/modeling.ipynb | 1511 +++--- notebooks/point_source/simulator.ipynb | 1783 ++++---- notebooks/point_source/simulator_sample.ipynb | 845 ++-- .../weak/features/strong_lensing/fit.ipynb | 307 ++ .../features/strong_lensing/modeling.ipynb | 392 ++ .../features/strong_lensing/simulator.ipynb | 319 ++ notebooks/weak/fit.ipynb | 561 ++- notebooks/weak/likelihood_function.ipynb | 456 ++ notebooks/weak/modeling.ipynb | 463 ++ notebooks/weak/real_data/a2744.ipynb | 440 ++ notebooks/weak/simulator.ipynb | 502 +- workspace_index.json | 20 +- 260 files changed, 158248 insertions(+), 144776 deletions(-) create mode 100644 notebooks/cluster/lenstool/README.md create mode 100644 notebooks/cluster/lenstool/data.ipynb create mode 100644 notebooks/cluster/lenstool/modeling.ipynb create mode 100644 notebooks/guides/point_source_pairing.ipynb create mode 100644 notebooks/weak/features/strong_lensing/fit.ipynb create mode 100644 notebooks/weak/features/strong_lensing/modeling.ipynb create mode 100644 notebooks/weak/features/strong_lensing/simulator.ipynb create mode 100644 notebooks/weak/likelihood_function.ipynb create mode 100644 notebooks/weak/modeling.ipynb create mode 100644 notebooks/weak/real_data/a2744.ipynb diff --git a/llms-full.txt b/llms-full.txt index f79e17a50..0d81ea184 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -389,7 +389,7 @@ AUTO-GENERATED by PyAutoBuild — do not edit by hand; regenerate with generate. - [Simulator: Cluster](scripts/cluster/simulator.py): This script simulates an example strong lens on the 'cluster' scale: a small cluster consisting of 2 main lens galaxies (a brightest cluster galaxy + a single satellite), 10 lower-mass cluster member galaxies on a luminosity-mass scaling relation, a single host dark matter halo not tied to any individual galaxy, and 2 multiply-imaged background source galaxies sitting at *different* redshifts (``z = 1.0`` and ``z = 2.0``) — making this a genuine multi-plane lens. - Contents: Multi-Plane Setup, Main Lens vs Scaling Members vs Host Halo vs Source Galaxies, Dataset Paths, Imaging and Visualization Grids, Galaxy Centres, Over Sampling, Main Lens Galaxies, Scaling Member Galaxies, Host Dark Matter Halo, Source Galaxies, Ray Tracing, JAX JIT, Point Solver, Point Datasets, Combined CSV, Manual CSV Editing, Scaling Galaxies CSV, Model CSVs, Tracer JSON, Imaging, Visualize - [Start Here: Cluster](scripts/cluster/start_here.py): Cluster-scale strong lenses are made of: - - Contents: JAX, Beta Feature, Google Colab Setup, Imports, Dataset, Model CSVs, Scaling Galaxies Table, Point Solver, Cluster Components, Model, Analysis + Factor Graph, Search, Model Fit, Live Visual Update, Result, Wrap Up + - Contents: JAX, Capabilities, Google Colab Setup, Imports, Dataset, Model CSVs, Scaling Galaxies Table, Point Solver, Cluster Components, Model, Analysis + Factor Graph, Search, Model Fit, Live Visual Update, Result, Wrap Up ## multi diff --git a/notebooks/cluster/csv_api.ipynb b/notebooks/cluster/csv_api.ipynb index 2b1730b46..19638d6f0 100644 --- a/notebooks/cluster/csv_api.ipynb +++ b/notebooks/cluster/csv_api.ipynb @@ -1,603 +1,667 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "CSV API: Cluster\n", - "================\n", - "\n", - "Cluster lens modelling uses a small family of CSV files as its canonical\n", - "input format. Rather than composing a tracer and a model inline in Python,\n", - "the cluster workflow stores every component \u2014 main galaxies, dark-matter\n", - "halo, source light, source point components, scaling-tier members \u2014 in\n", - "spreadsheet-editable CSVs that are then loaded into PyAutoLens objects.\n", - "\n", - "This guide walks through every CSV in the cluster surface, in the order\n", - "they are usually consumed:\n", - "\n", - " 1. **Point dataset CSV** \u2014 ``point_datasets.csv`` \u2014 the observational data\n", - " (multiple-image positions per source). Loaded *before* modelling, as is\n", - " the autolens convention. Produced by ``simulator.py`` for the bundled\n", - " cluster dataset; this guide builds an illustrative one to show the\n", - " schema.\n", - " 2. **Model CSVs (this PR)** \u2014 ``mass.csv``, ``light.csv``, ``point.csv``\n", - " \u2014 one CSV per profile family, each row carrying a galaxy name, an\n", - " attribute name, a profile class, and the constructor parameters of\n", - " that profile. The lens galaxies, the host dark-matter halo, and the\n", - " background sources all share the same three CSVs (grouping is by\n", - " profile family, not by tier). Joined across the family CSVs to\n", - " produce a full ``Galaxy`` (or ``af.Model[Galaxy]``).\n", - " 3. **Scaling galaxies CSV** \u2014 ``scaling_galaxies.csv`` \u2014 the legacy\n", - " 3-column ``y, x, luminosity`` schema for the scaling-relation tier\n", - " that drove the original CSV-first cluster work. Kept as-is because\n", - " naming each scaling-tier member would add overhead with no signal.\n", - "\n", - "Everything this guide writes goes into ``dataset/cluster/csv_api_example/``\n", - "(a scratch folder dedicated to this script). The canonical cluster dataset\n", - "at ``dataset/cluster/simple/`` is produced by ``simulator.py`` and is what\n", - "``modeling.py`` and ``start_here.py`` consume.\n", - "\n", - "__Why CSVs?__\n", - "\n", - "At cluster scale you might have tens of main galaxies, hundreds of member\n", - "galaxies, and a dozen background sources. Building all of that inline in\n", - "Python becomes unwieldy and error-prone. CSVs make every component:\n", - "\n", - " - **Spreadsheet-editable** \u2014 open in Excel / LibreOffice / any text\n", - " editor, tweak a value, save, re-run. No Python edit needed.\n", - " - **Diff-friendly** \u2014 git diff on a CSV is line-by-line, so a change to\n", - " one galaxy's mass is one line in the diff.\n", - " - **Easy to scale up** \u2014 adding more galaxies is a row append, not a\n", - " Python loop edit. The number of free parameters in the model does not\n", - " grow with the population size (see scaling galaxies below).\n", - " - **Round-trippable** \u2014 write a Python model out, read it back, get the\n", - " same Python model. This guide demonstrates that explicitly.\n", - "\n", - "__Contents__\n", - "\n", - "- **Imports & Output Path** \u2014 set up the scratch folder.\n", - "- **Point Dataset CSV** \u2014 write and load ``point_datasets.csv``.\n", - "- **Galaxy Model CSV API** \u2014 write, read, and round-trip family CSVs.\n", - "- **Scaling Galaxies CSV** \u2014 the legacy scaling-tier schema.\n", - "- **af.Model[Galaxy] from CSVs** \u2014 building modelling-ready objects.\n", - "- **Wrap Up** \u2014 where each CSV lives in the canonical cluster workflow." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Imports & Output Path__\n", - "\n", - "Everything this guide writes goes into ``dataset/cluster/csv_api_example/``,\n", - "a scratch folder dedicated to this script. The canonical cluster dataset\n", - "lives at ``dataset/cluster/simple/`` and is produced by ``simulator.py``;\n", - "this guide is illustrative and does not pair to that dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "output_path = Path(\"dataset\") / \"cluster\" / \"csv_api_example\"\n", - "output_path.mkdir(exist_ok=True, parents=True)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Dataset CSV__\n", - "\n", - "We begin with the observational data CSV because convention is \"load data\n", - "before modelling\". For point-source cluster lensing the observational data\n", - "is the image-plane positions of each multiply-imaged source.\n", - "\n", - "``al.PointDataset`` carries a source's ``name`` (which pairs to a ``Point``\n", - "component in the model), its ``positions`` (image-plane (y, x) coordinates\n", - "of the multiple images), a ``positions_noise_map`` (per-position\n", - "positional uncertainty), and an optional ``redshift``. A real cluster\n", - "typically has one ``PointDataset`` per background source.\n", - "\n", - "The CSV schema is one row per observed image, grouped by ``name``:\n", - "\n", - " ``name, y, x, positions_noise, redshift?``\n", - "\n", - "All rows sharing a ``name`` belong to the same source. The ``redshift``\n", - "column is consistent within a group (validated on load)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "point_dataset_0 = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=al.Grid2DIrregular(\n", - " [\n", - " (-9.0, -19.3),\n", - " (0.04, -0.08),\n", - " (1.84, 23.3),\n", - " ]\n", - " ),\n", - " positions_noise_map=0.005,\n", - " redshift=1.0,\n", - ")\n", - "\n", - "point_dataset_1 = al.PointDataset(\n", - " name=\"point_1\",\n", - " positions=al.Grid2DIrregular(\n", - " [\n", - " (-15.96, 18.99),\n", - " (0.68, -0.51),\n", - " (13.65, -14.32),\n", - " ]\n", - " ),\n", - " positions_noise_map=0.005,\n", - " redshift=2.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Write the combined dataset to a single CSV." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_csv(\n", - " datasets=[point_dataset_0, point_dataset_1],\n", - " file_path=output_path / \"point_datasets.csv\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Read it back. ``al.list_from_csv`` returns a ``List[PointDataset]`` with\n", - "each entry carrying the per-source positions, noise, and redshift." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "loaded_point_datasets = al.list_from_csv(file_path=output_path / \"point_datasets.csv\")\n", - "\n", - "print(\"=\" * 64)\n", - "print(\"Point datasets loaded from point_datasets.csv:\")\n", - "print(\"=\" * 64)\n", - "for dataset in loaded_point_datasets:\n", - " print(f\" name={dataset.name!r} redshift={dataset.redshift}\")\n", - " print(f\" positions={dataset.positions.in_list}\")\n", - " print(f\" noise={dataset.positions_noise_map}\")\n", - " print()\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Model CSV API__\n", - "\n", - "The model CSV API stores every galaxy component (main galaxies, the dark\n", - "matter halo, sources) in three family-level CSVs:\n", - "\n", - " - ``mass.csv`` \u2014 all mass profiles (``dPIEMassSph``, ``NFWMCRLudlowSph``, etc.)\n", - " - ``light.csv`` \u2014 all light profiles (``SersicSph``, ``SersicCore``, etc.)\n", - " - ``point.csv`` \u2014 all point-source components (``Point``)\n", - "\n", - "Each row carries:\n", - "\n", - " - ``galaxy`` \u2014 galaxy name. Rows sharing this name compose into one\n", - " ``Galaxy``.\n", - " - ``attr_name`` \u2014 attribute name to bind under on the Galaxy (e.g.\n", - " ``mass``, ``bulge``, ``point_0``).\n", - " - ``profile_class``\u2014 the concrete class name (looked up via ``getattr``\n", - " against the family namespace ``al.mp`` / ``al.lp`` /\n", - " ``al.ps``).\n", - " - ```` \u2014 the profile's constructor arguments. Tuples like\n", - " ``centre`` split into ``y`` / ``x`` columns; other\n", - " tuples (e.g. ``ell_comps``) into ``_0`` /\n", - " ``_1``.\n", - " - ``redshift`` \u2014 optional per-row Galaxy redshift (consistent across\n", - " every row for the same galaxy or all blank).\n", - "\n", - "**Sparse columns are supported.** Different profile classes inside the\n", - "same family CSV can use disjoint parameter columns; cells unused by a\n", - "row's class are simply blank. This is what lets one ``main_lens_mass.csv``\n", - "carry both 2 ``dPIEMassSph`` cluster members and 1 ``NFWMCRLudlowSph``\n", - "host halo even though their parameter sets don't overlap.\n", - "\n", - "We start by building a small illustrative cluster in Python." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "source_redshifts = [1.0, 2.0]\n", - "\n", - "# One dict per profile family. Each maps {galaxy_name: {attr_name: profile_instance}}.\n", - "# Mass family carries the 2 main-lens dPIE profiles + the host-halo NFW; light family\n", - "# carries the 2 main-lens Sersic bulges + the 2 source-galaxy SersicCore bulges; point\n", - "# family carries the 2 source-galaxy Point components.\n", - "\n", - "mass_profiles = {\n", - " \"lens_0\": {\"mass\": al.mp.dPIEMassSph(centre=(0.0, 0.0), ra=8.0, rs=20.0, b0=3.0)},\n", - " \"lens_1\": {\"mass\": al.mp.dPIEMassSph(centre=(10.0, 8.0), ra=5.0, rs=12.0, b0=1.2)},\n", - " \"host_halo\": {\n", - " \"dark\": al.mp.NFWMCRLudlowSph(\n", - " centre=(0.0, 0.0),\n", - " mass_at_200=10**15.3,\n", - " redshift_object=redshift_lens,\n", - " redshift_source=max(source_redshifts),\n", - " )\n", - " },\n", - "}\n", - "\n", - "light_profiles = {\n", - " \"lens_0\": {\n", - " \"bulge\": al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=1.5, effective_radius=3.0, sersic_index=4.0\n", - " )\n", - " },\n", - " \"lens_1\": {\n", - " \"bulge\": al.lp.SersicSph(\n", - " centre=(10.0, 8.0), intensity=0.8, effective_radius=1.5, sersic_index=3.5\n", - " )\n", - " },\n", - " \"source_0\": {\n", - " \"bulge\": al.lp.SersicCore(\n", - " centre=(0.3, 0.5),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=2.0,\n", - " effective_radius=0.3,\n", - " sersic_index=1.0,\n", - " )\n", - " },\n", - " \"source_1\": {\n", - " \"bulge\": al.lp.SersicCore(\n", - " centre=(-0.8, 1.2),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=90.0),\n", - " intensity=2.0,\n", - " effective_radius=0.3,\n", - " sersic_index=1.0,\n", - " )\n", - " },\n", - "}\n", - "\n", - "point_profiles = {\n", - " \"source_0\": {\"point_0\": al.ps.Point(centre=(0.3, 0.5))},\n", - " \"source_1\": {\"point_1\": al.ps.Point(centre=(-0.8, 1.2))},\n", - "}" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now have a complete cluster model expressed as nested dicts. Time to\n", - "write each profile family to its own CSV.\n", - "\n", - "The ``redshifts`` argument is a flat ``{galaxy_name: redshift}`` mapping\n", - "so the writer can emit the per-row ``redshift`` column. Mains + halo sit\n", - "at the lens redshift; sources carry their per-source redshift." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass_csv = output_path / \"mass.csv\"\n", - "light_csv = output_path / \"light.csv\"\n", - "point_csv = output_path / \"point.csv\"\n", - "\n", - "redshifts_by_galaxy = {\n", - " \"lens_0\": redshift_lens,\n", - " \"lens_1\": redshift_lens,\n", - " \"host_halo\": redshift_lens,\n", - " \"source_0\": source_redshifts[0],\n", - " \"source_1\": source_redshifts[1],\n", - "}\n", - "\n", - "al.galaxy_models_to_csv(\n", - " profiles_by_galaxy=mass_profiles,\n", - " file_path=mass_csv,\n", - " family=\"mass\",\n", - " redshifts=redshifts_by_galaxy,\n", - ")\n", - "\n", - "al.galaxy_models_to_csv(\n", - " profiles_by_galaxy=light_profiles,\n", - " file_path=light_csv,\n", - " family=\"light\",\n", - " redshifts=redshifts_by_galaxy,\n", - ")\n", - "\n", - "al.galaxy_models_to_csv(\n", - " profiles_by_galaxy=point_profiles,\n", - " file_path=point_csv,\n", - " family=\"point\",\n", - " redshifts=redshifts_by_galaxy,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Read each family CSV back. ``al.galaxy_models_from_csv`` returns a typed\n", - "``GalaxyModelTable`` with one ``GalaxyModelRow`` per CSV row: each row\n", - "carries the resolved ``profile_class`` (now a Python class, not a string)\n", - "and the parameter dict it needs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass_table = al.galaxy_models_from_csv(mass_csv, family=\"mass\")\n", - "light_table = al.galaxy_models_from_csv(light_csv, family=\"light\")\n", - "point_table = al.galaxy_models_from_csv(point_csv, family=\"point\")\n", - "\n", - "for label, table in [\n", - " (\"mass.csv\", mass_table),\n", - " (\"light.csv\", light_table),\n", - " (\"point.csv\", point_table),\n", - "]:\n", - " print(\"=\" * 64)\n", - " print(f\"{label}:\")\n", - " print(\"=\" * 64)\n", - " for row in table.rows:\n", - " print(\n", - " f\" galaxy={row.galaxy!r} attr_name={row.attr_name!r} \"\n", - " f\"class={row.profile_class.__name__} redshift={row.redshift}\"\n", - " )\n", - " print(f\" params={row.params}\")\n", - " print()\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Join the family tables into ``Galaxy`` instances. ``al.galaxies_from_csv_tables``\n", - "takes one or more ``GalaxyModelTable``s and groups rows by the ``galaxy``\n", - "column, attaching each profile under its ``attr_name``. Per-galaxy redshift\n", - "consistency is enforced across every family CSV that mentions the galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxies = al.galaxies_from_csv_tables(\n", - " mass_table,\n", - " light_table,\n", - " point_table,\n", - ")\n", - "\n", - "print(\"=\" * 64)\n", - "print(\"Galaxies built from CSVs (al.galaxies_from_csv_tables):\")\n", - "print(\"=\" * 64)\n", - "for name, galaxy in galaxies.items():\n", - " attr_summary = \", \".join(\n", - " f\"{a}={type(getattr(galaxy, a)).__name__}\"\n", - " for a in vars(galaxy)\n", - " if not a.startswith(\"_\") and a != \"redshift\"\n", - " )\n", - " print(f\" {name}: redshift={galaxy.redshift} attrs=[{attr_summary}]\")\n", - "print()\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__af.Model[Galaxy] from CSVs__\n", - "\n", - "The concrete ``Galaxy`` instances above are useful for visualisation and\n", - "for the simulator (which needs concrete parameter values to produce a\n", - "truth tracer). For *modelling*, what you want is an ``af.Model[Galaxy]``\n", - "where some parameters are free (priors) and others are fixed.\n", - "\n", - "``al.galaxy_af_models_from_csv_tables`` returns the same dict keyed by\n", - "galaxy name, but each value is an ``af.Model[Galaxy]`` with concrete CSV\n", - "values as fixed defaults. To free a parameter for the non-linear search,\n", - "mutate the returned model:\n", - "\n", - "```python\n", - "galaxy_models[\"lens_0\"].mass.ra = af.UniformPrior(lower_limit=1.0, upper_limit=15.0)\n", - "galaxy_models[\"lens_0\"].mass.rs = af.UniformPrior(lower_limit=5.0, upper_limit=40.0)\n", - "galaxy_models[\"lens_0\"].mass.b0 = af.UniformPrior(lower_limit=0.1, upper_limit=10.0)\n", - "```\n", - "\n", - "This is the same composition pattern as ``af.Model(al.Galaxy, mass=...)``\n", - "elsewhere in the workspace \u2014 only the construction step is replaced by\n", - "loading from CSVs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxy_models = al.galaxy_af_models_from_csv_tables(\n", - " mass_table,\n", - " light_table,\n", - " point_table,\n", - ")\n", - "\n", - "# Mutate selected params on the main-lens mass profiles into priors.\n", - "for name in (\"lens_0\", \"lens_1\"):\n", - " galaxy_models[name].mass.ra = af.UniformPrior(lower_limit=1.0, upper_limit=15.0)\n", - " galaxy_models[name].mass.rs = af.UniformPrior(lower_limit=5.0, upper_limit=40.0)\n", - " galaxy_models[name].mass.b0 = af.UniformPrior(lower_limit=0.1, upper_limit=10.0)\n", - "\n", - "# Host halo: free mass_at_200, keep centre + redshifts fixed.\n", - "galaxy_models[\"host_halo\"].dark.mass_at_200 = af.LogUniformPrior(\n", - " lower_limit=10**14.5, upper_limit=10**16.0\n", - ")\n", - "\n", - "model = af.Collection(galaxies=af.Collection(**galaxy_models))\n", - "\n", - "print(\"=\" * 64)\n", - "print(\"af.Collection built from CSVs (modelling-ready):\")\n", - "print(\"=\" * 64)\n", - "print(model.info)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling Galaxies CSV__\n", - "\n", - "The scaling-tier CSV format predates the named-galaxy model CSV API and\n", - "keeps its narrow 3-column schema: ``y, x, luminosity, redshift?``. The\n", - "scaling tier is implicitly one profile class per member, so naming each\n", - "member and emitting an ``attr_name`` column would be more overhead than\n", - "signal \u2014 every row uses the same ``dPIEMassSph`` mass profile with\n", - "parameters derived from the shared ``scaling_factor`` and\n", - "``scaling_exponent`` modelling parameters.\n", - "\n", - "``al.galaxy_table_to_csv`` and ``al.galaxy_table_from_csv`` are the\n", - "schema-specific writers/readers. The simulator emits 10 scaling members\n", - "into ``dataset/cluster/simple/scaling_galaxies.csv``; we write an\n", - "illustrative 3-member version here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scaling_centres = [(5.5, -6.5), (-7.5, 3.0), (12.0, -5.0)]\n", - "scaling_luminosities = [0.40, 0.32, 0.25]\n", - "\n", - "scaling_csv = output_path / \"scaling_galaxies.csv\"\n", - "\n", - "al.galaxy_table_to_csv(\n", - " centres=scaling_centres,\n", - " luminosities=scaling_luminosities,\n", - " file_path=scaling_csv,\n", - ")\n", - "\n", - "scaling_table = al.galaxy_table_from_csv(scaling_csv)\n", - "\n", - "print(\"=\" * 64)\n", - "print(\"scaling_galaxies.csv (legacy schema):\")\n", - "print(\"=\" * 64)\n", - "print(f\" centres={[tuple(c) for c in scaling_table.centres.in_list]}\")\n", - "print(f\" luminosities={scaling_table.luminosities}\")\n", - "print(f\" redshifts={scaling_table.redshifts}\")\n", - "print()\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The four CSVs above are the canonical input format for cluster lens\n", - "modelling in PyAutoLens. Where they live in the workflow:\n", - "\n", - " - ``dataset/cluster//point_datasets.csv`` \u2014 produced by\n", - " ``simulator.py``; consumed by ``modeling.py`` and ``start_here.py``\n", - " via ``al.list_from_csv``.\n", - " - ``dataset/cluster//mass.csv`` + ``light.csv`` +\n", - " ``point.csv`` \u2014 produced by ``simulator.py`` (truth values); consumed\n", - " by ``modeling.py`` and ``start_here.py`` via\n", - " ``al.galaxy_models_from_csv`` + ``al.galaxy_af_models_from_csv_tables``.\n", - " - ``dataset/cluster//scaling_galaxies.csv`` \u2014 produced by\n", - " ``simulator.py``; consumed by ``modeling.py`` and ``start_here.py``\n", - " via ``al.galaxy_table_from_csv``.\n", - "\n", - "To start modelling your own cluster:\n", - "\n", - " 1. Edit (or generate from a light-only fit) the model CSVs and\n", - " ``scaling_galaxies.csv`` for your cluster's main + scaling tiers.\n", - " 2. Edit ``point_datasets.csv`` with your measured multiple-image\n", - " positions per source.\n", - " 3. Drop the CSVs into ``dataset/cluster//``.\n", - " 4. Set ``dataset_name = \"\"`` in ``modeling.py`` and run.\n", - "\n", - "For a deeper view of what ``simulator.py`` and ``modeling.py`` do with\n", - "these CSVs end-to-end, follow the next files in this folder." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "CSV API: Cluster\n", + "================\n", + "\n", + "Cluster lens modelling uses a small family of CSV files as its canonical\n", + "input format. Rather than composing a tracer and a model inline in Python,\n", + "the cluster workflow stores every component \u2014 main galaxies, dark-matter\n", + "halo, source light, source point components, scaling-tier members \u2014 in\n", + "spreadsheet-editable CSVs that are then loaded into PyAutoLens objects.\n", + "\n", + "This guide walks through every CSV in the cluster surface, in the order\n", + "they are usually consumed:\n", + "\n", + " 1. **Point dataset CSV** \u2014 ``point_datasets.csv`` \u2014 the observational data\n", + " (multiple-image positions per source). Loaded *before* modelling, as is\n", + " the autolens convention. Produced by ``simulator.py`` for the bundled\n", + " cluster dataset; this guide builds an illustrative one to show the\n", + " schema.\n", + " 2. **Model CSVs (this PR)** \u2014 ``mass.csv``, ``light.csv``, ``point.csv``\n", + " \u2014 one CSV per profile family, each row carrying a galaxy name, an\n", + " attribute name, a profile class, and the constructor parameters of\n", + " that profile. The lens galaxies, the host dark-matter halo, and the\n", + " background sources all share the same three CSVs (grouping is by\n", + " profile family, not by tier). Joined across the family CSVs to\n", + " produce a full ``Galaxy`` (or ``af.Model[Galaxy]``).\n", + " 3. **Scaling galaxies CSV** \u2014 ``scaling_galaxies.csv`` \u2014 the legacy\n", + " 3-column ``y, x, luminosity`` schema for the scaling-relation tier\n", + " that drove the original CSV-first cluster work. Kept as-is because\n", + " naming each scaling-tier member would add overhead with no signal.\n", + "\n", + "Everything this guide writes goes into ``dataset/cluster/csv_api_example/``\n", + "(a scratch folder dedicated to this script). The canonical cluster dataset\n", + "at ``dataset/cluster/simple/`` is produced by ``simulator.py`` and is what\n", + "``modeling.py`` and ``start_here.py`` consume.\n", + "\n", + "__Why CSVs?__\n", + "\n", + "At cluster scale you might have tens of main galaxies, hundreds of member\n", + "galaxies, and a dozen background sources. Building all of that inline in\n", + "Python becomes unwieldy and error-prone. CSVs make every component:\n", + "\n", + " - **Spreadsheet-editable** \u2014 open in Excel / LibreOffice / any text\n", + " editor, tweak a value, save, re-run. No Python edit needed.\n", + " - **Diff-friendly** \u2014 git diff on a CSV is line-by-line, so a change to\n", + " one galaxy's mass is one line in the diff.\n", + " - **Easy to scale up** \u2014 adding more galaxies is a row append, not a\n", + " Python loop edit. The number of free parameters in the model does not\n", + " grow with the population size (see scaling galaxies below).\n", + " - **Round-trippable** \u2014 write a Python model out, read it back, get the\n", + " same Python model. This guide demonstrates that explicitly.\n", + "\n", + "__Contents__\n", + "\n", + "- **Imports & Output Path** \u2014 set up the scratch folder.\n", + "- **Point Dataset CSV** \u2014 write and load ``point_datasets.csv``.\n", + "- **Galaxy Model CSV API** \u2014 write, read, and round-trip family CSVs.\n", + "- **Scaling Galaxies CSV** \u2014 the legacy scaling-tier schema.\n", + "- **af.Model[Galaxy] from CSVs** \u2014 building modelling-ready objects.\n", + "- **Wrap Up** \u2014 where each CSV lives in the canonical cluster workflow." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imports & Output Path__\n", + "\n", + "Everything this guide writes goes into ``dataset/cluster/csv_api_example/``,\n", + "a scratch folder dedicated to this script. The canonical cluster dataset\n", + "lives at ``dataset/cluster/simple/`` and is produced by ``simulator.py``;\n", + "this guide is illustrative and does not pair to that dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "output_path = Path(\"dataset\") / \"cluster\" / \"csv_api_example\"\n", + "output_path.mkdir(exist_ok=True, parents=True)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Dataset CSV__\n", + "\n", + "We begin with the observational data CSV because convention is \"load data\n", + "before modelling\". For point-source cluster lensing the observational data\n", + "is the image-plane positions of each multiply-imaged source.\n", + "\n", + "``al.PointDataset`` carries a source's ``name`` (which pairs to a ``Point``\n", + "component in the model), its ``positions`` (image-plane (y, x) coordinates\n", + "of the multiple images), a ``positions_noise_map`` (per-position\n", + "positional uncertainty), and an optional ``redshift``. A real cluster\n", + "typically has one ``PointDataset`` per background source.\n", + "\n", + "The CSV schema is one row per observed image, grouped by ``name``:\n", + "\n", + " ``name, y, x, positions_noise, redshift?``\n", + "\n", + "All rows sharing a ``name`` belong to the same source. The ``redshift``\n", + "column is consistent within a group (validated on load)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "point_dataset_0 = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=al.Grid2DIrregular(\n", + " [\n", + " (-9.0, -19.3),\n", + " (0.04, -0.08),\n", + " (1.84, 23.3),\n", + " ]\n", + " ),\n", + " positions_noise_map=0.005,\n", + " redshift=1.0,\n", + ")\n", + "\n", + "point_dataset_1 = al.PointDataset(\n", + " name=\"point_1\",\n", + " positions=al.Grid2DIrregular(\n", + " [\n", + " (-15.96, 18.99),\n", + " (0.68, -0.51),\n", + " (13.65, -14.32),\n", + " ]\n", + " ),\n", + " positions_noise_map=0.005,\n", + " redshift=2.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Write the combined dataset to a single CSV." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_csv(\n", + " datasets=[point_dataset_0, point_dataset_1],\n", + " file_path=output_path / \"point_datasets.csv\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Read it back. ``al.list_from_csv`` returns a ``List[PointDataset]`` with\n", + "each entry carrying the per-source positions, noise, and redshift." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "loaded_point_datasets = al.list_from_csv(file_path=output_path / \"point_datasets.csv\")\n", + "\n", + "print(\"=\" * 64)\n", + "print(\"Point datasets loaded from point_datasets.csv:\")\n", + "print(\"=\" * 64)\n", + "for dataset in loaded_point_datasets:\n", + " print(f\" name={dataset.name!r} redshift={dataset.redshift}\")\n", + " print(f\" positions={dataset.positions.in_list}\")\n", + " print(f\" noise={dataset.positions_noise_map}\")\n", + " print()\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Model CSV API__\n", + "\n", + "The model CSV API stores every galaxy component (main galaxies, the dark\n", + "matter halo, sources) in three family-level CSVs:\n", + "\n", + " - ``mass.csv`` \u2014 all mass profiles (``dPIEMassSph``, ``NFWMCRLudlowSph``, etc.)\n", + " - ``light.csv`` \u2014 all light profiles (``SersicSph``, ``SersicCore``, etc.)\n", + " - ``point.csv`` \u2014 all point-source components (``Point``)\n", + "\n", + "Each row carries:\n", + "\n", + " - ``galaxy`` \u2014 galaxy name. Rows sharing this name compose into one\n", + " ``Galaxy``.\n", + " - ``attr_name`` \u2014 attribute name to bind under on the Galaxy (e.g.\n", + " ``mass``, ``bulge``, ``point_0``).\n", + " - ``profile_class``\u2014 the concrete class name (looked up via ``getattr``\n", + " against the family namespace ``al.mp`` / ``al.lp`` /\n", + " ``al.ps``).\n", + " - ```` \u2014 the profile's constructor arguments. Tuples like\n", + " ``centre`` split into ``y`` / ``x`` columns; other\n", + " tuples (e.g. ``ell_comps``) into ``_0`` /\n", + " ``_1``.\n", + " - ``redshift`` \u2014 optional per-row Galaxy redshift (consistent across\n", + " every row for the same galaxy or all blank).\n", + "\n", + "**Sparse columns are supported.** Different profile classes inside the\n", + "same family CSV can use disjoint parameter columns; cells unused by a\n", + "row's class are simply blank. This is what lets one ``main_lens_mass.csv``\n", + "carry both 2 ``dPIEMassSph`` cluster members and 1 ``NFWMCRLudlowSph``\n", + "host halo even though their parameter sets don't overlap.\n", + "\n", + "We start by building a small illustrative cluster in Python." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "source_redshifts = [1.0, 2.0]\n", + "\n", + "# One dict per profile family. Each maps {galaxy_name: {attr_name: profile_instance}}.\n", + "# Mass family carries the 2 main-lens dPIE profiles + the host-halo NFW; light family\n", + "# carries the 2 main-lens Sersic bulges + the 2 source-galaxy SersicCore bulges; point\n", + "# family carries the 2 source-galaxy Point components.\n", + "\n", + "mass_profiles = {\n", + " \"lens_0\": {\"mass\": al.mp.dPIEMassSph(centre=(0.0, 0.0), ra=8.0, rs=20.0, b0=3.0)},\n", + " \"lens_1\": {\"mass\": al.mp.dPIEMassSph(centre=(10.0, 8.0), ra=5.0, rs=12.0, b0=1.2)},\n", + " \"host_halo\": {\n", + " \"dark\": al.mp.NFWMCRLudlowSph(\n", + " centre=(0.0, 0.0),\n", + " mass_at_200=10**15.3,\n", + " redshift_object=redshift_lens,\n", + " redshift_source=max(source_redshifts),\n", + " )\n", + " },\n", + "}\n", + "\n", + "light_profiles = {\n", + " \"lens_0\": {\n", + " \"bulge\": al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=1.5, effective_radius=3.0, sersic_index=4.0\n", + " )\n", + " },\n", + " \"lens_1\": {\n", + " \"bulge\": al.lp.SersicSph(\n", + " centre=(10.0, 8.0), intensity=0.8, effective_radius=1.5, sersic_index=3.5\n", + " )\n", + " },\n", + " \"source_0\": {\n", + " \"bulge\": al.lp.SersicCore(\n", + " centre=(0.3, 0.5),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=2.0,\n", + " effective_radius=0.3,\n", + " sersic_index=1.0,\n", + " )\n", + " },\n", + " \"source_1\": {\n", + " \"bulge\": al.lp.SersicCore(\n", + " centre=(-0.8, 1.2),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=90.0),\n", + " intensity=2.0,\n", + " effective_radius=0.3,\n", + " sersic_index=1.0,\n", + " )\n", + " },\n", + "}\n", + "\n", + "point_profiles = {\n", + " \"source_0\": {\"point_0\": al.ps.Point(centre=(0.3, 0.5))},\n", + " \"source_1\": {\"point_1\": al.ps.Point(centre=(-0.8, 1.2))},\n", + "}" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now have a complete cluster model expressed as nested dicts. Time to\n", + "write each profile family to its own CSV.\n", + "\n", + "The ``redshifts`` argument is a flat ``{galaxy_name: redshift}`` mapping\n", + "so the writer can emit the per-row ``redshift`` column. Mains + halo sit\n", + "at the lens redshift; sources carry their per-source redshift." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass_csv = output_path / \"mass.csv\"\n", + "light_csv = output_path / \"light.csv\"\n", + "point_csv = output_path / \"point.csv\"\n", + "\n", + "redshifts_by_galaxy = {\n", + " \"lens_0\": redshift_lens,\n", + " \"lens_1\": redshift_lens,\n", + " \"host_halo\": redshift_lens,\n", + " \"source_0\": source_redshifts[0],\n", + " \"source_1\": source_redshifts[1],\n", + "}\n", + "\n", + "al.galaxy_models_to_csv(\n", + " profiles_by_galaxy=mass_profiles,\n", + " file_path=mass_csv,\n", + " family=\"mass\",\n", + " redshifts=redshifts_by_galaxy,\n", + ")\n", + "\n", + "al.galaxy_models_to_csv(\n", + " profiles_by_galaxy=light_profiles,\n", + " file_path=light_csv,\n", + " family=\"light\",\n", + " redshifts=redshifts_by_galaxy,\n", + ")\n", + "\n", + "al.galaxy_models_to_csv(\n", + " profiles_by_galaxy=point_profiles,\n", + " file_path=point_csv,\n", + " family=\"point\",\n", + " redshifts=redshifts_by_galaxy,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Read each family CSV back. ``al.galaxy_models_from_csv`` returns a typed\n", + "``GalaxyModelTable`` with one ``GalaxyModelRow`` per CSV row: each row\n", + "carries the resolved ``profile_class`` (now a Python class, not a string)\n", + "and the parameter dict it needs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass_table = al.galaxy_models_from_csv(mass_csv, family=\"mass\")\n", + "light_table = al.galaxy_models_from_csv(light_csv, family=\"light\")\n", + "point_table = al.galaxy_models_from_csv(point_csv, family=\"point\")\n", + "\n", + "for label, table in [\n", + " (\"mass.csv\", mass_table),\n", + " (\"light.csv\", light_table),\n", + " (\"point.csv\", point_table),\n", + "]:\n", + " print(\"=\" * 64)\n", + " print(f\"{label}:\")\n", + " print(\"=\" * 64)\n", + " for row in table.rows:\n", + " print(\n", + " f\" galaxy={row.galaxy!r} attr_name={row.attr_name!r} \"\n", + " f\"class={row.profile_class.__name__} redshift={row.redshift}\"\n", + " )\n", + " print(f\" params={row.params}\")\n", + " print()\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Join the family tables into ``Galaxy`` instances. ``al.galaxies_from_csv_tables``\n", + "takes one or more ``GalaxyModelTable``s and groups rows by the ``galaxy``\n", + "column, attaching each profile under its ``attr_name``. Per-galaxy redshift\n", + "consistency is enforced across every family CSV that mentions the galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxies = al.galaxies_from_csv_tables(\n", + " mass_table,\n", + " light_table,\n", + " point_table,\n", + ")\n", + "\n", + "print(\"=\" * 64)\n", + "print(\"Galaxies built from CSVs (al.galaxies_from_csv_tables):\")\n", + "print(\"=\" * 64)\n", + "for name, galaxy in galaxies.items():\n", + " attr_summary = \", \".join(\n", + " f\"{a}={type(getattr(galaxy, a)).__name__}\"\n", + " for a in vars(galaxy)\n", + " if not a.startswith(\"_\") and a != \"redshift\"\n", + " )\n", + " print(f\" {name}: redshift={galaxy.redshift} attrs=[{attr_summary}]\")\n", + "print()\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__af.Model[Galaxy] from CSVs__\n", + "\n", + "The concrete ``Galaxy`` instances above are useful for visualisation and\n", + "for the simulator (which needs concrete parameter values to produce a\n", + "truth tracer). For *modelling*, what you want is an ``af.Model[Galaxy]``\n", + "where some parameters are free (priors) and others are fixed.\n", + "\n", + "``al.galaxy_af_models_from_csv_tables`` returns the same dict keyed by\n", + "galaxy name, but each value is an ``af.Model[Galaxy]`` with concrete CSV\n", + "values as fixed defaults. To free a parameter for the non-linear search,\n", + "mutate the returned model:\n", + "\n", + "```python\n", + "galaxy_models[\"lens_0\"].mass.ra = af.UniformPrior(lower_limit=1.0, upper_limit=15.0)\n", + "galaxy_models[\"lens_0\"].mass.rs = af.UniformPrior(lower_limit=5.0, upper_limit=40.0)\n", + "galaxy_models[\"lens_0\"].mass.b0 = af.UniformPrior(lower_limit=0.1, upper_limit=10.0)\n", + "```\n", + "\n", + "This is the same composition pattern as ``af.Model(al.Galaxy, mass=...)``\n", + "elsewhere in the workspace \u2014 only the construction step is replaced by\n", + "loading from CSVs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxy_models = al.galaxy_af_models_from_csv_tables(\n", + " mass_table,\n", + " light_table,\n", + " point_table,\n", + ")\n", + "\n", + "# Mutate selected params on the main-lens mass profiles into priors.\n", + "for name in (\"lens_0\", \"lens_1\"):\n", + " galaxy_models[name].mass.ra = af.UniformPrior(lower_limit=1.0, upper_limit=15.0)\n", + " galaxy_models[name].mass.rs = af.UniformPrior(lower_limit=5.0, upper_limit=40.0)\n", + " galaxy_models[name].mass.b0 = af.UniformPrior(lower_limit=0.1, upper_limit=10.0)\n", + "\n", + "# Host halo: free mass_at_200, keep centre + redshifts fixed.\n", + "galaxy_models[\"host_halo\"].dark.mass_at_200 = af.LogUniformPrior(\n", + " lower_limit=10**14.5, upper_limit=10**16.0\n", + ")\n", + "\n", + "model = af.Collection(galaxies=af.Collection(**galaxy_models))\n", + "\n", + "print(\"=\" * 64)\n", + "print(\"af.Collection built from CSVs (modelling-ready):\")\n", + "print(\"=\" * 64)\n", + "print(model.info)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling Galaxies CSV__\n", + "\n", + "The scaling-tier CSV format predates the named-galaxy model CSV API and\n", + "keeps its narrow 3-column schema: ``y, x, luminosity, redshift?``. The\n", + "scaling tier is implicitly one profile class per member, so naming each\n", + "member and emitting an ``attr_name`` column would be more overhead than\n", + "signal \u2014 every row uses the same ``dPIEMassSph`` mass profile with\n", + "parameters derived from the reference-anchored scaling relation's shared\n", + "``b0_ref`` normalization (see ``modeling.py``).\n", + "\n", + "``al.galaxy_table_to_csv`` and ``al.galaxy_table_from_csv`` are the\n", + "schema-specific writers/readers. The simulator emits 10 scaling members\n", + "into ``dataset/cluster/simple/scaling_galaxies.csv``; we write an\n", + "illustrative 3-member version here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scaling_centres = [(5.5, -6.5), (-7.5, 3.0), (12.0, -5.0)]\n", + "scaling_luminosities = [0.40, 0.32, 0.25]\n", + "\n", + "scaling_csv = output_path / \"scaling_galaxies.csv\"\n", + "\n", + "al.galaxy_table_to_csv(\n", + " centres=scaling_centres,\n", + " luminosities=scaling_luminosities,\n", + " file_path=scaling_csv,\n", + ")\n", + "\n", + "scaling_table = al.galaxy_table_from_csv(scaling_csv)\n", + "\n", + "print(\"=\" * 64)\n", + "print(\"scaling_galaxies.csv (legacy schema):\")\n", + "print(\"=\" * 64)\n", + "print(f\" centres={[tuple(c) for c in scaling_table.centres.in_list]}\")\n", + "print(f\" luminosities={scaling_table.luminosities}\")\n", + "print(f\" redshifts={scaling_table.redshifts}\")\n", + "print()\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The four CSVs above are the canonical input format for cluster lens\n", + "modelling in PyAutoLens. Where they live in the workflow:\n", + "\n", + " - ``dataset/cluster//point_datasets.csv`` \u2014 produced by\n", + " ``simulator.py``; consumed by ``modeling.py`` and ``start_here.py``\n", + " via ``al.list_from_csv``.\n", + " - ``dataset/cluster//mass.csv`` + ``light.csv`` +\n", + " ``point.csv`` \u2014 produced by ``simulator.py`` (truth values); consumed\n", + " by ``modeling.py`` and ``start_here.py`` via\n", + " ``al.galaxy_models_from_csv`` + ``al.galaxy_af_models_from_csv_tables``.\n", + " - ``dataset/cluster//scaling_galaxies.csv`` \u2014 produced by\n", + " ``simulator.py``; consumed by ``modeling.py`` and ``start_here.py``\n", + " via ``al.galaxy_table_from_csv``.\n", + "\n", + "__Lenstool-Parameterized Rows__\n", + "\n", + "Because the ``profile_class`` column dispatches against the full ``al.mp`` namespace, a\n", + "``mass.csv`` can carry rows in **Lenstool's native parameterization** via ``dPIEMassLenstool``\n", + "\u2014 the columns become the ``.par``-file keywords verbatim::\n", + "\n", + " galaxy,attr_name,profile_class,y,x,ellipticity,angle_pos,sigma,r_core,r_cut,redshift_object,redshift_source,H0,Om0,redshift\n", + " O1,mass,dPIEMassLenstool,1.479,-2.997,0.678,8.971,987.34,18.96,283.54,0.39,11.76,70.0,0.3,0.39\n", + "\n", + "``sigma`` is Lenstool's fiducial ``v_disp`` (sigma_LT), radii are in arcsec, and the run's own\n", + "cosmology travels as the flat ``H0`` / ``Om0`` columns. ``scripts/cluster/lenstool/`` builds its\n", + "entire 149-component published model this way \u2014 the ``.par`` file becomes one canonical CSV. Note\n", + "the multi-plane convention: ``redshift_source`` must be the tracer's *final* (highest) source\n", + "plane.\n", + "\n", + "Light-profile CSVs (``light.csv``) support the linear / operated variants with qualified class\n", + "names (``linear.Sersic``, ``operated.Gaussian``); plain names resolve to the standard profiles.\n", + "\n", + "__Member Catalogues With Properties__\n", + "\n", + "``scaling_galaxies.csv`` / ``al.galaxy_table_from_csv`` accept any extra per-galaxy columns\n", + "beyond ``y, x, luminosity[, redshift]`` \u2014 numeric columns (``ellipticity``, ``angle_pos``,\n", + "``mag``) load as float lists in ``GalaxyTable.properties``, strings (names, notes) as string\n", + "lists. Nothing is silently dropped, and two loud guards protect the model CSVs: a typo'd\n", + "parameter column raises (instead of silently leaving the profile at its default), as does a\n", + "duplicate ``(galaxy, attr_name)`` row pair.\n", + "\n", + "To start modelling your own cluster:\n", + "\n", + " 1. Edit (or generate from a light-only fit) the model CSVs and\n", + " ``scaling_galaxies.csv`` for your cluster's main + scaling tiers.\n", + " 2. Edit ``point_datasets.csv`` with your measured multiple-image\n", + " positions per source.\n", + " 3. Drop the CSVs into ``dataset/cluster//``.\n", + " 4. Set ``dataset_name = \"\"`` in ``modeling.py`` and run.\n", + "\n", + "For a deeper view of what ``simulator.py`` and ``modeling.py`` do with\n", + "these CSVs end-to-end, follow the next files in this folder." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/cluster/lenstool/README.md b/notebooks/cluster/lenstool/README.md new file mode 100644 index 000000000..6d5c5193c --- /dev/null +++ b/notebooks/cluster/lenstool/README.md @@ -0,0 +1,67 @@ +# PyAutoLens for Lenstool Users + +If you model galaxy clusters with **Lenstool**, this folder shows you — with a real, published +cluster — how to do the same analysis in **PyAutoLens**, and what you gain on the other side. + +The worked example is **SMACS J0723.3−7327**, the first JWST cluster, using the Lenstool model of +[Mahler et al. 2023](https://arxiv.org/abs/2207.07101) whose complete workflow (`best.par`, +`input.par`, `arcs.dat`, `galcat.cat`) is +[public](https://github.com/guillaumemahler/SMACS0723-mahler2022). Two scripts: + +| Script | What it does | +|---|---| +| `data.py` | Downloads the published model files + a RELICS HST image; converts them to the PyAutoLens cluster CSV formats, recording every unit and sign convention. | +| `modeling.py` | **Reconstructs** the published best-fit (149 dPIE potentials, read straight from `best.par`), **verifies** it reproduces the observed multiple images, and composes the **refit** — the same free parameters and priors as `input.par` — for a from-scratch PyAutoLens fit you can compare with the paper's Table 3. | + +## The dictionary + +| Lenstool | PyAutoLens | Notes | +|---|---|---| +| `potentiel` with `profil 81` | `al.mp.dPIEMass` / `al.mp.dPIEMassSph` | Same code lineage: PyAutoLens's dPIE is ported from Lenstool's C source and validated against it numerically (`autolens_workspace_test/scripts/cluster/lenstool_parity.py`). | +| `x_centre`, `y_centre` | `centre=(y, x)` | Same relative-arcsec frame; **x is positive toward West** (see below). | +| `ellipticite` | `ellipticity` | Lenstool's mass ellipticity (a²−b²)/(a²+b²); PyAutoLens converts internally to its (1−q)/(1+q) exactly as `set_lens.c` does. | +| `angle_pos` | `angle_pos` | Degrees, counter-clockwise. | +| `core_radius`, `cut_radius` | `r_core`, `r_cut` | Arcsec. For the `_kpc` variants divide by `cosmology.kpc_per_arcsec_from(redshift=z_lens)`. | +| `v_disp` | `sigma` | **The fiducial σ_LT, not the physical central dispersion** — σ₀ = √(3/2)·σ_LT (Elíasdóttir et al. 2007, App. A). Quote `.par` values unchanged; never feed a measured stellar dispersion here. | +| `potfile` (`mag0`, `sigma`, `cutkpc`, `vdslope 4`, `slope 4`) | shared priors + derived parameters | σᵢ = σ\*·(L/L₀)^0.25, r_cut,ᵢ = r_cut\*·(L/L₀)^0.5, L/L₀ = 10^(0.4(mag0−m)) — the reference-anchored scaling relation, which is also the default in `scripts/cluster/` and `scripts/group/`. | +| `arcs.dat` | `point_datasets.csv` → `al.list_from_csv` | One `PointDataset` per system; redshifts per system. | +| `sigposArcsec` | `positions_noise` column | Identical positional chi-squared definition. | +| source-plane optimization | `al.FitPositionsSource` | Lenstool's default likelihood. | +| image-plane optimization | `al.FitPositionsImagePair*` | The rigorous (and expensive) version; see `scripts/cluster/likelihood_function.py`. | +| `best.par` | the max-likelihood instance of the fit | And `bayes.dat` ↔ the Nautilus samples. | +| cosmology block (`H0 70`, `omegaM 0.3`) | `al.cosmo.FlatLambdaCDM(H0=70.0, Om0=0.3)` | Pass it everywhere — PyAutoLens defaults to Planck15. | + +## Conventions verified against the data (not just documented) + +- **Coordinate frame**: Lenstool's relative frame has x = −ΔRA·cos δ₀·3600 (positive West), + y = +Δδ·3600. `data.py` verifies this by matching `galcat.cat` positions to their `potential` + sections in `best.par` to milliarcseconds. +- **σ convention and ellipticity conversion**: unit-tested in PyAutoGalaxy against the Lenstool C + source, and re-checked numerically by the 6-leg parity script in `autolens_workspace_test`. +- **End-to-end**: `modeling.py`'s Verification I ray-traces the 60 observed images through the + reconstructed 149-profile tracer — every system collapses to a compact source-plane group + (median RMS 0.07″, all below 0.29″ — consistent with the published image-plane RMS of 0.32″ + through typical magnifications), and the optional Verification II forward-solve checks the + image-plane RMS directly. One convention this exercise pinned down the honest way: PyAutoLens's + multi-plane tracer normalizes profile deflections to its **final** plane, so `from_lenstool` + must be given the highest source redshift in the system as `redshift_source`. + +## What's different (and why you might care) + +- **Sources are sampled, not eliminated.** Lenstool removes source positions analytically inside + its optimizer; PyAutoLens samples them (or, with `FitPositionsImagePair*`, solves the forward + problem per likelihood call). You see the joint posterior, including source–mass covariances. +- **The likelihood is yours to choose.** Source-plane χ² reproduces Lenstool; image-plane χ² with + proper handling of image multiplicity is one line away, and the full likelihood is documented + step by step (`scripts/cluster/likelihood_function.py`). +- **Everything after the point-source fit is in the same framework**: extended-source (arc) + modeling with pixelized sources, JAX acceleration of the point solver, Bayesian evidence for + model comparison, and the cluster visualization tools (per-source-plane critical curves and + caustics — `aplt.plot_critical_curves`) used throughout these scripts. + +## Attribution + +The SMACS J0723 model files are by Mahler et al. (2023), downloaded at runtime from their public +repository — cite them (and RELICS, Coe et al. 2019, for the imaging) in any work that uses this +example. The dPIE profile implementation in PyAutoLens derives from Lenstool's C code (Kneib, +Jullo et al.); see the profile docstrings for references. diff --git a/notebooks/cluster/lenstool/data.ipynb b/notebooks/cluster/lenstool/data.ipynb new file mode 100644 index 000000000..4edcd1ff7 --- /dev/null +++ b/notebooks/cluster/lenstool/data.ipynb @@ -0,0 +1,520 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lenstool Users: Data Preparation (SMACS J0723)\n", + "==============================================\n", + "\n", + "This script prepares everything needed to repeat a *published Lenstool cluster analysis* in\n", + "**PyAutoLens**: the strong-lensing model of the first JWST cluster, SMACS J0723.3-7327, by\n", + "Mahler et al. 2023 (ApJ 945, 49; arXiv:2207.07101).\n", + "\n", + "Mahler et al. released their complete Lenstool workflow \u2014 the optimized parameter file\n", + "(``best.par``), the model definition (``input.par``), the multiple-image constraints\n", + "(``arcs.dat``) and the cluster-member catalogue (``galcat.cat``) \u2014 at\n", + "https://github.com/guillaumemahler/SMACS0723-mahler2022 (ICLv2 release). This script downloads\n", + "those files (plus a public RELICS HST image for visualization), translates them into the CSV\n", + "formats the PyAutoLens cluster workflow reads, and records every unit / sign convention involved,\n", + "so that ``modeling.py`` can (a) reconstruct the published best-fit model exactly and (b) refit it\n", + "from scratch.\n", + "\n", + "If you are a Lenstool user: each output CSV corresponds to one input you already maintain \u2014\n", + "\n", + " - ``arcs.dat`` \u2192 ``point_datasets.csv`` (multiple-image positions + redshifts + noise)\n", + " - ``galcat.cat`` \u2192 ``members.csv`` (member catalogue: centres + shape/mag properties,\n", + " the ``al.galaxy_table_from_csv`` schema)\n", + " - ``best.par`` \u2192 ``mass.csv`` (every optimized ``potential`` section as one\n", + " ``dPIEMassLenstool`` row of the canonical named-galaxy model CSV \u2014 **the .par file as a\n", + " table**, read back with ``al.galaxy_models_from_csv`` like every other cluster dataset)\n", + "\n", + "__Attribution__\n", + "\n", + "The model files are \u00a9 Mahler et al. (GPL-licensed repository) and are downloaded at runtime with\n", + "attribution rather than redistributed here. Cite Mahler et al. 2023 (and RELICS, Coe et al. 2019,\n", + "for the imaging) in any work that uses them.\n", + "\n", + "__Contents__\n", + "\n", + "- **Paths + URLs:** where the Mahler files and the RELICS image live.\n", + "- **Downloads:** cached downloads of the four Lenstool files and the F814W image.\n", + "- **Coordinate Convention:** Lenstool's relative-arcsec frame, verified against the data.\n", + "- **Parse arcs.dat:** the multiple-image constraints, with per-system redshifts.\n", + "- **Parse galcat.cat:** the 144 red-sequence cluster members.\n", + "- **Parse best.par:** all 149 optimized dPIE potentials + model-optimized source redshifts.\n", + "- **Image Cutout:** a field-sized cutout of the RELICS F814W image in the Lenstool frame.\n", + "- **Write CSVs:** the PyAutoLens-side data products.\n", + "\n", + "__Coordinate Convention__\n", + "\n", + "Lenstool works in a relative tangent-plane frame centred on the ``runmode`` ``reference``\n", + "coordinate (here RA0, Dec0 = 110.826989, -73.454723):\n", + "\n", + " x_lt = -(RA - RA0) * cos(Dec0) * 3600 [arcsec, positive toward WEST]\n", + " y_lt = +(Dec - Dec0) * 3600 [arcsec, positive toward North]\n", + "\n", + "This is verified against the data itself: member #1 of ``galcat.cat`` (RA = 110.8355945) maps to\n", + "x = -8.820\", matching its ``potential`` section in ``best.par`` (x_centre = -8.822). Everything\n", + "this script writes uses (y, x) = (y_lt, x_lt), so all PyAutoLens positions, model centres and\n", + "critical curves live in Lenstool's own frame and can be compared number-for-number with the\n", + "``.par`` file. Only the *sign of x vs RA* differs from the usual sky convention \u2014 remember it if\n", + "you overlay results on a WCS-aligned image.\n", + "\n", + "__Position Uncertainties__\n", + "\n", + "Lenstool's positional likelihood uses a single ``sigposArcsec`` from ``input.par`` \u2014 here\n", + "0.44271887\" \u2014 not the per-image shape columns of ``arcs.dat``. We propagate that value into the\n", + "``positions_noise`` column of ``point_datasets.csv`` so the PyAutoLens chi-squared is defined\n", + "identically.\n", + "\n", + "__Source Redshifts__\n", + "\n", + "Five systems have spectroscopic redshifts carried by ``arcs.dat`` (systems 1, 2, 5: MUSE grism\n", + "values 1.449, 1.3779, 1.425; system 3: MUSE 1.9914; system 7: JWST/NIRSpec 5.17). The other 16\n", + "systems were optimized by Lenstool (``z_m_limit``); their best-fit values are read from the\n", + "``best.par`` header and treated as *fixed* inputs here \u2014 the refit in ``modeling.py`` frees the\n", + "mass model but not the source redshifts, which keeps the comparison to the published mass\n", + "parameters clean (and the multi-plane solve tractable)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "import re\n", + "import urllib.request\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths + URLs__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "DATASET_PATH = Path(\"dataset\") / \"cluster\" / \"smacs0723\"\n", + "LENSTOOL_PATH = DATASET_PATH / \"lenstool_model\"\n", + "\n", + "MAHLER_BASE = (\n", + " \"https://raw.githubusercontent.com/guillaumemahler/SMACS0723-mahler2022/main/ICLv2\"\n", + ")\n", + "MAHLER_FILES = [\"arcs.dat\", \"galcat.cat\", \"input.par\", \"best.par\", \"README.txt\"]\n", + "\n", + "RELICS_F814W_URL = (\n", + " \"https://archive.stsci.edu/hlsps/relics/smacs0723m73/images/60mas-resolution/\"\n", + " \"hlsp_relics_hst_acs-60mas_smacs0723m73_f814w_v1_drc.fits\"\n", + ")\n", + "\n", + "REFERENCE_RA = 110.826989\n", + "REFERENCE_DEC = -73.454723\n", + "\n", + "MAG0 = 19.12 # potfile reference magnitude (input.par)\n", + "SIGPOS_ARCSEC = 0.44271887242357316 # input.par sigposArcsec\n", + "\n", + "CUTOUT_ARCSEC = 220.0 # full width of the image cutout\n", + "PIXEL_SCALE = 0.06 # RELICS 60mas imaging" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Downloads__\n", + "\n", + "Each file is downloaded once and cached on disk; delete a file to force a re-download. The 96 MB\n", + "F814W mosaic is only fetched if the cutout does not already exist (and can be deleted afterwards \u2014\n", + "set ``KEEP_FULL_MOSAIC = False``)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "KEEP_FULL_MOSAIC = False\n", + "\n", + "LENSTOOL_PATH.mkdir(parents=True, exist_ok=True)\n", + "\n", + "for name in MAHLER_FILES:\n", + " path = LENSTOOL_PATH / name\n", + " if not path.exists():\n", + " print(f\"Downloading {name} from the Mahler et al. repository...\")\n", + " urllib.request.urlretrieve(f\"{MAHLER_BASE}/{name}\", path)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Coordinate Convention__ (see module docstring)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def lenstool_yx_from(ra: float, dec: float) -> tuple:\n", + " \"\"\"Convert (RA, Dec) in degrees to Lenstool's relative (y, x) in arcsec.\"\"\"\n", + " x = -(ra - REFERENCE_RA) * np.cos(np.deg2rad(REFERENCE_DEC)) * 3600.0\n", + " y = (dec - REFERENCE_DEC) * 3600.0\n", + " return (y, x)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parse best.par__\n", + "\n", + "``best.par`` is Lenstool's optimized output: a header of comment lines carrying the per-system\n", + "optimized redshifts (``# z: dlsds:<...>``), then one ``potential`` section per mass\n", + "component. This cluster has 149: five individually-optimized halos (labelled O1-O5: the\n", + "cluster-scale halo, the BCG, two light-concentration clumps and one galaxy modelled outside the\n", + "scaling relation) and 144 cluster members whose parameters Lenstool derived from the scaling\n", + "relation. Every section carries the *same five numbers you would type into PyAutoLens*:\n", + "\n", + " Lenstool ``potential`` \u2192 ``dPIEMass.from_lenstool`` argument\n", + " ---------------------------------------------------------------------\n", + " x_centre / y_centre [arcsec] \u2192 centre=(y, x)\n", + " ellipticity (a\u00b2-b\u00b2)/(a\u00b2+b\u00b2) \u2192 ellipticity\n", + " angle_pos [deg] \u2192 angle_pos\n", + " core_radius / cut_radius [arcsec] \u2192 r_core / r_cut\n", + " v_disp [km/s] (sigma_LT!) \u2192 sigma\n", + " z_lens \u2192 redshift_object" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def parse_best_par(path: Path) -> tuple:\n", + " \"\"\"Return (halo_list, z_by_system) parsed from a Lenstool best.par file.\"\"\"\n", + " text = path.read_text()\n", + "\n", + " z_by_system = {}\n", + " for match in re.finditer(r\"^#(\\d+)\\.0 z:([\\d.]+)\", text, re.M):\n", + " z_by_system[int(match.group(1))] = float(match.group(2))\n", + "\n", + " halos = []\n", + " for section in re.finditer(r\"^potential\\s+(\\S+)\\n(.*?)^\\s*end\", text, re.M | re.S):\n", + " label, body = section.group(1), section.group(2)\n", + " params = dict(re.findall(r\"^\\s*(\\S+)\\s+(-?[\\d.]+)\", body, re.M))\n", + " halos.append(\n", + " {\n", + " \"label\": label,\n", + " \"profile\": int(float(params[\"profile\"])),\n", + " \"x\": float(params[\"x_centre\"]),\n", + " \"y\": float(params[\"y_centre\"]),\n", + " \"ellipticity\": float(params[\"ellipticity\"]),\n", + " \"angle_pos\": float(params[\"angle_pos\"]),\n", + " \"r_core\": float(params[\"core_radius\"]),\n", + " \"r_cut\": float(params[\"cut_radius\"]),\n", + " \"sigma\": float(params[\"v_disp\"]),\n", + " \"z_lens\": float(params[\"z_lens\"]),\n", + " }\n", + " )\n", + " return halos, z_by_system\n", + "\n", + "\n", + "halos, z_by_system = parse_best_par(LENSTOOL_PATH / \"best.par\")\n", + "\n", + "named_halos = [h for h in halos if h[\"label\"].startswith(\"O\")]\n", + "member_halos = [h for h in halos if not h[\"label\"].startswith(\"O\")]\n", + "\n", + "print(\n", + " f\"best.par: {len(halos)} potentials \"\n", + " f\"({len(named_halos)} individually-optimized + {len(member_halos)} scaling members); \"\n", + " f\"{len(z_by_system)} model-optimized source redshifts.\"\n", + ")\n", + "assert all(h[\"profile\"] == 81 for h in halos), \"expected dPIE (profile 81) throughout\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parse arcs.dat__\n", + "\n", + "One row per multiple image: ``id RA Dec a b theta z mag``. The id encodes ``system.image``. A z of\n", + "0.0 means \"no spectroscopic redshift\" \u2014 those systems take their model-optimized value from the\n", + "best.par header instead." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_rows = []\n", + "for line in (LENSTOOL_PATH / \"arcs.dat\").read_text().splitlines():\n", + " parts = line.split()\n", + " if len(parts) < 7 or parts[0].startswith(\"#\"):\n", + " continue\n", + " system = int(parts[0].split(\".\")[0])\n", + " y, x = lenstool_yx_from(ra=float(parts[1]), dec=float(parts[2]))\n", + " z_spec = float(parts[6])\n", + " redshift = z_spec if z_spec > 0 else z_by_system[system]\n", + " image_rows.append((system, y, x, redshift, z_spec > 0))\n", + "\n", + "systems = sorted({row[0] for row in image_rows})\n", + "n_spec = len({r[0] for r in image_rows if r[4]})\n", + "print(\n", + " f\"arcs.dat: {len(image_rows)} images across {len(systems)} systems \"\n", + " f\"({n_spec} systems with spectroscopic redshifts).\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parse galcat.cat__\n", + "\n", + "One row per cluster member: ``id RA Dec a b theta mag [z]``. Commented rows (leading ``#``) were\n", + "excluded by Mahler et al. and are skipped here too. The shape columns give the light ellipse:\n", + "Lenstool converts (a, b) to its mass ellipticity as e = (a\u00b2 - b\u00b2) / (a\u00b2 + b\u00b2), and we store that\n", + "value so the refit can fix each member's geometry exactly as Lenstool did. The luminosity that\n", + "enters the scaling relation is *relative* to the potfile pivot mag0 = 19.12:\n", + "\n", + " L / L0 = 10 ** (0.4 * (mag0 - mag))" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "member_rows = []\n", + "for line in (LENSTOOL_PATH / \"galcat.cat\").read_text().splitlines():\n", + " parts = line.split()\n", + " if len(parts) < 7 or parts[0].startswith(\"#\"):\n", + " continue\n", + " y, x = lenstool_yx_from(ra=float(parts[1]), dec=float(parts[2]))\n", + " a, b = float(parts[3]), float(parts[4])\n", + " ellipticity = (a**2 - b**2) / (a**2 + b**2)\n", + " angle_pos = float(parts[5])\n", + " mag = float(parts[6])\n", + " luminosity = 10.0 ** (0.4 * (MAG0 - mag))\n", + " member_rows.append((y, x, ellipticity, angle_pos, mag, luminosity))\n", + "\n", + "print(f\"galcat.cat: {len(member_rows)} cluster members (relative to mag0 = {MAG0}).\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Write CSVs__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "with open(DATASET_PATH / \"point_datasets.csv\", \"w\") as f:\n", + " f.write(\"name,y,x,positions_noise,redshift\\n\")\n", + " for system, y, x, redshift, _ in image_rows:\n", + " f.write(f\"point_{system},{y:.6f},{x:.6f},{SIGPOS_ARCSEC},{redshift}\\n\")\n", + "\n", + "al.galaxy_table_to_csv(\n", + " centres=[(r[0], r[1]) for r in member_rows],\n", + " luminosities=[r[5] for r in member_rows],\n", + " file_path=DATASET_PATH / \"members.csv\",\n", + " properties={\n", + " \"ellipticity\": [r[2] for r in member_rows],\n", + " \"angle_pos\": [r[3] for r in member_rows],\n", + " \"mag\": [r[4] for r in member_rows],\n", + " },\n", + ")\n", + "\n", + "# The whole optimized model \u2014 5 named halos + 144 scaling members \u2014 becomes ONE canonical\n", + "# ``mass.csv``: each ``potential`` section is a ``dPIEMassLenstool`` row whose columns are the\n", + "# ``.par`` keywords verbatim (sigma, r_core, r_cut, ellipticity, angle_pos) plus the run's\n", + "# redshifts and cosmology as flat values. ``modeling.py`` reads it back with the same\n", + "# ``al.galaxy_models_from_csv`` call used throughout ``scripts/cluster/``.\n", + "#\n", + "# Every profile is normalized against the tracer's FINAL source plane (the multi-plane\n", + "# convention modeling.py explains) using the run's own cosmology (H0=70, Om0=0.3).\n", + "Z_FINAL_PLANE = max(z_by_system.values())\n", + "\n", + "profiles_by_galaxy = {}\n", + "for h in halos:\n", + " name = h[\"label\"] if h[\"label\"].startswith(\"O\") else f\"member_{h['label']}\"\n", + " profiles_by_galaxy[name] = {\n", + " \"mass\": al.mp.dPIEMassLenstool(\n", + " centre=(h[\"y\"], h[\"x\"]),\n", + " ellipticity=h[\"ellipticity\"],\n", + " angle_pos=h[\"angle_pos\"],\n", + " sigma=h[\"sigma\"],\n", + " r_core=h[\"r_core\"],\n", + " r_cut=h[\"r_cut\"],\n", + " redshift_object=h[\"z_lens\"],\n", + " redshift_source=Z_FINAL_PLANE,\n", + " H0=70.0,\n", + " Om0=0.3,\n", + " )\n", + " }\n", + "\n", + "al.galaxy_models_to_csv(\n", + " profiles_by_galaxy,\n", + " DATASET_PATH / \"mass.csv\",\n", + " family=\"mass\",\n", + " redshifts={name: 0.39 for name in profiles_by_galaxy},\n", + ")\n", + "\n", + "print(\"Wrote point_datasets.csv, members.csv, mass.csv (149 dPIEMassLenstool rows).\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image Cutout__\n", + "\n", + "A 220\" cutout of the RELICS F814W mosaic centred on the Lenstool reference coordinate, stored\n", + "**flipped in x** so that array coordinates match the Lenstool frame defined above (x positive\n", + "toward West). The cutout exists purely for visualization \u2014 none of the modeling uses pixel data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cutout_path = DATASET_PATH / \"data.fits\"\n", + "\n", + "import os\n", + "\n", + "if os.environ.get(\"PYAUTO_SMALL_DATASETS\") == \"1\":\n", + " print(\n", + " \"PYAUTO_SMALL_DATASETS=1: skipping the 96 MB RELICS mosaic download / cutout \"\n", + " \"(visualization-only product; the modeling data products above are complete).\"\n", + " )\n", + "elif not cutout_path.exists():\n", + " from astropy.io import fits\n", + " from astropy.wcs import WCS\n", + "\n", + " mosaic_path = DATASET_PATH / \"f814w_mosaic.fits\"\n", + " if not mosaic_path.exists():\n", + " print(\"Downloading the RELICS F814W mosaic (96 MB, one-off)...\")\n", + " urllib.request.urlretrieve(RELICS_F814W_URL, mosaic_path)\n", + "\n", + " with fits.open(mosaic_path) as hdul:\n", + " data = hdul[0].data\n", + " wcs = WCS(hdul[0].header)\n", + "\n", + " x_pix, y_pix = wcs.world_to_pixel_values(REFERENCE_RA, REFERENCE_DEC)\n", + " half = int(CUTOUT_ARCSEC / 2.0 / PIXEL_SCALE)\n", + " y0, x0 = int(round(float(y_pix))), int(round(float(x_pix)))\n", + " cutout = data[y0 - half : y0 + half, x0 - half : x0 + half]\n", + "\n", + " # Flip x so +x points West, matching the Lenstool relative frame.\n", + " cutout = cutout[:, ::-1]\n", + "\n", + " header = fits.Header()\n", + " header[\"PIXSCALE\"] = PIXEL_SCALE\n", + " header[\"COMMENT\"] = \"RELICS F814W cutout in the Lenstool relative frame (+x West)\"\n", + " fits.PrimaryHDU(cutout.astype(np.float32), header=header).writeto(\n", + " cutout_path, overwrite=True\n", + " )\n", + " print(\n", + " f'Wrote {cutout_path} ({cutout.shape[0]}x{cutout.shape[1]} @ {PIXEL_SCALE}\"/px).'\n", + " )\n", + "\n", + " if not KEEP_FULL_MOSAIC:\n", + " mosaic_path.unlink()\n", + " print(\"Deleted the full mosaic (set KEEP_FULL_MOSAIC = True to keep it).\")\n", + "\n", + "print(\"\\nData preparation complete.\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/cluster/lenstool/modeling.ipynb b/notebooks/cluster/lenstool/modeling.ipynb new file mode 100644 index 000000000..6c7993568 --- /dev/null +++ b/notebooks/cluster/lenstool/modeling.ipynb @@ -0,0 +1,679 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lenstool Users: Cluster Modeling (SMACS J0723)\n", + "==============================================\n", + "\n", + "**If you model galaxy clusters with Lenstool, this script is for you.** It repeats a published\n", + "Lenstool analysis \u2014 Mahler et al. 2023's model of SMACS J0723, the first JWST cluster \u2014 in\n", + "**PyAutoLens**, in three steps:\n", + "\n", + " 1. **Reconstruct** the published best-fit model, reading the 149 dPIE potentials of ``best.par``\n", + " directly into PyAutoLens profiles (no re-fitting, no conversion by hand).\n", + " 2. **Verify** that PyAutoLens reproduces the published lensing: the observed multiple images\n", + " trace back to compact source-plane groups, and (optionally) forward-solving the lens equation\n", + " recovers the observed image positions at the published RMS (0.32\").\n", + " 3. **Refit** the cluster from scratch, composing the same model Lenstool optimized \u2014 same free\n", + " parameters, same priors as ``input.par``, same positional likelihood \u2014 so the posterior can be\n", + " compared with Table 3 of the paper number for number.\n", + "\n", + "Run ``data.py`` first; it downloads the public model files and writes the CSVs read here.\n", + "\n", + "__The .par \u2192 PyAutoLens dictionary__\n", + "\n", + " Lenstool PyAutoLens\n", + " ------------------------------------------------------------------------------\n", + " ``potentiel`` / profil 81 ``al.mp.dPIEMass`` (elliptical dPIE)\n", + " x_centre, y_centre [arcsec] ``centre=(y, x)`` \u2014 same relative frame (data.py)\n", + " ellipticite = (a\u00b2-b\u00b2)/(a\u00b2+b\u00b2) ``ellipticity`` (converted internally to (1-q)/(1+q))\n", + " angle_pos [deg] ``angle_pos``\n", + " core_radius / cut_radius [arcsec] ``r_core`` / ``r_cut``\n", + " v_disp [km/s] ``sigma`` \u2014 Lenstool's *fiducial* sigma_LT, see below\n", + " z_lens ``redshift_object``\n", + " ``potfile`` (member scaling) shared priors + derived per-member parameters (below)\n", + " ``arcs.dat`` ``point_datasets.csv`` \u2192 ``al.PointDataset`` list\n", + " sigposArcsec ``positions_noise`` column (same chi-squared)\n", + " source-plane optimization ``al.FitPositionsSource`` (this script)\n", + " image-plane optimization ``al.FitPositionsImagePair*`` (heavier; see guides)\n", + " ``best.par`` the max-likelihood instance of the PyAutoLens fit\n", + "\n", + "__Three traps to know about__\n", + "\n", + " - **sigma is the fiducial velocity dispersion sigma_LT**, not the physical central velocity\n", + " dispersion: sigma_0 = sqrt(3/2) * sigma_LT (Eliasdottir et al. 2007, App. A). PyAutoLens's\n", + " ``from_lenstool`` / ``dPIEMassLenstool`` take sigma_LT \u2014 quote ``v_disp`` values unchanged.\n", + " Feeding a *measured* stellar velocity dispersion here overestimates the mass by 50%.\n", + " - **The x axis points West** in Lenstool's relative frame (data.py verifies this against the\n", + " data). All coordinates in this script live in that frame, so numbers compare directly to the\n", + " ``.par`` file; flip x when overlaying on a WCS-aligned image.\n", + " - **Radii in .par files come in arcsec and kpc variants** (``core_radius`` vs\n", + " ``core_radius_kpc``). PyAutoLens profiles work in arcsec; the kpc \u2192 arcsec conversion uses the\n", + " Lenstool run's own cosmology (below), 1\" = 5.313 kpc at z = 0.39.\n", + "\n", + "__Cosmology__\n", + "\n", + "Mahler et al. use flat LCDM with H0 = 70, Omega_m = 0.3 \u2014 the Lenstool-typical choice, *not*\n", + "PyAutoLens's Planck15 default. It is passed explicitly everywhere below; forgetting it shifts\n", + "D_LS/D_S and hence every mass normalization at the percent level.\n", + "\n", + "__Contents__\n", + "\n", + "- **Load Data:** the CSVs written by ``data.py``.\n", + "- **The Published Model, Reconstructed:** 149 ``from_lenstool`` profiles + 21 point sources.\n", + "- **Verification I \u2014 source-plane compactness:** observed images trace to tight source groups.\n", + "- **Verification II \u2014 image-plane RMS (optional):** forward-solve vs the published 0.32\".\n", + "- **Critical Curves (optional):** the per-source-plane critical curves over the HST image.\n", + "- **The Refit:** the same free parameters and priors as ``input.par``, fit with Nautilus.\n", + "\n", + "__Runtime__\n", + "\n", + "Reconstruction + Verification I run in a couple of minutes. Verification II forward-solves the\n", + "lens equation for 21 sources through a 149-profile multi-plane tracer (one-off JAX compile of a\n", + "few minutes; enable with ``RUN_IMAGE_PLANE = True``), and the critical-curve figure costs ~10+\n", + "minutes per plane in numpy (``MAKE_FIGURES = True``). The full\n", + "refit is a production job (tens of hours on a workstation; Mahler et al. report k = 46 free\n", + "parameters \u2014 ours is 30 mass + 42 source-position parameters with the 16 photometric source\n", + "redshifts held at their published values). ``PYAUTO_TEST_MODE=2`` exercises the full composition\n", + "without running the sampler." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "RUN_IMAGE_PLANE = False" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cosmology__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cosmology = al.cosmo.FlatLambdaCDM(H0=70.0, Om0=0.3)\n", + "\n", + "Z_LENS = 0.39\n", + "\n", + "kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=Z_LENS)\n", + "print(f'1\" = {kpc_per_arcsec:.3f} kpc at z = {Z_LENS} (Lenstool-run cosmology)')" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Load Data__\n", + "\n", + "``data.py`` wrote one CSV per Lenstool input (see its docstring for the full conventions):\n", + "\n", + " - ``point_datasets.csv`` \u2014 the 60 multiple images of 21 sources (``arcs.dat``), with Lenstool's\n", + " ``sigposArcsec`` as the position noise and per-system redshifts (spectroscopic where they\n", + " exist, the model-optimized values of ``best.par`` otherwise).\n", + " - ``mass.csv`` \u2014 the complete optimized mass model in the **canonical named-galaxy CSV**\n", + " (the same ``al.galaxy_models_from_csv`` format every cluster script uses): 149 rows of\n", + " ``profile_class = dPIEMassLenstool``, one per ``potential`` section of ``best.par``, whose\n", + " columns are the ``.par`` keywords verbatim. The five individually-optimized halos are named\n", + " O1 (cluster-scale), O2 (BCG), O3 (\"dNW\"), O4 (\"ICL\"), O5 (\"eCM\"); the 144 scaling members\n", + " are ``member_``.\n", + " - ``members.csv`` \u2014 the member *catalogue* (``galcat.cat``) in the ``al.galaxy_table_from_csv``\n", + " schema: centres + luminosities plus ``ellipticity`` / ``angle_pos`` / ``mag`` property\n", + " columns (the refit derives member masses from these + the shared scaling parameters, as\n", + " ``potfile`` does)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / \"cluster\" / \"smacs0723\"\n", + "\n", + "dataset_list = al.list_from_csv(file_path=dataset_path / \"point_datasets.csv\")\n", + "print(f\"{len(dataset_list)} point-source systems loaded.\")\n", + "\n", + "\n", + "mass_table = al.galaxy_models_from_csv(dataset_path / \"mass.csv\", family=\"mass\")\n", + "members_table = al.galaxy_table_from_csv(file_path=dataset_path / \"members.csv\")\n", + "\n", + "print(\n", + " f\"mass.csv: {len(mass_table.rows)} dPIEMassLenstool rows | \"\n", + " f\"members.csv: {len(members_table.luminosities)} catalogue members\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__The Published Model, Reconstructed__\n", + "\n", + "Every ``potential`` section of ``best.par`` becomes one ``al.mp.dPIEMass`` via ``from_lenstool``\n", + "\u2014 the arguments are the ``.par`` keywords, verbatim. This is the whole point of the Lenstool-native\n", + "API: nothing is transcribed by hand, and the sqrt(3/2) sigma convention, the ellipticity\n", + "conversion and the D_LS/D_S normalization are handled (and unit-tested) inside PyAutoLens.\n", + "\n", + "One subtlety that *will* bite anyone porting a model by hand: Lenstool normalizes each halo's\n", + "deflection per source plane at ray-trace time, whereas a PyAutoLens profile is normalized once, at\n", + "construction, against a chosen source redshift. In a multi-plane ``Tracer`` the cosmological\n", + "scaling factors beta are computed **relative to the tracer's final plane** \u2014 so every profile must\n", + "be normalized to the *highest* source redshift in the system (here the z = 11.76 candidate).\n", + "Normalize to any other plane and every deflection is off by a constant D-ratio (for z = 2 it is\n", + "~15%, which smears each source group by ~2\" \u2014 we found out the honest way). Both codes then reduce\n", + "to the same beta = D_LS/D_S ratios in the same cosmology." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "Z_REF_SOURCE = max(float(dataset.redshift) for dataset in dataset_list)\n", + "\n", + "# One call: every mass.csv row instantiates its dPIEMassLenstool with the .par values \u2014\n", + "# the redshift_source (final-plane) normalization and the run's H0/Om0 travel inside the\n", + "# CSV columns, so nothing here needs to remember them.\n", + "lens_galaxies = list(al.galaxies_from_csv_tables(mass_table).values())\n", + "\n", + "print(f\"Reconstructed {len(lens_galaxies)} dPIE mass components from mass.csv.\")\n", + "\n", + "source_galaxies = [\n", + " al.Galaxy(\n", + " redshift=dataset.redshift, **{dataset.name: al.ps.Point(centre=(0.0, 0.0))}\n", + " )\n", + " for dataset in dataset_list\n", + "]\n", + "\n", + "tracer = al.Tracer(galaxies=lens_galaxies + source_galaxies, cosmology=cosmology)\n", + "print(\n", + " f\"Tracer has {len(tracer.planes)} planes (1 lens + {len(tracer.planes)-1} source).\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Verification I \u2014 Source-Plane Compactness__\n", + "\n", + "Lenstool's model was optimized so the observed images of each system meet at a single source\n", + "position. Ray-tracing the observed image positions through the reconstructed tracer must therefore\n", + "produce *compact* per-system source-plane groups \u2014 this is the fast, solver-free check that the\n", + "model transferred correctly (a wrong sigma convention, ellipticity definition or cosmology blows\n", + "these numbers up immediately).\n", + "\n", + "For each system, every observed image is traced back to that system's source plane and the RMS\n", + "scatter about the group's centroid is reported. With the published image-plane RMS of 0.32\" and typical\n", + "magnifications of a few, per-system values of ~0.1\" (we measure a median of 0.07\", every system\n", + "below 0.29\") confirm the model transferred exactly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scatters = []\n", + "for dataset in dataset_list:\n", + " positions = al.Grid2DIrregular(np.atleast_2d(np.asarray(dataset.positions)))\n", + " plane_index = tracer.plane_index_via_redshift_from(redshift=dataset.redshift)\n", + " traced = tracer.traced_grid_2d_list_from(grid=positions)[plane_index]\n", + " traced = np.asarray(traced)\n", + " centroid = traced.mean(axis=0)\n", + " rms = float(np.sqrt(np.mean(np.sum((traced - centroid) ** 2, axis=1))))\n", + " scatters.append(rms)\n", + " print(\n", + " f\" {dataset.name:>9} (z={float(dataset.redshift):6.3f}, \"\n", + " f'{len(traced)} images): source-plane rms = {rms:6.4f}\"'\n", + " )\n", + "\n", + "print(f'Median source-plane rms: {np.median(scatters):.4f}\"')" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Verification II \u2014 Image-Plane RMS (optional)__\n", + "\n", + "The definitive check forward-solves the lens equation: from each system's model source position\n", + "(the traced centroid above), find every image position the reconstructed mass model predicts, pair\n", + "predictions with observations, and measure the image-plane RMS. Mahler et al. report 0.32\". This\n", + "uses the same JAX-compiled ``PointSolver`` as the cluster simulator; the one-off compile takes\n", + "several minutes for a 149-profile tracer, so it is gated behind ``RUN_IMAGE_PLANE``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if RUN_IMAGE_PLANE:\n", + " import jax\n", + "\n", + " al.jax.register_tracer_classes(tracer)\n", + "\n", + " grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.7)\n", + " solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.01, use_jax=True\n", + " )\n", + "\n", + " offsets = []\n", + " for dataset in dataset_list:\n", + " positions = al.Grid2DIrregular(np.atleast_2d(np.asarray(dataset.positions)))\n", + " plane_index = tracer.plane_index_via_redshift_from(redshift=dataset.redshift)\n", + " traced = np.asarray(\n", + " tracer.traced_grid_2d_list_from(grid=positions)[plane_index]\n", + " )\n", + " source_centre = tuple(traced.mean(axis=0))\n", + "\n", + " predicted = solver.solve(\n", + " tracer=tracer,\n", + " source_plane_coordinate=source_centre,\n", + " plane_redshift=float(dataset.redshift),\n", + " )\n", + " predicted = np.atleast_2d(np.asarray(predicted))\n", + "\n", + " for obs in np.atleast_2d(np.asarray(dataset.positions)):\n", + " offsets.append(np.min(np.linalg.norm(predicted - obs, axis=1)))\n", + " print(f\" {dataset.name}: {len(predicted)} predicted images\")\n", + "\n", + " rms_image_plane = float(np.sqrt(np.mean(np.array(offsets) ** 2)))\n", + " print(f'Image-plane rms = {rms_image_plane:.3f}\" (published: 0.32\")')" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Critical Curves (optional)__\n", + "\n", + "The cluster plotters draw each source plane's critical curves over the HST cutout \u2014 with 21\n", + "source planes we plot a representative subset (the lowest-redshift plane and the famous z = 11.8\n", + "candidate's plane). At cluster scale each plane's curves differ visibly; this is the multi-plane\n", + "structure a single-plane \"the critical curve\" plot hides.\n", + "\n", + "Fair warning on cost: a critical-curve evaluation walks the whole grid through all 149 dPIE\n", + "profiles several times \u2014 ~10+ minutes per figure on a laptop \u2014 so it is gated behind\n", + "``MAKE_FIGURES`` like the forward-solve. Enable both when producing final figures on real\n", + "hardware." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "MAKE_FIGURES = False\n", + "\n", + "if MAKE_FIGURES:\n", + " image = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=0.06\n", + " )\n", + "\n", + " # 150x150 @ 1\" resolves the cluster-scale curves; refine when producing final figures.\n", + " viz_grid = al.Grid2D.uniform(shape_native=(150, 150), pixel_scales=1.0)\n", + "\n", + " plane_redshifts = [float(p.redshift) for p in tracer.planes]\n", + " plane_indices = [\n", + " plane_redshifts.index(min(r for r in plane_redshifts if r > Z_LENS)),\n", + " len(plane_redshifts) - 1,\n", + " ]\n", + "\n", + " aplt.plot_critical_curves(\n", + " tracer,\n", + " grid=viz_grid,\n", + " image=image,\n", + " plane_indices=plane_indices,\n", + " output_path=str(dataset_path),\n", + " output_filename=\"critical_curves_reconstruction\",\n", + " output_format=\"png\",\n", + " )\n", + " print(\"Critical-curve figure written next to the dataset.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__The Refit__\n", + "\n", + "Everything above used the published answer. A real analysis *fits*: the model below reproduces the\n", + "composition Lenstool optimized (``input.par``), using the Lenstool-parameterized profile\n", + "``al.mp.dPIEMassLenstool`` so every free parameter, prior bound and posterior number is in\n", + "Lenstool units:\n", + "\n", + " - **O1, cluster halo**: centre U(-5,5)\" (both axes), ellipticity U(0,0.8), angle U(-90,90),\n", + " r_core U(10,150) kpc, sigma U(300,1200) km/s; r_cut fixed at 1500 kpc. [6 free]\n", + " - **O2, BCG**: geometry fixed to the light; r_core U(0.1,10) kpc, r_cut U(10,500) kpc,\n", + " sigma U(100,500). [3 free]\n", + " - **O3 \"dNW\"**: centre U(best \u00b1 2.6)\", ellipticity U(0,0.8), angle U(-90,90), r_core U(0,10) kpc,\n", + " r_cut U(10,300) kpc, sigma U(0,200). [7 free]\n", + " - **O4 \"ICL\"**: centre U(best \u00b1 10)\", ellipticity U(0,0.8), angle U(-90,90), r_core U(0,50) kpc,\n", + " r_cut U(50,1000) kpc, sigma U(0,700). [7 free]\n", + " - **O5 \"eCM\"**: centre fixed; ellipticity U(0,0.6), angle U(-90,90), r_core U(0,10) kpc,\n", + " r_cut U(10,200) kpc, sigma U(0.1,300). [5 free]\n", + " - **potfile members**: every catalogue member gets a ``dPIEMassLenstool`` with centre, shape and\n", + " angle *fixed to the light* (``galcat.cat``) and its sigma / r_cut derived from two shared free\n", + " parameters exactly as ``potfile`` defines \u2014\n", + "\n", + " sigma_i = sigma_star * (L_i/L0)^0.25 sigma_star ~ U(50, 500) km/s\n", + " r_cut_i = r_cut_star * (L_i/L0)^0.5 r_cut_star ~ U(1, 100) kpc\n", + " r_core_i = 0.15 kpc (fixed)\n", + "\n", + " (``vdslope 4`` and ``slope 4`` in ``input.par`` are these fixed exponents; the same\n", + " reference-anchored convention is the default throughout the PyAutoLens cluster workflow.)\n", + " [2 free for all 146 catalogue members]\n", + "\n", + " - **Sources**: one ``al.ps.Point`` per system with a free centre initialised from the traced\n", + " centroid of its observed images; redshifts fixed (spectroscopic or published model values).\n", + " [42 free]\n", + "\n", + "Total: 30 mass + 42 source-position parameters. Mahler et al.'s k = 46 counts 30 mass + 16 free\n", + "photometric redshifts \u2014 Lenstool eliminates source positions analytically in its optimization,\n", + "PyAutoLens samples them; the guides discuss this difference and the image-plane likelihood that\n", + "removes it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sigma_star = af.UniformPrior(lower_limit=50.0, upper_limit=500.0)\n", + "r_cut_star_kpc = af.UniformPrior(lower_limit=1.0, upper_limit=100.0)\n", + "\n", + "R_CORE_MEMBER = 0.15 / kpc_per_arcsec # potfile corekpc, converted once\n", + "\n", + "member_models = []\n", + "for centre, luminosity, ellipticity, angle_pos in zip(\n", + " members_table.centres,\n", + " members_table.luminosities,\n", + " members_table.properties[\"ellipticity\"],\n", + " members_table.properties[\"angle_pos\"],\n", + "):\n", + " mass = af.Model(al.mp.dPIEMassLenstool)\n", + " mass.centre = tuple(centre)\n", + " mass.ellipticity = ellipticity\n", + " mass.angle_pos = angle_pos\n", + " mass.r_core = R_CORE_MEMBER\n", + " mass.r_cut = (r_cut_star_kpc / float(kpc_per_arcsec)) * luminosity**0.5\n", + " mass.sigma = sigma_star * luminosity**0.25\n", + " mass.redshift_object = Z_LENS\n", + " mass.redshift_source = Z_REF_SOURCE\n", + " mass.H0 = 70.0\n", + " mass.Om0 = 0.3\n", + " member_models.append(af.Model(al.Galaxy, redshift=Z_LENS, mass=mass))\n", + "\n", + "\n", + "# The named halos start as af.Models straight from mass.csv \u2014 the canonical\n", + "# al.galaxy_af_models_from_csv_tables call gives every row's values as fixed\n", + "# defaults, and we promote exactly the parameters input.par optimized to priors\n", + "# (in Lenstool units). Redshifts and H0/Om0 ride in from the CSV already fixed.\n", + "halo_af_models = al.galaxy_af_models_from_csv_tables(mass_table)\n", + "\n", + "\n", + "def halo_model_from(label, limits):\n", + " galaxy_model = halo_af_models[label]\n", + " mass = galaxy_model.mass\n", + " row = next(r for r in mass_table.rows if r.galaxy == label)\n", + " best_y, best_x = row.params[\"centre\"]\n", + "\n", + " if \"centre\" in limits:\n", + " half = limits[\"centre\"]\n", + " mass.centre_0 = af.UniformPrior(\n", + " lower_limit=best_y - half, upper_limit=best_y + half\n", + " )\n", + " mass.centre_1 = af.UniformPrior(\n", + " lower_limit=best_x - half, upper_limit=best_x + half\n", + " )\n", + "\n", + " if \"ellipticity\" in limits:\n", + " mass.ellipticity = af.UniformPrior(0.0, limits[\"ellipticity\"])\n", + " mass.angle_pos = af.UniformPrior(-90.0, 90.0)\n", + "\n", + " lo, hi = limits[\"r_core_kpc\"]\n", + " mass.r_core = af.UniformPrior(\n", + " lo / float(kpc_per_arcsec), hi / float(kpc_per_arcsec)\n", + " )\n", + "\n", + " if \"r_cut_kpc\" in limits:\n", + " lo, hi = limits[\"r_cut_kpc\"]\n", + " mass.r_cut = af.UniformPrior(\n", + " lo / float(kpc_per_arcsec), hi / float(kpc_per_arcsec)\n", + " )\n", + " # else: r_cut stays at its CSV value (O1: fixed 1500 kpc, already in arcsec).\n", + "\n", + " lo, hi = limits[\"sigma\"]\n", + " mass.sigma = af.UniformPrior(lo, hi)\n", + " return galaxy_model\n", + "\n", + "\n", + "halo_models = [\n", + " halo_model_from(\n", + " \"O1\", dict(centre=5.0, ellipticity=0.8, r_core_kpc=(10, 150), sigma=(300, 1200))\n", + " ),\n", + " halo_model_from(\n", + " \"O2\", dict(r_core_kpc=(0.1, 10), r_cut_kpc=(10, 500), sigma=(100, 500))\n", + " ),\n", + " halo_model_from(\n", + " \"O3\",\n", + " dict(\n", + " centre=2.6,\n", + " ellipticity=0.8,\n", + " r_core_kpc=(0, 10),\n", + " r_cut_kpc=(10, 300),\n", + " sigma=(0, 200),\n", + " ),\n", + " ),\n", + " halo_model_from(\n", + " \"O4\",\n", + " dict(\n", + " centre=10.0,\n", + " ellipticity=0.8,\n", + " r_core_kpc=(0, 50),\n", + " r_cut_kpc=(50, 1000),\n", + " sigma=(0, 700),\n", + " ),\n", + " ),\n", + " halo_model_from(\n", + " \"O5\",\n", + " dict(\n", + " ellipticity=0.6, r_core_kpc=(0, 10), r_cut_kpc=(10, 200), sigma=(0.1, 300)\n", + " ),\n", + " ),\n", + "]\n", + "\n", + "source_models = []\n", + "for dataset in dataset_list:\n", + " positions = np.atleast_2d(np.asarray(dataset.positions))\n", + " plane_index = tracer.plane_index_via_redshift_from(redshift=dataset.redshift)\n", + " traced = np.asarray(\n", + " tracer.traced_grid_2d_list_from(grid=al.Grid2DIrregular(positions))[plane_index]\n", + " )\n", + " point = af.Model(al.ps.Point)\n", + " point.centre_0 = af.GaussianPrior(mean=float(traced[:, 0].mean()), sigma=2.0)\n", + " point.centre_1 = af.GaussianPrior(mean=float(traced[:, 1].mean()), sigma=2.0)\n", + " source_models.append(\n", + " af.Model(al.Galaxy, redshift=float(dataset.redshift), **{dataset.name: point})\n", + " )\n", + "\n", + "model = af.Collection(\n", + " halos=af.Collection(halo_models),\n", + " members=af.Collection(member_models),\n", + " sources=af.Collection(source_models),\n", + ")\n", + "\n", + "print(\n", + " f\"Refit model composed: {model.prior_count} free parameters \"\n", + " f\"(30 mass + {2 * len(dataset_list)} source positions).\"\n", + ")\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "Source-plane chi-squared (Lenstool's default likelihood) via ``AnalysisPoint``; the factor-graph\n", + "pattern is identical to ``scripts/cluster/modeling.py``. This is a production-scale job \u2014 run it\n", + "on real hardware, not in a tutorial session. The result's max-likelihood instance is the PyAutoLens\n", + "equivalent of ``best.par``; compare it to Table 3 of Mahler et al. 2023 and to the reconstruction\n", + "above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"cluster\") / \"smacs0723\",\n", + " name=\"lenstool_refit\",\n", + " unique_tag=\"mahler2023\",\n", + " n_live=400,\n", + " number_of_cores=4,\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=1.4),\n", + " pixel_scale_precision=0.025,\n", + ")\n", + "\n", + "analysis_list = [\n", + " al.AnalysisPoint(\n", + " dataset=dataset,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsSource,\n", + " cosmology=cosmology,\n", + " )\n", + " for dataset in dataset_list\n", + "]\n", + "\n", + "import os\n", + "\n", + "if os.environ.get(\"LENSTOOL_EXAMPLE_RUN_FIT\"):\n", + " analysis_factor_list = [\n", + " af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " for analysis in analysis_list\n", + " ]\n", + " factor_graph = af.FactorGraphModel(*analysis_factor_list)\n", + " result_list = search.fit(\n", + " model=factor_graph.global_prior_model, analysis=factor_graph\n", + " )\n", + " print(\"Refit complete \u2014 compare result_list max-likelihood values with Table 3.\")\n", + "else:\n", + " print(\n", + " \"Refit composition validated (the model.info above is the structural pass); set \"\n", + " \"LENSTOOL_EXAMPLE_RUN_FIT=1 to execute the production-scale search \u2014 a 72-parameter \"\n", + " \"factor-graph fit is never smoke-mode material.\"\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished. The README in this folder is the narrative companion: the full .par \u2194 PyAutoLens\n", + "dictionary, the conventions verified here, and where to go next (image-plane likelihood, JAX,\n", + "extended-source modeling of the arcs \u2014 everything Lenstool cannot do is on the other side of\n", + "this door)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/cluster/likelihood_function.ipynb b/notebooks/cluster/likelihood_function.ipynb index 3290c3b81..ba85f4bfe 100644 --- a/notebooks/cluster/likelihood_function.ipynb +++ b/notebooks/cluster/likelihood_function.ipynb @@ -1,1111 +1,1153 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Log Likelihood Function: Cluster Point Source\n", - "==============================================\n", - "\n", - "This script provides a step-by-step guide of the cluster point-source ``log_likelihood_function``,\n", - "the figure-of-merit Nautilus optimises when fitting a cluster lens model to ``point_datasets.csv``.\n", - "\n", - "Cluster point-source modelling has two distinct likelihood flavours:\n", - "\n", - " 1. **Source-plane chi\u00b2** (``FitPositionsSource``) \u2014 ray-trace every observed image-plane position\n", - " back through the lens, and measure the magnification-weighted scatter of the back-traced\n", - " positions around the source-plane centre. Cheap (no forward solve) and JAX-friendly.\n", - "\n", - " 2. **Image-plane chi\u00b2** (``FitPositionsImagePair``) \u2014 forward-solve the model's image-plane\n", - " positions for each source via the ``PointSolver``, pair each model position to the closest\n", - " observed position, and measure the image-plane residuals. More intuitive (residuals in\n", - " arc-seconds), but slower and carries pairing pathologies the source-plane variant avoids.\n", - "\n", - "We walk through both, end to end, with the actual library formulae. The standard cluster model is\n", - "assumed: all lens-plane galaxies at ``z = 0.5``, two background sources at *different* redshifts\n", - "(``z = 1.0`` and ``z = 2.0``). Multi-plane ray tracing therefore applies and we explain how the\n", - "recursive lens equation handles it.\n", - "\n", - "This is a companion to ``scripts/imaging/likelihood_function.py`` and\n", - "``scripts/group/likelihood_function.py``, both of which cover extended-imaging likelihoods. The\n", - "cluster surface is the first one in the workspace tutorial set that uses *point sources*\n", - "exclusively + multi-plane ray tracing + many lens galaxies in the deflection sum.\n", - "\n", - "__Aims__\n", - "\n", - " - Provide a resource that authors can include in papers using cluster modelling, so readers can\n", - " understand the likelihood (including references to the prior literature it builds on) without\n", - " needing to re-derive equations or trace the PyAutoLens source.\n", - " - Make the difference between source-plane and image-plane chi\u00b2 concrete so users can decide\n", - " which flavour fits their dataset + compute budget.\n", - " - Document the multi-plane ray-tracing convention used at cluster scale.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load the CCD imaging (used only for visualisation) and the per-source point datasets.\n", - "- **Truth Model:** Load the truth model from the family CSVs (mass / point / scaling_galaxies).\n", - "- **Tracer:** Build the ``Tracer`` carrying every lens-plane galaxy + the two source-plane galaxies.\n", - "\n", - "- **Source-Plane Chi Squared: Concept:** Why \"back-trace the observed images\" is a likelihood.\n", - "- **Multi-Plane Ray Tracing:** Recursive lens equation; scaling factors per source.\n", - "- **Back-Traced Source-Plane Positions:** Computing source-plane positions per multiple image.\n", - "- **Source-Plane Centroid:** Truth Point centre vs barycenter (which to use when).\n", - "- **Residual Map:** Source-plane distance per multiple image.\n", - "- **Magnifications at Positions:** Hessian-derived magnification per image.\n", - "- **Chi Squared Map (Source):** ``residual\u00b2 \u00d7 magnification\u00b2 / noise\u00b2``.\n", - "- **Per-Source Chi Squared:** Sum within one source's multiple-image set.\n", - "- **Total Chi Squared (Source):** Sum across sources.\n", - "- **Noise Normalization (Source):** ``sum log(2\u03c0 \u00d7 magnification^-2 \u00d7 noise\u00b2)``.\n", - "- **Source-Plane Log Likelihood:** ``-0.5 \u00d7 (chi\u00b2 + noise_normalization)``.\n", - "- **Source-Plane Validation:** ``al.FitPositionsSource`` matches the step-by-step result.\n", - "\n", - "- **Image-Plane Chi Squared: Concept:** Forward-solve, pair, measure.\n", - "- **Point Solver Setup:** Starting grid, precision, magnification threshold.\n", - "- **Forward Solving Model Positions:** Per source, via the multi-plane-aware solver.\n", - "- **Pairing Model to Observed:** Three pairing schemes + the too-many / too-few pathology.\n", - "- **Image-Plane Residual Map:** Image-plane distances per pair.\n", - "- **Chi Squared Map (Image):** ``residual\u00b2 / noise\u00b2`` \u2014 no magnification weighting here.\n", - "- **Per-Source / Total Chi Squared:** Sum across pairs and sources.\n", - "- **Noise Normalization (Image):** ``sum log(2\u03c0 \u00d7 noise\u00b2)``.\n", - "- **Image-Plane Log Likelihood:** ``-0.5 \u00d7 (chi\u00b2 + noise_normalization)``.\n", - "- **Image-Plane Validation:** ``al.FitPositionsImagePair`` matches the step-by-step result.\n", - "\n", - "- **Source-Plane vs Image-Plane: When to Use Which:** Practical comparison.\n", - "- **Wrap Up:** Pointers to ``modeling.py``, ``csv_api.py``, and the test-workspace sanity diagnostic." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "import subprocess\n", - "import sys\n", - "from pathlib import Path\n", - "\n", - "import numpy as np\n", - "\n", - "import autogalaxy as ag\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "The cluster point-source dataset lives in ``dataset/cluster/simple/``. The CCD image (``data.fits``)\n", - "is loaded for visualisation only \u2014 point-source modelling does NOT fit the imaging directly. The\n", - "positions of the multiple images of each source, plus per-position uncertainties, drive the\n", - "likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"cluster\" / dataset_name\n", - "\n", - "if (\n", - " not (dataset_path / \"data.fits\").exists()\n", - " or not (dataset_path / \"mass.csv\").exists()\n", - "):\n", - " subprocess.run([sys.executable, \"scripts/cluster/simulator.py\"], check=True)\n", - "\n", - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Datasets__\n", - "\n", - "``point_datasets.csv`` carries one row per observed multiple image, grouped by source ``name``.\n", - "``al.list_from_csv`` returns a ``List[PointDataset]`` where each entry carries:\n", - "\n", - " - ``name`` \u2014 source identifier (e.g. ``point_0``), used for name pairing with the model.\n", - " - ``positions`` \u2014 image-plane (y, x) positions of every multiple image of that source.\n", - " - ``positions_noise_map`` \u2014 per-position positional uncertainty in arc-seconds.\n", - " - ``redshift`` \u2014 the source redshift (different per source \u2014 this is a multi-plane system)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = al.list_from_csv(file_path=dataset_path / \"point_datasets.csv\")\n", - "\n", - "for dataset in dataset_list:\n", - " print(f\" {dataset.name}: z={dataset.redshift} {len(dataset.positions)} images\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Truth Model__\n", - "\n", - "The simulator wrote the truth lens model into the family CSVs (``mass.csv`` for dPIE + NFW mass\n", - "profiles; ``point.csv`` for source-galaxy ``Point`` components). We load both and build the\n", - "truth ``Tracer``. The scaling tier (10 low-mass members) has its own legacy ``scaling_galaxies.csv``\n", - "schema with shared scaling-relation parameters.\n", - "\n", - "The likelihood walkthrough below evaluates the chi\u00b2 at the *truth* model \u2014 the chi\u00b2 ought to be at\n", - "or near its minimum here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass_table = al.galaxy_models_from_csv(dataset_path / \"mass.csv\", family=\"mass\")\n", - "point_table = al.galaxy_models_from_csv(dataset_path / \"point.csv\", family=\"point\")\n", - "scaling_table = al.galaxy_table_from_csv(dataset_path / \"scaling_galaxies.csv\")\n", - "\n", - "galaxies_by_name = al.galaxies_from_csv_tables(mass_table, point_table)\n", - "\n", - "redshift_lens = 0.5\n", - "source_redshifts = sorted({float(d.redshift) for d in dataset_list})\n", - "\n", - "main_lens_galaxies = [galaxies_by_name[\"lens_0\"], galaxies_by_name[\"lens_1\"]]\n", - "host_halo_galaxy = galaxies_by_name[\"host_halo\"]\n", - "source_galaxies = [galaxies_by_name[\"source_0\"], galaxies_by_name[\"source_1\"]]\n", - "\n", - "# Scaling tier: per-member dPIE built from the legacy CSV + shared scaling relation.\n", - "scaling_galaxies = []\n", - "SCALING_FACTOR_TRUTH = 0.3\n", - "SCALING_EXPONENT_TRUTH = 1.0\n", - "for centre, luminosity in zip(\n", - " scaling_table.centres.in_list, scaling_table.luminosities\n", - "):\n", - " b0 = SCALING_FACTOR_TRUTH * luminosity**SCALING_EXPONENT_TRUTH\n", - " scaling_galaxies.append(\n", - " al.Galaxy(\n", - " redshift=redshift_lens,\n", - " mass=al.mp.dPIEMassSph(centre=tuple(centre), ra=0.1, rs=10.0, b0=b0),\n", - " )\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "The tracer carries:\n", - "\n", - " - 2 main lens galaxies (BCG + satellite) \u2014 individually-modelled dPIE mass profiles.\n", - " - 10 scaling-tier member galaxies \u2014 dPIE mass profiles whose ``b0`` is derived from the shared\n", - " scaling relation ``b0 = scaling_factor \u00d7 luminosity^scaling_exponent``.\n", - " - 1 host dark matter halo \u2014 ``NFWMCRLudlowSph`` at the cluster centre.\n", - " - 2 source galaxies \u2014 ``Point`` profiles at distinct redshifts (multi-plane).\n", - "\n", - "PyAutoLens automatically groups galaxies by redshift into ``planes`` and ray-traces through them\n", - "in redshift order. For this cluster the planes are: ``z=0.5`` (lens-plane, holds 13 galaxies),\n", - "``z=1.0`` (source_0 plane), ``z=2.0`` (source_1 plane)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=main_lens_galaxies\n", - " + scaling_galaxies\n", - " + [host_halo_galaxy]\n", - " + source_galaxies\n", - ")\n", - "\n", - "print(\n", - " f\"Tracer has {len(tracer.planes)} planes at redshifts \"\n", - " f\"{[float(p.redshift) for p in tracer.planes]}\"\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Chi Squared: Concept__\n", - "\n", - "If the lens model is correct, every observed image position of a given source ought to ray-trace\n", - "back to (approximately) the same source-plane location: the source's true centre. The figure of\n", - "merit for the **source-plane chi\u00b2** is the scatter of those back-traced positions around the\n", - "source-plane reference, weighted by magnification (so source-plane residuals are converted back to\n", - "image-plane scale where the noise is defined).\n", - "\n", - "This contrasts with the more intuitive image-plane chi\u00b2 which compares observed image positions to\n", - "the model's *forward-solved* multiple-image positions. Source-plane chi\u00b2 is cheaper (no forward\n", - "solve required, just one back-projection per image), and the math is JAX-friendly. The trade-off\n", - "is that the magnification weighting amplifies any positional error in the back-projection: at\n", - "cluster scale where magnifications near multi-image positions are ~100\u00d7, even sub-arcsecond\n", - "source-plane residuals can drive chi\u00b2 to 10\u2076\u201310\u2078 at the truth.\n", - "\n", - "This is the figure of merit used by ``al.FitPositionsSource`` and is described in detail in\n", - "Jullo et al. 2007 (\"A Bayesian approach to strong lensing modelling of galaxy clusters\") and the\n", - "references therein.\n", - "\n", - "__Multi-Plane Ray Tracing__\n", - "\n", - "The cluster has two sources at *different* redshifts (``z = 1.0`` and ``z = 2.0``), so ray-tracing\n", - "each source's positions back to its source plane goes through every earlier plane along the way.\n", - "This is **multi-plane ray tracing**, formalised by the recursive lens equation:\n", - "\n", - ".. math::\n", - "\n", - " \\\\theta_{j} = \\\\theta_{0} - \\\\sum_{i=1}^{j-1} \\\\beta_{ij} \\\\alpha_{i}(\\\\theta_{i})\n", - "\n", - "where:\n", - "\n", - " - :math:`\\\\theta_{0}` is the image-plane position,\n", - " - :math:`\\\\theta_{i}` is the position of the ray at plane ``i``,\n", - " - :math:`\\\\alpha_{i}(\\\\theta_{i})` is the deflection at plane ``i`` (sum of every galaxy at that\n", - " plane's redshift),\n", - " - :math:`\\\\beta_{ij} = (D_{ij} \\\\, D_{s}) / (D_{j} \\\\, D_{is})` is a cosmological scaling factor\n", - " that scales the deflection from plane ``i``'s angular-diameter convention to plane ``j``'s.\n", - "\n", - "For our cluster the relevant cases per source are:\n", - "\n", - " - **Source 0 (z=1.0)** \u2014 back-tracing its image positions just hits the lens plane (z=0.5).\n", - " Two planes total, so the scaling factor reduces to 1.0 and the equation is the familiar\n", - " :math:`\\\\theta_{1} = \\\\theta_{0} - \\\\alpha_{0}(\\\\theta_{0})`.\n", - "\n", - " - **Source 1 (z=2.0)** \u2014 back-tracing hits the lens plane (z=0.5) AND the source_0 plane\n", - " (z=1.0). Even though source_0 carries no mass, the recursive equation still walks through its\n", - " plane; the actual contribution depends on whether any galaxy at z=1.0 has a non-zero mass\n", - " profile (none here \u2014 the source_0 ``Point`` profile contributes nothing to the deflection\n", - " field).\n", - "\n", - "The library encapsulates all of this in ``Tracer.deflections_between_planes_from(grid, plane_i=0,\n", - "plane_j=)``. The function walks the planes from ``plane_i`` to ``plane_j``\n", - "in redshift order, applying the per-plane :math:`\\\\alpha_{i}(\\\\theta_{i})` sums and the scaling\n", - "factors :math:`\\\\beta_{ij}` automatically. The full derivation of the recursive equation, the\n", - "sign conventions, and the cosmological scaling factors live in the multi-plane guide at\n", - "``autolens_workspace/scripts/guides/lensing/multi_plane.py``.\n", - "\n", - "In practice the easiest entry point isn't ``deflections_between_planes_from`` (which returns the\n", - "*differences* between plane positions and requires you to apply the lens equation manually) but\n", - "``Tracer.traced_grid_2d_list_from(grid)``. That function returns the list of grids per plane,\n", - "fully traced through the recursive lens equation, with the cosmological scaling factors already\n", - "applied. For our 3-plane cluster:\n", - "\n", - " - ``traced_grids[0]`` is the input image-plane grid (unchanged).\n", - " - ``traced_grids[1]`` is the source_0 plane (z=1.0) position of every input ray after passing\n", - " through the lens-plane deflection.\n", - " - ``traced_grids[2]`` is the source_1 plane (z=2.0) position of every input ray after passing\n", - " through the lens plane *and* the source_0 plane (the recursive step).\n", - "\n", - "Each source's back-traced positions are then just ``traced_grids[]``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_plane_positions_per_source = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " plane_index = tracer.plane_index_via_redshift_from(redshift=dataset.redshift)\n", - " traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.positions)\n", - " source_plane_positions_per_source.append(traced_grids[plane_index])\n", - "\n", - " print(\n", - " f\" {dataset.name}: plane_index={plane_index}, \"\n", - " f\"back-traced positions = {traced_grids[plane_index].in_list}\"\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Back-Traced Source-Plane Positions: Conceptual Recap__\n", - "\n", - "The traced grid at plane ``j`` is, by construction:\n", - "\n", - ".. math::\n", - "\n", - " \\\\theta_{j} = \\\\theta_{0} - \\\\alpha_{\\\\text{multi-plane}}(\\\\theta_{0}; j)\n", - "\n", - "where the multi-plane deflection :math:`\\\\alpha_{\\\\text{multi-plane}}` is the recursive sum\n", - "discussed in the previous section \u2014 exactly the result you would get from manually applying\n", - "``deflections_between_planes_from`` and then ``grid_2d_via_deflection_grid_from``. Using\n", - "``traced_grid_2d_list_from`` is the same thing in one call.\n", - "\n", - "__Source-Plane Centroid__\n", - "\n", - "The \"reference point\" against which we measure the source-plane scatter has two options:\n", - "\n", - " 1. **Truth ``Point`` centre.** Each source carries a ``Point`` profile in the model whose\n", - " ``centre`` is a free parameter (or fixed to the truth here). At a model fit, ``Point.centre``\n", - " *is* the source-plane (y, x) the multiple images should converge to.\n", - "\n", - " 2. **Barycenter of back-traced positions.** Pretend you don't know the truth centre. Compute the\n", - " centroid of the back-traced positions \u2014 at the right model the centroid sits where the source\n", - " actually is, and the residuals are scatter around it.\n", - "\n", - "Option (1) is what ``al.FitPositionsSource(profile=point_profile)`` uses. Option (2) is what\n", - "``al.FitPositionsSource(profile=None)`` uses (the default during model fits, because at search\n", - "time the model doesn't yet know the truth).\n", - "\n", - "For this walkthrough we use option (1) so the residuals have a well-defined physical meaning\n", - "(distance from truth, not from a derived centroid)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_plane_centroids = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " # source_i's Point profile is attached to source_galaxies[i] under attr name \"point_i\".\n", - " point_profile = getattr(source_galaxies[i], dataset.name)\n", - " source_plane_centroids.append(point_profile.centre)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Residual Map__\n", - "\n", - "The per-image residual is just the source-plane distance between each back-traced position and the\n", - "source-plane centroid:\n", - "\n", - ".. math::\n", - "\n", - " r_{i} = |\\\\theta_{j,i} - \\\\theta_{j,\\\\text{centre}}|" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "residuals_per_source = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " sp_positions = source_plane_positions_per_source[i]\n", - " centre = source_plane_centroids[i]\n", - " residuals = sp_positions.distances_to_coordinate_from(coordinate=centre)\n", - " residuals_per_source.append(residuals)\n", - "\n", - " print(f\" {dataset.name}: residuals = {np.asarray(residuals)}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Magnifications at Positions__\n", - "\n", - "The source-plane chi\u00b2 weights each residual by the local **magnification** at that image position.\n", - "Why: position noise is defined in the image plane (arcsec), but the residual is in the source\n", - "plane. The magnification converts a source-plane distance back to its image-plane scale.\n", - "Specifically the formula below uses magnification squared as a multiplicative weight.\n", - "\n", - "PyAutoLens computes magnification from the Hessian of the deflection field via\n", - "``Tracer.magnification_2d_via_hessian_from(grid)``. Physically the Hessian gives the linear\n", - "distortion matrix at each image position; its determinant is the magnification (signed: negative\n", - "for parity-flipped images, but ``magnification_2d_via_hessian_from`` returns the absolute value).\n", - "The conceptual chain is **deflection field \u2192 Hessian via finite-difference \u2192 eigenvalues\n", - "(convergence \u03ba and shear \u03b3) \u2192 magnification :math:`\\\\mu = 1 / |(1-\\\\kappa)^2 - \\\\gamma^2|`**.\n", - "\n", - "Full derivation and the eigenvalue / critical-curve discussion live in the lens-calc guide at\n", - "``autolens_workspace/scripts/guides/lensing/lens_calc.py``.\n", - "\n", - "For multi-plane lenses the function uses the full multi-plane deflection in the Hessian, so the\n", - "magnification correctly reflects the cumulative distortion through every intervening plane. The\n", - "``LensCalc`` helper bundles the Hessian, magnification, and critical-curve calculations; we build\n", - "one per source so each magnification is evaluated against the correct multi-plane chain (the\n", - "source's plane index).\n", - "\n", - "The library wraps this with an ``abs(...)`` so the returned magnification is always positive (raw\n", - "magnification can be negative when the image is parity-flipped; squaring it in the chi\u00b2 formula\n", - "makes the sign irrelevant anyway)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnifications_per_source = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " plane_index = tracer.plane_index_via_redshift_from(redshift=dataset.redshift)\n", - " od = ag.LensCalc.from_tracer(\n", - " tracer=tracer, use_multi_plane=True, plane_j=plane_index\n", - " )\n", - " mag = abs(od.magnification_2d_via_hessian_from(grid=dataset.positions))\n", - " magnifications_per_source.append(mag)\n", - "\n", - " print(f\" {dataset.name}: magnifications = {np.asarray(mag)}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Chi Squared Map (Source)__\n", - "\n", - "The per-image source-plane chi\u00b2 combines the residual, the magnification, and the position noise:\n", - "\n", - ".. math::\n", - "\n", - " \\\\chi^{2}_{i} = \\\\frac{r_{i}^{2} \\\\, \\\\mu_{i}^{2}}{\\\\sigma_{i}^{2}}\n", - "\n", - "This is the formula in ``autolens/point/fit/positions/source/separations.py`` ::\n", - "\n", - " chi_squared_map = residual_map**2 / (magnifications_at_positions.array**-2 * noise_map.array**2)\n", - "\n", - "which expands to ``residual\u00b2 \u00d7 magnification\u00b2 / noise\u00b2``. (The :math:`\\\\mu^{-2}` in the denominator\n", - "of the source-code form cancels the :math:`\\\\mu^{2}` in the numerator and leaves the form above.)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_maps_per_source = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " r = np.asarray(residuals_per_source[i])\n", - " mu = np.asarray(magnifications_per_source[i])\n", - " sigma = np.asarray(dataset.positions_noise_map)\n", - " chi_sq_map = r**2 * mu**2 / sigma**2\n", - " chi_squared_maps_per_source.append(chi_sq_map)\n", - "\n", - " print(f\" {dataset.name}: chi\u00b2 per image = {chi_sq_map}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Source Chi Squared__\n", - "\n", - "Per-source: sum the chi\u00b2 map across that source's multiple images." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_per_source = [float(np.sum(m)) for m in chi_squared_maps_per_source]\n", - "for i, dataset in enumerate(dataset_list):\n", - " print(f\" {dataset.name}: chi\u00b2 = {chi_squared_per_source[i]:.4e}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Total Chi Squared (Source)__\n", - "\n", - "Each source is independent (different redshift, different multiple-image set), so the total chi\u00b2\n", - "is just the sum." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_chi_squared_source = float(sum(chi_squared_per_source))\n", - "print(f\"Total source-plane chi\u00b2 = {total_chi_squared_source:.4e}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization (Source)__\n", - "\n", - "The full Gaussian log-likelihood carries a normalisation term that depends on the noise:\n", - "\n", - ".. math::\n", - "\n", - " \\\\mathcal{N} = \\\\sum_{i} \\\\log \\\\left( 2\\\\pi \\\\, \\\\mu_{i}^{-2} \\\\, \\\\sigma_{i}^{2} \\\\right)\n", - "\n", - "The magnification factor appears here too because the chi\u00b2 formula effectively treats each image\n", - "position as having an *effective* source-plane noise of :math:`\\\\sigma_{i} / \\\\mu_{i}`. The\n", - "normalisation matches that interpretation so the resulting expression is a well-defined Gaussian\n", - "log-likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalizations_per_source = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " mu = np.asarray(magnifications_per_source[i])\n", - " sigma = np.asarray(dataset.positions_noise_map)\n", - " nn = float(np.sum(np.log(2 * np.pi * (mu**-2) * sigma**2)))\n", - " noise_normalizations_per_source.append(nn)\n", - "\n", - "total_noise_normalization_source = float(sum(noise_normalizations_per_source))\n", - "print(\n", - " f\"Total source-plane noise normalization = {total_noise_normalization_source:.4e}\"\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Log Likelihood__\n", - "\n", - "The standard Gaussian log-likelihood:\n", - "\n", - ".. math::\n", - "\n", - " \\\\log L = -\\\\frac{1}{2} \\\\left( \\\\chi^{2} + \\\\mathcal{N} \\\\right)\n", - "\n", - "This is what Nautilus maximises during a cluster point-source model fit. The chi\u00b2 term encodes\n", - "goodness-of-fit; the normalisation absorbs the constant noise-dependent prefactor of the Gaussian." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_likelihood_source = -0.5 * (\n", - " total_chi_squared_source + total_noise_normalization_source\n", - ")\n", - "print(f\"Source-plane log likelihood = {log_likelihood_source:.4e}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Validation__\n", - "\n", - "To verify the step-by-step derivation, we instantiate ``al.FitPositionsSource`` and confirm its\n", - "``log_likelihood`` matches. ``FitPositionsSource`` does exactly the calculation above internally\n", - "(see ``autolens/point/fit/positions/source/separations.py``)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sum_library_log_likelihood_source = 0.0\n", - "for i, dataset in enumerate(dataset_list):\n", - " point_profile = getattr(source_galaxies[i], dataset.name)\n", - " fit = al.FitPositionsSource(\n", - " name=dataset.name,\n", - " data=dataset.positions,\n", - " noise_map=dataset.positions_noise_map,\n", - " tracer=tracer,\n", - " solver=None,\n", - " profile=point_profile,\n", - " )\n", - " sum_library_log_likelihood_source += float(fit.log_likelihood)\n", - "\n", - "print(f\"Library source-plane log likelihood = {sum_library_log_likelihood_source:.4e}\")\n", - "print(\n", - " f\"Match: {np.isclose(log_likelihood_source, sum_library_log_likelihood_source, rtol=1e-6)}\"\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image-Plane Chi Squared: Concept__\n", - "\n", - "The image-plane chi\u00b2 takes the opposite approach. Instead of ray-tracing observed positions to\n", - "the source plane and measuring source-plane scatter, it **forward-solves** the model's image-plane\n", - "positions for each source and compares them to the observed positions in the image plane.\n", - "\n", - "Mechanically:\n", - "\n", - " 1. Take the source-plane centre (model's ``Point.centre`` per source).\n", - " 2. Use a ``PointSolver`` to ray-trace triangles from a fine image-plane grid forward through the\n", - " lens until they converge to that source-plane centre \u2014 these are the **model positions**.\n", - " 3. For each model position, find the closest observed position and pair them.\n", - " 4. The residual per pair is just the image-plane Euclidean distance.\n", - "\n", - "The chi\u00b2 is then ``residual\u00b2 / noise\u00b2`` \u2014 no magnification weighting, because the residual is\n", - "already in the same units (arc-seconds) as the noise. This makes the absolute chi\u00b2 much smaller\n", - "than the source-plane variant at the same model: at the truth, residuals are bounded by the\n", - "``PointSolver`` precision (~0.001\"), so chi\u00b2 \u2248 (0.001 / 0.005)\u00b2 \u00d7 N \u2248 N \u00d7 0.04, where N is the\n", - "number of paired images. Far below the source-plane chi\u00b2 which is dominated by magnification \u00d7 the\n", - "same precision floor.\n", - "\n", - "The trade-off: image-plane chi\u00b2 requires a **forward solve per evaluation**, which is slow,\n", - "introduces solver precision as a noise floor, and carries pairing pathologies discussed below.\n", - "The source-plane chi\u00b2 has none of those costs.\n", - "\n", - "__Point Solver Setup__\n", - "\n", - "The ``PointSolver`` searches for image-plane positions whose forward ray-trace lands at the\n", - "source-plane target. Conceptually it does this by:\n", - "\n", - " 1. Starting from a coarse image-plane grid.\n", - " 2. Tessellating into triangles. For each triangle, ray-trace its three vertices to the source\n", - " plane. If the source-plane target lies inside the traced triangle, that image-plane triangle\n", - " contains a model position.\n", - " 3. Refining each containing triangle \u2014 subdivide, ray-trace the sub-triangles, repeat \u2014 until\n", - " the triangle size drops below a configurable ``pixel_scale_precision``.\n", - " 4. The triangle centroid at convergence is the model position.\n", - "\n", - "This is iterative and the runtime scales with the source-plane precision target. The solver also\n", - "filters out central images via ``magnification_threshold``: highly demagnified images (often the\n", - "unobservable central image of a strong-lens configuration) are discarded so the model doesn't pair\n", - "unphysical demagnified solutions to the observed multi-image set.\n", - "\n", - "A standalone walkthrough of the triangle-refinement algorithm (sub-grid traversal, magnification\n", - "filtering, multi-plane handling, JAX-compatibility) lives in\n", - "``autolens_workspace/scripts/guides/point_source/triangle_solver.py`` \u2014 TODO, currently not yet\n", - "written.\n", - "\n", - "In code, the solver is constructed once and reused per evaluation:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=1.0),\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.1,\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Forward Solving Model Positions__\n", - "\n", - "For each source, call ``solver.solve(tracer, source_plane_coordinate=...)`` with the source's\n", - "truth source-plane centre. The solver returns the image-plane (y, x) positions of every\n", - "multiple image whose magnification is above ``magnification_threshold``.\n", - "\n", - "Multi-plane ray tracing happens *inside* ``solve()`` automatically. For source_1 (z=2.0) the\n", - "triangles are forward-traced through the z=1.0 source_0 plane before reaching z=2.0 \u2014 the same\n", - "recursive lens equation discussed above is applied during the forward solve. The solver picks the\n", - "right target plane via the ``plane_redshift`` argument: this is the source galaxy's redshift, and\n", - "the solver maps it to the corresponding plane index in the tracer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_positions_per_source = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " point_profile = getattr(source_galaxies[i], dataset.name)\n", - " model_positions = solver.solve(\n", - " tracer=tracer,\n", - " source_plane_coordinate=point_profile.centre,\n", - " plane_redshift=dataset.redshift,\n", - " )\n", - " model_positions_per_source.append(model_positions)\n", - " print(f\" {dataset.name}: model positions = {model_positions.in_list}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pairing Model to Observed__\n", - "\n", - "We have two sets of image-plane positions per source: the **model positions** (forward-solved) and\n", - "the **observed positions** (from ``point_datasets.csv``). The chi\u00b2 requires a one-to-one mapping\n", - "between them \u2014 but the two sets don't come pre-paired, and may not even have the same length.\n", - "\n", - "**Three pairing schemes** exist in PyAutoLens, each with different behaviour when the counts\n", - "don't match:\n", - "\n", - " 1. **``FitPositionsImagePair`` (Hungarian, no-repeat).** Pairs each observed to its nearest\n", - " predicted image via the **Hungarian algorithm** (also called linear sum assignment). The\n", - " algorithm finds the unique 1-to-1 pairing that *globally minimises the total distance* \u2014\n", - " not greedy. Two consequences: (a) it gives the optimal assignment even when greedy would\n", - " fail; (b) when counts differ, unmatched positions still get unpaired and contribute nothing\n", - " to chi\u00b2.\n", - "\n", - " 2. **``FitPositionsImagePairAll`` (closest, with replacement).** Each model position pairs to its\n", - " nearest observed position independently. A single observed position may be paired to multiple\n", - " model positions. Every model position contributes to chi\u00b2.\n", - "\n", - " 3. **``FitPositionsImagePairRepeat`` (repeats allowed).** Similar to (2) but with explicit\n", - " handling of cases where the model genuinely predicts multiple images at nearly the same\n", - " location (e.g. near a caustic crossing).\n", - "\n", - "This walkthrough uses scheme (1) \u2014 ``FitPositionsImagePair`` \u2014 which is the historical default and\n", - "the one ``al.FitPointDataset`` constructs by default. We follow that convention.\n", - "\n", - "**The known pathology** (documented in the PyAutoLens source for ``FitPositionsImagePair``):\n", - "\n", - " - **Model predicts too many images** (more model than observed). The pairing leaves some model\n", - " positions unpaired. These model positions contribute *nothing* to chi\u00b2. A model that\n", - " spuriously generates extra demagnified images can therefore have its chi\u00b2 artificially\n", - " *reduced* compared to a model that produces exactly the observed count.\n", - " - **Model predicts too few images** (fewer model than observed). Some observed positions go\n", - " unpaired and also contribute nothing. The chi\u00b2 is similarly artificially reduced.\n", - "\n", - "In both cases the optimiser can prefer pathological solutions with the wrong image count over the\n", - "correct solution. ``magnification_threshold`` partially mitigates this by discarding very\n", - "demagnified spurious images, but it is not a complete fix. ``FitPositionsImagePairAll`` /\n", - "``...Repeat`` address some of the cases, at the cost of other failure modes. Selecting the right\n", - "scheme for a real cluster fit is a per-dataset judgement call.\n", - "\n", - "In code we do the same pairing the library does \u2014 Hungarian / linear sum assignment \u2014 by\n", - "computing the pairwise distance matrix and handing it to ``scipy.optimize.linear_sum_assignment``:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.optimize import linear_sum_assignment\n", - "\n", - "\n", - "def _pair_hungarian(model_positions, observed_positions):\n", - " \"\"\"Optimal 1-1 pairing via Hungarian algorithm.\n", - "\n", - " Returns a list of (model_position, observed_position, distance) tuples. When the number of\n", - " model and observed positions differ, only ``min(N_model, N_observed)`` pairs are returned\n", - " \u2014 the excess positions on the larger side go unpaired and contribute nothing to chi\u00b2.\n", - " \"\"\"\n", - " model_arr = np.asarray(model_positions)\n", - " observed_arr = np.asarray(observed_positions)\n", - "\n", - " # Pairwise distance matrix: distances[i, j] = ||model_i - observed_j||.\n", - " distances = np.linalg.norm(\n", - " model_arr[:, np.newaxis, :] - observed_arr[np.newaxis, :, :], axis=2\n", - " )\n", - "\n", - " row_ind, col_ind = linear_sum_assignment(distances)\n", - "\n", - " return [\n", - " (model_arr[i], observed_arr[j], float(distances[i, j]))\n", - " for i, j in zip(row_ind, col_ind)\n", - " ]\n", - "\n", - "\n", - "pairs_per_source = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " pairs = _pair_hungarian(model_positions_per_source[i], dataset.positions)\n", - " pairs_per_source.append(pairs)\n", - " print(\n", - " f\" {dataset.name}: {len(pairs)} pairs / {len(dataset.positions)} observed / \"\n", - " f\"{len(model_positions_per_source[i])} model\"\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image-Plane Residual Map__\n", - "\n", - "Residual per pair = image-plane Euclidean distance between model and observed (in arc-seconds).\n", - "This is what ``_pair_closest_no_repeat`` already returned as the third element of each tuple." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "residual_maps_image_per_source = []\n", - "for i, pairs in enumerate(pairs_per_source):\n", - " residuals = np.array([p[2] for p in pairs])\n", - " residual_maps_image_per_source.append(residuals)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Chi Squared Map (Image)__\n", - "\n", - "Image-plane chi\u00b2 per pair:\n", - "\n", - ".. math::\n", - "\n", - " \\\\chi^{2}_{i} = \\\\frac{r_{i}^{2}}{\\\\sigma_{i}^{2}}\n", - "\n", - "No magnification weighting \u2014 the residual is already in image-plane arc-seconds, the same units\n", - "as the noise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_maps_image_per_source = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " r = residual_maps_image_per_source[i]\n", - " # Use the noise of each paired observed position. For simplicity we use the dataset's mean\n", - " # noise since all positions in the cluster simulator share the same \u03c3; in a real analysis\n", - " # this should index per-observed-position.\n", - " sigma = float(np.mean(np.asarray(dataset.positions_noise_map)))\n", - " chi_sq_map = r**2 / sigma**2\n", - " chi_squared_maps_image_per_source.append(chi_sq_map)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Source / Total Chi Squared (Image)__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_per_source_image = [\n", - " float(np.sum(m)) for m in chi_squared_maps_image_per_source\n", - "]\n", - "total_chi_squared_image = float(sum(chi_squared_per_source_image))\n", - "print(f\"Total image-plane chi\u00b2 = {total_chi_squared_image:.4e}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization (Image)__\n", - "\n", - "The Gaussian-log normalisation for image-plane chi\u00b2 is the standard form:\n", - "\n", - ".. math::\n", - "\n", - " \\\\mathcal{N} = \\\\sum_{i} \\\\log \\\\left( 2\\\\pi \\\\, \\\\sigma_{i}^{2} \\\\right)\n", - "\n", - "No magnification factor (residuals are already in image-plane units)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalizations_per_source_image = []\n", - "for i, dataset in enumerate(dataset_list):\n", - " n_pairs = len(residual_maps_image_per_source[i])\n", - " sigma = float(np.mean(np.asarray(dataset.positions_noise_map)))\n", - " nn = float(n_pairs * np.log(2 * np.pi * sigma**2))\n", - " noise_normalizations_per_source_image.append(nn)\n", - "\n", - "total_noise_normalization_image = float(sum(noise_normalizations_per_source_image))\n", - "print(f\"Total image-plane noise normalization = {total_noise_normalization_image:.4e}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image-Plane Log Likelihood__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_likelihood_image = -0.5 * (\n", - " total_chi_squared_image + total_noise_normalization_image\n", - ")\n", - "print(f\"Image-plane log likelihood = {log_likelihood_image:.4e}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image-Plane Validation__\n", - "\n", - "Instantiate ``al.FitPositionsImagePair`` and confirm match. Note: ``FitPositionsImagePair`` uses\n", - "the same closest-no-repeat pairing scheme as our manual implementation above, so the chi\u00b2 values\n", - "should agree to within numerical precision (the library's solver may use slightly different\n", - "triangle precision at refinement edge cases, producing sub-percent residual differences)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sum_library_log_likelihood_image = 0.0\n", - "for i, dataset in enumerate(dataset_list):\n", - " point_profile = getattr(source_galaxies[i], dataset.name)\n", - " fit = al.FitPositionsImagePair(\n", - " name=dataset.name,\n", - " data=dataset.positions,\n", - " noise_map=dataset.positions_noise_map,\n", - " tracer=tracer,\n", - " solver=solver,\n", - " profile=point_profile,\n", - " )\n", - " sum_library_log_likelihood_image += float(fit.log_likelihood)\n", - "\n", - "print(f\"Library image-plane log likelihood = {sum_library_log_likelihood_image:.4e}\")\n", - "print(\n", - " f\"Match (rtol=1e-2): \"\n", - " f\"{np.isclose(log_likelihood_image, sum_library_log_likelihood_image, rtol=1e-2)}\"\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane vs Image-Plane: When to Use Which__\n", - "\n", - "| Aspect | Source-plane (``FitPositionsSource``) | Image-plane (``FitPositionsImagePair``) |\n", - "|---|---|---|\n", - "| **Forward solve required?** | No | Yes (one per evaluation) |\n", - "| **Per-evaluation cost** | Cheap (just back-projection) | Expensive (triangle refinement) |\n", - "| **JAX-friendly?** | Yes | The solver is JAX-jitted, but each evaluation triggers a forward solve regardless |\n", - "| **Magnification weighting** | Yes (``\u03bc\u00b2`` in chi\u00b2) | No |\n", - "| **Sensitive to ``PointSolver`` precision?** | Amplified by ``\u03bc\u00b2`` at multi-image positions | Bounded by precision directly |\n", - "| **Pairing pathology** | None (uses all images) | Yes (too-many / too-few; see above) |\n", - "| **Absolute chi\u00b2 scale (at truth)** | Large (~``\u03bc\u00b2`` \u00d7 precision\u00b2) | Small (~precision\u00b2) |\n", - "| **Best for** | Fast Nautilus fits; JAX-jit'd parameter estimation | Final residual visualisation; cases where pairing is unambiguous |\n", - "\n", - "For most cluster fits the source-plane chi\u00b2 is the right default: it's faster, JAX-compatible,\n", - "and doesn't suffer pairing pathologies. The image-plane chi\u00b2 is most useful for diagnostic\n", - "visualisation of where the model's predicted images sit relative to the observed ones, and for\n", - "cases where the source-plane chi\u00b2 has bias issues that need cross-checking.\n", - "\n", - "The cluster modelling script at ``scripts/cluster/modeling.py`` uses ``AnalysisPoint`` which\n", - "selects the chi\u00b2 flavour via its constructor; consult the ``AnalysisPoint`` docstring for the\n", - "current default.\n", - "\n", - "__Wrap Up__\n", - "\n", - "This script has walked through the cluster point-source likelihood end to end, in both source-plane\n", - "and image-plane flavours, for a multi-plane (2-source) cluster.\n", - "\n", - "Next steps:\n", - "\n", - "- ``scripts/cluster/csv_api.py`` \u2014 the family CSV schema this script reads.\n", - "- ``scripts/cluster/simulator.py`` \u2014 how the truth dataset was generated.\n", - "- ``scripts/cluster/modeling.py`` \u2014 full Nautilus fit using the same likelihood function discussed\n", - " here, with priors and ``AnalysisPoint`` configuration.\n", - "- ``autolens_workspace_test/scripts/cluster/likelihood_sanity.py`` \u2014 perturbation sweep that\n", - " surfaces the ``PointSolver`` precision-floor pathology described in the source-plane section.\n", - "- ``autolens_workspace/scripts/guides/lensing/multi_plane.py`` \u2014 full derivation of the recursive\n", - " lens equation.\n", - "- ``autolens_workspace/scripts/guides/lensing/lens_calc.py`` \u2014 Hessian / magnification /\n", - " critical-curve mechanics.\n", - "\n", - "For a deeper understanding of cluster lens modelling and point-source likelihoods, the\n", - "**HowToLens** Jupyter notebook lectures cover both topics in detail.\n", - "\n", - "__JAX__\n", - "\n", - "The chi-squared walkthrough above is pure NumPy. To JAX-accelerate it,\n", - "wrap construction in `@jax.jit` with the post-Phase-2 `PointSolver`\n", - "pattern:\n", - "\n", - "```python\n", - "import jax\n", - "import jax.numpy as jnp\n", - "from autolens.jax import register_tracer_classes\n", - "\n", - "register_tracer_classes(tracer) # one-time\n", - "\n", - "@jax.jit\n", - "def cluster_log_likelihood(tracer, dataset, source_galaxy):\n", - " fit = al.FitPositionsSource(\n", - " name=dataset.name,\n", - " data=dataset.positions,\n", - " noise_map=dataset.positions_noise_map,\n", - " tracer=tracer,\n", - " solver=None, # source-plane chi\u00b2 doesn't need the solver\n", - " profile=getattr(source_galaxy, dataset.name),\n", - " )\n", - " return fit.log_likelihood\n", - "```\n", - "\n", - "For the canonical search-driven path (`AnalysisPoint(use_jax=True)`),\n", - "see `modeling.py`. For JIT-ing library methods directly without going\n", - "through `FitPositionsSource`, see `scripts/guides/lens_calc.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Log Likelihood Function: Cluster Point Source\n", + "==============================================\n", + "\n", + "This script provides a step-by-step guide of the cluster point-source ``log_likelihood_function``,\n", + "the figure-of-merit Nautilus optimises when fitting a cluster lens model to ``point_datasets.csv``.\n", + "\n", + "Cluster point-source modelling has two distinct likelihood flavours:\n", + "\n", + " 1. **Source-plane chi\u00b2** (``FitPositionsSource``) \u2014 ray-trace every observed image-plane position\n", + " back through the lens, and measure the magnification-weighted scatter of the back-traced\n", + " positions around the source-plane centre. Cheap (no forward solve) and JAX-friendly.\n", + "\n", + " 2. **Image-plane chi\u00b2** (``FitPositionsImagePair``) \u2014 forward-solve the model's image-plane\n", + " positions for each source via the ``PointSolver``, pair each model position to the closest\n", + " observed position, and measure the image-plane residuals. More intuitive (residuals in\n", + " arc-seconds), but slower and carries pairing pathologies the source-plane variant avoids.\n", + "\n", + "We walk through both, end to end, with the actual library formulae. The standard cluster model is\n", + "assumed: all lens-plane galaxies at ``z = 0.5``, two background sources at *different* redshifts\n", + "(``z = 1.0`` and ``z = 2.0``). Multi-plane ray tracing therefore applies and we explain how the\n", + "recursive lens equation handles it.\n", + "\n", + "This is a companion to ``scripts/imaging/likelihood_function.py`` and\n", + "``scripts/group/likelihood_function.py``, both of which cover extended-imaging likelihoods. The\n", + "cluster surface is the first one in the workspace tutorial set that uses *point sources*\n", + "exclusively + multi-plane ray tracing + many lens galaxies in the deflection sum.\n", + "\n", + "__Aims__\n", + "\n", + " - Provide a resource that authors can include in papers using cluster modelling, so readers can\n", + " understand the likelihood (including references to the prior literature it builds on) without\n", + " needing to re-derive equations or trace the PyAutoLens source.\n", + " - Make the difference between source-plane and image-plane chi\u00b2 concrete so users can decide\n", + " which flavour fits their dataset + compute budget.\n", + " - Document the multi-plane ray-tracing convention used at cluster scale.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load the CCD imaging (used only for visualisation) and the per-source point datasets.\n", + "- **Truth Model:** Load the truth model from the family CSVs (mass / point / scaling_galaxies).\n", + "- **Tracer:** Build the ``Tracer`` carrying every lens-plane galaxy + the two source-plane galaxies.\n", + "\n", + "- **Source-Plane Chi Squared: Concept:** Why \"back-trace the observed images\" is a likelihood.\n", + "- **Multi-Plane Ray Tracing:** Recursive lens equation; scaling factors per source.\n", + "- **Back-Traced Source-Plane Positions:** Computing source-plane positions per multiple image.\n", + "- **Source-Plane Centroid:** Truth Point centre vs barycenter (which to use when).\n", + "- **Residual Map:** Source-plane distance per multiple image.\n", + "- **Magnifications at Positions:** Hessian-derived magnification per image.\n", + "- **Chi Squared Map (Source):** ``residual\u00b2 \u00d7 magnification\u00b2 / noise\u00b2``.\n", + "- **Per-Source Chi Squared:** Sum within one source's multiple-image set.\n", + "- **Total Chi Squared (Source):** Sum across sources.\n", + "- **Noise Normalization (Source):** ``sum log(2\u03c0 \u00d7 magnification^-2 \u00d7 noise\u00b2)``.\n", + "- **Source-Plane Log Likelihood:** ``-0.5 \u00d7 (chi\u00b2 + noise_normalization)``.\n", + "- **Source-Plane Validation:** ``al.FitPositionsSource`` matches the step-by-step result.\n", + "\n", + "- **Image-Plane Chi Squared: Concept:** Forward-solve, pair, measure.\n", + "- **Point Solver Setup:** Starting grid, precision, magnification threshold.\n", + "- **Forward Solving Model Positions:** Per source, via the multi-plane-aware solver.\n", + "- **Pairing Model to Observed:** Three pairing schemes + the too-many / too-few pathology.\n", + "- **Image-Plane Residual Map:** Image-plane distances per pair.\n", + "- **Chi Squared Map (Image):** ``residual\u00b2 / noise\u00b2`` \u2014 no magnification weighting here.\n", + "- **Per-Source / Total Chi Squared:** Sum across pairs and sources.\n", + "- **Noise Normalization (Image):** ``sum log(2\u03c0 \u00d7 noise\u00b2)``.\n", + "- **Image-Plane Log Likelihood:** ``-0.5 \u00d7 (chi\u00b2 + noise_normalization)``.\n", + "- **Image-Plane Validation:** ``al.FitPositionsImagePair`` matches the step-by-step result.\n", + "\n", + "- **Source-Plane vs Image-Plane: When to Use Which:** Practical comparison.\n", + "- **Wrap Up:** Pointers to ``modeling.py``, ``csv_api.py``, and the test-workspace sanity diagnostic." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "import subprocess\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "\n", + "import autogalaxy as ag\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "The cluster point-source dataset lives in ``dataset/cluster/simple/``. The CCD image (``data.fits``)\n", + "is loaded for visualisation only \u2014 point-source modelling does NOT fit the imaging directly. The\n", + "positions of the multiple images of each source, plus per-position uncertainties, drive the\n", + "likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"cluster\" / dataset_name\n", + "\n", + "if (\n", + " not (dataset_path / \"data.fits\").exists()\n", + " or not (dataset_path / \"mass.csv\").exists()\n", + "):\n", + " subprocess.run([sys.executable, \"scripts/cluster/simulator.py\"], check=True)\n", + "\n", + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Datasets__\n", + "\n", + "``point_datasets.csv`` carries one row per observed multiple image, grouped by source ``name``.\n", + "``al.list_from_csv`` returns a ``List[PointDataset]`` where each entry carries:\n", + "\n", + " - ``name`` \u2014 source identifier (e.g. ``point_0``), used for name pairing with the model.\n", + " - ``positions`` \u2014 image-plane (y, x) positions of every multiple image of that source.\n", + " - ``positions_noise_map`` \u2014 per-position positional uncertainty in arc-seconds.\n", + " - ``redshift`` \u2014 the source redshift (different per source \u2014 this is a multi-plane system)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = al.list_from_csv(file_path=dataset_path / \"point_datasets.csv\")\n", + "\n", + "for dataset in dataset_list:\n", + " print(f\" {dataset.name}: z={dataset.redshift} {len(dataset.positions)} images\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Truth Model__\n", + "\n", + "The simulator wrote the truth lens model into the family CSVs (``mass.csv`` for dPIE + NFW mass\n", + "profiles; ``point.csv`` for source-galaxy ``Point`` components). We load both and build the\n", + "truth ``Tracer``. The scaling tier (10 low-mass members) has its own legacy ``scaling_galaxies.csv``\n", + "schema with shared scaling-relation parameters.\n", + "\n", + "The likelihood walkthrough below evaluates the chi\u00b2 at the *truth* model \u2014 the chi\u00b2 ought to be at\n", + "or near its minimum here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass_table = al.galaxy_models_from_csv(dataset_path / \"mass.csv\", family=\"mass\")\n", + "point_table = al.galaxy_models_from_csv(dataset_path / \"point.csv\", family=\"point\")\n", + "scaling_table = al.galaxy_table_from_csv(dataset_path / \"scaling_galaxies.csv\")\n", + "\n", + "galaxies_by_name = al.galaxies_from_csv_tables(mass_table, point_table)\n", + "\n", + "redshift_lens = 0.5\n", + "source_redshifts = sorted({float(d.redshift) for d in dataset_list})\n", + "\n", + "main_lens_galaxies = [galaxies_by_name[\"lens_0\"], galaxies_by_name[\"lens_1\"]]\n", + "host_halo_galaxy = galaxies_by_name[\"host_halo\"]\n", + "source_galaxies = [galaxies_by_name[\"source_0\"], galaxies_by_name[\"source_1\"]]\n", + "\n", + "# Scaling tier: per-member dPIE built from the legacy CSV + the reference-anchored\n", + "# scaling relation (Lenstool convention; see modeling.py for the full rationale).\n", + "scaling_galaxies = []\n", + "SCALING_B0_REF_TRUTH = 0.12\n", + "SCALING_EXPONENT = 0.5\n", + "SCALING_RS_REF = 10.0\n", + "luminosity_ref = max(scaling_table.luminosities)\n", + "for centre, luminosity in zip(\n", + " scaling_table.centres.in_list, scaling_table.luminosities\n", + "):\n", + " luminosity_ratio = luminosity / luminosity_ref\n", + " b0 = SCALING_B0_REF_TRUTH * luminosity_ratio**SCALING_EXPONENT\n", + " rs = SCALING_RS_REF * luminosity_ratio**SCALING_EXPONENT\n", + " scaling_galaxies.append(\n", + " al.Galaxy(\n", + " redshift=redshift_lens,\n", + " mass=al.mp.dPIEMassSph(centre=tuple(centre), ra=0.1, rs=rs, b0=b0),\n", + " )\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "The tracer carries:\n", + "\n", + " - 2 main lens galaxies (BCG + satellite) \u2014 individually-modelled dPIE mass profiles.\n", + " - 10 scaling-tier member galaxies \u2014 dPIE mass profiles whose ``b0`` and ``rs`` derive from the\n", + " reference-anchored scaling relation ``b0 = b0_ref \u00d7 (L/L_ref)^0.5`` (Lenstool convention).\n", + " - 1 host dark matter halo \u2014 ``NFWMCRLudlowSph`` at the cluster centre.\n", + " - 2 source galaxies \u2014 ``Point`` profiles at distinct redshifts (multi-plane).\n", + "\n", + "PyAutoLens automatically groups galaxies by redshift into ``planes`` and ray-traces through them\n", + "in redshift order. For this cluster the planes are: ``z=0.5`` (lens-plane, holds 13 galaxies),\n", + "``z=1.0`` (source_0 plane), ``z=2.0`` (source_1 plane)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=main_lens_galaxies\n", + " + scaling_galaxies\n", + " + [host_halo_galaxy]\n", + " + source_galaxies\n", + ")\n", + "\n", + "print(\n", + " f\"Tracer has {len(tracer.planes)} planes at redshifts \"\n", + " f\"{[float(p.redshift) for p in tracer.planes]}\"\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Chi Squared: Concept__\n", + "\n", + "If the lens model is correct, every observed image position of a given source ought to ray-trace\n", + "back to (approximately) the same source-plane location: the source's true centre. The figure of\n", + "merit for the **source-plane chi\u00b2** is the scatter of those back-traced positions around the\n", + "source-plane reference, weighted by magnification (so source-plane residuals are converted back to\n", + "image-plane scale where the noise is defined).\n", + "\n", + "This contrasts with the more intuitive image-plane chi\u00b2 which compares observed image positions to\n", + "the model's *forward-solved* multiple-image positions. Source-plane chi\u00b2 is cheaper (no forward\n", + "solve required, just one back-projection per image), and the math is JAX-friendly. The trade-off\n", + "is that the magnification weighting amplifies any positional error in the back-projection: at\n", + "cluster scale where magnifications near multi-image positions are ~100\u00d7, even sub-arcsecond\n", + "source-plane residuals can drive chi\u00b2 to 10\u2076\u201310\u2078 at the truth.\n", + "\n", + "This is the figure of merit used by ``al.FitPositionsSource`` and is described in detail in\n", + "Jullo et al. 2007 (\"A Bayesian approach to strong lensing modelling of galaxy clusters\") and the\n", + "references therein.\n", + "\n", + "__Multi-Plane Ray Tracing__\n", + "\n", + "The cluster has two sources at *different* redshifts (``z = 1.0`` and ``z = 2.0``), so ray-tracing\n", + "each source's positions back to its source plane goes through every earlier plane along the way.\n", + "This is **multi-plane ray tracing**, formalised by the recursive lens equation:\n", + "\n", + ".. math::\n", + "\n", + " \\\\theta_{j} = \\\\theta_{0} - \\\\sum_{i=1}^{j-1} \\\\beta_{ij} \\\\alpha_{i}(\\\\theta_{i})\n", + "\n", + "where:\n", + "\n", + " - :math:`\\\\theta_{0}` is the image-plane position,\n", + " - :math:`\\\\theta_{i}` is the position of the ray at plane ``i``,\n", + " - :math:`\\\\alpha_{i}(\\\\theta_{i})` is the deflection at plane ``i`` (sum of every galaxy at that\n", + " plane's redshift),\n", + " - :math:`\\\\beta_{ij} = (D_{ij} \\\\, D_{s}) / (D_{j} \\\\, D_{is})` is a cosmological scaling factor\n", + " that scales the deflection from plane ``i``'s angular-diameter convention to plane ``j``'s.\n", + "\n", + "For our cluster the relevant cases per source are:\n", + "\n", + " - **Source 0 (z=1.0)** \u2014 back-tracing its image positions just hits the lens plane (z=0.5).\n", + " Two planes total, so the scaling factor reduces to 1.0 and the equation is the familiar\n", + " :math:`\\\\theta_{1} = \\\\theta_{0} - \\\\alpha_{0}(\\\\theta_{0})`.\n", + "\n", + " - **Source 1 (z=2.0)** \u2014 back-tracing hits the lens plane (z=0.5) AND the source_0 plane\n", + " (z=1.0). Even though source_0 carries no mass, the recursive equation still walks through its\n", + " plane; the actual contribution depends on whether any galaxy at z=1.0 has a non-zero mass\n", + " profile (none here \u2014 the source_0 ``Point`` profile contributes nothing to the deflection\n", + " field).\n", + "\n", + "The library encapsulates all of this in ``Tracer.deflections_between_planes_from(grid, plane_i=0,\n", + "plane_j=)``. The function walks the planes from ``plane_i`` to ``plane_j``\n", + "in redshift order, applying the per-plane :math:`\\\\alpha_{i}(\\\\theta_{i})` sums and the scaling\n", + "factors :math:`\\\\beta_{ij}` automatically. The full derivation of the recursive equation, the\n", + "sign conventions, and the cosmological scaling factors live in the multi-plane guide at\n", + "``autolens_workspace/scripts/guides/lensing/multi_plane.py``.\n", + "\n", + "In practice the easiest entry point isn't ``deflections_between_planes_from`` (which returns the\n", + "*differences* between plane positions and requires you to apply the lens equation manually) but\n", + "``Tracer.traced_grid_2d_list_from(grid)``. That function returns the list of grids per plane,\n", + "fully traced through the recursive lens equation, with the cosmological scaling factors already\n", + "applied. For our 3-plane cluster:\n", + "\n", + " - ``traced_grids[0]`` is the input image-plane grid (unchanged).\n", + " - ``traced_grids[1]`` is the source_0 plane (z=1.0) position of every input ray after passing\n", + " through the lens-plane deflection.\n", + " - ``traced_grids[2]`` is the source_1 plane (z=2.0) position of every input ray after passing\n", + " through the lens plane *and* the source_0 plane (the recursive step).\n", + "\n", + "Each source's back-traced positions are then just ``traced_grids[]``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_plane_positions_per_source = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " plane_index = tracer.plane_index_via_redshift_from(redshift=dataset.redshift)\n", + " traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.positions)\n", + " source_plane_positions_per_source.append(traced_grids[plane_index])\n", + "\n", + " print(\n", + " f\" {dataset.name}: plane_index={plane_index}, \"\n", + " f\"back-traced positions = {traced_grids[plane_index].in_list}\"\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Back-Traced Source-Plane Positions: Conceptual Recap__\n", + "\n", + "The traced grid at plane ``j`` is, by construction:\n", + "\n", + ".. math::\n", + "\n", + " \\\\theta_{j} = \\\\theta_{0} - \\\\alpha_{\\\\text{multi-plane}}(\\\\theta_{0}; j)\n", + "\n", + "where the multi-plane deflection :math:`\\\\alpha_{\\\\text{multi-plane}}` is the recursive sum\n", + "discussed in the previous section \u2014 exactly the result you would get from manually applying\n", + "``deflections_between_planes_from`` and then ``grid_2d_via_deflection_grid_from``. Using\n", + "``traced_grid_2d_list_from`` is the same thing in one call.\n", + "\n", + "__Source-Plane Centroid__\n", + "\n", + "The \"reference point\" against which we measure the source-plane scatter has two options:\n", + "\n", + " 1. **Truth ``Point`` centre.** Each source carries a ``Point`` profile in the model whose\n", + " ``centre`` is a free parameter (or fixed to the truth here). At a model fit, ``Point.centre``\n", + " *is* the source-plane (y, x) the multiple images should converge to.\n", + "\n", + " 2. **Barycenter of back-traced positions.** Pretend you don't know the truth centre. Compute the\n", + " centroid of the back-traced positions \u2014 at the right model the centroid sits where the source\n", + " actually is, and the residuals are scatter around it.\n", + "\n", + "Option (1) is what ``al.FitPositionsSource(profile=point_profile)`` uses. Option (2) is what\n", + "``al.FitPositionsSource(profile=None)`` uses (the default during model fits, because at search\n", + "time the model doesn't yet know the truth).\n", + "\n", + "For this walkthrough we use option (1) so the residuals have a well-defined physical meaning\n", + "(distance from truth, not from a derived centroid)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_plane_centroids = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " # source_i's Point profile is attached to source_galaxies[i] under attr name \"point_i\".\n", + " point_profile = getattr(source_galaxies[i], dataset.name)\n", + " source_plane_centroids.append(point_profile.centre)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Residual Map__\n", + "\n", + "The per-image residual is just the source-plane distance between each back-traced position and the\n", + "source-plane centroid:\n", + "\n", + ".. math::\n", + "\n", + " r_{i} = |\\\\theta_{j,i} - \\\\theta_{j,\\\\text{centre}}|" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "residuals_per_source = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " sp_positions = source_plane_positions_per_source[i]\n", + " centre = source_plane_centroids[i]\n", + " residuals = sp_positions.distances_to_coordinate_from(coordinate=centre)\n", + " residuals_per_source.append(residuals)\n", + "\n", + " print(f\" {dataset.name}: residuals = {np.asarray(residuals)}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Magnifications at Positions__\n", + "\n", + "The source-plane chi\u00b2 weights each residual by the local **magnification** at that image position.\n", + "Why: position noise is defined in the image plane (arcsec), but the residual is in the source\n", + "plane. The magnification converts a source-plane distance back to its image-plane scale.\n", + "Specifically the formula below uses magnification squared as a multiplicative weight.\n", + "\n", + "PyAutoLens computes magnification from the Hessian of the deflection field via\n", + "``Tracer.magnification_2d_via_hessian_from(grid)``. Physically the Hessian gives the linear\n", + "distortion matrix at each image position; its determinant is the magnification (signed: negative\n", + "for parity-flipped images, but ``magnification_2d_via_hessian_from`` returns the absolute value).\n", + "The conceptual chain is **deflection field \u2192 Hessian via finite-difference \u2192 eigenvalues\n", + "(convergence \u03ba and shear \u03b3) \u2192 magnification :math:`\\\\mu = 1 / |(1-\\\\kappa)^2 - \\\\gamma^2|`**.\n", + "\n", + "Full derivation and the eigenvalue / critical-curve discussion live in the lens-calc guide at\n", + "``autolens_workspace/scripts/guides/lensing/lens_calc.py``.\n", + "\n", + "For multi-plane lenses the function uses the full multi-plane deflection in the Hessian, so the\n", + "magnification correctly reflects the cumulative distortion through every intervening plane. The\n", + "``LensCalc`` helper bundles the Hessian, magnification, and critical-curve calculations; we build\n", + "one per source so each magnification is evaluated against the correct multi-plane chain (the\n", + "source's plane index).\n", + "\n", + "The library wraps this with an ``abs(...)`` so the returned magnification is always positive (raw\n", + "magnification can be negative when the image is parity-flipped; squaring it in the chi\u00b2 formula\n", + "makes the sign irrelevant anyway)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnifications_per_source = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " plane_index = tracer.plane_index_via_redshift_from(redshift=dataset.redshift)\n", + " od = ag.LensCalc.from_tracer(\n", + " tracer=tracer, use_multi_plane=True, plane_j=plane_index\n", + " )\n", + " mag = abs(od.magnification_2d_via_hessian_from(grid=dataset.positions))\n", + " magnifications_per_source.append(mag)\n", + "\n", + " print(f\" {dataset.name}: magnifications = {np.asarray(mag)}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Chi Squared Map (Source)__\n", + "\n", + "The per-image source-plane chi\u00b2 combines the residual, the magnification, and the position noise:\n", + "\n", + ".. math::\n", + "\n", + " \\\\chi^{2}_{i} = \\\\frac{r_{i}^{2} \\\\, \\\\mu_{i}^{2}}{\\\\sigma_{i}^{2}}\n", + "\n", + "This is the formula in ``autolens/point/fit/positions/source/separations.py`` ::\n", + "\n", + " chi_squared_map = residual_map**2 / (magnifications_at_positions.array**-2 * noise_map.array**2)\n", + "\n", + "which expands to ``residual\u00b2 \u00d7 magnification\u00b2 / noise\u00b2``. (The :math:`\\\\mu^{-2}` in the denominator\n", + "of the source-code form cancels the :math:`\\\\mu^{2}` in the numerator and leaves the form above.)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_maps_per_source = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " r = np.asarray(residuals_per_source[i])\n", + " mu = np.asarray(magnifications_per_source[i])\n", + " sigma = np.asarray(dataset.positions_noise_map)\n", + " chi_sq_map = r**2 * mu**2 / sigma**2\n", + " chi_squared_maps_per_source.append(chi_sq_map)\n", + "\n", + " print(f\" {dataset.name}: chi\u00b2 per image = {chi_sq_map}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Source Chi Squared__\n", + "\n", + "Per-source: sum the chi\u00b2 map across that source's multiple images." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_per_source = [float(np.sum(m)) for m in chi_squared_maps_per_source]\n", + "for i, dataset in enumerate(dataset_list):\n", + " print(f\" {dataset.name}: chi\u00b2 = {chi_squared_per_source[i]:.4e}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Total Chi Squared (Source)__\n", + "\n", + "Each source is independent (different redshift, different multiple-image set), so the total chi\u00b2\n", + "is just the sum." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_chi_squared_source = float(sum(chi_squared_per_source))\n", + "print(f\"Total source-plane chi\u00b2 = {total_chi_squared_source:.4e}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization (Source)__\n", + "\n", + "The full Gaussian log-likelihood carries a normalisation term that depends on the noise:\n", + "\n", + ".. math::\n", + "\n", + " \\\\mathcal{N} = \\\\sum_{i} \\\\log \\\\left( 2\\\\pi \\\\, \\\\mu_{i}^{-2} \\\\, \\\\sigma_{i}^{2} \\\\right)\n", + "\n", + "The magnification factor appears here too because the chi\u00b2 formula effectively treats each image\n", + "position as having an *effective* source-plane noise of :math:`\\\\sigma_{i} / \\\\mu_{i}`. The\n", + "normalisation matches that interpretation so the resulting expression is a well-defined Gaussian\n", + "log-likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalizations_per_source = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " mu = np.asarray(magnifications_per_source[i])\n", + " sigma = np.asarray(dataset.positions_noise_map)\n", + " nn = float(np.sum(np.log(2 * np.pi * (mu**-2) * sigma**2)))\n", + " noise_normalizations_per_source.append(nn)\n", + "\n", + "total_noise_normalization_source = float(sum(noise_normalizations_per_source))\n", + "print(\n", + " f\"Total source-plane noise normalization = {total_noise_normalization_source:.4e}\"\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Log Likelihood__\n", + "\n", + "The standard Gaussian log-likelihood:\n", + "\n", + ".. math::\n", + "\n", + " \\\\log L = -\\\\frac{1}{2} \\\\left( \\\\chi^{2} + \\\\mathcal{N} \\\\right)\n", + "\n", + "This is what Nautilus maximises during a cluster point-source model fit. The chi\u00b2 term encodes\n", + "goodness-of-fit; the normalisation absorbs the constant noise-dependent prefactor of the Gaussian." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_likelihood_source = -0.5 * (\n", + " total_chi_squared_source + total_noise_normalization_source\n", + ")\n", + "print(f\"Source-plane log likelihood = {log_likelihood_source:.4e}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Validation__\n", + "\n", + "To verify the step-by-step derivation, we instantiate ``al.FitPositionsSource`` and confirm its\n", + "``log_likelihood`` matches. ``FitPositionsSource`` does exactly the calculation above internally\n", + "(see ``autolens/point/fit/positions/source/separations.py``)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sum_library_log_likelihood_source = 0.0\n", + "for i, dataset in enumerate(dataset_list):\n", + " point_profile = getattr(source_galaxies[i], dataset.name)\n", + " fit = al.FitPositionsSource(\n", + " name=dataset.name,\n", + " data=dataset.positions,\n", + " noise_map=dataset.positions_noise_map,\n", + " tracer=tracer,\n", + " solver=None,\n", + " profile=point_profile,\n", + " )\n", + " sum_library_log_likelihood_source += float(fit.log_likelihood)\n", + "\n", + "print(f\"Library source-plane log likelihood = {sum_library_log_likelihood_source:.4e}\")\n", + "print(\n", + " f\"Match: {np.isclose(log_likelihood_source, sum_library_log_likelihood_source, rtol=1e-6)}\"\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image-Plane Chi Squared: Concept__\n", + "\n", + "The image-plane chi\u00b2 takes the opposite approach. Instead of ray-tracing observed positions to\n", + "the source plane and measuring source-plane scatter, it **forward-solves** the model's image-plane\n", + "positions for each source and compares them to the observed positions in the image plane.\n", + "\n", + "Mechanically:\n", + "\n", + " 1. Take the source-plane centre (model's ``Point.centre`` per source).\n", + " 2. Use a ``PointSolver`` to ray-trace triangles from a fine image-plane grid forward through the\n", + " lens until they converge to that source-plane centre \u2014 these are the **model positions**.\n", + " 3. For each model position, find the closest observed position and pair them.\n", + " 4. The residual per pair is just the image-plane Euclidean distance.\n", + "\n", + "The chi\u00b2 is then ``residual\u00b2 / noise\u00b2`` \u2014 no magnification weighting, because the residual is\n", + "already in the same units (arc-seconds) as the noise. This makes the absolute chi\u00b2 much smaller\n", + "than the source-plane variant at the same model: at the truth, residuals are bounded by the\n", + "``PointSolver`` precision (~0.001\"), so chi\u00b2 \u2248 (0.001 / 0.005)\u00b2 \u00d7 N \u2248 N \u00d7 0.04, where N is the\n", + "number of paired images. Far below the source-plane chi\u00b2 which is dominated by magnification \u00d7 the\n", + "same precision floor.\n", + "\n", + "The trade-off: image-plane chi\u00b2 requires a **forward solve per evaluation**, which is slow,\n", + "introduces solver precision as a noise floor, and carries pairing pathologies discussed below.\n", + "The source-plane chi\u00b2 has none of those costs.\n", + "\n", + "__Point Solver Setup__\n", + "\n", + "The ``PointSolver`` searches for image-plane positions whose forward ray-trace lands at the\n", + "source-plane target. Conceptually it does this by:\n", + "\n", + " 1. Starting from a coarse image-plane grid.\n", + " 2. Tessellating into triangles. For each triangle, ray-trace its three vertices to the source\n", + " plane. If the source-plane target lies inside the traced triangle, that image-plane triangle\n", + " contains a model position.\n", + " 3. Refining each containing triangle \u2014 subdivide, ray-trace the sub-triangles, repeat \u2014 until\n", + " the triangle size drops below a configurable ``pixel_scale_precision``.\n", + " 4. The triangle centroid at convergence is the model position.\n", + "\n", + "This is iterative and the runtime scales with the source-plane precision target. The solver also\n", + "filters out central images via ``magnification_threshold``: highly demagnified images (often the\n", + "unobservable central image of a strong-lens configuration) are discarded so the model doesn't pair\n", + "unphysical demagnified solutions to the observed multi-image set.\n", + "\n", + "A standalone walkthrough of the triangle-refinement algorithm (sub-grid traversal, magnification\n", + "filtering, multi-plane handling, JAX-compatibility) lives in\n", + "``autolens_workspace/scripts/guides/point_source/triangle_solver.py`` \u2014 TODO, currently not yet\n", + "written.\n", + "\n", + "In code, the solver is constructed once and reused per evaluation:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=1.0),\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.1,\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Forward Solving Model Positions__\n", + "\n", + "For each source, call ``solver.solve(tracer, source_plane_coordinate=...)`` with the source's\n", + "truth source-plane centre. The solver returns the image-plane (y, x) positions of every\n", + "multiple image whose magnification is above ``magnification_threshold``.\n", + "\n", + "Multi-plane ray tracing happens *inside* ``solve()`` automatically. For source_1 (z=2.0) the\n", + "triangles are forward-traced through the z=1.0 source_0 plane before reaching z=2.0 \u2014 the same\n", + "recursive lens equation discussed above is applied during the forward solve. The solver picks the\n", + "right target plane via the ``plane_redshift`` argument: this is the source galaxy's redshift, and\n", + "the solver maps it to the corresponding plane index in the tracer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_positions_per_source = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " point_profile = getattr(source_galaxies[i], dataset.name)\n", + " model_positions = solver.solve(\n", + " tracer=tracer,\n", + " source_plane_coordinate=point_profile.centre,\n", + " plane_redshift=dataset.redshift,\n", + " )\n", + " model_positions_per_source.append(model_positions)\n", + " print(f\" {dataset.name}: model positions = {model_positions.in_list}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pairing Model to Observed__\n", + "\n", + "We have two sets of image-plane positions per source: the **model positions** (forward-solved) and\n", + "the **observed positions** (from ``point_datasets.csv``). The chi\u00b2 requires a one-to-one mapping\n", + "between them \u2014 but the two sets don't come pre-paired, and may not even have the same length.\n", + "\n", + "**Three pairing schemes** exist in PyAutoLens, each with different behaviour when the counts\n", + "don't match:\n", + "\n", + " 1. **``FitPositionsImagePair`` (Hungarian, no-repeat).** Pairs each observed to its nearest\n", + " predicted image via the **Hungarian algorithm** (also called linear sum assignment). The\n", + " algorithm finds the unique 1-to-1 pairing that *globally minimises the total distance* \u2014\n", + " not greedy. Two consequences: (a) it gives the optimal assignment even when greedy would\n", + " fail; (b) when counts differ, unmatched positions still get unpaired and contribute nothing\n", + " to chi\u00b2.\n", + "\n", + " 2. **``FitPositionsImagePairAll`` (closest, with replacement).** Each model position pairs to its\n", + " nearest observed position independently. A single observed position may be paired to multiple\n", + " model positions. Every model position contributes to chi\u00b2.\n", + "\n", + " 3. **``FitPositionsImagePairRepeat`` (repeats allowed).** Similar to (2) but with explicit\n", + " handling of cases where the model genuinely predicts multiple images at nearly the same\n", + " location (e.g. near a caustic crossing).\n", + "\n", + "This walkthrough uses scheme (1) \u2014 ``FitPositionsImagePair`` \u2014 which is the historical default and\n", + "the one ``al.FitPointDataset`` constructs by default. We follow that convention.\n", + "\n", + "**The known pathology** (documented in the PyAutoLens source for ``FitPositionsImagePair``):\n", + "\n", + " - **Model predicts too many images** (more model than observed). The pairing leaves some model\n", + " positions unpaired. These model positions contribute *nothing* to chi\u00b2. A model that\n", + " spuriously generates extra demagnified images can therefore have its chi\u00b2 artificially\n", + " *reduced* compared to a model that produces exactly the observed count.\n", + " - **Model predicts too few images** (fewer model than observed). Some observed positions go\n", + " unpaired and also contribute nothing. The chi\u00b2 is similarly artificially reduced.\n", + "\n", + "In both cases the optimiser can prefer pathological solutions with the wrong image count over the\n", + "correct solution. ``magnification_threshold`` partially mitigates this by discarding very\n", + "demagnified spurious images, but it is not a complete fix. ``FitPositionsImagePairAll`` /\n", + "``...Repeat`` address some of the cases, at the cost of other failure modes. Selecting the right\n", + "scheme for a real cluster fit is a per-dataset judgement call.\n", + "\n", + "In code we do the same pairing the library does \u2014 Hungarian / linear sum assignment \u2014 by\n", + "computing the pairwise distance matrix and handing it to ``scipy.optimize.linear_sum_assignment``:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.optimize import linear_sum_assignment\n", + "\n", + "\n", + "def _pair_hungarian(model_positions, observed_positions):\n", + " \"\"\"Optimal 1-1 pairing via Hungarian algorithm.\n", + "\n", + " Returns a list of (model_position, observed_position, distance) tuples. When the number of\n", + " model and observed positions differ, only ``min(N_model, N_observed)`` pairs are returned\n", + " \u2014 the excess positions on the larger side go unpaired and contribute nothing to chi\u00b2.\n", + " \"\"\"\n", + " model_arr = np.asarray(model_positions)\n", + " observed_arr = np.asarray(observed_positions)\n", + "\n", + " # Pairwise distance matrix: distances[i, j] = ||model_i - observed_j||.\n", + " distances = np.linalg.norm(\n", + " model_arr[:, np.newaxis, :] - observed_arr[np.newaxis, :, :], axis=2\n", + " )\n", + "\n", + " row_ind, col_ind = linear_sum_assignment(distances)\n", + "\n", + " return [\n", + " (model_arr[i], observed_arr[j], float(distances[i, j]))\n", + " for i, j in zip(row_ind, col_ind)\n", + " ]\n", + "\n", + "\n", + "pairs_per_source = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " pairs = _pair_hungarian(model_positions_per_source[i], dataset.positions)\n", + " pairs_per_source.append(pairs)\n", + " print(\n", + " f\" {dataset.name}: {len(pairs)} pairs / {len(dataset.positions)} observed / \"\n", + " f\"{len(model_positions_per_source[i])} model\"\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image-Plane Residual Map__\n", + "\n", + "Residual per pair = image-plane Euclidean distance between model and observed (in arc-seconds).\n", + "This is what ``_pair_closest_no_repeat`` already returned as the third element of each tuple." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "residual_maps_image_per_source = []\n", + "for i, pairs in enumerate(pairs_per_source):\n", + " residuals = np.array([p[2] for p in pairs])\n", + " residual_maps_image_per_source.append(residuals)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Chi Squared Map (Image)__\n", + "\n", + "Image-plane chi\u00b2 per pair:\n", + "\n", + ".. math::\n", + "\n", + " \\\\chi^{2}_{i} = \\\\frac{r_{i}^{2}}{\\\\sigma_{i}^{2}}\n", + "\n", + "No magnification weighting \u2014 the residual is already in image-plane arc-seconds, the same units\n", + "as the noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_maps_image_per_source = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " r = residual_maps_image_per_source[i]\n", + " # Use the noise of each paired observed position. For simplicity we use the dataset's mean\n", + " # noise since all positions in the cluster simulator share the same \u03c3; in a real analysis\n", + " # this should index per-observed-position.\n", + " sigma = float(np.mean(np.asarray(dataset.positions_noise_map)))\n", + " chi_sq_map = r**2 / sigma**2\n", + " chi_squared_maps_image_per_source.append(chi_sq_map)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Source / Total Chi Squared (Image)__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_per_source_image = [\n", + " float(np.sum(m)) for m in chi_squared_maps_image_per_source\n", + "]\n", + "total_chi_squared_image = float(sum(chi_squared_per_source_image))\n", + "print(f\"Total image-plane chi\u00b2 = {total_chi_squared_image:.4e}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization (Image)__\n", + "\n", + "The Gaussian-log normalisation for image-plane chi\u00b2 is the standard form:\n", + "\n", + ".. math::\n", + "\n", + " \\\\mathcal{N} = \\\\sum_{i} \\\\log \\\\left( 2\\\\pi \\\\, \\\\sigma_{i}^{2} \\\\right)\n", + "\n", + "No magnification factor (residuals are already in image-plane units)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalizations_per_source_image = []\n", + "for i, dataset in enumerate(dataset_list):\n", + " n_pairs = len(residual_maps_image_per_source[i])\n", + " sigma = float(np.mean(np.asarray(dataset.positions_noise_map)))\n", + " nn = float(n_pairs * np.log(2 * np.pi * sigma**2))\n", + " noise_normalizations_per_source_image.append(nn)\n", + "\n", + "total_noise_normalization_image = float(sum(noise_normalizations_per_source_image))\n", + "print(f\"Total image-plane noise normalization = {total_noise_normalization_image:.4e}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image-Plane Log Likelihood__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_likelihood_image = -0.5 * (\n", + " total_chi_squared_image + total_noise_normalization_image\n", + ")\n", + "print(f\"Image-plane log likelihood = {log_likelihood_image:.4e}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image-Plane Validation__\n", + "\n", + "Instantiate ``al.FitPositionsImagePair`` and confirm match. Note: ``FitPositionsImagePair`` uses\n", + "the same closest-no-repeat pairing scheme as our manual implementation above, so the chi\u00b2 values\n", + "should agree to within numerical precision (the library's solver may use slightly different\n", + "triangle precision at refinement edge cases, producing sub-percent residual differences)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sum_library_log_likelihood_image = 0.0\n", + "for i, dataset in enumerate(dataset_list):\n", + " point_profile = getattr(source_galaxies[i], dataset.name)\n", + " fit = al.FitPositionsImagePair(\n", + " name=dataset.name,\n", + " data=dataset.positions,\n", + " noise_map=dataset.positions_noise_map,\n", + " tracer=tracer,\n", + " solver=solver,\n", + " profile=point_profile,\n", + " )\n", + " sum_library_log_likelihood_image += float(fit.log_likelihood)\n", + "\n", + "print(f\"Library image-plane log likelihood = {sum_library_log_likelihood_image:.4e}\")\n", + "print(\n", + " f\"Match (rtol=1e-2): \"\n", + " f\"{np.isclose(log_likelihood_image, sum_library_log_likelihood_image, rtol=1e-2)}\"\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane vs Image-Plane: When to Use Which__\n", + "\n", + "| Aspect | Source-plane (``FitPositionsSource``) | Image-plane (``FitPositionsImagePair``) |\n", + "|---|---|---|\n", + "| **Forward solve required?** | No | Yes (one per evaluation) |\n", + "| **Per-evaluation cost** | Cheap (just back-projection) | Expensive (triangle refinement) |\n", + "| **JAX-friendly?** | Yes | The solver is JAX-jitted, but each evaluation triggers a forward solve regardless |\n", + "| **Magnification weighting** | Yes (``\u03bc\u00b2`` in chi\u00b2) | No |\n", + "| **Sensitive to ``PointSolver`` precision?** | Amplified by ``\u03bc\u00b2`` at multi-image positions | Bounded by precision directly |\n", + "| **Pairing pathology** | None (uses all images) | Yes (too-many / too-few; see above) |\n", + "| **Absolute chi\u00b2 scale (at truth)** | Large (~``\u03bc\u00b2`` \u00d7 precision\u00b2) | Small (~precision\u00b2) |\n", + "| **Best for** | Fast Nautilus fits; JAX-jit'd parameter estimation | Final residual visualisation; cases where pairing is unambiguous |\n", + "\n", + "For most cluster fits the source-plane chi\u00b2 is the right default: it's faster, JAX-compatible,\n", + "and doesn't suffer pairing pathologies. The image-plane chi\u00b2 is most useful for diagnostic\n", + "visualisation of where the model's predicted images sit relative to the observed ones, and for\n", + "cases where the source-plane chi\u00b2 has bias issues that need cross-checking.\n", + "\n", + "The cluster modelling script at ``scripts/cluster/modeling.py`` uses ``AnalysisPoint`` which\n", + "selects the chi\u00b2 flavour via its constructor; consult the ``AnalysisPoint`` docstring for the\n", + "current default.\n", + "\n", + "__Wrap Up__\n", + "\n", + "This script has walked through the cluster point-source likelihood end to end, in both source-plane\n", + "and image-plane flavours, for a multi-plane (2-source) cluster.\n", + "\n", + "Next steps:\n", + "\n", + "- ``scripts/cluster/csv_api.py`` \u2014 the family CSV schema this script reads.\n", + "- ``scripts/cluster/simulator.py`` \u2014 how the truth dataset was generated.\n", + "- ``scripts/cluster/modeling.py`` \u2014 full Nautilus fit using the same likelihood function discussed\n", + " here, with priors and ``AnalysisPoint`` configuration.\n", + "- ``autolens_workspace_test/scripts/cluster/likelihood_sanity.py`` \u2014 perturbation sweep that\n", + " surfaces the ``PointSolver`` precision-floor pathology described in the source-plane section.\n", + "- ``autolens_workspace/scripts/guides/lensing/multi_plane.py`` \u2014 full derivation of the recursive\n", + " lens equation.\n", + "- ``autolens_workspace/scripts/guides/lensing/lens_calc.py`` \u2014 Hessian / magnification /\n", + " critical-curve mechanics.\n", + "\n", + "For a deeper understanding of cluster lens modelling and point-source likelihoods, the\n", + "**HowToLens** Jupyter notebook lectures cover both topics in detail.\n", + "\n", + "__JAX__\n", + "\n", + "The chi-squared walkthrough above is pure NumPy. To JAX-accelerate it,\n", + "wrap construction in `@jax.jit` with the post-Phase-2 `PointSolver`\n", + "pattern:\n", + "\n", + "```python\n", + "import jax\n", + "import jax.numpy as jnp\n", + "from autolens.jax import register_tracer_classes\n", + "\n", + "register_tracer_classes(tracer) # one-time\n", + "\n", + "@jax.jit\n", + "def cluster_log_likelihood(tracer, dataset, source_galaxy):\n", + " fit = al.FitPositionsSource(\n", + " name=dataset.name,\n", + " data=dataset.positions,\n", + " noise_map=dataset.positions_noise_map,\n", + " tracer=tracer,\n", + " solver=None, # source-plane chi\u00b2 doesn't need the solver\n", + " profile=getattr(source_galaxy, dataset.name),\n", + " )\n", + " return fit.log_likelihood\n", + "```\n", + "\n", + "For the canonical search-driven path (`AnalysisPoint(use_jax=True)`),\n", + "see `modeling.py`. For JIT-ing library methods directly without going\n", + "through `FitPositionsSource`, see `scripts/guides/lens_calc.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/cluster/modeling.ipynb b/notebooks/cluster/modeling.ipynb index 13c17ede5..3fcea5e88 100644 --- a/notebooks/cluster/modeling.ipynb +++ b/notebooks/cluster/modeling.ipynb @@ -1,855 +1,919 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Cluster\n", - "=================\n", - "\n", - "This script models the small multi-plane cluster lens simulated by ``cluster/simulator.py``: a Brightest\n", - "Cluster Galaxy (BCG) plus one satellite member at the lens redshift ``z = 0.5``, 10 lower-mass cluster\n", - "members modelled collectively via a luminosity-mass scaling relation, a standalone NFW host dark matter\n", - "halo *not* tied to any individual galaxy, and 2 background sources at *different* redshifts (``z = 1.0``\n", - "and ``z = 2.0``) which the cluster lenses into multiple images.\n", - "\n", - "Cluster modeling almost always uses the *point source* API: rather than fitting the extended arc light\n", - "of each lensed source, only the image-plane positions of the brightest pixels of each multiple image are\n", - "fitted. Per-source positions, noise maps, and redshifts are loaded from a single hand-editable CSV that\n", - "the simulator writes (``point_datasets.csv``) via ``al.list_from_csv``. Cluster-member centres and\n", - "luminosities for the scaling tier are loaded from a second CSV (``scaling_galaxies.csv``) via\n", - "``al.galaxy_table_from_csv``.\n", - "\n", - "__Contents__\n", - "\n", - "- **Example:** What this script fits and the underlying simulator.\n", - "- **Simulation:** Overview of how the simulated dataset was generated.\n", - "- **Dataset:** Load the CCD image and the per-source point datasets from the combined CSV.\n", - "- **Model CSVs:** Load the named-galaxy mass + point CSVs written by the simulator.\n", - "- **Scaling Galaxies Table:** Load the scaling-tier centres + luminosities from the CSV.\n", - "- **Point Solver:** Set up the image-plane multiple-image solver.\n", - "- **Chi Squared:** Why this script uses an image-plane chi-squared.\n", - "- **Cluster Components:** The four categories of lensing object \u2014 main lens galaxies, scaling members, host halo, sources.\n", - "- **Redshifts:** Multi-plane redshift handling and the source-redshift / dataset-redshift pairing.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Scaling Relation:** The shared two-parameter relation that ties every scaling member's mass to its luminosity.\n", - "- **Name Pairing:** How each ``Point`` model component is paired to its ``PointDataset``.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", - "- **Analysis:** Create the ``AnalysisPoint`` objects, one per dataset.\n", - "- **Factor Graph:** Combine per-dataset analyses into one global ``FactorGraphModel``.\n", - "- **Run Times:** Profiling the expected run time of the model-fit.\n", - "- **Output Folder Layout:** Description of the ``output`` folder structure.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Example__\n", - "\n", - "This script fits a ``PointDataset`` of a small multi-plane cluster where:\n", - "\n", - " - There are 2 main lens galaxies with ``dPIEMassSph`` total mass distributions, each with their centre\n", - " fixed to the values written out by the simulator [6 parameters].\n", - " - There are 10 scaling-tier member galaxies. Each carries a ``dPIEMassSph`` mass with centre fixed,\n", - " ``ra`` and ``rs`` fixed at the simulator truth values, and ``b0`` derived from the *shared*\n", - " scaling-relation parameters and the per-member luminosity [2 parameters total for the entire tier].\n", - " - There is 1 standalone ``NFWMCRLudlowSph`` host dark matter halo with its centre fixed and a free\n", - " ``mass_at_200`` [1 parameter].\n", - " - There are 2 source galaxies modeled as ``Point`` sources, each with its redshift pinned to the value\n", - " in its ``PointDataset`` row [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=13.\n", - "\n", - "The defining feature of cluster modeling is the scaling tier: 10 lower-mass members are fit jointly\n", - "with just *2 free parameters* (``scaling_factor`` and ``scaling_exponent``). Adding more members to\n", - "``scaling_galaxies.csv`` in the future does not grow the dimensionality of parameter space.\n", - "\n", - "__Simulation__\n", - "\n", - "This script fits the simulated cluster dataset produced by ``autolens_workspace/*/cluster/simulator.py``.\n", - "That simulator writes:\n", - "\n", - " - ``data.fits`` / ``noise_map.fits`` / ``psf.fits`` \u2014 CCD imaging of the cluster (used for visualization).\n", - " - ``point_datasets.csv`` \u2014 one row per observed multiple image, grouped by source ``name``, with a\n", - " ``redshift`` column per group. Loaded here with ``al.list_from_csv``.\n", - " - ``mass.csv`` + ``light.csv`` + ``point.csv`` \u2014 named-galaxy CSVs carrying the full truth model\n", - " (main galaxies + host halo + sources). Loaded here with ``al.galaxy_models_from_csv``. See\n", - " ``scripts/cluster/csv_api.py`` for the schema walkthrough.\n", - " - ``scaling_galaxies.csv`` \u2014 one row per scaling-tier member with columns ``y, x, luminosity``. Loaded\n", - " here with ``al.galaxy_table_from_csv``.\n", - " - ``tracer.json`` \u2014 true ``Tracer`` (used by visualization, not modeling)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens dataset ``cluster``, which is the dataset we will use to perform lens modeling.\n", - "\n", - "We begin by loading a CCD image of the dataset. Although we perform point-source modeling and will not\n", - "use the imaging data in the model-fit, it is useful to load it for visualization.\n", - "\n", - "The ``pixel_scales`` define the arc-second to pixel conversion factor of the image, which for the\n", - "dataset we are using is 0.1\" / pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", \"cluster\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if (\n", - " not (dataset_path / \"data.fits\").exists()\n", - " or not (dataset_path / \"scaling_galaxies.csv\").exists()\n", - " or not (dataset_path / \"mass.csv\").exists()\n", - "):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/cluster/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Datasets__\n", - "\n", - "We load the point datasets from the combined ``point_datasets.csv`` written by the simulator. This\n", - "returns a ``List[PointDataset]`` where each entry carries:\n", - "\n", - " - ``positions``: the image-plane (y, x) positions of that source's multiple images.\n", - " - ``positions_noise_map``: the per-position positional uncertainty.\n", - " - ``redshift``: the source redshift (different for each source \u2014 this is a multi-plane system).\n", - "\n", - "The CSV is the recommended hand-editable input for cluster datasets: a user can edit positions, noise\n", - "values, or per-source redshifts directly in a spreadsheet rather than by writing Python." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = al.list_from_csv(file_path=dataset_path / \"point_datasets.csv\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can print each dataset's name, positions, noise, and redshift." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for dataset in dataset_list:\n", - "\n", - " print(\"Point Dataset Info:\")\n", - " print(dataset.info)\n", - " print(f\"Redshift: {dataset.redshift}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot the cluster image with each source's positions overlaid.\n", - "\n", - "Cluster-scale visualization (per-source colouring, per-image-group zoom panels, kpc scale bars) is\n", - "prototyped in ``autolens_workspace_test/scripts/imaging/visualization_cluster.py``; the default\n", - "``aplt`` helpers below are sufficient for this script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.plot_grid(\n", - " grid=al.Grid2DIrregular(np.atleast_2d(dataset.positions)),\n", - " title=dataset.name,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model CSVs__\n", - "\n", - "The simulator writes the truth model into three family-level CSVs \u2014 ``mass.csv`` (lens + halo mass\n", - "profiles), ``light.csv`` (lens + source light profiles), ``point.csv`` (source point components) \u2014 keyed\n", - "by a ``galaxy`` column with ``profile_class`` dispatch (see ``scripts/cluster/csv_api.py`` for the full\n", - "schema walkthrough). We load the mass and point families here; the light family is not needed for\n", - "point-source modeling (light profiles do not affect the lensing).\n", - "\n", - "In a real analysis these CSVs come from upstream measurement: light-profile fits to the imaging data\n", - "populate ``light.csv``; the lens-galaxy centres in ``mass.csv`` are typically pinned to the light\n", - "centres (with their values then taken as ground truth). For cluster-scale point-source modeling the\n", - "observed centres remove a large block of degenerate parameters that the multiple-image positions alone\n", - "cannot constrain." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass_table = al.galaxy_models_from_csv(\n", - " file_path=dataset_path / \"mass.csv\", family=\"mass\"\n", - ")\n", - "point_table = al.galaxy_models_from_csv(\n", - " file_path=dataset_path / \"point.csv\", family=\"point\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling Galaxies Table__\n", - "\n", - "The 10 scaling-tier cluster members live in ``scaling_galaxies.csv`` \u2014 one row per member with columns\n", - "``y, x, luminosity``. ``al.galaxy_table_from_csv`` returns a typed ``GalaxyTable`` carrying:\n", - "\n", - " - ``.centres`` \u2014 a ``Grid2DIrregular`` of per-member centres (y, x).\n", - " - ``.luminosities`` \u2014 a list of per-member luminosities.\n", - "\n", - "Both arrive in the same order as the CSV rows; the model loop below zips them together. Adding more\n", - "scaling members to a real cluster amounts to extending the CSV \u2014 no Python edits required.\n", - "\n", - "In a real analysis the luminosities come from a prior light-only fit (e.g. an MGE bulge fit to the\n", - "imaging data, or a SLaM ``source_lp_0`` stage). See\n", - "``scripts/group/features/scaling_relation/modeling_for_luminosities.py`` for the standalone-fit\n", - "pattern, and ``scripts/group/features/scaling_relation/modeling.py`` for the full prose discussion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scaling_galaxies_table = al.galaxy_table_from_csv(\n", - " file_path=dataset_path / \"scaling_galaxies.csv\"\n", - ")\n", - "scaling_galaxies_centres = scaling_galaxies_table.centres\n", - "scaling_galaxies_luminosity_list = scaling_galaxies_table.luminosities\n", - "\n", - "print(f\"Scaling galaxies centres: {scaling_galaxies_centres}\")\n", - "print(f\"Scaling galaxies luminosities: {scaling_galaxies_luminosity_list}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "For point-source modeling we require a ``PointSolver``, which determines the multiple images of the mass\n", - "model for a point source at (y, x) in the source plane.\n", - "\n", - "It does this by ray-tracing triangles from the image plane to the source plane and checking whether each\n", - "source-plane (y, x) point lies inside the traced triangle. The method gradually refines smaller and\n", - "smaller triangles so multiple images are computed with sub-pixel precision.\n", - "\n", - "The ``PointSolver`` needs an initial image-plane grid (defined below) and a ``pixel_scale_precision``\n", - "controlling resolution \u2014 smaller values are more accurate but slower; 0.001 balances the two.\n", - "\n", - "Strong-lens mass models have a \"central image\" which is nearly always so demagnified it cannot be\n", - "observed. We discard it via ``magnification_threshold=0.1``; raise/lower this if your dataset does/does\n", - "not include a central image.\n", - "\n", - "__Chi Squared__\n", - "\n", - "For point-source modeling, the likelihood can be defined in the *image plane* (compare model\n", - "multiple-image positions to observed positions) or the *source plane* (collapse observed positions back\n", - "to a common source-plane location). This script uses the image-plane chi-squared via the ``PointSolver``;\n", - "see ``autolens_workspace/*/point_source/log_likelihood_function`` for a full walkthrough." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=1.0, # The pixel-scale converts pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cluster Components__\n", - "\n", - "We organise the lensing objects into four distinct categories that map directly onto the simulator:\n", - "\n", - " - ``main_lens_galaxies``: The 2 individually-modelled cluster members (BCG + satellite). Each is fitted\n", - " with a ``dPIEMassSph`` total mass profile whose centre is fixed to ``main_lens_centres[i]``.\n", - "\n", - " - ``scaling_galaxies``: The 10 scaling-tier cluster members. Each carries a ``dPIEMassSph`` mass with\n", - " centre fixed (from the CSV), ``ra`` and ``rs`` fixed at the simulator truth values, and ``b0``\n", - " derived from the *shared* ``scaling_factor`` and ``scaling_exponent`` parameters plus the per-member\n", - " luminosity. The whole tier contributes just 2 free parameters to the model regardless of how many\n", - " members are in the CSV.\n", - "\n", - " - ``host_halo``: A single standalone ``Galaxy`` carrying the cluster's ``NFWMCRLudlowSph`` dark matter\n", - " halo. The halo is *not* tied to any individual member \u2014 it sits \"on top of\" the members and\n", - " dominates the large-scale lensing.\n", - "\n", - " - ``source_galaxies``: 2 background sources, *at different redshifts* (so this is a genuine multi-plane\n", - " lens). Each is modeled as a ``Point`` source whose redshift is pinned to the value in its\n", - " ``PointDataset``.\n", - "\n", - "The galaxy-scale analogue of the scaling-relation tier (with extended-light imaging modeling rather than\n", - "point-source) is demonstrated at\n", - "``scripts/group/features/scaling_relation/modeling.py``.\n", - "\n", - "__Redshifts__\n", - "\n", - "The two sources sit at *different* redshifts (``z = 1.0`` and ``z = 2.0``), so the ``Tracer`` ray-traces\n", - "through both planes when solving for the multiple images of the further source. The lens galaxies and\n", - "host halo all sit at ``z = 0.5``.\n", - "\n", - "We pin each source's ``Galaxy.redshift`` to the redshift carried by its ``PointDataset`` \u2014 that's the\n", - "whole point of the CSV ``redshift`` column. Hardcoding ``redshift=1.0`` here would silently produce the\n", - "wrong multi-plane geometry.\n", - "\n", - "For the host halo, ``NFWMCRLudlowSph`` requires ``redshift_object`` and ``redshift_source`` to evaluate\n", - "the Ludlow et al. (2016) concentration-mass relation. We anchor ``redshift_source`` to the *furthest*\n", - "source redshift (matching the convention used by the simulator), so the concentration is computed\n", - "against the deepest light cone in the system.\n", - "\n", - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The 2 main lens galaxies each have a ``dPIEMassSph`` mass profile with centre fixed and free\n", - " ``ra``, ``rs``, ``b0`` \u2014 3 free parameters per galaxy [6 parameters].\n", - " - The 10 scaling-tier members share two free parameters: ``scaling_factor`` and ``scaling_exponent``.\n", - " Each member's ``b0`` is computed as ``scaling_factor * luminosity ** scaling_exponent``; ``ra`` and\n", - " ``rs`` are held fixed at the simulator truth values (0.1\" and 10.0\") [2 parameters].\n", - " - The host halo has an ``NFWMCRLudlowSph`` mass profile with centre fixed and a free ``mass_at_200``\n", - " [1 parameter].\n", - " - Each source has a ``Point`` model with free ``centre_0`` / ``centre_1`` priors initialised from the\n", - " mean of that source's observed positions [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=13.\n", - "\n", - "__Scaling Relation__\n", - "\n", - "The scaling relation has the form ``b0 = scaling_factor * luminosity ** scaling_exponent``. With 10\n", - "members and shared (``scaling_factor``, ``scaling_exponent``) the relation is well-constrained: as long\n", - "as the per-member luminosities span a meaningful dynamic range, the multi-image positions pull the two\n", - "relation parameters tightly. The simulator's truth values are ``scaling_factor = 0.3`` and\n", - "``scaling_exponent = 1.0``. Priors below are wider than the truth to give the search room." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "source_redshifts = [dataset.redshift for dataset in dataset_list]\n", - "\n", - "# Build af.Model[Galaxy] instances from the family CSVs. Concrete CSV values\n", - "# become fixed af.Model defaults; we then promote selected params to priors\n", - "# below. This dict is keyed by galaxy name (lens_0, lens_1, host_halo,\n", - "# source_0, source_1) \u2014 the same naming convention the simulator uses.\n", - "\n", - "galaxy_models = al.galaxy_af_models_from_csv_tables(mass_table, point_table)\n", - "\n", - "# Main Lens Galaxies: free dPIE ra / rs / b0 on each; centre stays fixed at the CSV value.\n", - "for name in (\"lens_0\", \"lens_1\"):\n", - " galaxy_models[name].mass.ra = af.UniformPrior(lower_limit=1.0, upper_limit=15.0)\n", - " galaxy_models[name].mass.rs = af.UniformPrior(lower_limit=5.0, upper_limit=40.0)\n", - " galaxy_models[name].mass.b0 = af.UniformPrior(lower_limit=0.1, upper_limit=10.0)\n", - "\n", - "# Host Halo: free mass_at_200; centre + redshift_object + redshift_source stay fixed.\n", - "galaxy_models[\"host_halo\"].dark.mass_at_200 = af.LogUniformPrior(\n", - " lower_limit=10**14.5, upper_limit=10**16.0\n", - ")\n", - "\n", - "# Source Galaxies: free Point centres with GaussianPrior initialised from the\n", - "# mean of each source's observed multiple-image positions in its PointDataset.\n", - "# This deliberately ignores the truth centre stored in point.csv \u2014 in a real\n", - "# analysis you don't know the source's true source-plane position, you only\n", - "# have the image-plane positions of its multiple images.\n", - "for i, dataset in enumerate(dataset_list):\n", - " positions = np.atleast_2d(dataset.positions)\n", - " point_attr = getattr(galaxy_models[f\"source_{i}\"], f\"point_{i}\")\n", - " point_attr.centre_0 = af.GaussianPrior(\n", - " mean=float(np.mean(positions[:, 0])), sigma=3.0\n", - " )\n", - " point_attr.centre_1 = af.GaussianPrior(\n", - " mean=float(np.mean(positions[:, 1])), sigma=3.0\n", - " )\n", - "\n", - "# Scaling Tier Members (dPIEMassSph, b0 derived from shared scaling relation).\n", - "#\n", - "# scaling_factor and scaling_exponent are defined ONCE outside the loop. Every\n", - "# member's b0 is a derived prior of these two shared parameters plus its own\n", - "# (fixed) luminosity, so the entire tier contributes 2 free parameters regardless\n", - "# of how many members are in scaling_galaxies.csv.\n", - "\n", - "scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=1.0)\n", - "scaling_exponent = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - "scaling_ra_fixed = 0.1\n", - "scaling_rs_fixed = 10.0\n", - "\n", - "scaling_galaxies_list = []\n", - "for centre, luminosity in zip(\n", - " scaling_galaxies_centres, scaling_galaxies_luminosity_list\n", - "):\n", - " mass = af.Model(al.mp.dPIEMassSph)\n", - " mass.centre = tuple(centre)\n", - " mass.ra = scaling_ra_fixed\n", - " mass.rs = scaling_rs_fixed\n", - " mass.b0 = scaling_factor * luminosity**scaling_exponent\n", - "\n", - " scaling_galaxies_list.append(af.Model(al.Galaxy, redshift=redshift_lens, mass=mass))\n", - "\n", - "scaling_galaxies = af.Collection(scaling_galaxies_list)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**galaxy_models),\n", - " scaling_galaxies=scaling_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ``info`` attribute shows the model in a readable format. This prints the main lens galaxies, the\n", - "host halo, and the source galaxies, each with its free / fixed parameters.\n", - "\n", - "The ``info`` below may not display optimally on your computer screen \u2014 for example whitespace between\n", - "parameter names on the left and parameter priors on the right may break across multiple lines. The\n", - "``info_whitespace_length`` parameter in ``config/general.yaml`` controls this; reset the Jupyter\n", - "kernel after changing it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Name Pairing__\n", - "\n", - "Every ``PointDataset`` has a ``name`` (e.g. ``point_0``, ``point_1``). This pairs the dataset to the\n", - "``Point`` model component with the same name. Above, the ``af.Model(al.ps.Point)`` for source ``i`` is\n", - "attached to its ``af.Model(al.Galaxy)`` under the key ``point_i`` \u2014 that's what the ``**{f\"point_{i}\":\n", - "point}`` expansion does.\n", - "\n", - "If a dataset has no matching ``Point`` in the model, that dataset is ignored. If a ``Point`` exists with\n", - "no matching dataset, **PyAutoLens** raises an error.\n", - "\n", - "In multi-source cluster lenses, this name pairing is what ensures every source's positions are fitted by\n", - "the correct model component." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The lens model is fitted to the data using the nested sampling algorithm Nautilus (see\n", - "``point_source/start_here.py`` for a full description).\n", - "\n", - "The folders ``autolens_workspace/*/guides/modeling/searches`` and ``customize`` give overviews of the\n", - "non-linear searches PyAutoLens supports and how to customize the fit, including the priors.\n", - "\n", - "Results are output to::\n", - "\n", - " /autolens_workspace/output/cluster/simple/modeling//\n", - "\n", - "__Unique Identifier__\n", - "\n", - "The ``unique_identifier`` is generated from the model, search, and dataset, so re-running with the same\n", - "configuration resumes the existing fit. Changing any of the three regenerates the identifier.\n", - "\n", - "__Iterations Per Update__\n", - "\n", - "Every N iterations the search prints the max-likelihood model and best-fit image. On GPU ~2500 keeps the\n", - "output cadence around once per minute; on CPU a similar cadence is reached at lower N.\n", - "\n", - "__Live Visual Update__\n", - "\n", - "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", - "live display surface:\n", - "\n", - "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", - " watch the fit converge without leaving your terminal.\n", - "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", - " refreshes in place every `iterations_per_quick_update`.\n", - "\n", - "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", - "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"cluster\"),\n", - " name=\"modeling\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We create one ``AnalysisPoint`` per dataset. Each defines the ``log_likelihood_function`` Nautilus uses\n", - "to fit the model to that dataset's multiple-image positions.\n", - "\n", - "We then wrap each analysis in an ``AnalysisFactor`` pairing it to the *shared* lens model, and combine\n", - "all factors into a single ``FactorGraphModel``. The total log likelihood is the sum of the per-dataset\n", - "log likelihoods; each dataset gets its own output subdirectory for visualization.\n", - "\n", - "__JAX__\n", - "\n", - "`AnalysisPoint(use_jax=True)` per-dataset; the search driver wraps the\n", - "joint likelihood in `jax.vmap(jax.jit(...))`. Cluster point-source fits\n", - "get the largest speedup from JAX on GPU (triangle refinement + multi-\n", - "plane deflection sum dominate runtime). Force NumPy with `use_jax=False`\n", - "when debugging." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [\n", - " al.AnalysisPoint(dataset=dataset, solver=solver, use_jax=True)\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis Factor__\n", - "\n", - "Each analysis is wrapped in an ``AnalysisFactor`` paired with the shared model. The factor-graph API is\n", - "used heavily for advanced multi-dataset lens modeling \u2014 multi-wavelength imaging, joint\n", - "imaging+interferometer fits, multiple-source cluster fits like this one." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = [\n", - " af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " for analysis in analysis_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Factor Graph__\n", - "\n", - "All ``AnalysisFactor`` objects combine into one ``FactorGraphModel``. The per-dataset log likelihoods\n", - "are summed; results land in a unified directory with per-dataset visualization subdirs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Print the global model the factor graph fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "Cluster lens modeling is computationally expensive \u2014 full Nautilus runs are typically hours on CPU,\n", - "minutes on GPU. Run times scale with (a) the log-likelihood evaluation time of a single sample and\n", - "(b) the number of iterations Nautilus needs to converge.\n", - "\n", - "For this 2-main + halo + 2-source model the log-likelihood evaluation is < 1 second on CPU and < 0.02 s\n", - "on GPU. A converged fit typically takes a few thousand iterations.\n", - "\n", - "__Model-Fit__\n", - "\n", - "Pass the factor-graph model and the factor graph itself (as the analysis) to ``search.fit``. Watch\n", - "``autolens_workspace/output`` for on-the-fly visualization while the fit runs.\n", - "\n", - "**Run Time Error:** On certain operating systems and Python versions, the code below may produce an\n", - "error. If this occurs, see ``autolens_workspace/guides/modeling/bug_fix``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell will progress once the search has completed \u2014 this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "print(\"The search has finished run \u2014 you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output Folder Layout__\n", - "\n", - "Now the fit is running you should checkout the ``autolens_workspace/output`` folder. Results are written\n", - "to hard-disk on the fly in human-readable formats \u2014 ``.json``, ``.csv``, ``.fits``, ``.png`` and plain\n", - "text \u2014 using the highest-likelihood model found so far.\n", - "\n", - "Each completed fit lives at::\n", - "\n", - " output/cluster//modeling//\n", - " files/ <- JSON + CSV: loadable Python objects\n", - " tracer.json <- max log likelihood Tracer\n", - " model.json <- fitted af.Collection model\n", - " samples.csv <- full Nautilus samples\n", - " samples_summary.json <- max log likelihood parameter values + errors\n", - " samples_info.json <- metadata about the samples\n", - " search.json <- non-linear search configuration\n", - " settings.json <- search settings\n", - " cosmology.json <- cosmology used for the fit\n", - " covariance.csv <- parameter covariance matrix\n", - " image/ <- FITS + PNG: imaging + point-source products\n", - " dataset.fits <- data, noise-map and PSF\n", - " fit.fits <- model image, residuals, chi-squared map\n", - " tracer.fits <- per-galaxy image-plane images\n", - " source_plane_images.fits <- source-plane reconstructions\n", - " positions.png <- observed vs model multiple-image positions\n", - " dataset.png, fit.png, tracer.png <- visualisations\n", - " model.info <- human-readable model summary\n", - " model.results <- human-readable fit summary\n", - " search.summary <- search run summary\n", - " search_internal/ <- files used to resume / visualise the search\n", - " metadata <- run metadata\n", - "\n", - "The ```` is a 32-character identifier derived from the model, search and dataset.\n", - "\n", - "__Result__\n", - "\n", - "``search.fit`` on a factor-graph returns a list of ``Result`` objects, one per ``AnalysisFactor`` (i.e.\n", - "one per dataset). Each carries the same ``max_log_likelihood_instance`` (since they share the global\n", - "model) but its own per-dataset visualization and ``FitPoint`` object.\n", - "\n", - "[The ``info_whitespace_length`` config setting also controls whitespace in ``result.info``.]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " print(result.max_log_likelihood_instance)\n", - "\n", - " aplt.subplot_tracer(\n", - " tracer=result.max_log_likelihood_tracer,\n", - " grid=grid,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ``Samples`` are identical across results (the model is global), so a single corner plot from the\n", - "first result is enough." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result_list[0].samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This script gives a concise overview of the basic cluster modeling API for a small multi-plane cluster.\n", - "\n", - "__Data Preparation__\n", - "\n", - "If you are looking to fit your own point-source cluster data, see\n", - "``autolens_workspace/*/data_preparation/point_source/README.md`` for the input-data standards.\n", - "\n", - "__HowToLens__\n", - "\n", - "For a deeper understanding of how lens modeling, ray-tracing, and non-linear searches actually work, see\n", - "the **HowToLens** Jupyter notebook lectures at https://github.com/PyAutoLabs/HowToLens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Cluster\n", + "=================\n", + "\n", + "This script models the small multi-plane cluster lens simulated by ``cluster/simulator.py``: a Brightest\n", + "Cluster Galaxy (BCG) plus one satellite member at the lens redshift ``z = 0.5``, 10 lower-mass cluster\n", + "members modelled collectively via a luminosity-mass scaling relation, a standalone NFW host dark matter\n", + "halo *not* tied to any individual galaxy, and 2 background sources at *different* redshifts (``z = 1.0``\n", + "and ``z = 2.0``) which the cluster lenses into multiple images.\n", + "\n", + "Cluster modeling almost always uses the *point source* API: rather than fitting the extended arc light\n", + "of each lensed source, only the image-plane positions of the brightest pixels of each multiple image are\n", + "fitted. Per-source positions, noise maps, and redshifts are loaded from a single hand-editable CSV that\n", + "the simulator writes (``point_datasets.csv``) via ``al.list_from_csv``. Cluster-member centres and\n", + "luminosities for the scaling tier are loaded from a second CSV (``scaling_galaxies.csv``) via\n", + "``al.galaxy_table_from_csv``.\n", + "\n", + "__Contents__\n", + "\n", + "- **Example:** What this script fits and the underlying simulator.\n", + "- **Simulation:** Overview of how the simulated dataset was generated.\n", + "- **Dataset:** Load the CCD image and the per-source point datasets from the combined CSV.\n", + "- **Model CSVs:** Load the named-galaxy mass + point CSVs written by the simulator.\n", + "- **Scaling Galaxies Table:** Load the scaling-tier centres + luminosities from the CSV.\n", + "- **Point Solver:** Set up the image-plane multiple-image solver.\n", + "- **Chi Squared:** Why this script uses an image-plane chi-squared.\n", + "- **Cluster Components:** The four categories of lensing object \u2014 main lens galaxies, scaling members, host halo, sources.\n", + "- **Redshifts:** Multi-plane redshift handling and the source-redshift / dataset-redshift pairing.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Scaling Relation:** The shared two-parameter relation that ties every scaling member's mass to its luminosity.\n", + "- **Name Pairing:** How each ``Point`` model component is paired to its ``PointDataset``.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", + "- **Analysis:** Create the ``AnalysisPoint`` objects, one per dataset.\n", + "- **Factor Graph:** Combine per-dataset analyses into one global ``FactorGraphModel``.\n", + "- **Run Times:** Profiling the expected run time of the model-fit.\n", + "- **Output Folder Layout:** Description of the ``output`` folder structure.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Example__\n", + "\n", + "This script fits a ``PointDataset`` of a small multi-plane cluster where:\n", + "\n", + " - There are 2 main lens galaxies with ``dPIEMassSph`` total mass distributions, each with their centre\n", + " fixed to the values written out by the simulator [6 parameters].\n", + " - There are 10 scaling-tier member galaxies. Each carries a ``dPIEMassSph`` mass with centre fixed,\n", + " ``ra`` fixed, and ``b0`` / ``rs`` derived from the shared reference-anchored scaling relation and the\n", + " per-member luminosity [1 parameter total for the entire tier].\n", + " - There is 1 standalone ``NFWMCRLudlowSph`` host dark matter halo with its centre fixed and a free\n", + " ``mass_at_200`` [1 parameter].\n", + " - There are 2 source galaxies modeled as ``Point`` sources, each with its redshift pinned to the value\n", + " in its ``PointDataset`` row [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=12.\n", + "\n", + "The defining feature of cluster modeling is the scaling tier: 10 lower-mass members are fit jointly\n", + "with a *single free parameter* (``b0_ref``, the lens strength of the brightest member; the relation's\n", + "exponent is fixed at the Faber-Jackson value of 0.5, following the Lenstool convention). Adding more\n", + "members to ``scaling_galaxies.csv`` in the future does not grow the dimensionality of parameter space.\n", + "\n", + "__Simulation__\n", + "\n", + "This script fits the simulated cluster dataset produced by ``autolens_workspace/*/cluster/simulator.py``.\n", + "That simulator writes:\n", + "\n", + " - ``data.fits`` / ``noise_map.fits`` / ``psf.fits`` \u2014 CCD imaging of the cluster (used for visualization).\n", + " - ``point_datasets.csv`` \u2014 one row per observed multiple image, grouped by source ``name``, with a\n", + " ``redshift`` column per group. Loaded here with ``al.list_from_csv``.\n", + " - ``mass.csv`` + ``light.csv`` + ``point.csv`` \u2014 named-galaxy CSVs carrying the full truth model\n", + " (main galaxies + host halo + sources). Loaded here with ``al.galaxy_models_from_csv``. See\n", + " ``scripts/cluster/csv_api.py`` for the schema walkthrough.\n", + " - ``scaling_galaxies.csv`` \u2014 one row per scaling-tier member with columns ``y, x, luminosity``. Loaded\n", + " here with ``al.galaxy_table_from_csv``.\n", + " - ``tracer.json`` \u2014 true ``Tracer`` (used by visualization, not modeling)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens dataset ``cluster``, which is the dataset we will use to perform lens modeling.\n", + "\n", + "We begin by loading a CCD image of the dataset. Although we perform point-source modeling and will not\n", + "use the imaging data in the model-fit, it is useful to load it for visualization.\n", + "\n", + "The ``pixel_scales`` define the arc-second to pixel conversion factor of the image, which for the\n", + "dataset we are using is 0.1\" / pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", \"cluster\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/cluster/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Datasets__\n", + "\n", + "We load the point datasets from the combined ``point_datasets.csv`` written by the simulator. This\n", + "returns a ``List[PointDataset]`` where each entry carries:\n", + "\n", + " - ``positions``: the image-plane (y, x) positions of that source's multiple images.\n", + " - ``positions_noise_map``: the per-position positional uncertainty.\n", + " - ``redshift``: the source redshift (different for each source \u2014 this is a multi-plane system).\n", + "\n", + "The CSV is the recommended hand-editable input for cluster datasets: a user can edit positions, noise\n", + "values, or per-source redshifts directly in a spreadsheet rather than by writing Python." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = al.list_from_csv(file_path=dataset_path / \"point_datasets.csv\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can print each dataset's name, positions, noise, and redshift." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for dataset in dataset_list:\n", + "\n", + " print(\"Point Dataset Info:\")\n", + " print(dataset.info)\n", + " print(f\"Redshift: {dataset.redshift}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot the cluster image with each source's positions overlaid.\n", + "\n", + "Cluster-scale visualization (per-source colouring, per-image-group zoom panels, kpc scale bars) is\n", + "prototyped in ``autolens_workspace_test/scripts/imaging/visualization_cluster.py``; the default\n", + "``aplt`` helpers below are sufficient for this script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.plot_grid(\n", + " grid=al.Grid2DIrregular(np.atleast_2d(dataset.positions)),\n", + " title=dataset.name,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model CSVs__\n", + "\n", + "The simulator writes the truth model into three family-level CSVs \u2014 ``mass.csv`` (lens + halo mass\n", + "profiles), ``light.csv`` (lens + source light profiles), ``point.csv`` (source point components) \u2014 keyed\n", + "by a ``galaxy`` column with ``profile_class`` dispatch (see ``scripts/cluster/csv_api.py`` for the full\n", + "schema walkthrough). We load the mass and point families here; the light family is not needed for\n", + "point-source modeling (light profiles do not affect the lensing).\n", + "\n", + "In a real analysis these CSVs come from upstream measurement: light-profile fits to the imaging data\n", + "populate ``light.csv``; the lens-galaxy centres in ``mass.csv`` are typically pinned to the light\n", + "centres (with their values then taken as ground truth). For cluster-scale point-source modeling the\n", + "observed centres remove a large block of degenerate parameters that the multiple-image positions alone\n", + "cannot constrain." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass_table = al.galaxy_models_from_csv(\n", + " file_path=dataset_path / \"mass.csv\", family=\"mass\"\n", + ")\n", + "point_table = al.galaxy_models_from_csv(\n", + " file_path=dataset_path / \"point.csv\", family=\"point\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling Galaxies Table__\n", + "\n", + "The 10 scaling-tier cluster members live in ``scaling_galaxies.csv`` \u2014 one row per member with columns\n", + "``y, x, luminosity``. ``al.galaxy_table_from_csv`` returns a typed ``GalaxyTable`` carrying:\n", + "\n", + " - ``.centres`` \u2014 a ``Grid2DIrregular`` of per-member centres (y, x).\n", + " - ``.luminosities`` \u2014 a list of per-member luminosities.\n", + "\n", + "Both arrive in the same order as the CSV rows; the model loop below zips them together. Adding more\n", + "scaling members to a real cluster amounts to extending the CSV \u2014 no Python edits required.\n", + "\n", + "In a real analysis the luminosities come from a prior light-only fit (e.g. an MGE bulge fit to the\n", + "imaging data, or a SLaM ``source_lp_0`` stage). See\n", + "``scripts/group/features/scaling_relation/modeling_for_luminosities.py`` for the standalone-fit\n", + "pattern, and ``scripts/group/features/scaling_relation/modeling.py`` for the full prose discussion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scaling_galaxies_table = al.galaxy_table_from_csv(\n", + " file_path=dataset_path / \"scaling_galaxies.csv\"\n", + ")\n", + "scaling_galaxies_centres = scaling_galaxies_table.centres\n", + "scaling_galaxies_luminosity_list = scaling_galaxies_table.luminosities\n", + "\n", + "print(f\"Scaling galaxies centres: {scaling_galaxies_centres}\")\n", + "print(f\"Scaling galaxies luminosities: {scaling_galaxies_luminosity_list}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "For point-source modeling we require a ``PointSolver``, which determines the multiple images of the mass\n", + "model for a point source at (y, x) in the source plane.\n", + "\n", + "It does this by ray-tracing triangles from the image plane to the source plane and checking whether each\n", + "source-plane (y, x) point lies inside the traced triangle. The method gradually refines smaller and\n", + "smaller triangles so multiple images are computed with sub-pixel precision.\n", + "\n", + "The ``PointSolver`` needs an initial image-plane grid (defined below) and a ``pixel_scale_precision``\n", + "controlling resolution \u2014 smaller values are more accurate but slower; 0.001 balances the two.\n", + "\n", + "Strong-lens mass models have a \"central image\" which is nearly always so demagnified it cannot be\n", + "observed. We discard it via ``magnification_threshold=0.1``; raise/lower this if your dataset does/does\n", + "not include a central image.\n", + "\n", + "__Chi Squared__\n", + "\n", + "For point-source modeling, the likelihood can be defined in the *image plane* (compare model\n", + "multiple-image positions to observed positions) or the *source plane* (collapse observed positions back\n", + "to a common source-plane location). This script uses the image-plane chi-squared via the ``PointSolver``;\n", + "see ``autolens_workspace/*/point_source/log_likelihood_function`` for a full walkthrough." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=1.0, # The pixel-scale converts pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cluster Components__\n", + "\n", + "We organise the lensing objects into four distinct categories that map directly onto the simulator:\n", + "\n", + " - ``main_lens_galaxies``: The 2 individually-modelled cluster members (BCG + satellite). Each is fitted\n", + " with a ``dPIEMassSph`` total mass profile whose centre is fixed to ``main_lens_centres[i]``.\n", + "\n", + " - ``scaling_galaxies``: The 10 scaling-tier cluster members. Each carries a ``dPIEMassSph`` mass with\n", + " centre fixed (from the CSV), ``ra`` fixed, and ``b0`` / ``rs`` derived from the shared\n", + " reference-anchored scaling relation (single free normalization ``b0_ref``, exponents fixed at 0.5)\n", + " plus the per-member luminosity. The whole tier contributes 1 free parameter to the model regardless\n", + " of how many members are in the CSV.\n", + "\n", + " - ``host_halo``: A single standalone ``Galaxy`` carrying the cluster's ``NFWMCRLudlowSph`` dark matter\n", + " halo. The halo is *not* tied to any individual member \u2014 it sits \"on top of\" the members and\n", + " dominates the large-scale lensing.\n", + "\n", + " - ``source_galaxies``: 2 background sources, *at different redshifts* (so this is a genuine multi-plane\n", + " lens). Each is modeled as a ``Point`` source whose redshift is pinned to the value in its\n", + " ``PointDataset``.\n", + "\n", + "The galaxy-scale analogue of the scaling-relation tier (with extended-light imaging modeling rather than\n", + "point-source) is demonstrated at\n", + "``scripts/group/features/scaling_relation/modeling.py``.\n", + "\n", + "__Redshifts__\n", + "\n", + "The two sources sit at *different* redshifts (``z = 1.0`` and ``z = 2.0``), so the ``Tracer`` ray-traces\n", + "through both planes when solving for the multiple images of the further source. The lens galaxies and\n", + "host halo all sit at ``z = 0.5``.\n", + "\n", + "We pin each source's ``Galaxy.redshift`` to the redshift carried by its ``PointDataset`` \u2014 that's the\n", + "whole point of the CSV ``redshift`` column. Hardcoding ``redshift=1.0`` here would silently produce the\n", + "wrong multi-plane geometry.\n", + "\n", + "For the host halo, ``NFWMCRLudlowSph`` requires ``redshift_object`` and ``redshift_source`` to evaluate\n", + "the Ludlow et al. (2016) concentration-mass relation. We anchor ``redshift_source`` to the *furthest*\n", + "source redshift (matching the convention used by the simulator), so the concentration is computed\n", + "against the deepest light cone in the system.\n", + "\n", + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The 2 main lens galaxies each have a ``dPIEMassSph`` mass profile with centre fixed and free\n", + " ``ra``, ``rs``, ``b0`` \u2014 3 free parameters per galaxy [6 parameters].\n", + " - The 10 scaling-tier members share a single free parameter: ``b0_ref``, the lens strength of the\n", + " *brightest* member (the reference galaxy). Each member's ``b0`` and ``rs`` are computed as\n", + " ``b0_ref * (L / L_ref) ** 0.5`` and ``rs_ref * (L / L_ref) ** 0.5`` with the exponents fixed;\n", + " ``ra`` (0.1\") and ``rs_ref`` (10.0\") are held fixed at the simulator truth values [1 parameter].\n", + " - The host halo has an ``NFWMCRLudlowSph`` mass profile with centre fixed and a free ``mass_at_200``\n", + " [1 parameter].\n", + " - Each source has a ``Point`` model with free ``centre_0`` / ``centre_1`` priors initialised from the\n", + " mean of that source's observed positions [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=12.\n", + "\n", + "__Scaling Relation__\n", + "\n", + "The scaling relation is reference-anchored, the convention used by Lenstool and essentially every\n", + "published cluster strong-lensing analysis (Limousin et al. 2005; Eliasdottir et al. 2007; Bergamini et\n", + "al. 2019):\n", + "\n", + " b0_i = b0_ref * (L_i / L_ref) ** 0.5\n", + " rs_i = rs_ref * (L_i / L_ref) ** 0.5\n", + "\n", + "The single free parameter ``b0_ref`` is the lens strength of the brightest scaling member \u2014 a physically\n", + "interpretable quantity (roughly that galaxy's Einstein radius) for which a prior range is easy to\n", + "motivate, unlike an abstract multiplicative factor whose units depend on the (arbitrary) luminosity\n", + "normalization. The exponent is *fixed* at 0.5 rather than fitted: b0 \u221d sigma\u00b2 for the dPIE, and\n", + "Faber-Jackson (sigma \u221d L^(1/4)) gives b0 \u221d L^(1/2). The truncation radius scales with the same fixed\n", + "exponent, mirroring Lenstool's r_cut \u221d L^(1/2). Only luminosity *ratios* enter, so the CSV's luminosity\n", + "units are irrelevant; magnitude catalogues convert via ``L_i / L_ref = 10 ** (0.4 * (m_ref - m_i))``.\n", + "\n", + "Freeing the exponent (or ``rs_ref``) is a one-line change shown in the code comment below \u2014 useful as a\n", + "systematics test, at the cost of the degeneracy between normalization and slope that the fixed-exponent\n", + "convention exists to avoid. Kinematic calibrations (Bergamini et al. 2019: sigma \u221d L^0.27-0.28 from MUSE\n", + "member kinematics, i.e. a b0 exponent \u2248 0.55, with the r_cut exponent from the fundamental plane) are\n", + "the standard refinement when member velocity dispersions are available.\n", + "\n", + "The simulator's truth value is ``b0_ref = 0.12`` arcsec. The prior below is much wider than the truth to\n", + "give the search room." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "source_redshifts = [dataset.redshift for dataset in dataset_list]\n", + "\n", + "# Build af.Model[Galaxy] instances from the family CSVs. Concrete CSV values\n", + "# become fixed af.Model defaults; we then promote selected params to priors\n", + "# below. This dict is keyed by galaxy name (lens_0, lens_1, host_halo,\n", + "# source_0, source_1) \u2014 the same naming convention the simulator uses.\n", + "\n", + "galaxy_models = al.galaxy_af_models_from_csv_tables(mass_table, point_table)\n", + "\n", + "# Main Lens Galaxies: free dPIE ra / rs / b0 on each; centre stays fixed at the CSV value.\n", + "for name in (\"lens_0\", \"lens_1\"):\n", + " galaxy_models[name].mass.ra = af.UniformPrior(lower_limit=1.0, upper_limit=15.0)\n", + " galaxy_models[name].mass.rs = af.UniformPrior(lower_limit=5.0, upper_limit=40.0)\n", + " galaxy_models[name].mass.b0 = af.UniformPrior(lower_limit=0.1, upper_limit=10.0)\n", + "\n", + "# Host Halo: free mass_at_200; centre + redshift_object + redshift_source stay fixed.\n", + "galaxy_models[\"host_halo\"].dark.mass_at_200 = af.LogUniformPrior(\n", + " lower_limit=10**14.5, upper_limit=10**16.0\n", + ")\n", + "\n", + "# Source Galaxies: free Point centres with GaussianPrior initialised from the\n", + "# mean of each source's observed multiple-image positions in its PointDataset.\n", + "# This deliberately ignores the truth centre stored in point.csv \u2014 in a real\n", + "# analysis you don't know the source's true source-plane position, you only\n", + "# have the image-plane positions of its multiple images.\n", + "for i, dataset in enumerate(dataset_list):\n", + " positions = np.atleast_2d(dataset.positions)\n", + " point_attr = getattr(galaxy_models[f\"source_{i}\"], f\"point_{i}\")\n", + " point_attr.centre_0 = af.GaussianPrior(\n", + " mean=float(np.mean(positions[:, 0])), sigma=3.0\n", + " )\n", + " point_attr.centre_1 = af.GaussianPrior(\n", + " mean=float(np.mean(positions[:, 1])), sigma=3.0\n", + " )\n", + "\n", + "# Scaling Tier Members (dPIEMassSph, b0 and rs derived from the reference-anchored\n", + "# scaling relation).\n", + "#\n", + "# b0_ref is defined ONCE outside the loop \u2014 it is the tier's only free parameter,\n", + "# the lens strength of the brightest (reference) member. Every member's b0 and rs\n", + "# are derived priors of b0_ref (or plain fixed values, for rs) scaled by its\n", + "# luminosity ratio to the reference, with the exponents fixed at the Faber-Jackson\n", + "# value of 0.5. The entire tier therefore contributes 1 free parameter regardless\n", + "# of how many members are in scaling_galaxies.csv.\n", + "#\n", + "# To free the exponent as a systematics test, replace the fixed value with e.g.\n", + "# `scaling_exponent = af.UniformPrior(lower_limit=0.0, upper_limit=1.0)` \u2014 every\n", + "# member's b0 then derives from two shared parameters, as in older versions of\n", + "# this example.\n", + "\n", + "scaling_b0_ref = af.UniformPrior(lower_limit=0.0, upper_limit=1.0)\n", + "scaling_exponent = 0.5\n", + "\n", + "scaling_luminosity_ref = max(scaling_galaxies_luminosity_list)\n", + "scaling_ra_fixed = 0.1\n", + "scaling_rs_ref_fixed = 10.0\n", + "\n", + "scaling_galaxies_list = []\n", + "for centre, luminosity in zip(\n", + " scaling_galaxies_centres, scaling_galaxies_luminosity_list\n", + "):\n", + " luminosity_ratio = luminosity / scaling_luminosity_ref\n", + "\n", + " mass = af.Model(al.mp.dPIEMassSph)\n", + " mass.centre = tuple(centre)\n", + " mass.ra = scaling_ra_fixed\n", + " mass.rs = scaling_rs_ref_fixed * luminosity_ratio**scaling_exponent\n", + " mass.b0 = scaling_b0_ref * luminosity_ratio**scaling_exponent\n", + "\n", + " scaling_galaxies_list.append(af.Model(al.Galaxy, redshift=redshift_lens, mass=mass))\n", + "\n", + "scaling_galaxies = af.Collection(scaling_galaxies_list)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**galaxy_models),\n", + " scaling_galaxies=scaling_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The ``info`` attribute shows the model in a readable format. This prints the main lens galaxies, the\n", + "host halo, and the source galaxies, each with its free / fixed parameters.\n", + "\n", + "The ``info`` below may not display optimally on your computer screen \u2014 for example whitespace between\n", + "parameter names on the left and parameter priors on the right may break across multiple lines. The\n", + "``info_whitespace_length`` parameter in ``config/general.yaml`` controls this; reset the Jupyter\n", + "kernel after changing it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Name Pairing__\n", + "\n", + "Every ``PointDataset`` has a ``name`` (e.g. ``point_0``, ``point_1``). This pairs the dataset to the\n", + "``Point`` model component with the same name. Above, the ``af.Model(al.ps.Point)`` for source ``i`` is\n", + "attached to its ``af.Model(al.Galaxy)`` under the key ``point_i`` \u2014 that's what the ``**{f\"point_{i}\":\n", + "point}`` expansion does.\n", + "\n", + "If a dataset has no matching ``Point`` in the model, that dataset is ignored. If a ``Point`` exists with\n", + "no matching dataset, **PyAutoLens** raises an error.\n", + "\n", + "In multi-source cluster lenses, this name pairing is what ensures every source's positions are fitted by\n", + "the correct model component." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The lens model is fitted to the data using the nested sampling algorithm Nautilus (see\n", + "``point_source/start_here.py`` for a full description).\n", + "\n", + "The folders ``autolens_workspace/*/guides/modeling/searches`` and ``customize`` give overviews of the\n", + "non-linear searches PyAutoLens supports and how to customize the fit, including the priors.\n", + "\n", + "Results are output to::\n", + "\n", + " /autolens_workspace/output/cluster/simple/modeling//\n", + "\n", + "__Unique Identifier__\n", + "\n", + "The ``unique_identifier`` is generated from the model, search, and dataset, so re-running with the same\n", + "configuration resumes the existing fit. Changing any of the three regenerates the identifier.\n", + "\n", + "__Iterations Per Update__\n", + "\n", + "Every N iterations the search prints the max-likelihood model and best-fit image. On GPU ~2500 keeps the\n", + "output cadence around once per minute; on CPU a similar cadence is reached at lower N.\n", + "\n", + "__Live Visual Update__\n", + "\n", + "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", + "live display surface:\n", + "\n", + "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", + " watch the fit converge without leaving your terminal.\n", + "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", + " refreshes in place every `iterations_per_quick_update`.\n", + "\n", + "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", + "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"cluster\"),\n", + " name=\"modeling\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We create one ``AnalysisPoint`` per dataset. Each defines the ``log_likelihood_function`` Nautilus uses\n", + "to fit the model to that dataset's multiple-image positions.\n", + "\n", + "We then wrap each analysis in an ``AnalysisFactor`` pairing it to the *shared* lens model, and combine\n", + "all factors into a single ``FactorGraphModel``. The total log likelihood is the sum of the per-dataset\n", + "log likelihoods; each dataset gets its own output subdirectory for visualization.\n", + "\n", + "__JAX__\n", + "\n", + "`AnalysisPoint(use_jax=True)` per-dataset; the search driver wraps the\n", + "joint likelihood in `jax.vmap(jax.jit(...))`. Cluster point-source fits\n", + "get the largest speedup from JAX on GPU (triangle refinement + multi-\n", + "plane deflection sum dominate runtime). Force NumPy with `use_jax=False`\n", + "when debugging." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [\n", + " al.AnalysisPoint(dataset=dataset, solver=solver, use_jax=True)\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis Factor__\n", + "\n", + "Each analysis is wrapped in an ``AnalysisFactor`` paired with the shared model. The factor-graph API is\n", + "used heavily for advanced multi-dataset lens modeling \u2014 multi-wavelength imaging, joint\n", + "imaging+interferometer fits, multiple-source cluster fits like this one." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = [\n", + " af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " for analysis in analysis_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Factor Graph__\n", + "\n", + "All ``AnalysisFactor`` objects combine into one ``FactorGraphModel``. The per-dataset log likelihoods\n", + "are summed; results land in a unified directory with per-dataset visualization subdirs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Print the global model the factor graph fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "Cluster lens modeling is computationally expensive \u2014 full Nautilus runs are typically hours on CPU,\n", + "minutes on GPU. Run times scale with (a) the log-likelihood evaluation time of a single sample and\n", + "(b) the number of iterations Nautilus needs to converge.\n", + "\n", + "For this 2-main + halo + 2-source model the log-likelihood evaluation is < 1 second on CPU and < 0.02 s\n", + "on GPU. A converged fit typically takes a few thousand iterations.\n", + "\n", + "__Model-Fit__\n", + "\n", + "Pass the factor-graph model and the factor graph itself (as the analysis) to ``search.fit``. Watch\n", + "``autolens_workspace/output`` for on-the-fly visualization while the fit runs.\n", + "\n", + "**Run Time Error:** On certain operating systems and Python versions, the code below may produce an\n", + "error. If this occurs, see ``autolens_workspace/guides/modeling/bug_fix``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell will progress once the search has completed \u2014 this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "print(\"The search has finished run \u2014 you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output Folder Layout__\n", + "\n", + "Now the fit is running you should checkout the ``autolens_workspace/output`` folder. Results are written\n", + "to hard-disk on the fly in human-readable formats \u2014 ``.json``, ``.csv``, ``.fits``, ``.png`` and plain\n", + "text \u2014 using the highest-likelihood model found so far.\n", + "\n", + "Each completed fit lives at::\n", + "\n", + " output/cluster//modeling//\n", + " files/ <- JSON + CSV: loadable Python objects\n", + " tracer.json <- max log likelihood Tracer\n", + " model.json <- fitted af.Collection model\n", + " samples.csv <- full Nautilus samples\n", + " samples_summary.json <- max log likelihood parameter values + errors\n", + " samples_info.json <- metadata about the samples\n", + " search.json <- non-linear search configuration\n", + " settings.json <- search settings\n", + " cosmology.json <- cosmology used for the fit\n", + " covariance.csv <- parameter covariance matrix\n", + " image/ <- FITS + PNG: imaging + point-source products\n", + " dataset.fits <- data, noise-map and PSF\n", + " fit.fits <- model image, residuals, chi-squared map\n", + " tracer.fits <- per-galaxy image-plane images\n", + " source_plane_images.fits <- source-plane reconstructions\n", + " positions.png <- observed vs model multiple-image positions\n", + " dataset.png, fit.png, tracer.png <- visualisations\n", + " model.info <- human-readable model summary\n", + " model.results <- human-readable fit summary\n", + " search.summary <- search run summary\n", + " search_internal/ <- files used to resume / visualise the search\n", + " metadata <- run metadata\n", + "\n", + "The ```` is a 32-character identifier derived from the model, search and dataset.\n", + "\n", + "__Result__\n", + "\n", + "``search.fit`` on a factor-graph returns a list of ``Result`` objects, one per ``AnalysisFactor`` (i.e.\n", + "one per dataset). Each carries the same ``max_log_likelihood_instance`` (since they share the global\n", + "model) but its own per-dataset visualization and ``FitPoint`` object.\n", + "\n", + "[The ``info_whitespace_length`` config setting also controls whitespace in ``result.info``.]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " print(result.max_log_likelihood_instance)\n", + "\n", + " aplt.subplot_tracer(\n", + " tracer=result.max_log_likelihood_tracer,\n", + " grid=grid,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The ``Samples`` are identical across results (the model is global), so a single corner plot from the\n", + "first result is enough." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result_list[0].samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This script gives a concise overview of the basic cluster modeling API for a small multi-plane cluster.\n", + "\n", + "__Data Preparation__\n", + "\n", + "If you are looking to fit your own point-source cluster data, see\n", + "``autolens_workspace/*/data_preparation/point_source/README.md`` for the input-data standards.\n", + "\n", + "__HowToLens__\n", + "\n", + "For a deeper understanding of how lens modeling, ray-tracing, and non-linear searches actually work, see\n", + "the **HowToLens** Jupyter notebook lectures at https://github.com/PyAutoLabs/HowToLens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/cluster/simulator.ipynb b/notebooks/cluster/simulator.ipynb index fcf2fdd88..365d0316c 100644 --- a/notebooks/cluster/simulator.ipynb +++ b/notebooks/cluster/simulator.ipynb @@ -1,907 +1,965 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Cluster\n", - "==================\n", - "\n", - "This script simulates an example strong lens on the 'cluster' scale: a small cluster consisting of 2 main\n", - "lens galaxies (a brightest cluster galaxy + a single satellite), 10 lower-mass cluster member galaxies on\n", - "a luminosity-mass scaling relation, a single host dark matter halo not tied to any individual galaxy, and\n", - "2 multiply-imaged background source galaxies sitting at *different* redshifts (``z = 1.0`` and ``z = 2.0``)\n", - "\u2014 making this a genuine multi-plane lens.\n", - "\n", - "Real clusters can have tens or hundreds of member galaxies and several background sources. The example\n", - "keeps the main-lens tier minimal (2 individually-modelled galaxies) but is paired with a population of 10\n", - "scaling members so the dataset already exercises the full cluster workflow \u2014 the scaling-relation tier is\n", - "the cluster default rather than an opt-in feature, because every real cluster carries a population of\n", - "lower-mass members that must be modelled collectively. Scaling up to a larger cluster amounts to adding\n", - "rows to ``scaling_galaxies.csv`` (and, optionally, more main galaxies).\n", - "\n", - "Modeling at cluster scale almost always uses the *point source* API: rather than fitting the extended arc\n", - "light of a lensed source, we fit only the image-plane positions of the brightest pixels of each multiple\n", - "image. This script simulates that point-source data alongside CCD imaging \u2014 the imaging is used to\n", - "*measure* the point positions in real datasets and to visually confirm the lens configuration.\n", - "\n", - "__Contents__\n", - "\n", - "- **Multi-Plane Setup:** Why the two sources sit at different redshifts and what that buys the example.\n", - "- **Main Lens vs Scaling Members vs Host Halo vs Source Galaxies:** Galaxies are organized into four categories.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a name.\n", - "- **Imaging and Visualization Grids:** Define the high-res rendering grid and a coarse viz grid.\n", - "- **Galaxy Centres:** Define the centres of the main lens galaxies, scaling members, and sources; used for over-sampling and CSV/JSON output.\n", - "- **Over Sampling:** Adaptive over-sampling grid for accurate light profile evaluation near galaxy centres.\n", - "- **Main Lens Galaxies:** The 2 individually-modelled cluster members \u2014 each has a `SersicSph` light profile and a `dPIEMassSph` mass.\n", - "- **Scaling Member Galaxies:** 10 lower-mass members on a luminosity-mass relation \u2014 collectively important, individually weak.\n", - "- **Host Dark Matter Halo:** A standalone `NFWMCRLudlowSph` halo with `mass_at_200 = 10^15.3` at z=0.5.\n", - "- **Source Galaxies:** The 2 multi-plane background sources, each a `SersicCore` light + a `Point` model.\n", - "- **Ray Tracing:** Combine all galaxies into a single `Tracer` capable of multi-plane ray tracing.\n", - "- **JAX JIT:** Register the tracer's underlying classes as JAX pytrees and compile the point solver.\n", - "- **Point Solver:** Solve for image-plane multiple-image positions of each source.\n", - "- **Point Datasets:** Collect per-source image positions (with noise) into `PointDataset` objects, one per source.\n", - "- **Combined CSV:** Write *all* point datasets to a single CSV so a user can hand-edit positions and noise in a spreadsheet.\n", - "- **Manual CSV Editing:** Instructions for editing the combined CSV by hand, which is the preferred cluster workflow.\n", - "- **Scaling Galaxies CSV:** Write the scaling-member centres and luminosities to ``scaling_galaxies.csv``.\n", - "- **Model CSVs:** Write the truth model to ``mass.csv`` + ``light.csv`` + ``point.csv`` (the named-galaxy CSV API).\n", - "- **Tracer JSON:** Save the true `Tracer` for future inspection.\n", - "- **Imaging:** Simulate CCD imaging of the cluster (used to measure positions in real datasets and for visualization).\n", - "- **Visualize:** Plot the point-source dataset, tracer, and imaging.\n", - "\n", - "__Multi-Plane Setup__\n", - "\n", - "The two background sources sit at distinct redshifts: source 0 at ``z = 1.0`` and source 1 at ``z = 2.0``.\n", - "A real cluster lenses many sources at many different redshifts simultaneously; restricting an example to\n", - "a single source plane is a galaxy-scale approximation that hides multi-plane ray-tracing entirely. By\n", - "choosing two distinct redshifts here we get a concrete multi-plane testbed with the smallest possible\n", - "configuration \u2014 the `Tracer` ray-traces through *both* source planes when solving for the image positions\n", - "of the further source, exercising the multi-plane code path.\n", - "\n", - "The host halo's ``redshift_source`` parameter is anchored to the *furthest* source (``z = 2.0``) so its\n", - "``NFWMCRLudlow`` concentration is set against the deepest light cone in the system. The halo mass\n", - "``10^15.3 M_sun`` is large enough that *both* sources end up multiply-imaged.\n", - "\n", - "__Main Lens vs Scaling Members vs Host Halo vs Source Galaxies__\n", - "\n", - "- `main_lens_galaxies`: The 2 individually-modelled cluster members that dominate the light and contribute\n", - " the brightest galaxy-scale lensing. Each carries its own `SersicSph` light profile and `dPIEMassSph`\n", - " mass; their parameters are free in the modeling script.\n", - "\n", - "- `scaling_galaxies`: 10 lower-mass cluster members modelled collectively via a luminosity-mass scaling\n", - " relation. Each member is individually weak compared to the main galaxies or the host halo, but the\n", - " population together perturbs the deflection field non-trivially \u2014 exactly the regime in which the\n", - " scaling-relation tier of the modeling API earns its keep. The number of free parameters does not grow\n", - " with the number of scaling members; the shared `scaling_factor` and `scaling_exponent` (2 parameters\n", - " total) determine every member's mass from its luminosity.\n", - "\n", - "- `host_halo_galaxy`: A standalone `Galaxy` holding the cluster's `NFWMCRLudlowSph` dark matter halo. It\n", - " is not tied to any individual member galaxy \u2014 the halo is a separate mass component sitting \"on top of\"\n", - " the members.\n", - "\n", - "- `source_galaxies`: The 2 background sources at *different* redshifts. Each carries both a `SersicCore`\n", - " light profile (for visualization of the lensed arcs) and a `Point` model component (used during\n", - " point-source modeling).\n", - "\n", - "Main lens, host halo, and source centres are saved to JSON files. Scaling-member centres and luminosities\n", - "are saved to ``scaling_galaxies.csv`` (the canonical input for the scaling tier).\n", - "\n", - "__dPIE Mass Profile__\n", - "\n", - "The cluster member galaxies use the dual Pseudo-Isothermal Elliptical (dPIE) mass profile introduced in\n", - "Eliasdottir 2007 (https://arxiv.org/abs/0710.5636), the de facto standard for cluster strong lens modeling.\n", - "In spherical form (`dPIEMassSph`), its parameters are:\n", - "\n", - " - `ra` (arcsec): the core radius, below which the density profile flattens (kept small, ~0.05\u20130.1\" at z=0.5).\n", - " - `rs` (arcsec): the truncation radius, above which the density falls as R^-4 (kept ~10\u201330\" for cluster members).\n", - " - `b0` (arcsec): the mass normalization, roughly setting the galaxy-scale Einstein radius.\n", - "\n", - "Per-galaxy values for the 2 main-tier galaxies are hand-tuned below; for the 10 scaling-tier members they\n", - "are derived from each member's luminosity via the relation described next.\n", - "\n", - "__Luminosity-Mass Scaling Relation__\n", - "\n", - "The 10 scaling members share a single one-parameter relation for the dPIE mass normalization:\n", - "\n", - " b0 = scaling_factor * luminosity ** scaling_exponent\n", - "\n", - "Truth values used in this simulator are ``scaling_factor = 0.3`` arcsec / unit luminosity and\n", - "``scaling_exponent = 1.0`` (i.e. linear in luminosity). The core radius ``ra`` and truncation radius\n", - "``rs`` are held fixed across all scaling members at ``ra = 0.1\"`` and ``rs = 10.0\"``; only ``b0`` varies\n", - "per member. Luminosities are log-spaced across roughly 0.05\u20130.40, so per-member ``b0`` values run from\n", - "~0.015 to ~0.12 arcsec \u2014 each member is individually well below the BCG (``b0 = 3.0``) but the 10 of\n", - "them sum to a few-tenths of an arcsec of effective mass, perturbing the deflection field by ~10\u201315%.\n", - "\n", - "The modeling script promotes ``scaling_factor`` and ``scaling_exponent`` to free parameters; their truth\n", - "values above are recovered when the model is fit to the simulated point datasets. Adding more scaling\n", - "members in the future amounts to adding rows to ``scaling_galaxies.csv`` \u2014 the number of free parameters\n", - "in the model stays at 2 for the entire tier.\n", - "\n", - "__NFWMCRLudlow Host Halo__\n", - "\n", - "The host dark matter halo uses `NFWMCRLudlowSph`, which parameterises an NFW profile by the physical mass\n", - "within r_200 (`mass_at_200`) and the lens and source redshifts. Internally the concentration-mass\n", - "relation of Ludlow et al. (2016) sets the concentration, which together with the cosmology determines\n", - "``kappa_s`` and ``scale_radius``. ``mass_at_200 = 10^15.3`` (~2e15 M_sun) is chosen so the combined\n", - "halo + member lensing produces genuinely multiply-imaged sources within the field \u2014 lighter halos\n", - "(``10^14.5``) would only weakly lens these source positions and give a single image each, which is not\n", - "useful as a modeling testbed.\n", - "\n", - "__JAX JIT__\n", - "\n", - "Solving the lens equation for the image-plane positions of a point source is iterative and numerically\n", - "expensive \u2014 at cluster scale (many lens galaxies, multi-plane ray tracing) it dominates the simulator's\n", - "runtime by an order of magnitude. We accelerate it with JAX via ``al.PointSolver(use_jax=True)`` and a\n", - "``@jax.jit`` wrapper around the solve call.\n", - "\n", - "The library handles pytree registration of ``Tracer`` + every Galaxy / profile class internally via\n", - "``autolens.jax.register_tracer_classes(tracer)`` (one user-visible setup line, called before the first\n", - "``@jax.jit`` invocation). The compiled triangle-refinement kernel is cached and reused across both\n", - "sources \u2014 turning what was ~5 minutes of Python-loop overhead into a few seconds of compiled JAX\n", - "execution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import jax\n", - "import jax.numpy as jnp\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "# Pytree registration is now handled by autolens.jax.register_tracer_classes,\n", - "# called once before the @jax.jit'd PointSolver call further below." - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive\n", - "name. They define the folder the dataset is output to on your hard-disk:\n", - "\n", - " - The image will be output to `/autolens_workspace/dataset/cluster/simple/data.fits`.\n", - " - The point datasets will be written to `/autolens_workspace/dataset/cluster/simple/point_datasets.csv`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"cluster\"\n", - "dataset_name = \"simple\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "All main lens galaxies and the host dark matter halo sit at the same lens redshift ``z = 0.5``. The two\n", - "sources sit at *different* redshifts (``z = 1.0`` and ``z = 2.0``); see the ``__Multi-Plane Setup__``\n", - "section in the module docstring for the rationale." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "source_redshifts = [1.0, 2.0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Define the centres of the main lens galaxies, the 10 scaling-tier members, the host halo, and the\n", - "sources. The host halo is anchored at the cluster centre (the origin); the two main galaxies are placed\n", - "at the centre and a single satellite location offset to the upper-right. Scaling-member centres are\n", - "hand-tuned to sit at radii of 5\u201315\" from the centre \u2014 well inside the strongly-lensed region of the\n", - "host halo but clear of the cores of the two main galaxies. Source centres are chosen so that both\n", - "sources land in the strongly-lensed region, producing genuine multiple images." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [\n", - " (0.0, 0.0), # BCG at cluster centre\n", - " (10.0, 8.0), # satellite member\n", - "]\n", - "\n", - "scaling_galaxies_centres = [\n", - " (5.5, -6.5),\n", - " (-7.5, 3.0),\n", - " (12.0, -5.0),\n", - " (-4.0, -9.0),\n", - " (3.0, 13.0),\n", - " (-14.0, 4.0),\n", - " (15.0, 9.0),\n", - " (-9.0, -12.0),\n", - " (8.5, 5.5),\n", - " (-6.5, 11.0),\n", - "]\n", - "\n", - "scaling_galaxies_luminosities = [\n", - " 0.40,\n", - " 0.32,\n", - " 0.25,\n", - " 0.20,\n", - " 0.16,\n", - " 0.13,\n", - " 0.10,\n", - " 0.08,\n", - " 0.06,\n", - " 0.05,\n", - "]\n", - "\n", - "host_halo_centre = (0.0, 0.0)\n", - "\n", - "source_centres = [\n", - " (0.3, 0.5),\n", - " (-0.8, 1.2),\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Imaging and Visualization Grids__\n", - "\n", - "Two grids are used for image rendering and one for visualization plotting:\n", - "\n", - " - ``imaging_grid``: a high-resolution (1000x1000 @ 0.1\"/px) grid with adaptive over-sampling around each\n", - " cluster member. This is the grid passed to ``SimulatorImaging.via_tracer_from`` and gives an accurate\n", - " simulated CCD image.\n", - " - ``viz_grid``: a coarse (200x200 @ 0.5\"/px), un-over-sampled grid passed only to the visualization\n", - " plotters at the end of the script. Visualization plots are illustrative \u2014 they don't need the same\n", - " resolution or sub-sampling as the rendered data, and using the imaging grid for them dominated the\n", - " simulator's runtime in earlier versions of this script.\n", - "\n", - "Both grids span the same 100\"x100\" field \u2014 the typical Einstein radius of a ``10^15`` M_sun halo is\n", - "~20\u201330\" and the member galaxies span ~30\" across, so the field has to be large to capture the multiple\n", - "images and extended arc light. The PointSolver builds *its own* internal grid for triangle root-finding\n", - "(see the ``__Point Solver__`` section below); that grid is independent of these rendering grids by design." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "imaging_grid = al.Grid2D.uniform(\n", - " shape_native=(1000, 1000),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "viz_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.5)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling evaluates light profiles on a higher-resolution sub-grid in bright central regions, trading\n", - "compute for accuracy. For cluster lenses we over-sample around the centre of every cluster member \u2014\n", - "both the 2 main galaxies and the 10 scaling members \u2014 so each galaxy's Sersic profile is rendered\n", - "accurately even at the smaller effective radii of the scaling-tier members.\n", - "\n", - "The source galaxies use a cored `SersicCore` profile so that lensed arcs can be evaluated without\n", - "explicit source-plane over-sampling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "imaging_over_sample = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=imaging_grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=main_lens_centres + scaling_galaxies_centres,\n", - ")\n", - "\n", - "imaging_grid = imaging_grid.apply_over_sampling(over_sample_size=imaging_over_sample)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "The 2 cluster member galaxies. Each is given a `SersicSph` light profile (used only for visualization \u2014\n", - "the imaging data is not used in point-source modeling) and a `dPIEMassSph` mass profile with hand-tuned\n", - "parameters representative of cluster members: a larger central BCG and one smaller satellite galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_dpie_params = [\n", - " # (ra, rs, b0) per galaxy \u2014 arcsec\n", - " (8.0, 20.0, 3.0), # BCG \u2014 strongest\n", - " (5.0, 12.0, 1.2), # satellite\n", - "]\n", - "\n", - "main_lens_sersic_params = [\n", - " # (intensity, effective_radius, sersic_index)\n", - " (1.5, 3.0, 4.0), # BCG \u2014 bright and extended\n", - " (0.8, 1.5, 3.5), # satellite\n", - "]\n", - "\n", - "main_lens_galaxies = []\n", - "for centre, (ra, rs, b0), (intensity, effective_radius, sersic_index) in zip(\n", - " main_lens_centres, main_lens_dpie_params, main_lens_sersic_params\n", - "):\n", - " bulge = al.lp.SersicSph(\n", - " centre=centre,\n", - " intensity=intensity,\n", - " effective_radius=effective_radius,\n", - " sersic_index=sersic_index,\n", - " )\n", - " mass = al.mp.dPIEMassSph(centre=centre, ra=ra, rs=rs, b0=b0)\n", - " main_lens_galaxies.append(al.Galaxy(redshift=redshift_lens, bulge=bulge, mass=mass))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling Member Galaxies__\n", - "\n", - "The 10 cluster members modelled collectively via the luminosity-mass scaling relation (see the\n", - "``__Luminosity-Mass Scaling Relation__`` section of the module docstring). The simulator hardcodes the\n", - "truth values of ``scaling_factor`` and ``scaling_exponent`` here and derives each member's `b0` from its\n", - "luminosity. ``ra`` and ``rs`` are held fixed across all scaling members \u2014 only ``b0`` varies. Light\n", - "profiles use the per-member luminosity as the central intensity so the rendered image visibly traces the\n", - "scaling-tier population." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scaling_factor_truth = 0.3\n", - "scaling_exponent_truth = 1.0\n", - "scaling_ra = 0.1\n", - "scaling_rs = 10.0\n", - "\n", - "scaling_galaxies = []\n", - "for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities):\n", - " bulge = al.lp.SersicSph(\n", - " centre=centre,\n", - " intensity=luminosity,\n", - " effective_radius=0.8,\n", - " sersic_index=3.0,\n", - " )\n", - " b0 = scaling_factor_truth * luminosity**scaling_exponent_truth\n", - " mass = al.mp.dPIEMassSph(centre=centre, ra=scaling_ra, rs=scaling_rs, b0=b0)\n", - " scaling_galaxies.append(al.Galaxy(redshift=redshift_lens, bulge=bulge, mass=mass))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Host Dark Matter Halo__\n", - "\n", - "A standalone galaxy holding the cluster's NFW dark matter halo. It has no light profile \u2014 it sits in the\n", - "tracer solely to contribute mass. `NFWMCRLudlowSph` is parameterised by the physical halo mass within\n", - "r_200 and the redshifts; the concentration is set by the Ludlow et al. (2016) concentration-mass relation.\n", - "The ``redshift_source`` argument is anchored to the *furthest* source (``z = 2.0``) so the concentration\n", - "is computed against the deepest light cone in the multi-plane system." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "host_halo = al.mp.NFWMCRLudlowSph(\n", - " centre=host_halo_centre,\n", - " mass_at_200=10**15.3,\n", - " redshift_object=redshift_lens,\n", - " redshift_source=max(source_redshifts),\n", - ")\n", - "\n", - "host_halo_galaxy = al.Galaxy(redshift=redshift_lens, dark=host_halo)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxies__\n", - "\n", - "The 2 background sources at *different* redshifts. Each carries a `SersicCore` light profile (used only\n", - "for visual confirmation of the lensed arcs \u2014 the cored profile changes gradually in the centre so explicit\n", - "source-plane over-sampling is unnecessary) and a `Point` model component whose multiple-image positions\n", - "we solve for and use as the modeling data.\n", - "\n", - "Each source's redshift is taken from ``source_redshifts``, so source 0 sits at ``z = 1.0`` and source 1\n", - "at ``z = 2.0``. The `Tracer` ray-traces multi-plane through both planes automatically." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxies = []\n", - "for i, (centre, src_z) in enumerate(zip(source_centres, source_redshifts)):\n", - " bulge = al.lp.SersicCore(\n", - " centre=centre,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0 + 30.0 * i),\n", - " intensity=2.0,\n", - " effective_radius=0.3,\n", - " sersic_index=1.0,\n", - " )\n", - " point = al.ps.Point(centre=centre)\n", - " source_galaxies.append(\n", - " al.Galaxy(redshift=src_z, bulge=bulge, **{f\"point_{i}\": point})\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Combine main lens galaxies, the scaling-tier members, the host halo galaxy, and the source galaxies into\n", - "a single tracer that produces the simulated image. With sources at distinct redshifts, the tracer\n", - "automatically handles multi-plane ray tracing. The scaling members share the lens redshift, so they\n", - "contribute to the single lens-plane deflection alongside the main galaxies and the halo." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=main_lens_galaxies\n", - " + scaling_galaxies\n", - " + [host_halo_galaxy]\n", - " + source_galaxies\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__JAX JIT \u2014 Point Solver__\n", - "\n", - "Solving the lens equation for the image-plane positions of a point source is iterative and numerically\n", - "expensive \u2014 at cluster scale it dominates the simulator's runtime. We use ``al.PointSolver(use_jax=True)``\n", - "and wrap the solve in ``@jax.jit`` for the speedup.\n", - "\n", - "The library handles pytree registration of ``Tracer`` + every reachable galaxy / profile class via\n", - "the one-time ``autolens.jax.register_tracer_classes(tracer)`` call below. Before\n", - "PR PyAutoLens#538 + PyAutoArray#335 (Phase 2 of ``z_features/jax_user_intro.md``), this section was\n", - "a ~60-line manual ceremony with an ``af.Model`` mirror, ``register_model``, and ``register_instance_pytree``;\n", - "the new API collapses it to a single import + single registration call + ``use_jax=True`` flag.\n", - "\n", - "``PointSolver(use_jax=True).solve`` defaults ``remove_infinities=False`` to honour the JAX static-shape\n", - "contract \u2014 the returned positions are padded with ``inf`` where no image was found, which is JIT-safe.\n", - "We strip them outside the jit below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autolens.jax import register_tracer_classes\n", - "\n", - "register_tracer_classes(tracer)\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.1),\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.1,\n", - " use_jax=True,\n", - ")\n", - "\n", - "\n", - "@jax.jit\n", - "def jitted_solve(tracer, source_plane_coordinate):\n", - " return solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_plane_coordinate\n", - " ).array\n", - "\n", - "\n", - "positions_list = []\n", - "for i, src_centre in enumerate(source_centres):\n", - " coord = jnp.asarray(src_centre)\n", - " raw = np.asarray(jitted_solve(tracer, coord))\n", - " finite = ~(np.isinf(raw).any(axis=1) | np.isnan(raw).any(axis=1))\n", - " positions_list.append(al.Grid2DIrregular(raw[finite]))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Datasets__\n", - "\n", - "One `PointDataset` per source. The `name` (e.g. `point_0`, `point_1`) pairs each dataset with the\n", - "matching `Point` component in the lens model during modeling.\n", - "\n", - "`redshift` is populated on each `PointDataset` so that the per-source redshifts round-trip through the\n", - "combined CSV below. This is the piece that makes the CSV self-describing for cluster modeling \u2014 position,\n", - "noise, and redshift live in a single spreadsheet.\n", - "\n", - "The position uncertainty is set to 0.005\" (5 mas), reflecting the centroid precision achievable by PSF\n", - "fitting on HST or adaptive-optics imaging \u2014 *not* the imaging pixel scale, which is the detector's\n", - "sampling rather than its centroiding precision. See `scripts/point_source/simulator.py` for a full\n", - "discussion of this value." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "position_noise = 0.005\n", - "\n", - "dataset_list = []\n", - "for i, positions in enumerate(positions_list):\n", - " dataset = al.PointDataset(\n", - " name=f\"point_{i}\",\n", - " positions=positions,\n", - " positions_noise_map=position_noise,\n", - " redshift=source_redshifts[i],\n", - " )\n", - " dataset_list.append(dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output one .json file per dataset (exact round-trip; this is the canonical modeling input)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for i, dataset in enumerate(dataset_list):\n", - " al.output_to_json(\n", - " obj=dataset,\n", - " file_path=dataset_path / f\"point_dataset_{i}.json\",\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Combined CSV__\n", - "\n", - "For cluster-scale workflows with tens or hundreds of sources, a single CSV with one row per observed\n", - "image \u2014 grouped by ``name`` \u2014 is far easier to edit in a spreadsheet than many per-source JSON files.\n", - "``al.output_to_csv`` writes every dataset into one file. The `redshift` column is emitted automatically\n", - "because each dataset has its redshift set above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_csv(\n", - " datasets=dataset_list,\n", - " file_path=dataset_path / \"point_datasets.csv\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Manual CSV Editing__\n", - "\n", - "The combined CSV is the preferred cluster input: it is human-readable, editable in Excel / LibreOffice /\n", - "any text editor, and round-trips cleanly back into `list_from_csv`. The expected format is one row per\n", - "observed multiple image with the following columns:\n", - "\n", - " - `name` \u2014 the source identifier (e.g. `point_0`). All rows sharing a `name` belong to the same source.\n", - " - `y`, `x` \u2014 the image-plane position of the multiple image, in arc-seconds.\n", - " - `positions_noise` \u2014 the positional uncertainty in arc-seconds. This is the PSF-fit centroid\n", - " uncertainty on each multiple image, *not* the imaging pixel scale. For HST or adaptive-optics\n", - " data on bright cluster member images, ~0.005\" (5 mas) is a defensible default; ground-based\n", - " seeing-limited data may warrant a few \u00d7 0.01\" depending on image SNR.\n", - " - `redshift` \u2014 the source redshift. Every row for a given `name` must share the *same* redshift\n", - " (validated on load; `list_from_csv` raises if a group's rows disagree). Leave the cell blank if the\n", - " redshift is unknown; blank is tolerated as long as *all* rows in a group are blank.\n", - "\n", - "Optional columns `flux`, `flux_noise`, `time_delay`, `time_delay_noise` are also round-tripped \u2014 populate\n", - "them when the observation provides those measurements, leave them blank otherwise.\n", - "\n", - "To build a cluster dataset by hand, simulate or manually collect one set of images per source, then edit\n", - "the CSV directly: add or remove rows, adjust positions or noises, and save. Reload the dataset in a\n", - "modeling script with ``al.list_from_csv(file_path=dataset_path / \"point_datasets.csv\")``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Scaling Galaxies CSV__\n", - "\n", - "The scaling-tier members are written to a separate CSV \u2014 ``scaling_galaxies.csv`` \u2014 with one row per\n", - "member carrying its centre and luminosity. ``al.galaxy_table_to_csv`` produces the canonical schema\n", - "(`y, x, luminosity, redshift?`) that the modeling script consumes via ``al.galaxy_table_from_csv``.\n", - "\n", - "Scaling up a real cluster to a larger member population is then a CSV-level edit: add a row per\n", - "additional member, fill in its centre and luminosity, save. The modeling script picks up the new rows\n", - "automatically and the number of free parameters in the scaling tier stays at 2 (`scaling_factor` and\n", - "`scaling_exponent`).\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.galaxy_table_to_csv(\n", - " centres=scaling_galaxies_centres,\n", - " luminosities=scaling_galaxies_luminosities,\n", - " file_path=dataset_path / \"scaling_galaxies.csv\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer JSON__\n", - "\n", - "Save the `Tracer` so the true light profiles, mass profiles and galaxies can be inspected after the fact.\n", - "This can be loaded via `tracer = al.from_json(file_path)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=dataset_path / \"tracer.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model CSVs__\n", - "\n", - "Write the truth model out as three family-level CSVs \u2014 ``mass.csv``, ``light.csv``, ``point.csv`` \u2014\n", - "keyed by galaxy name. The modeling and start_here scripts load these directly with\n", - "``al.galaxy_models_from_csv`` and compose them into ``af.Model[Galaxy]`` instances ready for non-linear\n", - "search. See ``scripts/cluster/csv_api.py`` for the full schema walkthrough.\n", - "\n", - "The scaling tier keeps its narrow 3-column ``scaling_galaxies.csv`` schema written above \u2014 naming each\n", - "scaling member and emitting an ``attr_name`` column would be more overhead than signal." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass_profiles = {\n", - " **{f\"lens_{i}\": {\"mass\": g.mass} for i, g in enumerate(main_lens_galaxies)},\n", - " \"host_halo\": {\"dark\": host_halo_galaxy.dark},\n", - "}\n", - "\n", - "light_profiles = {\n", - " **{f\"lens_{i}\": {\"bulge\": g.bulge} for i, g in enumerate(main_lens_galaxies)},\n", - " **{f\"source_{i}\": {\"bulge\": g.bulge} for i, g in enumerate(source_galaxies)},\n", - "}\n", - "\n", - "point_profiles = {\n", - " f\"source_{i}\": {f\"point_{i}\": getattr(g, f\"point_{i}\")}\n", - " for i, g in enumerate(source_galaxies)\n", - "}\n", - "\n", - "redshifts_by_galaxy = {\n", - " **{f\"lens_{i}\": redshift_lens for i in range(len(main_lens_galaxies))},\n", - " \"host_halo\": redshift_lens,\n", - " **{f\"source_{i}\": z for i, z in enumerate(source_redshifts)},\n", - "}\n", - "\n", - "al.galaxy_models_to_csv(\n", - " profiles_by_galaxy=mass_profiles,\n", - " file_path=dataset_path / \"mass.csv\",\n", - " family=\"mass\",\n", - " redshifts=redshifts_by_galaxy,\n", - ")\n", - "\n", - "al.galaxy_models_to_csv(\n", - " profiles_by_galaxy=light_profiles,\n", - " file_path=dataset_path / \"light.csv\",\n", - " family=\"light\",\n", - " redshifts=redshifts_by_galaxy,\n", - ")\n", - "\n", - "al.galaxy_models_to_csv(\n", - " profiles_by_galaxy=point_profiles,\n", - " file_path=dataset_path / \"point.csv\",\n", - " family=\"point\",\n", - " redshifts=redshifts_by_galaxy,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Imaging__\n", - "\n", - "Strong lens clusters typically come with imaging data \u2014 used to *measure* the point positions and to\n", - "visually confirm the lens configuration. Although modeling here is point-source only, we output CCD\n", - "imaging so the dataset looks like a realistic cluster observation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=imaging_grid.pixel_scales\n", - ")\n", - "\n", - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")\n", - "\n", - "dataset = simulator.via_tracer_from(tracer=tracer, grid=imaging_grid)\n", - "\n", - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output .png plots of the simulated dataset, the tracer, and the per-source point datasets.\n", - "\n", - "These use the default galaxy-scale plotters and are known to be suboptimal for cluster-scale systems \u2014\n", - "arcs span a much larger field, per-source images benefit from distinct colours, and multi-source overlays\n", - "are useful. A follow-up prompt (`PyAutoPrompt/cluster/1_visualization.md`) addresses these\n", - "visualization requirements." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for i, pd in enumerate(dataset_list):\n", - " aplt.subplot_point_dataset(\n", - " dataset=pd, output_path=dataset_path, output_format=\"png\"\n", - " )\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=viz_grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=viz_grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Cluster\n", + "==================\n", + "\n", + "This script simulates an example strong lens on the 'cluster' scale: a small cluster consisting of 2 main\n", + "lens galaxies (a brightest cluster galaxy + a single satellite), 10 lower-mass cluster member galaxies on\n", + "a luminosity-mass scaling relation, a single host dark matter halo not tied to any individual galaxy, and\n", + "2 multiply-imaged background source galaxies sitting at *different* redshifts (``z = 1.0`` and ``z = 2.0``)\n", + "\u2014 making this a genuine multi-plane lens.\n", + "\n", + "Real clusters can have tens or hundreds of member galaxies and several background sources. The example\n", + "keeps the main-lens tier minimal (2 individually-modelled galaxies) but is paired with a population of 10\n", + "scaling members so the dataset already exercises the full cluster workflow \u2014 the scaling-relation tier is\n", + "the cluster default rather than an opt-in feature, because every real cluster carries a population of\n", + "lower-mass members that must be modelled collectively. Scaling up to a larger cluster amounts to adding\n", + "rows to ``scaling_galaxies.csv`` (and, optionally, more main galaxies).\n", + "\n", + "Modeling at cluster scale almost always uses the *point source* API: rather than fitting the extended arc\n", + "light of a lensed source, we fit only the image-plane positions of the brightest pixels of each multiple\n", + "image. This script simulates that point-source data alongside CCD imaging \u2014 the imaging is used to\n", + "*measure* the point positions in real datasets and to visually confirm the lens configuration.\n", + "\n", + "__Contents__\n", + "\n", + "- **Multi-Plane Setup:** Why the two sources sit at different redshifts and what that buys the example.\n", + "- **Main Lens vs Scaling Members vs Host Halo vs Source Galaxies:** Galaxies are organized into four categories.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a name.\n", + "- **Imaging and Visualization Grids:** Define the high-res rendering grid and a coarse viz grid.\n", + "- **Galaxy Centres:** Define the centres of the main lens galaxies, scaling members, and sources; used for over-sampling and CSV/JSON output.\n", + "- **Over Sampling:** Adaptive over-sampling grid for accurate light profile evaluation near galaxy centres.\n", + "- **Main Lens Galaxies:** The 2 individually-modelled cluster members \u2014 each has a `SersicSph` light profile and a `dPIEMassSph` mass.\n", + "- **Scaling Member Galaxies:** 10 lower-mass members on a luminosity-mass relation \u2014 collectively important, individually weak.\n", + "- **Host Dark Matter Halo:** A standalone `NFWMCRLudlowSph` halo with `mass_at_200 = 10^15.3` at z=0.5.\n", + "- **Source Galaxies:** The 2 multi-plane background sources, each a `SersicCore` light + a `Point` model.\n", + "- **Ray Tracing:** Combine all galaxies into a single `Tracer` capable of multi-plane ray tracing.\n", + "- **JAX JIT:** Register the tracer's underlying classes as JAX pytrees and compile the point solver.\n", + "- **Point Solver:** Solve for image-plane multiple-image positions of each source.\n", + "- **Point Datasets:** Collect per-source image positions (with noise) into `PointDataset` objects, one per source.\n", + "- **Combined CSV:** Write *all* point datasets to a single CSV so a user can hand-edit positions and noise in a spreadsheet.\n", + "- **Manual CSV Editing:** Instructions for editing the combined CSV by hand, which is the preferred cluster workflow.\n", + "- **Scaling Galaxies CSV:** Write the scaling-member centres and luminosities to ``scaling_galaxies.csv``.\n", + "- **Model CSVs:** Write the truth model to ``mass.csv`` + ``light.csv`` + ``point.csv`` (the named-galaxy CSV API).\n", + "- **Tracer JSON:** Save the true `Tracer` for future inspection.\n", + "- **Imaging:** Simulate CCD imaging of the cluster (used to measure positions in real datasets and for visualization).\n", + "- **Visualize:** Plot the point-source dataset, tracer, and imaging.\n", + "\n", + "__Multi-Plane Setup__\n", + "\n", + "The two background sources sit at distinct redshifts: source 0 at ``z = 1.0`` and source 1 at ``z = 2.0``.\n", + "A real cluster lenses many sources at many different redshifts simultaneously; restricting an example to\n", + "a single source plane is a galaxy-scale approximation that hides multi-plane ray-tracing entirely. By\n", + "choosing two distinct redshifts here we get a concrete multi-plane testbed with the smallest possible\n", + "configuration \u2014 the `Tracer` ray-traces through *both* source planes when solving for the image positions\n", + "of the further source, exercising the multi-plane code path.\n", + "\n", + "The host halo's ``redshift_source`` parameter is anchored to the *furthest* source (``z = 2.0``) so its\n", + "``NFWMCRLudlow`` concentration is set against the deepest light cone in the system. The halo mass\n", + "``10^15.3 M_sun`` is large enough that *both* sources end up multiply-imaged.\n", + "\n", + "__Main Lens vs Scaling Members vs Host Halo vs Source Galaxies__\n", + "\n", + "- `main_lens_galaxies`: The 2 individually-modelled cluster members that dominate the light and contribute\n", + " the brightest galaxy-scale lensing. Each carries its own `SersicSph` light profile and `dPIEMassSph`\n", + " mass; their parameters are free in the modeling script.\n", + "\n", + "- `scaling_galaxies`: 10 lower-mass cluster members modelled collectively via a luminosity-mass scaling\n", + " relation. Each member is individually weak compared to the main galaxies or the host halo, but the\n", + " population together perturbs the deflection field non-trivially \u2014 exactly the regime in which the\n", + " scaling-relation tier of the modeling API earns its keep. The number of free parameters does not grow\n", + " with the number of scaling members; a single shared normalization `b0_ref` (the lens strength of the\n", + " brightest member, with the relation's exponent fixed at the Faber-Jackson value) determines every\n", + " member's mass from its luminosity.\n", + "\n", + "- `host_halo_galaxy`: A standalone `Galaxy` holding the cluster's `NFWMCRLudlowSph` dark matter halo. It\n", + " is not tied to any individual member galaxy \u2014 the halo is a separate mass component sitting \"on top of\"\n", + " the members.\n", + "\n", + "- `source_galaxies`: The 2 background sources at *different* redshifts. Each carries both a `SersicCore`\n", + " light profile (for visualization of the lensed arcs) and a `Point` model component (used during\n", + " point-source modeling).\n", + "\n", + "Main lens, host halo, and source truth parameters (including centres) are saved to the named-galaxy CSVs\n", + "(``mass.csv`` / ``light.csv`` / ``point.csv``). Scaling-member centres and luminosities are saved to\n", + "``scaling_galaxies.csv`` (the canonical input for the scaling tier).\n", + "\n", + "__dPIE Mass Profile__\n", + "\n", + "The cluster member galaxies use the dual Pseudo-Isothermal Elliptical (dPIE) mass profile introduced in\n", + "Eliasdottir 2007 (https://arxiv.org/abs/0710.5636), the de facto standard for cluster strong lens modeling.\n", + "In spherical form (`dPIEMassSph`), its parameters are:\n", + "\n", + " - `ra` (arcsec): the core radius, below which the density profile flattens (kept small, ~0.05\u20130.1\" at z=0.5).\n", + " - `rs` (arcsec): the truncation radius, above which the density falls as R^-4 (kept ~10\u201330\" for cluster members).\n", + " - `b0` (arcsec): the mass normalization, roughly setting the galaxy-scale Einstein radius.\n", + "\n", + "Per-galaxy values for the 2 main-tier galaxies are hand-tuned below; for the 10 scaling-tier members they\n", + "are derived from each member's luminosity via the relation described next.\n", + "\n", + "__Luminosity-Mass Scaling Relation__\n", + "\n", + "The 10 scaling members share a reference-anchored relation for the dPIE mass normalization \u2014 the\n", + "convention used by Lenstool and essentially every published cluster strong-lensing analysis\n", + "(Limousin et al. 2005; Eliasdottir et al. 2007; Bergamini et al. 2019):\n", + "\n", + " b0_i = b0_ref * (L_i / L_ref) ** 0.5\n", + " rs_i = rs_ref * (L_i / L_ref) ** 0.5\n", + "\n", + "where ``L_ref`` is the luminosity of the *brightest scaling member* (the reference galaxy) and\n", + "``b0_ref`` is that member's lens strength. Anchoring to a reference galaxy makes the normalization\n", + "physically interpretable \u2014 it is the Einstein-radius-like strength of a galaxy you can point at in the\n", + "image \u2014 which is what makes a sensible prior range easy to define. The exponent is **fixed at 0.5**\n", + "rather than fitted: for the dPIE, ``b0`` is proportional to the velocity dispersion squared, and the\n", + "Faber-Jackson relation (L \u221d sigma^4, i.e. sigma \u221d L^(1/4)) then gives b0 \u221d L^(1/2). Lenstool applies\n", + "the same fixed-exponent scaling to the truncation radius (r_cut \u221d L^(1/2)), which is why ``rs`` scales\n", + "here too; the core radius ``ra`` is held fixed at a small value across the tier (0.1\"), again following\n", + "standard practice, since strong lensing barely constrains it.\n", + "\n", + "Truth values used in this simulator are ``b0_ref = 0.12`` arcsec and ``rs_ref = 10.0`` arcsec, anchored\n", + "to the brightest member (``L_ref = 0.40``). Luminosities are log-spaced across roughly 0.05\u20130.40, so\n", + "per-member ``b0`` values run from ~0.042 to 0.12 arcsec \u2014 each member is individually well below the\n", + "BCG (``b0 = 3.0``) but the 10 of them together perturb the deflection field by ~10\u201315%.\n", + "\n", + "The modeling script promotes ``b0_ref`` to the tier's single free parameter and recovers the truth value\n", + "when fit to the simulated point datasets. Adding more scaling members amounts to adding rows to\n", + "``scaling_galaxies.csv`` \u2014 the tier's free-parameter count stays at 1. Note that only the luminosity\n", + "*ratios* ``L_i / L_ref`` enter the relation, so the units of the luminosity column are irrelevant;\n", + "observational catalogues quoting magnitudes convert via ``L_i / L_ref = 10 ** (0.4 * (m_ref - m_i))``.\n", + "Kinematic calibrations of the exponent exist for when higher fidelity is needed \u2014 Bergamini et al. 2019\n", + "measure sigma \u221d L^0.27-0.28 from MUSE member kinematics (b0 exponent \u2248 0.55) and derive the r_cut\n", + "exponent from the fundamental plane \u2014 but 0.5 is the standard default.\n", + "\n", + "__NFWMCRLudlow Host Halo__\n", + "\n", + "The host dark matter halo uses `NFWMCRLudlowSph`, which parameterises an NFW profile by the physical mass\n", + "within r_200 (`mass_at_200`) and the lens and source redshifts. Internally the concentration-mass\n", + "relation of Ludlow et al. (2016) sets the concentration, which together with the cosmology determines\n", + "``kappa_s`` and ``scale_radius``. ``mass_at_200 = 10^15.3`` (~2e15 M_sun) is chosen so the combined\n", + "halo + member lensing produces genuinely multiply-imaged sources within the field \u2014 lighter halos\n", + "(``10^14.5``) would only weakly lens these source positions and give a single image each, which is not\n", + "useful as a modeling testbed.\n", + "\n", + "__JAX JIT__\n", + "\n", + "Solving the lens equation for the image-plane positions of a point source is iterative and numerically\n", + "expensive \u2014 at cluster scale (many lens galaxies, multi-plane ray tracing) it dominates the simulator's\n", + "runtime by an order of magnitude. We accelerate it with JAX via ``al.PointSolver(use_jax=True)`` and a\n", + "``@jax.jit`` wrapper around the solve call.\n", + "\n", + "The library handles pytree registration of ``Tracer`` + every Galaxy / profile class internally via\n", + "``autolens.jax.register_tracer_classes(tracer)`` (one user-visible setup line, called before the first\n", + "``@jax.jit`` invocation). The compiled triangle-refinement kernel is cached and reused across both\n", + "sources \u2014 turning what was ~5 minutes of Python-loop overhead into a few seconds of compiled JAX\n", + "execution." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import jax\n", + "import jax.numpy as jnp\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "# Pytree registration is now handled by autolens.jax.register_tracer_classes,\n", + "# called once before the @jax.jit'd PointSolver call further below." + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive\n", + "name. They define the folder the dataset is output to on your hard-disk:\n", + "\n", + " - The image will be output to `/autolens_workspace/dataset/cluster/simple/data.fits`.\n", + " - The point datasets will be written to `/autolens_workspace/dataset/cluster/simple/point_datasets.csv`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"cluster\"\n", + "dataset_name = \"simple\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "All main lens galaxies and the host dark matter halo sit at the same lens redshift ``z = 0.5``. The two\n", + "sources sit at *different* redshifts (``z = 1.0`` and ``z = 2.0``); see the ``__Multi-Plane Setup__``\n", + "section in the module docstring for the rationale." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "source_redshifts = [1.0, 2.0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Define the centres of the main lens galaxies, the 10 scaling-tier members, the host halo, and the\n", + "sources. The host halo is anchored at the cluster centre (the origin); the two main galaxies are placed\n", + "at the centre and a single satellite location offset to the upper-right. Scaling-member centres are\n", + "hand-tuned to sit at radii of 5\u201315\" from the centre \u2014 well inside the strongly-lensed region of the\n", + "host halo but clear of the cores of the two main galaxies. Source centres are chosen so that both\n", + "sources land in the strongly-lensed region, producing genuine multiple images." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [\n", + " (0.0, 0.0), # BCG at cluster centre\n", + " (10.0, 8.0), # satellite member\n", + "]\n", + "\n", + "scaling_galaxies_centres = [\n", + " (5.5, -6.5),\n", + " (-7.5, 3.0),\n", + " (12.0, -5.0),\n", + " (-4.0, -9.0),\n", + " (3.0, 13.0),\n", + " (-14.0, 4.0),\n", + " (15.0, 9.0),\n", + " (-9.0, -12.0),\n", + " (8.5, 5.5),\n", + " (-6.5, 11.0),\n", + "]\n", + "\n", + "scaling_galaxies_luminosities = [\n", + " 0.40,\n", + " 0.32,\n", + " 0.25,\n", + " 0.20,\n", + " 0.16,\n", + " 0.13,\n", + " 0.10,\n", + " 0.08,\n", + " 0.06,\n", + " 0.05,\n", + "]\n", + "\n", + "host_halo_centre = (0.0, 0.0)\n", + "\n", + "source_centres = [\n", + " (0.3, 0.5),\n", + " (-0.8, 1.2),\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imaging and Visualization Grids__\n", + "\n", + "Two grids are used for image rendering and one for visualization plotting:\n", + "\n", + " - ``imaging_grid``: a high-resolution (1000x1000 @ 0.1\"/px) grid with adaptive over-sampling around each\n", + " cluster member. This is the grid passed to ``SimulatorImaging.via_tracer_from`` and gives an accurate\n", + " simulated CCD image.\n", + " - ``viz_grid``: a coarse (200x200 @ 0.5\"/px), un-over-sampled grid passed only to the visualization\n", + " plotters at the end of the script. Visualization plots are illustrative \u2014 they don't need the same\n", + " resolution or sub-sampling as the rendered data, and using the imaging grid for them dominated the\n", + " simulator's runtime in earlier versions of this script.\n", + "\n", + "Both grids span the same 100\"x100\" field \u2014 the typical Einstein radius of a ``10^15`` M_sun halo is\n", + "~20\u201330\" and the member galaxies span ~30\" across, so the field has to be large to capture the multiple\n", + "images and extended arc light. The PointSolver builds *its own* internal grid for triangle root-finding\n", + "(see the ``__Point Solver__`` section below); that grid is independent of these rendering grids by design." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "imaging_grid = al.Grid2D.uniform(\n", + " shape_native=(1000, 1000),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "viz_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.5)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling evaluates light profiles on a higher-resolution sub-grid in bright central regions, trading\n", + "compute for accuracy. For cluster lenses we over-sample around the centre of every cluster member \u2014\n", + "both the 2 main galaxies and the 10 scaling members \u2014 so each galaxy's Sersic profile is rendered\n", + "accurately even at the smaller effective radii of the scaling-tier members.\n", + "\n", + "The source galaxies use a cored `SersicCore` profile so that lensed arcs can be evaluated without\n", + "explicit source-plane over-sampling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "imaging_over_sample = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=imaging_grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=main_lens_centres + scaling_galaxies_centres,\n", + ")\n", + "\n", + "imaging_grid = imaging_grid.apply_over_sampling(over_sample_size=imaging_over_sample)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "The 2 cluster member galaxies. Each is given a `SersicSph` light profile (used only for visualization \u2014\n", + "the imaging data is not used in point-source modeling) and a `dPIEMassSph` mass profile with hand-tuned\n", + "parameters representative of cluster members: a larger central BCG and one smaller satellite galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_dpie_params = [\n", + " # (ra, rs, b0) per galaxy \u2014 arcsec\n", + " (8.0, 20.0, 3.0), # BCG \u2014 strongest\n", + " (5.0, 12.0, 1.2), # satellite\n", + "]\n", + "\n", + "main_lens_sersic_params = [\n", + " # (intensity, effective_radius, sersic_index)\n", + " (1.5, 3.0, 4.0), # BCG \u2014 bright and extended\n", + " (0.8, 1.5, 3.5), # satellite\n", + "]\n", + "\n", + "main_lens_galaxies = []\n", + "for centre, (ra, rs, b0), (intensity, effective_radius, sersic_index) in zip(\n", + " main_lens_centres, main_lens_dpie_params, main_lens_sersic_params\n", + "):\n", + " bulge = al.lp.SersicSph(\n", + " centre=centre,\n", + " intensity=intensity,\n", + " effective_radius=effective_radius,\n", + " sersic_index=sersic_index,\n", + " )\n", + " mass = al.mp.dPIEMassSph(centre=centre, ra=ra, rs=rs, b0=b0)\n", + " main_lens_galaxies.append(al.Galaxy(redshift=redshift_lens, bulge=bulge, mass=mass))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling Member Galaxies__\n", + "\n", + "The 10 cluster members modelled collectively via the luminosity-mass scaling relation (see the\n", + "``__Luminosity-Mass Scaling Relation__`` section of the module docstring). The simulator hardcodes the\n", + "truth value of ``b0_ref`` (the brightest member's lens strength) and derives each member's ``b0`` and\n", + "``rs`` from its luminosity ratio to the reference, with both exponents fixed at the Faber-Jackson value\n", + "of 0.5. ``ra`` is held fixed across all scaling members. Light profiles use the per-member luminosity\n", + "as the central intensity so the rendered image visibly traces the scaling-tier population." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scaling_b0_ref_truth = 0.12\n", + "scaling_exponent = 0.5\n", + "scaling_luminosity_ref = max(scaling_galaxies_luminosities)\n", + "scaling_ra = 0.1\n", + "scaling_rs_ref = 10.0\n", + "\n", + "scaling_galaxies = []\n", + "for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities):\n", + " bulge = al.lp.SersicSph(\n", + " centre=centre,\n", + " intensity=luminosity,\n", + " effective_radius=0.8,\n", + " sersic_index=3.0,\n", + " )\n", + " luminosity_ratio = luminosity / scaling_luminosity_ref\n", + " b0 = scaling_b0_ref_truth * luminosity_ratio**scaling_exponent\n", + " rs = scaling_rs_ref * luminosity_ratio**scaling_exponent\n", + " mass = al.mp.dPIEMassSph(centre=centre, ra=scaling_ra, rs=rs, b0=b0)\n", + " scaling_galaxies.append(al.Galaxy(redshift=redshift_lens, bulge=bulge, mass=mass))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Host Dark Matter Halo__\n", + "\n", + "A standalone galaxy holding the cluster's NFW dark matter halo. It has no light profile \u2014 it sits in the\n", + "tracer solely to contribute mass. `NFWMCRLudlowSph` is parameterised by the physical halo mass within\n", + "r_200 and the redshifts; the concentration is set by the Ludlow et al. (2016) concentration-mass relation.\n", + "The ``redshift_source`` argument is anchored to the *furthest* source (``z = 2.0``) so the concentration\n", + "is computed against the deepest light cone in the multi-plane system." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "host_halo = al.mp.NFWMCRLudlowSph(\n", + " centre=host_halo_centre,\n", + " mass_at_200=10**15.3,\n", + " redshift_object=redshift_lens,\n", + " redshift_source=max(source_redshifts),\n", + ")\n", + "\n", + "host_halo_galaxy = al.Galaxy(redshift=redshift_lens, dark=host_halo)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxies__\n", + "\n", + "The 2 background sources at *different* redshifts. Each carries a `SersicCore` light profile (used only\n", + "for visual confirmation of the lensed arcs \u2014 the cored profile changes gradually in the centre so explicit\n", + "source-plane over-sampling is unnecessary) and a `Point` model component whose multiple-image positions\n", + "we solve for and use as the modeling data.\n", + "\n", + "Each source's redshift is taken from ``source_redshifts``, so source 0 sits at ``z = 1.0`` and source 1\n", + "at ``z = 2.0``. The `Tracer` ray-traces multi-plane through both planes automatically." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxies = []\n", + "for i, (centre, src_z) in enumerate(zip(source_centres, source_redshifts)):\n", + " bulge = al.lp.SersicCore(\n", + " centre=centre,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0 + 30.0 * i),\n", + " intensity=2.0,\n", + " effective_radius=0.3,\n", + " sersic_index=1.0,\n", + " )\n", + " point = al.ps.Point(centre=centre)\n", + " source_galaxies.append(\n", + " al.Galaxy(redshift=src_z, bulge=bulge, **{f\"point_{i}\": point})\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Combine main lens galaxies, the scaling-tier members, the host halo galaxy, and the source galaxies into\n", + "a single tracer that produces the simulated image. With sources at distinct redshifts, the tracer\n", + "automatically handles multi-plane ray tracing. The scaling members share the lens redshift, so they\n", + "contribute to the single lens-plane deflection alongside the main galaxies and the halo." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=main_lens_galaxies\n", + " + scaling_galaxies\n", + " + [host_halo_galaxy]\n", + " + source_galaxies\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__JAX JIT \u2014 Point Solver__\n", + "\n", + "Solving the lens equation for the image-plane positions of a point source is iterative and numerically\n", + "expensive \u2014 at cluster scale it dominates the simulator's runtime. We use ``al.PointSolver(use_jax=True)``\n", + "and wrap the solve in ``@jax.jit`` for the speedup.\n", + "\n", + "The library handles pytree registration of ``Tracer`` + every reachable galaxy / profile class via\n", + "the one-time ``autolens.jax.register_tracer_classes(tracer)`` call below. Before\n", + "PR PyAutoLens#538 + PyAutoArray#335 (Phase 2 of ``z_features/jax_user_intro.md``), this section was\n", + "a ~60-line manual ceremony with an ``af.Model`` mirror, ``register_model``, and ``register_instance_pytree``;\n", + "the new API collapses it to a single import + single registration call + ``use_jax=True`` flag.\n", + "\n", + "``PointSolver(use_jax=True).solve`` defaults ``remove_infinities=False`` to honour the JAX static-shape\n", + "contract \u2014 the returned positions are padded with ``inf`` where no image was found, which is JIT-safe.\n", + "We strip them outside the jit below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autolens.jax import register_tracer_classes\n", + "\n", + "register_tracer_classes(tracer)\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.1),\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.1,\n", + " use_jax=True,\n", + ")\n", + "\n", + "\n", + "@jax.jit\n", + "def jitted_solve(tracer, source_plane_coordinate):\n", + " return solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_plane_coordinate\n", + " ).array\n", + "\n", + "\n", + "positions_list = []\n", + "for i, src_centre in enumerate(source_centres):\n", + " coord = jnp.asarray(src_centre)\n", + " raw = np.asarray(jitted_solve(tracer, coord))\n", + " finite = ~(np.isinf(raw).any(axis=1) | np.isnan(raw).any(axis=1))\n", + " positions_list.append(al.Grid2DIrregular(raw[finite]))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Datasets__\n", + "\n", + "One `PointDataset` per source. The `name` (e.g. `point_0`, `point_1`) pairs each dataset with the\n", + "matching `Point` component in the lens model during modeling.\n", + "\n", + "`redshift` is populated on each `PointDataset` so that the per-source redshifts round-trip through the\n", + "combined CSV below. This is the piece that makes the CSV self-describing for cluster modeling \u2014 position,\n", + "noise, and redshift live in a single spreadsheet.\n", + "\n", + "The position uncertainty is set to 0.005\" (5 mas), reflecting the centroid precision achievable by PSF\n", + "fitting on HST or adaptive-optics imaging \u2014 *not* the imaging pixel scale, which is the detector's\n", + "sampling rather than its centroiding precision. See `scripts/point_source/simulator.py` for a full\n", + "discussion of this value." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "position_noise = 0.005\n", + "\n", + "dataset_list = []\n", + "for i, positions in enumerate(positions_list):\n", + " dataset = al.PointDataset(\n", + " name=f\"point_{i}\",\n", + " positions=positions,\n", + " positions_noise_map=position_noise,\n", + " redshift=source_redshifts[i],\n", + " )\n", + " dataset_list.append(dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output one .json file per dataset (exact round-trip; this is the canonical modeling input)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for i, dataset in enumerate(dataset_list):\n", + " al.output_to_json(\n", + " obj=dataset,\n", + " file_path=dataset_path / f\"point_dataset_{i}.json\",\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Combined CSV__\n", + "\n", + "For cluster-scale workflows with tens or hundreds of sources, a single CSV with one row per observed\n", + "image \u2014 grouped by ``name`` \u2014 is far easier to edit in a spreadsheet than many per-source JSON files.\n", + "``al.output_to_csv`` writes every dataset into one file. The `redshift` column is emitted automatically\n", + "because each dataset has its redshift set above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_csv(\n", + " datasets=dataset_list,\n", + " file_path=dataset_path / \"point_datasets.csv\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Manual CSV Editing__\n", + "\n", + "The combined CSV is the preferred cluster input: it is human-readable, editable in Excel / LibreOffice /\n", + "any text editor, and round-trips cleanly back into `list_from_csv`. The expected format is one row per\n", + "observed multiple image with the following columns:\n", + "\n", + " - `name` \u2014 the source identifier (e.g. `point_0`). All rows sharing a `name` belong to the same source.\n", + " - `y`, `x` \u2014 the image-plane position of the multiple image, in arc-seconds.\n", + " - `positions_noise` \u2014 the positional uncertainty in arc-seconds. This is the PSF-fit centroid\n", + " uncertainty on each multiple image, *not* the imaging pixel scale. For HST or adaptive-optics\n", + " data on bright cluster member images, ~0.005\" (5 mas) is a defensible default; ground-based\n", + " seeing-limited data may warrant a few \u00d7 0.01\" depending on image SNR.\n", + " - `redshift` \u2014 the source redshift. Every row for a given `name` must share the *same* redshift\n", + " (validated on load; `list_from_csv` raises if a group's rows disagree). Leave the cell blank if the\n", + " redshift is unknown; blank is tolerated as long as *all* rows in a group are blank.\n", + "\n", + "Optional columns `flux`, `flux_noise`, `time_delay`, `time_delay_noise` are also round-tripped \u2014 populate\n", + "them when the observation provides those measurements, leave them blank otherwise.\n", + "\n", + "To build a cluster dataset by hand, simulate or manually collect one set of images per source, then edit\n", + "the CSV directly: add or remove rows, adjust positions or noises, and save. Reload the dataset in a\n", + "modeling script with ``al.list_from_csv(file_path=dataset_path / \"point_datasets.csv\")``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Scaling Galaxies CSV__\n", + "\n", + "The scaling-tier members are written to a separate CSV \u2014 ``scaling_galaxies.csv`` \u2014 with one row per\n", + "member carrying its centre and luminosity. ``al.galaxy_table_to_csv`` produces the canonical schema\n", + "(`y, x, luminosity, redshift?`) that the modeling script consumes via ``al.galaxy_table_from_csv``.\n", + "\n", + "Scaling up a real cluster to a larger member population is then a CSV-level edit: add a row per\n", + "additional member, fill in its centre and luminosity, save. The modeling script picks up the new rows\n", + "automatically and the scaling tier's free-parameter count stays at 1 (``b0_ref``, the reference\n", + "member's lens strength; the relation's exponents stay fixed at 0.5). Only luminosity *ratios* enter the\n", + "relation, so any consistent luminosity convention works \u2014 including converting from magnitudes.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.galaxy_table_to_csv(\n", + " centres=scaling_galaxies_centres,\n", + " luminosities=scaling_galaxies_luminosities,\n", + " file_path=dataset_path / \"scaling_galaxies.csv\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer JSON__\n", + "\n", + "Save the `Tracer` so the true light profiles, mass profiles and galaxies can be inspected after the fact.\n", + "This can be loaded via `tracer = al.from_json(file_path)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=dataset_path / \"tracer.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model CSVs__\n", + "\n", + "Write the truth model out as three family-level CSVs \u2014 ``mass.csv``, ``light.csv``, ``point.csv`` \u2014\n", + "keyed by galaxy name. The modeling and start_here scripts load these directly with\n", + "``al.galaxy_models_from_csv`` and compose them into ``af.Model[Galaxy]`` instances ready for non-linear\n", + "search. See ``scripts/cluster/csv_api.py`` for the full schema walkthrough.\n", + "\n", + "The scaling tier keeps its narrow 3-column ``scaling_galaxies.csv`` schema written above \u2014 naming each\n", + "scaling member and emitting an ``attr_name`` column would be more overhead than signal." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass_profiles = {\n", + " **{f\"lens_{i}\": {\"mass\": g.mass} for i, g in enumerate(main_lens_galaxies)},\n", + " \"host_halo\": {\"dark\": host_halo_galaxy.dark},\n", + "}\n", + "\n", + "light_profiles = {\n", + " **{f\"lens_{i}\": {\"bulge\": g.bulge} for i, g in enumerate(main_lens_galaxies)},\n", + " **{f\"source_{i}\": {\"bulge\": g.bulge} for i, g in enumerate(source_galaxies)},\n", + "}\n", + "\n", + "point_profiles = {\n", + " f\"source_{i}\": {f\"point_{i}\": getattr(g, f\"point_{i}\")}\n", + " for i, g in enumerate(source_galaxies)\n", + "}\n", + "\n", + "redshifts_by_galaxy = {\n", + " **{f\"lens_{i}\": redshift_lens for i in range(len(main_lens_galaxies))},\n", + " \"host_halo\": redshift_lens,\n", + " **{f\"source_{i}\": z for i, z in enumerate(source_redshifts)},\n", + "}\n", + "\n", + "al.galaxy_models_to_csv(\n", + " profiles_by_galaxy=mass_profiles,\n", + " file_path=dataset_path / \"mass.csv\",\n", + " family=\"mass\",\n", + " redshifts=redshifts_by_galaxy,\n", + ")\n", + "\n", + "al.galaxy_models_to_csv(\n", + " profiles_by_galaxy=light_profiles,\n", + " file_path=dataset_path / \"light.csv\",\n", + " family=\"light\",\n", + " redshifts=redshifts_by_galaxy,\n", + ")\n", + "\n", + "al.galaxy_models_to_csv(\n", + " profiles_by_galaxy=point_profiles,\n", + " file_path=dataset_path / \"point.csv\",\n", + " family=\"point\",\n", + " redshifts=redshifts_by_galaxy,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imaging__\n", + "\n", + "Strong lens clusters typically come with imaging data \u2014 used to *measure* the point positions and to\n", + "visually confirm the lens configuration. Although modeling here is point-source only, we output CCD\n", + "imaging so the dataset looks like a realistic cluster observation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=imaging_grid.pixel_scales\n", + ")\n", + "\n", + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")\n", + "\n", + "dataset = simulator.via_tracer_from(tracer=tracer, grid=imaging_grid)\n", + "\n", + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output .png plots of the simulated dataset, the tracer, and the per-source point datasets.\n", + "\n", + "These use the default galaxy-scale plotters and are known to be suboptimal for cluster-scale systems \u2014\n", + "arcs span a much larger field, per-source images benefit from distinct colours, and multi-source overlays\n", + "are useful. A follow-up prompt (`PyAutoPrompt/cluster/1_visualization.md`) addresses these\n", + "visualization requirements." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for i, pd in enumerate(dataset_list):\n", + " aplt.subplot_point_dataset(\n", + " dataset=pd, output_path=dataset_path, output_format=\"png\"\n", + " )\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=viz_grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=viz_grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/data_preparation/start_here.ipynb b/notebooks/group/data_preparation/start_here.ipynb index be4812d89..0d27afc8f 100644 --- a/notebooks/group/data_preparation/start_here.ipynb +++ b/notebooks/group/data_preparation/start_here.ipynb @@ -1,447 +1,484 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Group: Data Preparation\n", - "=======================\n", - "\n", - "When a group-scale strong lens CCD imaging dataset is analysed, it must conform to certain standards in order for the\n", - "analysis to be performed correctly. This tutorial describes these standards and links to more detailed scripts which\n", - "will help you prepare your dataset to adhere to them if it does not already.\n", - "\n", - "Group-scale lenses differ from galaxy-scale lenses in that there are multiple lens galaxies whose light and mass\n", - "must be modeled. This data preparation script is the group-scale equivalent of the imaging data preparation script,\n", - "with additional sections covering the specification of main lens galaxy centres, extra galaxy centres, and scaling\n", - "galaxy centres.\n", - "\n", - "__Contents__\n", - "\n", - "- **Pixel Scale:** The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units.\n", - "- **Image:** The image is the image of your group-scale strong lens, which comes from a telescope like the.\n", - "- **Noise Map:** The noise-map defines the uncertainty in every pixel of your strong lens image, where values are.\n", - "- **PSF:** The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope.\n", - "- **Data Processing Complete:** If your image, noise-map and PSF conform the standards above, you are ready to analyse your dataset!\n", - "\n", - "__Pixel Scale__\n", - "\n", - "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", - "your telescope. You should look up now if you are unsure of the value.\n", - "\n", - "The pixel scale of some common telescopes is as follows:\n", - "\n", - " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", - " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", - " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", - " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", - " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", - "\n", - "It is absolutely vital you use the correct pixel scale, so double check this value!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image__\n", - "\n", - "The image is the image of your group-scale strong lens, which comes from a telescope like the Hubble Space\n", - "telescope (HST).\n", - "\n", - "Lets inspect an image which conforms to **PyAutoLens** standards:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\") / \"group\" / \"simple\"\n", - "\n", - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This image conforms to **PyAutoLens** standards for the following reasons.\n", - "\n", - " - Units: The image flux is in units of electrons per second (as opposed to electrons, counts, ADU`s etc.).\n", - " Internal **PyAutoLens** functions for computing quantities like galaxy magnitudes assume the data and model\n", - " light profiles are in electrons per second.\n", - "\n", - " - Centering: For group-scale lenses, the primary lens galaxy should be roughly centred in the image. The other\n", - " galaxies in the group will be at various offsets from the centre. Default **PyAutoLens** parameter priors assume\n", - " the primary lens galaxy is at the centre of the image.\n", - "\n", - " - Stamp Size: For group-scale lenses, the postage stamp cut-out must be large enough to include all galaxies in\n", - " the group, not just the primary lens galaxy. This is typically larger than for a galaxy-scale lens. It is still\n", - " advised to cut out a postage stamp rather than using the entire image, as this reduces the amount of memory\n", - " **PyAutoLens** uses, speeds up the analysis and ensures visualization zooms around the group. However, conforming\n", - " to this standard is not necessary to ensure an accurate **PyAutoLens** analysis.\n", - "\n", - "If your image conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you noise-map and PSF conform to standards first!).\n", - "\n", - "**Links / Resources:**\n", - "\n", - " - `imaging/data_preparation/examples/data.ipynb`: tools to process the data to conform to these standards.\n", - "\n", - "__Noise Map__\n", - "\n", - "The noise-map defines the uncertainty in every pixel of your strong lens image, where values are defined as the\n", - "RMS standard deviation in every pixel (not the variances, HST WHT-map values, etc.).\n", - "\n", - "Lets inspect a noise-map which conforms to **PyAutoLens** standards:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_map = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"noise_map.fits\", pixel_scales=0.1\n", - ")\n", - "\n", - "aplt.plot_array(array=noise_map, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This noise-map conforms to **PyAutoLens** standards for the following reasons:\n", - "\n", - " - Units: Like its corresponding image, it is in units of electrons per second (as opposed to electrons, counts,\n", - " ADU`s etc.). Internal **PyAutoLens** functions for computing quantities like a galaxy magnitude assume the data and\n", - " model light profiles are in electrons per second.\n", - "\n", - " - Values: The noise-map values themselves are the RMS standard deviations of the noise in every pixel. When a model\n", - " is fitted to data in **PyAutoLens** and a likelihood is evaluated, this calculation assumes that this is the\n", - " corresponding definition of the noise-map. The noise map therefore should not be the variance of the noise, or\n", - " another definition of noise.\n", - "\n", - "If you are not certain what the definition of the noise-map you have available to you is, or do not know how to\n", - "compute a noise-map at all, you should refer to the instrument handbook of the telescope your data is from. It is\n", - "absolutely vital that the noise-map is correct, as it is the only way **PyAutoLens** can quantify the goodness-of-fit.\n", - "\n", - "A sanity check for a reliable noise map is that the signal-to-noise of the lens galaxy is somewhere between a value of\n", - "10 - 300 and source around 5 - 50, however this is not a definitive test.\n", - "\n", - "Given the image should be centred and cut-out around the primary lens galaxy, so should the noise-map.\n", - "\n", - "If your noise-map conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you image and PSF conform to standards first!).\n", - "\n", - "**Links / Resources:**\n", - "\n", - " - `imaging/data_preparation/examples/noise_map.ipynb`: tools to process the noise-map to conform to these standards.\n", - "\n", - "__PSF__\n", - "\n", - "The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope. It is used when fitting\n", - "a dataset to include these effects, such that does not bias the model.\n", - "\n", - "It should be estimated from a stack of stars in the image during data reduction or using a PSF simulator (e.g. TinyTim\n", - "for Hubble).\n", - "\n", - "Lets inspect a PSF which conforms to **PyAutoLens** standards:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_fits(\n", - " file_path=dataset_path / \"psf.fits\", hdu=0, pixel_scales=0.1\n", - ")\n", - "\n", - "aplt.plot_array(array=psf.kernel, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This psf conforms to **PyAutoLens** standards for the following reasons.\n", - "\n", - " - Size: The PSF has a shape 21 x 21 pixels, which is large enough to capture the PSF core and thus capture the\n", - " majority of the blurring effect, but not so large that the convolution slows down the analysis. Large\n", - " PSFs (e.g. 51 x 51) are supported, but will lead to much slower run times. The size of the PSF should be carefully\n", - " chosen to ensure it captures the majority of blurring due to the telescope optics, which for most instruments is\n", - " something around 11 x 11 to 21 x 21.\n", - "\n", - " - Oddness: The PSF has dimensions which are odd (an even PSF would for example have shape 20 x 20). The\n", - " convolution of an even PSF introduces a small shift in the model images and produces an offset in the inferred\n", - " lens model parameters. Inputting an even PSF will lead **PyAutoLens** to raise an error.\n", - "\n", - " - Normalization: The PSF has been normalized such that all values within the kernel sum to 1 (note how all values in\n", - " the example PSF are below zero with the majority below 0.01). This ensures that flux is conserved when convolution\n", - " is performed, ensuring that quantities like a galaxy's magnitude are computed accurately.\n", - "\n", - " - Centering: The PSF is at the centre of the array (as opposed to in a corner), ensuring that no shift is introduced\n", - " due to PSF blurring on the inferred model parameters.\n", - "\n", - "If your PSF conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you noise-map and image conform to standards first!).\n", - "\n", - "**Links / Resources:**\n", - "\n", - " - `imaging/data_preparation/examples/psf.ipynb`: tools to process the PSF to conform to these standards.\n", - "\n", - "__Data Processing Complete__\n", - "\n", - "If your image, noise-map and PSF conform the standards above, you are ready to analyse your dataset!\n", - "\n", - "Below, we provide an overview of additional data preparation steps specific to group-scale lenses, as well as optional\n", - "steps which prepare other aspects of the analysis.\n", - "\n", - "New users are recommended to read the \"Main Lens Galaxy Centres\" section carefully, as this is a required step for\n", - "group-scale lens modeling. The remaining optional steps can be skimmed and revisited later if needed.\n", - "\n", - "__Main Lens Galaxy Centres (Required)__\n", - "\n", - "For group-scale lenses, the centres of ALL main lens galaxies must be specified in a `main_lens_centres.json` file\n", - "in the dataset folder. This is a critical step that is required for the list-based model composition API used by\n", - "**PyAutoLens** to build group-scale lens models.\n", - "\n", - "Each main lens galaxy is modeled with full free light and mass models, meaning their light profiles, mass profiles\n", - "and centres are all free parameters in the model-fit (or initialized from these input centres).\n", - "\n", - "The centres can be determined using one of the following methods:\n", - "\n", - " - The GUI tool provided in `data_preparation/gui/main_lens_centres.py`, which allows you to click on the image\n", - " to mark the centres interactively.\n", - "\n", - " - Image processing software (e.g. ds9, IRAF) to identify the galaxy centres from the image.\n", - "\n", - " - By fitting each galaxy individually with a light profile model and using the inferred centres.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/gui/main_lens_centres.py`: use a Graphical User Interface (GUI) to mark the main lens galaxy centres.\n", - "\n", - "__Extra Galaxies Centres (Optional)__\n", - "\n", - "There may be galaxies nearby the lens and source galaxies, whose emission blends with that of the lens and source\n", - "and whose mass may contribute to the ray-tracing and lens model.\n", - "\n", - "We can include these galaxies in the lens model, either as light profiles, mass profiles, or both, using the\n", - "modeling API, where these nearby objects are denoted `extra_galaxies`.\n", - "\n", - "For group-scale lenses, the distinction between \"main\" and \"extra\" galaxies is important:\n", - "\n", - " - Main galaxies are the primary lens galaxies in the group. They are modeled with full free light and mass models,\n", - " meaning their profiles and parameters are individually well constrained by the data.\n", - "\n", - " - Extra galaxies are additional companions that are individually significant but are modeled with more restrictive\n", - " models (e.g. fixed centres, scaling relations). There may be more extra galaxies in a group-scale lens than in a\n", - " typical galaxy-scale lens.\n", - "\n", - "The script `extra_galaxies_centres.py` marks the (y,x) arcsecond locations of these extra galaxies, so that when they\n", - "are included in the lens model the centre of these extra galaxies light and / or mass profiles are fixed to these\n", - "values (or their priors are initialized surrounding these centres). The centres are stored in an\n", - "`extra_galaxies_centres.json` file.\n", - "\n", - "The example `mask_extra_galaxies.py` (see below) masks the regions of an image where extra galaxies are present.\n", - "This mask is used to remove their signal from the data and increase their noise to make them not impact the fit.\n", - "This means their luminous emission does not need to be included in the model, reducing the number of free parameters\n", - "and speeding up the analysis. It is still a choice whether their mass is included in the model.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/gui/extra_galaxies_centres.py`: use a Graphical User Interface (GUI) to mark the extra galaxy centres.\n", - "- `features/extra_galaxies.py` how to use extra galaxies in a model-fit, including loading the extra galaxy centres.\n", - "\n", - "__Scaling Galaxies Centres (Optional)__\n", - "\n", - "For group-scale lenses, there is often a larger ensemble of companion galaxies whose individual masses cannot be\n", - "well constrained by the data. These galaxies are modeled using a shared luminosity-to-mass scaling relation, where\n", - "the mass of each galaxy is determined by its luminosity and a common set of scaling parameters.\n", - "\n", - "The centres of these scaling galaxies are specified in a `scaling_galaxies_centres.json` file in the dataset folder.\n", - "\n", - "The key distinction between scaling galaxies and extra galaxies is:\n", - "\n", - " - Scaling galaxies: A larger ensemble of companions whose individual masses cannot be independently constrained.\n", - " Their masses are linked via a shared scaling relation (e.g. a luminosity-mass relation), reducing the number of\n", - " free parameters. Use scaling galaxies when there are many group members whose individual mass contributions are\n", - " small but collectively significant.\n", - "\n", - " - Extra galaxies: Fewer companions that are individually significant enough to warrant their own mass (and\n", - " possibly light) model parameters. Use extra galaxies when a companion is bright or massive enough that its\n", - " individual contribution to the lensing can be meaningfully constrained.\n", - "\n", - "In practice, a group-scale lens model may include both extra galaxies and scaling galaxies simultaneously.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/gui/scaling_galaxies_centres.py`: use a Graphical User Interface (GUI) to mark the scaling galaxy centres.\n", - "\n", - "__Mask (Optional)__\n", - "\n", - "The mask removes the regions of the image where the lens and source galaxy are not present, typically the edges of the\n", - "image.\n", - "\n", - "Example modeling scripts internally create a 3.0\" circular mask and therefore do not require that a mask has been\n", - "created externally via a data preparation script.\n", - "\n", - "For group-scale lenses, the mask should typically be larger than for galaxy-scale lenses, to ensure that all galaxies\n", - "in the group and their associated lensed images are included within the masked region.\n", - "\n", - "This script shows how to create customize masked (e.g. annular, ellipses) which are tailored to match the lens or\n", - "lensed source emission.\n", - "\n", - "If you have not analysed your dataset yet and do not know of a specific reason why you need the bespoke masks\n", - "created by this script, it is recommended that you simply use the default ~3.0\" circular mask internally made in each\n", - "script and omit this data preparation tutorial.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/mask.ipynb`: tools to create a bespoke mask for your dataset.\n", - "- `data_preparation/examples/gui/mask.ipynb`: use a Graphical User Interface (GUI) to create a bespoke mask.\n", - "\n", - "__Positions (Optional)__\n", - "\n", - "The script allows you to mark the (y,x) arc-second positions of the multiply imaged lensed source galaxy in\n", - "the image-plane, under the assumption that they originate from the same location in the source-plane.\n", - "\n", - "A non-linear search (e.g. Nautilus) can then use these positions to preferentially choose mass models where these\n", - "positions trace close to one another in the source-plane. This speeding up the initial fitting of lens models and\n", - "removes unwanted solutions from parameter space which have too much or too little mass in the lens galaxy.\n", - "\n", - "If you create positions for your dataset, you must also update your modeling script to use them by loading them\n", - "and passing them to the `Analysis` object via a `PositionsLH` object.\n", - "\n", - "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using positions\n", - "to help the non-linear search find a good lens model.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "Position-based lens model resampling is particularly important for fitting pixelized source models, for the\n", - "reasons disucssed in the following readthedocs\n", - "webapge https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", - "\n", - "- `data_preparation/examples/optional/positions.ipynb`: input the positions manually into a Python script.\n", - "- `data_preparation/gui/positions.ipynb` use a Graphical User Interface (GUI) to mark the positions.\n", - "- `guides/modeling/customize` for an example.of how to use positions in a `modeling` script.\n", - "\n", - "__Mask Extra Galaxies (Optional)__\n", - "\n", - "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated\n", - "with the strong lens we are studying. The emission from these images will impact our model fitting and needs to be\n", - "removed from the analysis.\n", - "\n", - "For group-scale lenses, there are typically more galaxies in the field that may need to be masked compared to\n", - "galaxy-scale lenses.\n", - "\n", - "This script creates a mask of these regions of the image, called the `mask_extra_galaxies`, which can be used to\n", - "prevent them from impacting a fit. This mask may also include emission from objects which are not technically galaxies,\n", - "but blend with the galaxy we are studying in a similar way. Common examples of such objects are foreground stars\n", - "or emission due to the data reduction process.\n", - "\n", - "The mask can be applied in different ways. For example, it could be applied such that the image pixels are discarded\n", - "from the fit entirely, Alternatively the mask could be used to set the image values to (near) zero and increase their\n", - "corresponding noise-map to large values.\n", - "\n", - "The exact method used depends on the nature of the model being fitted. For simple fits like a light profile a mask\n", - "is appropriate, as removing image pixels does not change how the model is fitted. However, for more complex models\n", - "fits, like those using a pixelization, masking regions of the image in a way that removes their image pixels entirely\n", - "from the fit can produce discontinuities in the pixelixation. In this case, scaling the data and noise-map values\n", - "may be a better approach.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/mask_extra_galaxies.py`: create the extra galaxies mask manually via a Python script.\n", - "- `data_preparation/gui/extra_galaxies_mask.ipynb` use a Graphical User Interface (GUI) to create the extra galaxies mask.\n", - "- `features/extra_galaxies.py` how to use the extra galaxies mask in a model-fit.\n", - "\n", - "__Info (Optional)__\n", - "\n", - "Auxiliary information about a strong lens dataset may used during an analysis or afterwards when interpreting the\n", - "modeling results. For example, the redshifts of the source and lens galaxy.\n", - "\n", - "For group-scale lenses, additional information may be particularly useful, such as:\n", - "\n", - " - The number of galaxies in the group.\n", - " - The individual redshifts of each group member galaxy.\n", - " - Spectroscopic or photometric redshift estimates for each galaxy.\n", - " - The velocity dispersion of the group.\n", - " - Previous modeling results from the literature.\n", - "\n", - "By storing these as an `info.json` file in the lens's dataset folder, it is straight forward to load the redshifts\n", - "in a modeling script and pass them to a fit, such that **PyAutoLens** can then output results in physical\n", - "units (e.g. kpc instead of arc-seconds).\n", - "\n", - "For analysing large quantities of modeling results, **PyAutoLens** has an sqlite database feature. The info file\n", - "may can also be loaded by the database after a model-fit has completed, such that when one is interpreting\n", - "the results of a model fit additional data on a lens can be used to.\n", - "\n", - "For example, to plot the model-results against other measurements of a lens not made by PyAutoLens. Examples of such\n", - "data might be:\n", - "\n", - "- The velocity dispersion of the lens galaxy.\n", - "- The stellar mass of the lens galaxy.\n", - "- The results of previous strong lens models to the lens performed in previous papers.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/info.py`: create the info file manually via a Python script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Group: Data Preparation\n", + "=======================\n", + "\n", + "When a group-scale strong lens CCD imaging dataset is analysed, it must conform to certain standards in order for the\n", + "analysis to be performed correctly. This tutorial describes these standards and links to more detailed scripts which\n", + "will help you prepare your dataset to adhere to them if it does not already.\n", + "\n", + "Group-scale lenses differ from galaxy-scale lenses in that there are multiple lens galaxies whose light and mass\n", + "must be modeled. This data preparation script is the group-scale equivalent of the imaging data preparation script,\n", + "with additional sections covering the specification of main lens galaxy centres, extra galaxy centres, and scaling\n", + "galaxy centres.\n", + "\n", + "__Contents__\n", + "\n", + "- **Pixel Scale:** The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units.\n", + "- **Image:** The image is the image of your group-scale strong lens, which comes from a telescope like the.\n", + "- **Noise Map:** The noise-map defines the uncertainty in every pixel of your strong lens image, where values are.\n", + "- **PSF:** The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope.\n", + "- **Data Processing Complete:** If your image, noise-map and PSF conform the standards above, you are ready to analyse your dataset!\n", + "\n", + "__Pixel Scale__\n", + "\n", + "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", + "your telescope. You should look up now if you are unsure of the value.\n", + "\n", + "The pixel scale of some common telescopes is as follows:\n", + "\n", + " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", + " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", + " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", + " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", + " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", + "\n", + "It is absolutely vital you use the correct pixel scale, so double check this value!" + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image__\n", + "\n", + "The image is the image of your group-scale strong lens, which comes from a telescope like the Hubble Space\n", + "telescope (HST).\n", + "\n", + "Lets inspect an image which conforms to **PyAutoLens** standards:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / \"group\" / \"simple\"\n", + "\n", + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This image conforms to **PyAutoLens** standards for the following reasons.\n", + "\n", + " - Units: The image flux is in units of electrons per second (as opposed to electrons, counts, ADU`s etc.).\n", + " Internal **PyAutoLens** functions for computing quantities like galaxy magnitudes assume the data and model\n", + " light profiles are in electrons per second.\n", + "\n", + " - Centering: For group-scale lenses, the primary lens galaxy should be roughly centred in the image. The other\n", + " galaxies in the group will be at various offsets from the centre. Default **PyAutoLens** parameter priors assume\n", + " the primary lens galaxy is at the centre of the image.\n", + "\n", + " - Stamp Size: For group-scale lenses, the postage stamp cut-out must be large enough to include all galaxies in\n", + " the group, not just the primary lens galaxy. This is typically larger than for a galaxy-scale lens. It is still\n", + " advised to cut out a postage stamp rather than using the entire image, as this reduces the amount of memory\n", + " **PyAutoLens** uses, speeds up the analysis and ensures visualization zooms around the group. However, conforming\n", + " to this standard is not necessary to ensure an accurate **PyAutoLens** analysis.\n", + "\n", + "If your image conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you noise-map and PSF conform to standards first!).\n", + "\n", + "**Links / Resources:**\n", + "\n", + " - `imaging/data_preparation/examples/data.ipynb`: tools to process the data to conform to these standards.\n", + "\n", + "__Noise Map__\n", + "\n", + "The noise-map defines the uncertainty in every pixel of your strong lens image, where values are defined as the\n", + "RMS standard deviation in every pixel (not the variances, HST WHT-map values, etc.).\n", + "\n", + "Lets inspect a noise-map which conforms to **PyAutoLens** standards:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_map = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"noise_map.fits\", pixel_scales=0.1\n", + ")\n", + "\n", + "aplt.plot_array(array=noise_map, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This noise-map conforms to **PyAutoLens** standards for the following reasons:\n", + "\n", + " - Units: Like its corresponding image, it is in units of electrons per second (as opposed to electrons, counts,\n", + " ADU`s etc.). Internal **PyAutoLens** functions for computing quantities like a galaxy magnitude assume the data and\n", + " model light profiles are in electrons per second.\n", + "\n", + " - Values: The noise-map values themselves are the RMS standard deviations of the noise in every pixel. When a model\n", + " is fitted to data in **PyAutoLens** and a likelihood is evaluated, this calculation assumes that this is the\n", + " corresponding definition of the noise-map. The noise map therefore should not be the variance of the noise, or\n", + " another definition of noise.\n", + "\n", + "If you are not certain what the definition of the noise-map you have available to you is, or do not know how to\n", + "compute a noise-map at all, you should refer to the instrument handbook of the telescope your data is from. It is\n", + "absolutely vital that the noise-map is correct, as it is the only way **PyAutoLens** can quantify the goodness-of-fit.\n", + "\n", + "A sanity check for a reliable noise map is that the signal-to-noise of the lens galaxy is somewhere between a value of\n", + "10 - 300 and source around 5 - 50, however this is not a definitive test.\n", + "\n", + "Given the image should be centred and cut-out around the primary lens galaxy, so should the noise-map.\n", + "\n", + "If your noise-map conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you image and PSF conform to standards first!).\n", + "\n", + "**Links / Resources:**\n", + "\n", + " - `imaging/data_preparation/examples/noise_map.ipynb`: tools to process the noise-map to conform to these standards.\n", + "\n", + "__PSF__\n", + "\n", + "The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope. It is used when fitting\n", + "a dataset to include these effects, such that does not bias the model.\n", + "\n", + "It should be estimated from a stack of stars in the image during data reduction or using a PSF simulator (e.g. TinyTim\n", + "for Hubble).\n", + "\n", + "Lets inspect a PSF which conforms to **PyAutoLens** standards:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_fits(\n", + " file_path=dataset_path / \"psf.fits\", hdu=0, pixel_scales=0.1\n", + ")\n", + "\n", + "aplt.plot_array(array=psf.kernel, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This psf conforms to **PyAutoLens** standards for the following reasons.\n", + "\n", + " - Size: The PSF has a shape 21 x 21 pixels, which is large enough to capture the PSF core and thus capture the\n", + " majority of the blurring effect, but not so large that the convolution slows down the analysis. Large\n", + " PSFs (e.g. 51 x 51) are supported, but will lead to much slower run times. The size of the PSF should be carefully\n", + " chosen to ensure it captures the majority of blurring due to the telescope optics, which for most instruments is\n", + " something around 11 x 11 to 21 x 21.\n", + "\n", + " - Oddness: The PSF has dimensions which are odd (an even PSF would for example have shape 20 x 20). The\n", + " convolution of an even PSF introduces a small shift in the model images and produces an offset in the inferred\n", + " lens model parameters. Inputting an even PSF will lead **PyAutoLens** to raise an error.\n", + "\n", + " - Normalization: The PSF has been normalized such that all values within the kernel sum to 1 (note how all values in\n", + " the example PSF are below zero with the majority below 0.01). This ensures that flux is conserved when convolution\n", + " is performed, ensuring that quantities like a galaxy's magnitude are computed accurately.\n", + "\n", + " - Centering: The PSF is at the centre of the array (as opposed to in a corner), ensuring that no shift is introduced\n", + " due to PSF blurring on the inferred model parameters.\n", + "\n", + "If your PSF conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you noise-map and image conform to standards first!).\n", + "\n", + "**Links / Resources:**\n", + "\n", + " - `imaging/data_preparation/examples/psf.ipynb`: tools to process the PSF to conform to these standards.\n", + "\n", + "__Data Processing Complete__\n", + "\n", + "If your image, noise-map and PSF conform the standards above, you are ready to analyse your dataset!\n", + "\n", + "Below, we provide an overview of additional data preparation steps specific to group-scale lenses, as well as optional\n", + "steps which prepare other aspects of the analysis.\n", + "\n", + "New users are recommended to read the \"Main Lens Galaxy Centres\" section carefully, as this is a required step for\n", + "group-scale lens modeling. The remaining optional steps can be skimmed and revisited later if needed.\n", + "\n", + "__Main Lens Galaxy Centres (Required)__\n", + "\n", + "For group-scale lenses, the centres of ALL main lens galaxies must be specified in a `main_lens_centres.json` file\n", + "in the dataset folder. This is a critical step that is required for the list-based model composition API used by\n", + "**PyAutoLens** to build group-scale lens models.\n", + "\n", + "Each main lens galaxy is modeled with full free light and mass models, meaning their light profiles, mass profiles\n", + "and centres are all free parameters in the model-fit (or initialized from these input centres).\n", + "\n", + "The centres can be determined using one of the following methods:\n", + "\n", + " - The GUI tool provided in `data_preparation/gui/main_lens_centres.py`, which allows you to click on the image\n", + " to mark the centres interactively.\n", + "\n", + " - Image processing software (e.g. ds9, IRAF) to identify the galaxy centres from the image.\n", + "\n", + " - By fitting each galaxy individually with a light profile model and using the inferred centres.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/gui/main_lens_centres.py`: use a Graphical User Interface (GUI) to mark the main lens galaxy centres.\n", + "\n", + "__Extra Galaxies Centres (Optional)__\n", + "\n", + "There may be galaxies nearby the lens and source galaxies, whose emission blends with that of the lens and source\n", + "and whose mass may contribute to the ray-tracing and lens model.\n", + "\n", + "We can include these galaxies in the lens model, either as light profiles, mass profiles, or both, using the\n", + "modeling API, where these nearby objects are denoted `extra_galaxies`.\n", + "\n", + "For group-scale lenses, the distinction between \"main\" and \"extra\" galaxies is important:\n", + "\n", + " - Main galaxies are the primary lens galaxies in the group. They are modeled with full free light and mass models,\n", + " meaning their profiles and parameters are individually well constrained by the data.\n", + "\n", + " - Extra galaxies are additional companions that are individually significant but are modeled with more restrictive\n", + " models (e.g. fixed centres, scaling relations). There may be more extra galaxies in a group-scale lens than in a\n", + " typical galaxy-scale lens.\n", + "\n", + "The script `extra_galaxies_centres.py` marks the (y,x) arcsecond locations of these extra galaxies, so that when they\n", + "are included in the lens model the centre of these extra galaxies light and / or mass profiles are fixed to these\n", + "values (or their priors are initialized surrounding these centres). The centres are stored in an\n", + "`extra_galaxies_centres.json` file.\n", + "\n", + "The example `mask_extra_galaxies.py` (see below) masks the regions of an image where extra galaxies are present.\n", + "This mask is used to remove their signal from the data and increase their noise to make them not impact the fit.\n", + "This means their luminous emission does not need to be included in the model, reducing the number of free parameters\n", + "and speeding up the analysis. It is still a choice whether their mass is included in the model.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/gui/extra_galaxies_centres.py`: use a Graphical User Interface (GUI) to mark the extra galaxy centres.\n", + "- `features/extra_galaxies.py` how to use extra galaxies in a model-fit, including loading the extra galaxy centres.\n", + "\n", + "__Scaling Galaxies Centres (Optional)__\n", + "\n", + "For group-scale lenses, there is often a larger ensemble of companion galaxies whose individual masses cannot be\n", + "well constrained by the data. These galaxies are modeled using a shared luminosity-to-mass scaling relation, where\n", + "the mass of each galaxy is determined by its luminosity and a common set of scaling parameters.\n", + "\n", + "The centres of these scaling galaxies are specified in a `scaling_galaxies_centres.json` file in the dataset folder.\n", + "\n", + "The key distinction between scaling galaxies and extra galaxies is:\n", + "\n", + " - Scaling galaxies: A larger ensemble of companions whose individual masses cannot be independently constrained.\n", + " Their masses are linked via a shared scaling relation (e.g. a luminosity-mass relation), reducing the number of\n", + " free parameters. Use scaling galaxies when there are many group members whose individual mass contributions are\n", + " small but collectively significant.\n", + "\n", + " - Extra galaxies: Fewer companions that are individually significant enough to warrant their own mass (and\n", + " possibly light) model parameters. Use extra galaxies when a companion is bright or massive enough that its\n", + " individual contribution to the lensing can be meaningfully constrained.\n", + "\n", + "In practice, a group-scale lens model may include both extra galaxies and scaling galaxies simultaneously.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/gui/scaling_galaxies_centres.py`: use a Graphical User Interface (GUI) to mark the scaling galaxy centres.\n", + "\n", + "__Mask (Optional)__\n", + "\n", + "The mask removes the regions of the image where the lens and source galaxy are not present, typically the edges of the\n", + "image.\n", + "\n", + "Example modeling scripts internally create a 3.0\" circular mask and therefore do not require that a mask has been\n", + "created externally via a data preparation script.\n", + "\n", + "For group-scale lenses, the mask should typically be larger than for galaxy-scale lenses, to ensure that all galaxies\n", + "in the group and their associated lensed images are included within the masked region.\n", + "\n", + "This script shows how to create customize masked (e.g. annular, ellipses) which are tailored to match the lens or\n", + "lensed source emission.\n", + "\n", + "If you have not analysed your dataset yet and do not know of a specific reason why you need the bespoke masks\n", + "created by this script, it is recommended that you simply use the default ~3.0\" circular mask internally made in each\n", + "script and omit this data preparation tutorial.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/mask.ipynb`: tools to create a bespoke mask for your dataset.\n", + "- `data_preparation/examples/gui/mask.ipynb`: use a Graphical User Interface (GUI) to create a bespoke mask.\n", + "\n", + "__Positions (Optional)__\n", + "\n", + "The script allows you to mark the (y,x) arc-second positions of the multiply imaged lensed source galaxy in\n", + "the image-plane, under the assumption that they originate from the same location in the source-plane.\n", + "\n", + "A non-linear search (e.g. Nautilus) can then use these positions to preferentially choose mass models where these\n", + "positions trace close to one another in the source-plane. This speeding up the initial fitting of lens models and\n", + "removes unwanted solutions from parameter space which have too much or too little mass in the lens galaxy.\n", + "\n", + "If you create positions for your dataset, you must also update your modeling script to use them by loading them\n", + "and passing them to the `Analysis` object via a `PositionsLH` object.\n", + "\n", + "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using positions\n", + "to help the non-linear search find a good lens model.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "Position-based lens model resampling is particularly important for fitting pixelized source models, for the\n", + "reasons disucssed in the following readthedocs\n", + "webapge https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", + "\n", + "- `data_preparation/examples/optional/positions.ipynb`: input the positions manually into a Python script.\n", + "- `data_preparation/gui/positions.ipynb` use a Graphical User Interface (GUI) to mark the positions.\n", + "- `guides/modeling/customize` for an example.of how to use positions in a `modeling` script.\n", + "\n", + "__Mask Extra Galaxies (Optional)__\n", + "\n", + "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated\n", + "with the strong lens we are studying. The emission from these images will impact our model fitting and needs to be\n", + "removed from the analysis.\n", + "\n", + "For group-scale lenses, there are typically more galaxies in the field that may need to be masked compared to\n", + "galaxy-scale lenses.\n", + "\n", + "This script creates a mask of these regions of the image, called the `mask_extra_galaxies`, which can be used to\n", + "prevent them from impacting a fit. This mask may also include emission from objects which are not technically galaxies,\n", + "but blend with the galaxy we are studying in a similar way. Common examples of such objects are foreground stars\n", + "or emission due to the data reduction process.\n", + "\n", + "The mask can be applied in different ways. For example, it could be applied such that the image pixels are discarded\n", + "from the fit entirely, Alternatively the mask could be used to set the image values to (near) zero and increase their\n", + "corresponding noise-map to large values.\n", + "\n", + "The exact method used depends on the nature of the model being fitted. For simple fits like a light profile a mask\n", + "is appropriate, as removing image pixels does not change how the model is fitted. However, for more complex models\n", + "fits, like those using a pixelization, masking regions of the image in a way that removes their image pixels entirely\n", + "from the fit can produce discontinuities in the pixelixation. In this case, scaling the data and noise-map values\n", + "may be a better approach.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/mask_extra_galaxies.py`: create the extra galaxies mask manually via a Python script.\n", + "- `data_preparation/gui/extra_galaxies_mask.ipynb` use a Graphical User Interface (GUI) to create the extra galaxies mask.\n", + "- `features/extra_galaxies.py` how to use the extra galaxies mask in a model-fit.\n", + "\n", + "__Info (Optional)__\n", + "\n", + "Auxiliary information about a strong lens dataset may used during an analysis or afterwards when interpreting the\n", + "modeling results. For example, the redshifts of the source and lens galaxy.\n", + "\n", + "For group-scale lenses, additional information may be particularly useful, such as:\n", + "\n", + " - The number of galaxies in the group.\n", + " - The individual redshifts of each group member galaxy.\n", + " - Spectroscopic or photometric redshift estimates for each galaxy.\n", + " - The velocity dispersion of the group.\n", + " - Previous modeling results from the literature.\n", + "\n", + "By storing these as an `info.json` file in the lens's dataset folder, it is straight forward to load the redshifts\n", + "in a modeling script and pass them to a fit, such that **PyAutoLens** can then output results in physical\n", + "units (e.g. kpc instead of arc-seconds).\n", + "\n", + "For analysing large quantities of modeling results, **PyAutoLens** has an sqlite database feature. The info file\n", + "may can also be loaded by the database after a model-fit has completed, such that when one is interpreting\n", + "the results of a model fit additional data on a lens can be used to.\n", + "\n", + "For example, to plot the model-results against other measurements of a lens not made by PyAutoLens. Examples of such\n", + "data might be:\n", + "\n", + "- The velocity dispersion of the lens galaxy.\n", + "- The stellar mass of the lens galaxy.\n", + "- The results of previous strong lens models to the lens performed in previous papers.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/info.py`: create the info file manually via a Python script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/double_einstein_ring/chaining.ipynb b/notebooks/group/features/advanced/double_einstein_ring/chaining.ipynb index 2a340a9ae..641ad793f 100644 --- a/notebooks/group/features/advanced/double_einstein_ring/chaining.ipynb +++ b/notebooks/group/features/advanced/double_einstein_ring/chaining.ipynb @@ -1,445 +1,482 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Chaining: Group Double Einstein Ring\n", - "====================================\n", - "\n", - "This script chains two non-linear searches to fit `Imaging` data of a group-scale double Einstein ring lens \u2014\n", - "two source galaxies at different redshifts lensed by multiple main lens galaxies at the lens-plane redshift.\n", - "\n", - "In the final model:\n", - "\n", - " - Each main lens galaxy at z=0.5 has an MGE bulge and an `Isothermal` mass profile, composed via the group\n", - " `lens_dict` convention (one entry per main lens centre loaded from JSON).\n", - " - `source_0` at z=1.0 has an MGE bulge and an `IsothermalSph` mass \u2014 it deflects light from `source_1`.\n", - " - `source_1` at z=2.0 has an MGE bulge only.\n", - "\n", - "The two searches break down as follows:\n", - "\n", - " 1) Fit each main lens galaxy's mass and bulge, and `source_0`'s MGE bulge. `source_0`'s mass and `source_1`\n", - " are omitted entirely. A smaller mask removes the second source's emission from the fit.\n", - " 2) Pass the search 1 results forward as instances (fixed values), then introduce `source_0`'s mass and\n", - " `source_1`'s MGE bulge as new free parameters. A larger mask includes both source galaxies' emission.\n", - "\n", - "__Why Chain?__\n", - "\n", - "A group-scale double Einstein ring has many more free parameters than either a single-plane group lens OR a\n", - "single-lens-galaxy double Einstein ring. A single Nautilus search on the combined model is impractical: the\n", - "parameter space is too high-dimensional and local maxima are abundant.\n", - "\n", - "Chaining exploits a key physical observation: ray-tracing of `source_0` is fully independent of `source_1`'s\n", - "properties, so we can initialise the lens model + `source_0` first, then introduce `source_1` as a (more\n", - "tractable) extension.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Paths:** Load data; choose the chained-search output path.\n", - "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", - "- **Masking (Search 1):** Smaller mask that removes `source_1`.\n", - "- **Model (Search 1):** `lens_dict` (MGE bulge + `Isothermal` mass each) plus `source_0` MGE bulge.\n", - "- **Search + Result (Search 1):** Run search 1.\n", - "- **Masking (Search 2):** Larger mask that includes `source_1`.\n", - "- **Model (Search 2):** Search 1 results passed as instances, plus new free `source_0` mass and `source_1` MGE.\n", - "- **Search + Result (Search 2):** Run search 2.\n", - "- **Wrap Up.**\n", - "\n", - "__Prerequisites__\n", - "\n", - "For background on the canonical group lens chaining workflow and the single-lens double Einstein ring chaining\n", - "workflow, see:\n", - "\n", - " - `autolens_workspace/scripts/group/start_here.py` \u2014 group `lens_dict` model composition.\n", - " - `autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/chaining.py` \u2014 single-lens\n", - " double Einstein ring chained search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the group double Einstein ring `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Centres__\n", - "\n", - "Load the two main lens galaxy centres saved by the simulator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__\n", - "\n", - "All chained-search results are written under this path prefix." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"group\") / \"chaining\" / \"double_einstein_ring\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masking (Search 1)__\n", - "\n", - "A smaller mask that removes the light of `source_1`. The radius is chosen to encompass both main lens galaxies\n", - "and the primary Einstein ring around `source_0`, but to exclude the more distant `source_1` arcs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 2.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Each main lens galaxy gets an MGE bulge (20 Gaussians) and an `Isothermal` mass profile, composed via the\n", - "group `lens_dict` convention. `source_0` gets an MGE bulge with 20 Gaussians (no mass yet). `source_1` is\n", - "omitted entirely." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict_1 = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " centre_prior_is_uniform=True,\n", - " centre=(centre[0], centre[1]),\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = (centre[0], centre[1])\n", - "\n", - " lens_dict_1[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " )\n", - "\n", - "source_0_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source_0 = af.Model(al.Galaxy, redshift=1.0, bulge=source_0_bulge)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(**lens_dict_1, source_0=source_0))\n", - "\n", - "print(model_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__lens_dict_source_0_parametric\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)\n", - "\n", - "print(result_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masking (Search 2)__\n", - "\n", - "Reload the dataset with a larger mask that includes `source_1`'s arcs around the lens system." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 4.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "Each main lens galaxy's mass and bulge are passed forward as `instance` (i.e. fixed at the search 1 results,\n", - "not refit). `source_0`'s MGE bulge is also fixed. We then add:\n", - "\n", - " - `source_0`'s `IsothermalSph` mass (3 free parameters).\n", - " - `source_1`'s MGE bulge centred near (0.0, 0.0), with a narrow Gaussian prior (6 free parameters)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict_2 = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " lens_dict_2[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=getattr(result_1.instance.galaxies, f\"lens_{i}\").bulge,\n", - " mass=getattr(result_1.instance.galaxies, f\"lens_{i}\").mass,\n", - " )\n", - "\n", - "source_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " bulge=result_1.instance.galaxies.source_0.bulge,\n", - " mass=al.mp.IsothermalSph,\n", - ")\n", - "\n", - "source_1_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=source_1_bulge)\n", - "\n", - "model_2 = af.Collection(\n", - " galaxies=af.Collection(**lens_dict_2, source_0=source_0, source_1=source_1),\n", - ")\n", - "\n", - "print(model_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 2)__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__source_1_parametric\",\n", - " unique_tag=dataset_name,\n", - " n_live=200,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)\n", - "\n", - "print(result_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Two chained searches initialised a model for a group-scale double Einstein ring system. The key idea is to\n", - "exploit the independence of `source_0`'s ray-tracing from `source_1`'s properties \u2014 search 1 nails down the\n", - "lens-plane and `source_0` light using a mask that excludes `source_1`, and search 2 fixes that result and adds\n", - "the second source plus `source_0`'s mass.\n", - "\n", - "For a fully production-quality fit on real data (including pixelized source reconstructions), see `slam.py` in\n", - "the same directory. The single-search `modeling.py` example is for tutorial purposes only \u2014 it \"cheats\" by\n", - "initialising priors at the simulator's true values, which is impossible on real data.\n", - "\n", - "__Advanced Chaining__\n", - "\n", - "A more elaborate pipeline (5+ searches, eventually swapping the parametric source models for pixelized\n", - "reconstructions) is shown for the single-lens double Einstein ring in\n", - "`autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/chaining.py`. The same logic\n", - "transfers to the group-scale case, but no canonical template is currently provided for group DSPL \u2014 we are\n", - "still working out the most effective way to model these systems." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Chaining: Group Double Einstein Ring\n", + "====================================\n", + "\n", + "This script chains two non-linear searches to fit `Imaging` data of a group-scale double Einstein ring lens \u2014\n", + "two source galaxies at different redshifts lensed by multiple main lens galaxies at the lens-plane redshift.\n", + "\n", + "In the final model:\n", + "\n", + " - Each main lens galaxy at z=0.5 has an MGE bulge and an `Isothermal` mass profile, composed via the group\n", + " `lens_dict` convention (one entry per main lens centre loaded from JSON).\n", + " - `source_0` at z=1.0 has an MGE bulge and an `IsothermalSph` mass \u2014 it deflects light from `source_1`.\n", + " - `source_1` at z=2.0 has an MGE bulge only.\n", + "\n", + "The two searches break down as follows:\n", + "\n", + " 1) Fit each main lens galaxy's mass and bulge, and `source_0`'s MGE bulge. `source_0`'s mass and `source_1`\n", + " are omitted entirely. A smaller mask removes the second source's emission from the fit.\n", + " 2) Pass the search 1 results forward as instances (fixed values), then introduce `source_0`'s mass and\n", + " `source_1`'s MGE bulge as new free parameters. A larger mask includes both source galaxies' emission.\n", + "\n", + "__Why Chain?__\n", + "\n", + "A group-scale double Einstein ring has many more free parameters than either a single-plane group lens OR a\n", + "single-lens-galaxy double Einstein ring. A single Nautilus search on the combined model is impractical: the\n", + "parameter space is too high-dimensional and local maxima are abundant.\n", + "\n", + "Chaining exploits a key physical observation: ray-tracing of `source_0` is fully independent of `source_1`'s\n", + "properties, so we can initialise the lens model + `source_0` first, then introduce `source_1` as a (more\n", + "tractable) extension.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Paths:** Load data; choose the chained-search output path.\n", + "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", + "- **Masking (Search 1):** Smaller mask that removes `source_1`.\n", + "- **Model (Search 1):** `lens_dict` (MGE bulge + `Isothermal` mass each) plus `source_0` MGE bulge.\n", + "- **Search + Result (Search 1):** Run search 1.\n", + "- **Masking (Search 2):** Larger mask that includes `source_1`.\n", + "- **Model (Search 2):** Search 1 results passed as instances, plus new free `source_0` mass and `source_1` MGE.\n", + "- **Search + Result (Search 2):** Run search 2.\n", + "- **Wrap Up.**\n", + "\n", + "__Prerequisites__\n", + "\n", + "For background on the canonical group lens chaining workflow and the single-lens double Einstein ring chaining\n", + "workflow, see:\n", + "\n", + " - `autolens_workspace/scripts/group/start_here.py` \u2014 group `lens_dict` model composition.\n", + " - `autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/chaining.py` \u2014 single-lens\n", + " double Einstein ring chained search." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the group double Einstein ring `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Centres__\n", + "\n", + "Load the two main lens galaxy centres saved by the simulator." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__\n", + "\n", + "All chained-search results are written under this path prefix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"group\") / \"chaining\" / \"double_einstein_ring\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masking (Search 1)__\n", + "\n", + "A smaller mask that removes the light of `source_1`. The radius is chosen to encompass both main lens galaxies\n", + "and the primary Einstein ring around `source_0`, but to exclude the more distant `source_1` arcs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 2.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Each main lens galaxy gets an MGE bulge (20 Gaussians) and an `Isothermal` mass profile, composed via the\n", + "group `lens_dict` convention. `source_0` gets an MGE bulge with 20 Gaussians (no mass yet). `source_1` is\n", + "omitted entirely." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict_1 = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " centre_prior_is_uniform=True,\n", + " centre=(centre[0], centre[1]),\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = (centre[0], centre[1])\n", + "\n", + " lens_dict_1[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " )\n", + "\n", + "source_0_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source_0 = af.Model(al.Galaxy, redshift=1.0, bulge=source_0_bulge)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(**lens_dict_1, source_0=source_0))\n", + "\n", + "print(model_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__lens_dict_source_0_parametric\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)\n", + "\n", + "print(result_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masking (Search 2)__\n", + "\n", + "Reload the dataset with a larger mask that includes `source_1`'s arcs around the lens system." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 4.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "Each main lens galaxy's mass and bulge are passed forward as `instance` (i.e. fixed at the search 1 results,\n", + "not refit). `source_0`'s MGE bulge is also fixed. We then add:\n", + "\n", + " - `source_0`'s `IsothermalSph` mass (3 free parameters).\n", + " - `source_1`'s MGE bulge centred near (0.0, 0.0), with a narrow Gaussian prior (6 free parameters)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict_2 = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " lens_dict_2[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=getattr(result_1.instance.galaxies, f\"lens_{i}\").bulge,\n", + " mass=getattr(result_1.instance.galaxies, f\"lens_{i}\").mass,\n", + " )\n", + "\n", + "source_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " bulge=result_1.instance.galaxies.source_0.bulge,\n", + " mass=al.mp.IsothermalSph,\n", + ")\n", + "\n", + "source_1_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=source_1_bulge)\n", + "\n", + "model_2 = af.Collection(\n", + " galaxies=af.Collection(**lens_dict_2, source_0=source_0, source_1=source_1),\n", + ")\n", + "\n", + "print(model_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 2)__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__source_1_parametric\",\n", + " unique_tag=dataset_name,\n", + " n_live=200,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)\n", + "\n", + "print(result_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Two chained searches initialised a model for a group-scale double Einstein ring system. The key idea is to\n", + "exploit the independence of `source_0`'s ray-tracing from `source_1`'s properties \u2014 search 1 nails down the\n", + "lens-plane and `source_0` light using a mask that excludes `source_1`, and search 2 fixes that result and adds\n", + "the second source plus `source_0`'s mass.\n", + "\n", + "For a fully production-quality fit on real data (including pixelized source reconstructions), see `slam.py` in\n", + "the same directory. The single-search `modeling.py` example is for tutorial purposes only \u2014 it \"cheats\" by\n", + "initialising priors at the simulator's true values, which is impossible on real data.\n", + "\n", + "__Advanced Chaining__\n", + "\n", + "A more elaborate pipeline (5+ searches, eventually swapping the parametric source models for pixelized\n", + "reconstructions) is shown for the single-lens double Einstein ring in\n", + "`autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/chaining.py`. The same logic\n", + "transfers to the group-scale case, but no canonical template is currently provided for group DSPL \u2014 we are\n", + "still working out the most effective way to model these systems." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/double_einstein_ring/fit.ipynb b/notebooks/group/features/advanced/double_einstein_ring/fit.ipynb index a092afce2..3940e42e0 100644 --- a/notebooks/group/features/advanced/double_einstein_ring/fit.ipynb +++ b/notebooks/group/features/advanced/double_einstein_ring/fit.ipynb @@ -1,405 +1,442 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Group Double Einstein Ring Fit\n", - "========================================\n", - "\n", - "A group-scale double Einstein ring lens, where two source galaxies at different redshifts are lensed by multiple\n", - "main lens galaxies at the lens-plane redshift.\n", - "\n", - "This script demonstrates the API for fitting a group-scale double Einstein ring system via the standard `Tracer`\n", - "and `FitImaging` objects, without invoking a non-linear search. The two main lens galaxies are composed using\n", - "the group `lens_dict` convention (loaded from a JSON file of centres), and the two source galaxies share the\n", - "3-plane ray-tracing structure used in the imaging double Einstein ring example.\n", - "\n", - "The source galaxies are both modelled with a Multi Gaussian Expansion (MGE), as is each main lens galaxy's bulge.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script.\n", - "- **Dataset, Mask, Over Sampling:** Standard set up.\n", - "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", - "- **MGE Bases:** Build linear-Gaussian bases for each main lens galaxy bulge and for each source galaxy bulge.\n", - "- **Galaxies:** Compose `lens_dict` (two main lens galaxies, each with MGE bulge + mass), `source_0` (MGE bulge\n", - " + mass), `source_1` (MGE bulge).\n", - "- **Tracer:** Build the three-plane `Tracer` from `list(lens_dict.values()) + [source_0, source_1]`.\n", - "- **Fit:** Run `FitImaging` and inspect the fit.\n", - "- **Multi-Plane Ray-Tracing:** A short tour confirming the deflection chain accumulates contributions from both\n", - " main lens galaxies and from `source_0`'s mass.\n", - "- **Wrap Up.**\n", - "\n", - "__Prerequisites__\n", - "\n", - "This script combines the group-scale `lens_dict` API and the double Einstein ring multi-plane API. Read first:\n", - "\n", - " - `autolens_workspace/scripts/group/fit.py` \u2014 the standard group-scale fit.\n", - " - `autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/fit.py` \u2014 the single-lens double\n", - " Einstein ring fit, including the MGE basis pattern used here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the group double Einstein ring dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Centres__\n", - "\n", - "Load the two main lens galaxy centres saved by the simulator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "A 4.0\" circular mask, centred at the origin (midpoint of the two main lens galaxies), large enough to enclose\n", - "both Einstein rings." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 4.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling at each main lens galaxy centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE Bases__\n", - "\n", - "We build a linear-Gaussian `Basis` for each main lens galaxy bulge and for each source galaxy bulge. The Gaussians\n", - "share log10-spaced `sigma` values and have spherical symmetry for simplicity. The `intensity` of each Gaussian\n", - "is solved for at fit time via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "log10_sigma_list_lens = np.linspace(-2, np.log10(2.0), total_gaussians)\n", - "log10_sigma_list_source = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_basis(centre, log10_sigma_list):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "The lens-plane (z=0.5) is composed via the group `lens_dict` convention. Each main lens galaxy has an MGE bulge\n", - "centred on its loaded position, and an `IsothermalSph` mass profile matching the simulator's true values.\n", - "\n", - "`source_0` at z=1.0 has an MGE bulge AND an `IsothermalSph` mass \u2014 it deflects light from `source_1`.\n", - "`source_1` at z=2.0 has an MGE bulge only." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " lens_dict[f\"lens_{i}\"] = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=build_basis(\n", - " centre=(centre[0], centre[1]), log10_sigma_list=log10_sigma_list_lens\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(centre[0], centre[1]), einstein_radius=1.2),\n", - " )\n", - "\n", - "source_0_bulge = build_basis(\n", - " centre=(0.0, 0.0), log10_sigma_list=log10_sigma_list_source\n", - ")\n", - "source_1_bulge = build_basis(\n", - " centre=(-0.3, 0.3), log10_sigma_list=log10_sigma_list_source\n", - ")\n", - "\n", - "source_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=source_0_bulge,\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.25),\n", - ")\n", - "\n", - "source_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " bulge=source_1_bulge,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "Build the multi-plane Tracer. Galaxies are ordered internally by redshift, so the deflection chain runs:\n", - "\n", - " Plane 0 (image) : theta\n", - " Plane 1 (source_0, z=1.0) : theta - sum_over_main_lens_galaxies alpha_lens_i(theta)\n", - " Plane 2 (source_1, z=2.0) : theta - sum_over_main_lens_galaxies alpha_lens_i(theta)\n", - " - beta_01 * alpha_source_0(plane_1_grid)\n", - "\n", - "Note the contribution to `source_1`'s grid from EACH main lens galaxy as well as from `source_0`'s mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=list(lens_dict.values()) + [source_0, source_1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "Run `FitImaging` to solve for every Gaussian's `intensity` via linear algebra and to compute the model image,\n", - "residuals and log likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multi-Plane Ray-Tracing__\n", - "\n", - "Confirm the deflection chain by extracting one grid per plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.grid)\n", - "\n", - "print(f\"Number of planes traced through: {len(traced_grids)}\")\n", - "print(f\"Plane 0 (image-plane) \u2014 first coord: {traced_grids[0][0]}\")\n", - "print(f\"Plane 1 (source_0 at z=1.0) \u2014 first coord: {traced_grids[1][0]}\")\n", - "print(f\"Plane 2 (source_1 at z=2.0) \u2014 first coord: {traced_grids[2][0]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "After the fit, the linear MGE Gaussians have been assigned intensities. The fit exposes a Tracer with the\n", - "linear profiles converted to standard light profiles via\n", - "`fit.model_obj_linear_light_profiles_to_light_profiles`, which we use to visualise each MGE basis with its\n", - "solved-for amplitudes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "# Galaxies are ordered by redshift, so the first len(lens_dict) entries are the main lens galaxies,\n", - "# followed by source_0 (index N) and source_1 (index N+1).\n", - "n_lens = len(lens_dict)\n", - "\n", - "for i in range(n_lens):\n", - " subplot_basis_image(basis=tracer_fitted.galaxies[i].bulge, grid=plot_grid)\n", - "\n", - "subplot_basis_image(basis=tracer_fitted.galaxies[n_lens].bulge, grid=plot_grid)\n", - "subplot_basis_image(basis=tracer_fitted.galaxies[n_lens + 1].bulge, grid=plot_grid)\n", - "\n", - "print(f\"\\nFit log_likelihood: {fit.log_likelihood}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated the API for fitting a group-scale double Einstein ring system without invoking a\n", - "non-linear search. The two new aspects relative to the imaging double Einstein ring example are:\n", - "\n", - " 1. The lens-plane is composed via the group `lens_dict` convention, scaling naturally to any number of main\n", - " lens galaxies.\n", - " 2. The deflection chain to `source_1`'s plane accumulates contributions from EVERY main lens galaxy at z=0.5\n", - " plus `source_0`'s mass.\n", - "\n", - "For a realistic group DSPL fit on real data, use `chaining.py` (two chained searches) or `slam.py` (full\n", - "SLaM pipeline). `modeling.py` provides a tutorial single-search example, but cheats by initialising priors at\n", - "the true simulator values and is not appropriate for real data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Group Double Einstein Ring Fit\n", + "========================================\n", + "\n", + "A group-scale double Einstein ring lens, where two source galaxies at different redshifts are lensed by multiple\n", + "main lens galaxies at the lens-plane redshift.\n", + "\n", + "This script demonstrates the API for fitting a group-scale double Einstein ring system via the standard `Tracer`\n", + "and `FitImaging` objects, without invoking a non-linear search. The two main lens galaxies are composed using\n", + "the group `lens_dict` convention (loaded from a JSON file of centres), and the two source galaxies share the\n", + "3-plane ray-tracing structure used in the imaging double Einstein ring example.\n", + "\n", + "The source galaxies are both modelled with a Multi Gaussian Expansion (MGE), as is each main lens galaxy's bulge.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script.\n", + "- **Dataset, Mask, Over Sampling:** Standard set up.\n", + "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", + "- **MGE Bases:** Build linear-Gaussian bases for each main lens galaxy bulge and for each source galaxy bulge.\n", + "- **Galaxies:** Compose `lens_dict` (two main lens galaxies, each with MGE bulge + mass), `source_0` (MGE bulge\n", + " + mass), `source_1` (MGE bulge).\n", + "- **Tracer:** Build the three-plane `Tracer` from `list(lens_dict.values()) + [source_0, source_1]`.\n", + "- **Fit:** Run `FitImaging` and inspect the fit.\n", + "- **Multi-Plane Ray-Tracing:** A short tour confirming the deflection chain accumulates contributions from both\n", + " main lens galaxies and from `source_0`'s mass.\n", + "- **Wrap Up.**\n", + "\n", + "__Prerequisites__\n", + "\n", + "This script combines the group-scale `lens_dict` API and the double Einstein ring multi-plane API. Read first:\n", + "\n", + " - `autolens_workspace/scripts/group/fit.py` \u2014 the standard group-scale fit.\n", + " - `autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/fit.py` \u2014 the single-lens double\n", + " Einstein ring fit, including the MGE basis pattern used here." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the group double Einstein ring dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Centres__\n", + "\n", + "Load the two main lens galaxy centres saved by the simulator." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "A 4.0\" circular mask, centred at the origin (midpoint of the two main lens galaxies), large enough to enclose\n", + "both Einstein rings." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 4.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling at each main lens galaxy centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE Bases__\n", + "\n", + "We build a linear-Gaussian `Basis` for each main lens galaxy bulge and for each source galaxy bulge. The Gaussians\n", + "share log10-spaced `sigma` values and have spherical symmetry for simplicity. The `intensity` of each Gaussian\n", + "is solved for at fit time via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "log10_sigma_list_lens = np.linspace(-2, np.log10(2.0), total_gaussians)\n", + "log10_sigma_list_source = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_basis(centre, log10_sigma_list):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "The lens-plane (z=0.5) is composed via the group `lens_dict` convention. Each main lens galaxy has an MGE bulge\n", + "centred on its loaded position, and an `IsothermalSph` mass profile matching the simulator's true values.\n", + "\n", + "`source_0` at z=1.0 has an MGE bulge AND an `IsothermalSph` mass \u2014 it deflects light from `source_1`.\n", + "`source_1` at z=2.0 has an MGE bulge only." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " lens_dict[f\"lens_{i}\"] = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=build_basis(\n", + " centre=(centre[0], centre[1]), log10_sigma_list=log10_sigma_list_lens\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(centre[0], centre[1]), einstein_radius=1.2),\n", + " )\n", + "\n", + "source_0_bulge = build_basis(\n", + " centre=(0.0, 0.0), log10_sigma_list=log10_sigma_list_source\n", + ")\n", + "source_1_bulge = build_basis(\n", + " centre=(-0.3, 0.3), log10_sigma_list=log10_sigma_list_source\n", + ")\n", + "\n", + "source_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=source_0_bulge,\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.25),\n", + ")\n", + "\n", + "source_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " bulge=source_1_bulge,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "Build the multi-plane Tracer. Galaxies are ordered internally by redshift, so the deflection chain runs:\n", + "\n", + " Plane 0 (image) : theta\n", + " Plane 1 (source_0, z=1.0) : theta - sum_over_main_lens_galaxies alpha_lens_i(theta)\n", + " Plane 2 (source_1, z=2.0) : theta - sum_over_main_lens_galaxies alpha_lens_i(theta)\n", + " - beta_01 * alpha_source_0(plane_1_grid)\n", + "\n", + "Note the contribution to `source_1`'s grid from EACH main lens galaxy as well as from `source_0`'s mass." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=list(lens_dict.values()) + [source_0, source_1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "Run `FitImaging` to solve for every Gaussian's `intensity` via linear algebra and to compute the model image,\n", + "residuals and log likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multi-Plane Ray-Tracing__\n", + "\n", + "Confirm the deflection chain by extracting one grid per plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.grid)\n", + "\n", + "print(f\"Number of planes traced through: {len(traced_grids)}\")\n", + "print(f\"Plane 0 (image-plane) \u2014 first coord: {traced_grids[0][0]}\")\n", + "print(f\"Plane 1 (source_0 at z=1.0) \u2014 first coord: {traced_grids[1][0]}\")\n", + "print(f\"Plane 2 (source_1 at z=2.0) \u2014 first coord: {traced_grids[2][0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "After the fit, the linear MGE Gaussians have been assigned intensities. The fit exposes a Tracer with the\n", + "linear profiles converted to standard light profiles via\n", + "`fit.model_obj_linear_light_profiles_to_light_profiles`, which we use to visualise each MGE basis with its\n", + "solved-for amplitudes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "# Galaxies are ordered by redshift, so the first len(lens_dict) entries are the main lens galaxies,\n", + "# followed by source_0 (index N) and source_1 (index N+1).\n", + "n_lens = len(lens_dict)\n", + "\n", + "for i in range(n_lens):\n", + " subplot_basis_image(basis=tracer_fitted.galaxies[i].bulge, grid=plot_grid)\n", + "\n", + "subplot_basis_image(basis=tracer_fitted.galaxies[n_lens].bulge, grid=plot_grid)\n", + "subplot_basis_image(basis=tracer_fitted.galaxies[n_lens + 1].bulge, grid=plot_grid)\n", + "\n", + "print(f\"\\nFit log_likelihood: {fit.log_likelihood}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated the API for fitting a group-scale double Einstein ring system without invoking a\n", + "non-linear search. The two new aspects relative to the imaging double Einstein ring example are:\n", + "\n", + " 1. The lens-plane is composed via the group `lens_dict` convention, scaling naturally to any number of main\n", + " lens galaxies.\n", + " 2. The deflection chain to `source_1`'s plane accumulates contributions from EVERY main lens galaxy at z=0.5\n", + " plus `source_0`'s mass.\n", + "\n", + "For a realistic group DSPL fit on real data, use `chaining.py` (two chained searches) or `slam.py` (full\n", + "SLaM pipeline). `modeling.py` provides a tutorial single-search example, but cheats by initialising priors at\n", + "the true simulator values and is not appropriate for real data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/double_einstein_ring/likelihood_function.ipynb b/notebooks/group/features/advanced/double_einstein_ring/likelihood_function.ipynb index f741ad1c4..6292111c5 100644 --- a/notebooks/group/features/advanced/double_einstein_ring/likelihood_function.ipynb +++ b/notebooks/group/features/advanced/double_einstein_ring/likelihood_function.ipynb @@ -1,345 +1,382 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Group Double Einstein Ring__\n", - "\n", - "This script describes the additional steps required to compute the `log_likelihood` for a group-scale double\n", - "Einstein ring lens \u2014 two source galaxies at different redshifts lensed by multiple main lens galaxies at the\n", - "lens-plane redshift.\n", - "\n", - "This script does NOT repeat the steps shared with single-plane or single-lens-galaxy lensing (mask, image-plane\n", - "grid, convolution, chi-squared, noise normalization, MGE linear algebra). It documents only what is new for the\n", - "combined group + double Einstein ring case.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function below builds on standard imaging, MGE and single-lens double Einstein ring likelihoods.\n", - "Read these first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 single-plane imaging log likelihood, covering\n", - " chi-squared and noise normalization.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 the MGE\n", - " `Basis` of linear Gaussians and the associated linear-algebra terms.\n", - " - `autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/likelihood_function.py` \u2014 the\n", - " multi-plane deflection chain for a single-lens-galaxy double Einstein ring.\n", - "\n", - "This script focuses entirely on what differs for the group-scale case.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script (see above).\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", - "- **Galaxies:** Multiple main lens galaxies at z=0.5, source_0 at z=1.0 (light+mass), source_1 at z=2.0 (light).\n", - "- **Multi-Plane Ray-Tracing:** The deflection chain to `source_1`'s plane accumulates contributions from EVERY\n", - " main lens galaxy at z=0.5, plus from `source_0`'s mass.\n", - "- **Source-Plane Images:** Both source galaxies are evaluated at their respective ray-traced grids.\n", - "- **Likelihood:** Reference up to canonical scripts for chi-squared / noise / linear algebra.\n", - "- **Fit Check:** Confirm the manual reconstruction matches `FitImaging.log_likelihood`.\n", - "- **Wrap Up.**\n", - "\n", - "__What Changes Relative To The Single-Lens Double Einstein Ring__\n", - "\n", - "In the single-lens double Einstein ring likelihood function, the deflection map applied to image-plane\n", - "coordinates was simply `alpha_lens(theta)` \u2014 the deflection field of the lone foreground lens galaxy.\n", - "\n", - "In the group case, the lens-plane contains MULTIPLE main lens galaxies (indexed `lens_0`, `lens_1`, ...), each\n", - "with its own mass distribution. The lens-plane deflection field is the SUM of every main lens galaxy's\n", - "deflection:\n", - "\n", - " alpha_total(theta) = sum_i alpha_lens_i(theta)\n", - "\n", - "The remainder of the multi-plane chain is unchanged from the single-lens DSPL case:\n", - "\n", - " Plane 0 (image-plane) : theta\n", - " Plane 1 (source_0 at z=1.0) : theta - alpha_total(theta)\n", - " Plane 2 (source_1 at z=2.0) : theta - alpha_total(theta) - beta_01 * alpha_source_0(plane_1_grid)\n", - "\n", - "PyAutoLens handles this summation automatically inside `Tracer.traced_grid_2d_list_from`, so no special code is\n", - "required \u2014 but it is the conceptual difference relative to a single-lens-galaxy double Einstein ring." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the group double Einstein ring dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 4.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "The galaxies that participate in the multi-plane ray-tracing:\n", - "\n", - " - Two main lens galaxies at z=0.5, each with an `IsothermalSph` mass and a simple MGE bulge. Their mass\n", - " profiles together set the deflection field for the lens-plane.\n", - " - `source_0` at z=1.0 with an MGE bulge AND an `IsothermalSph` mass \u2014 this source deflects light from\n", - " `source_1`.\n", - " - `source_1` at z=2.0 with an MGE bulge only.\n", - "\n", - "Each galaxy's parameters are set to the simulator's true values so the manual likelihood computation produces a\n", - "sensible model image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 10\n", - "log10_sigma_list_lens = np.linspace(-2, np.log10(2.0), total_gaussians)\n", - "log10_sigma_list_source = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_basis(centre, log10_sigma_list):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "lens_galaxies = []\n", - "for centre in main_lens_centres:\n", - " lens_galaxies.append(\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=build_basis(\n", - " centre=(centre[0], centre[1]), log10_sigma_list=log10_sigma_list_lens\n", - " ),\n", - " mass=al.mp.IsothermalSph(\n", - " centre=(centre[0], centre[1]), einstein_radius=1.2\n", - " ),\n", - " )\n", - " )\n", - "\n", - "source_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=build_basis(centre=(0.0, 0.0), log10_sigma_list=log10_sigma_list_source),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.25),\n", - ")\n", - "\n", - "source_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " bulge=build_basis(centre=(-0.3, 0.3), log10_sigma_list=log10_sigma_list_source),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=lens_galaxies + [source_0, source_1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multi-Plane Ray-Tracing__\n", - "\n", - "`traced_grid_2d_list_from` returns one grid per plane, in redshift order. For the group DSPL case there are\n", - "still THREE planes (image, source_0, source_1) \u2014 the multi-plane chain does not grow with the number of main\n", - "lens galaxies, because all main lens galaxies share the same redshift z=0.5.\n", - "\n", - "What changes is the deflection field applied at the lens-plane boundary: it is the SUM over all main lens\n", - "galaxies' deflection contributions, computed automatically by `Tracer`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=dataset.grid)\n", - "\n", - "grid_image_plane = traced_grid_list[0]\n", - "grid_source_0 = traced_grid_list[1]\n", - "grid_source_1 = traced_grid_list[2]\n", - "\n", - "print(f\"Number of planes traced: {len(traced_grid_list)}\")\n", - "print(f\"Number of main lens galaxies at z=0.5: {len(lens_galaxies)}\")\n", - "print(f\"Plane 1 (source_0) first coord: {grid_source_0[0]}\")\n", - "print(f\"Plane 2 (source_1) first coord: {grid_source_1[0]}\")\n", - "\n", - "aplt.plot_grid(grid=grid_source_0, title=\"Ray-traced grid at source_0 plane (z=1.0)\")\n", - "aplt.plot_grid(grid=grid_source_1, title=\"Ray-traced grid at source_1 plane (z=2.0)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Images__\n", - "\n", - "`tracer.image_2d_from` evaluates every galaxy's light at the correct plane and sums the contributions into a\n", - "single model image. Internally it performs:\n", - "\n", - " 1. Ray-traces the image-plane grid through the lens-plane (summing the deflection fields of all main lens\n", - " galaxies) to obtain `grid_source_0`.\n", - " 2. Continues to ray-trace through `source_0`'s mass to obtain `grid_source_1`.\n", - " 3. Evaluates `source_0`'s MGE basis at `grid_source_0` and `source_1`'s MGE basis at `grid_source_1`.\n", - " 4. Sums all source contributions, returning the model image.\n", - "\n", - "The only group-specific step is (1). Every other step is identical to the imaging double Einstein ring case." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image_unconvolved = tracer.image_2d_from(grid=dataset.grid)\n", - "\n", - "aplt.plot_array(\n", - " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood__\n", - "\n", - "PSF convolution, chi-squared, noise normalization, and the MGE linear-algebra terms are unchanged. We delegate\n", - "to `FitImaging`, which handles the linear solve that recovers each Gaussian's `intensity` and assembles the\n", - "full `log_likelihood`.\n", - "\n", - "For the form of these terms, refer to:\n", - "\n", - " - `imaging/likelihood_function.py` \u2014 chi-squared and noise normalization.\n", - " - `imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 linear-algebra terms." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(\n", - " f\"\\nLog likelihood of the manual group double Einstein ring fit: {fit.log_likelihood}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The group double Einstein ring `log_likelihood` differs from the single-lens DSPL case in exactly one place:\n", - "the lens-plane deflection field is the sum of contributions from EVERY main lens galaxy at z=0.5, rather than\n", - "from a lone lens galaxy. `Tracer` handles the summation automatically \u2014 no new code is required.\n", - "\n", - "The deflection scaling factor `beta_01` between source_0 and source_1 (and therefore the cosmological\n", - "sensitivity of the system) is unchanged from the imaging case. Group DSPLs are valuable for cosmology because\n", - "they tend to have larger Einstein radii, which improves astrometric precision and therefore the angular\n", - "diameter distance ratios that `beta_01` depends on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Group Double Einstein Ring__\n", + "\n", + "This script describes the additional steps required to compute the `log_likelihood` for a group-scale double\n", + "Einstein ring lens \u2014 two source galaxies at different redshifts lensed by multiple main lens galaxies at the\n", + "lens-plane redshift.\n", + "\n", + "This script does NOT repeat the steps shared with single-plane or single-lens-galaxy lensing (mask, image-plane\n", + "grid, convolution, chi-squared, noise normalization, MGE linear algebra). It documents only what is new for the\n", + "combined group + double Einstein ring case.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function below builds on standard imaging, MGE and single-lens double Einstein ring likelihoods.\n", + "Read these first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 single-plane imaging log likelihood, covering\n", + " chi-squared and noise normalization.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 the MGE\n", + " `Basis` of linear Gaussians and the associated linear-algebra terms.\n", + " - `autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/likelihood_function.py` \u2014 the\n", + " multi-plane deflection chain for a single-lens-galaxy double Einstein ring.\n", + "\n", + "This script focuses entirely on what differs for the group-scale case.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script (see above).\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", + "- **Galaxies:** Multiple main lens galaxies at z=0.5, source_0 at z=1.0 (light+mass), source_1 at z=2.0 (light).\n", + "- **Multi-Plane Ray-Tracing:** The deflection chain to `source_1`'s plane accumulates contributions from EVERY\n", + " main lens galaxy at z=0.5, plus from `source_0`'s mass.\n", + "- **Source-Plane Images:** Both source galaxies are evaluated at their respective ray-traced grids.\n", + "- **Likelihood:** Reference up to canonical scripts for chi-squared / noise / linear algebra.\n", + "- **Fit Check:** Confirm the manual reconstruction matches `FitImaging.log_likelihood`.\n", + "- **Wrap Up.**\n", + "\n", + "__What Changes Relative To The Single-Lens Double Einstein Ring__\n", + "\n", + "In the single-lens double Einstein ring likelihood function, the deflection map applied to image-plane\n", + "coordinates was simply `alpha_lens(theta)` \u2014 the deflection field of the lone foreground lens galaxy.\n", + "\n", + "In the group case, the lens-plane contains MULTIPLE main lens galaxies (indexed `lens_0`, `lens_1`, ...), each\n", + "with its own mass distribution. The lens-plane deflection field is the SUM of every main lens galaxy's\n", + "deflection:\n", + "\n", + " alpha_total(theta) = sum_i alpha_lens_i(theta)\n", + "\n", + "The remainder of the multi-plane chain is unchanged from the single-lens DSPL case:\n", + "\n", + " Plane 0 (image-plane) : theta\n", + " Plane 1 (source_0 at z=1.0) : theta - alpha_total(theta)\n", + " Plane 2 (source_1 at z=2.0) : theta - alpha_total(theta) - beta_01 * alpha_source_0(plane_1_grid)\n", + "\n", + "PyAutoLens handles this summation automatically inside `Tracer.traced_grid_2d_list_from`, so no special code is\n", + "required \u2014 but it is the conceptual difference relative to a single-lens-galaxy double Einstein ring." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the group double Einstein ring dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 4.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "The galaxies that participate in the multi-plane ray-tracing:\n", + "\n", + " - Two main lens galaxies at z=0.5, each with an `IsothermalSph` mass and a simple MGE bulge. Their mass\n", + " profiles together set the deflection field for the lens-plane.\n", + " - `source_0` at z=1.0 with an MGE bulge AND an `IsothermalSph` mass \u2014 this source deflects light from\n", + " `source_1`.\n", + " - `source_1` at z=2.0 with an MGE bulge only.\n", + "\n", + "Each galaxy's parameters are set to the simulator's true values so the manual likelihood computation produces a\n", + "sensible model image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 10\n", + "log10_sigma_list_lens = np.linspace(-2, np.log10(2.0), total_gaussians)\n", + "log10_sigma_list_source = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_basis(centre, log10_sigma_list):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "lens_galaxies = []\n", + "for centre in main_lens_centres:\n", + " lens_galaxies.append(\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=build_basis(\n", + " centre=(centre[0], centre[1]), log10_sigma_list=log10_sigma_list_lens\n", + " ),\n", + " mass=al.mp.IsothermalSph(\n", + " centre=(centre[0], centre[1]), einstein_radius=1.2\n", + " ),\n", + " )\n", + " )\n", + "\n", + "source_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=build_basis(centre=(0.0, 0.0), log10_sigma_list=log10_sigma_list_source),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.25),\n", + ")\n", + "\n", + "source_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " bulge=build_basis(centre=(-0.3, 0.3), log10_sigma_list=log10_sigma_list_source),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=lens_galaxies + [source_0, source_1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multi-Plane Ray-Tracing__\n", + "\n", + "`traced_grid_2d_list_from` returns one grid per plane, in redshift order. For the group DSPL case there are\n", + "still THREE planes (image, source_0, source_1) \u2014 the multi-plane chain does not grow with the number of main\n", + "lens galaxies, because all main lens galaxies share the same redshift z=0.5.\n", + "\n", + "What changes is the deflection field applied at the lens-plane boundary: it is the SUM over all main lens\n", + "galaxies' deflection contributions, computed automatically by `Tracer`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=dataset.grid)\n", + "\n", + "grid_image_plane = traced_grid_list[0]\n", + "grid_source_0 = traced_grid_list[1]\n", + "grid_source_1 = traced_grid_list[2]\n", + "\n", + "print(f\"Number of planes traced: {len(traced_grid_list)}\")\n", + "print(f\"Number of main lens galaxies at z=0.5: {len(lens_galaxies)}\")\n", + "print(f\"Plane 1 (source_0) first coord: {grid_source_0[0]}\")\n", + "print(f\"Plane 2 (source_1) first coord: {grid_source_1[0]}\")\n", + "\n", + "aplt.plot_grid(grid=grid_source_0, title=\"Ray-traced grid at source_0 plane (z=1.0)\")\n", + "aplt.plot_grid(grid=grid_source_1, title=\"Ray-traced grid at source_1 plane (z=2.0)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Images__\n", + "\n", + "`tracer.image_2d_from` evaluates every galaxy's light at the correct plane and sums the contributions into a\n", + "single model image. Internally it performs:\n", + "\n", + " 1. Ray-traces the image-plane grid through the lens-plane (summing the deflection fields of all main lens\n", + " galaxies) to obtain `grid_source_0`.\n", + " 2. Continues to ray-trace through `source_0`'s mass to obtain `grid_source_1`.\n", + " 3. Evaluates `source_0`'s MGE basis at `grid_source_0` and `source_1`'s MGE basis at `grid_source_1`.\n", + " 4. Sums all source contributions, returning the model image.\n", + "\n", + "The only group-specific step is (1). Every other step is identical to the imaging double Einstein ring case." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image_unconvolved = tracer.image_2d_from(grid=dataset.grid)\n", + "\n", + "aplt.plot_array(\n", + " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood__\n", + "\n", + "PSF convolution, chi-squared, noise normalization, and the MGE linear-algebra terms are unchanged. We delegate\n", + "to `FitImaging`, which handles the linear solve that recovers each Gaussian's `intensity` and assembles the\n", + "full `log_likelihood`.\n", + "\n", + "For the form of these terms, refer to:\n", + "\n", + " - `imaging/likelihood_function.py` \u2014 chi-squared and noise normalization.\n", + " - `imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 linear-algebra terms." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(\n", + " f\"\\nLog likelihood of the manual group double Einstein ring fit: {fit.log_likelihood}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The group double Einstein ring `log_likelihood` differs from the single-lens DSPL case in exactly one place:\n", + "the lens-plane deflection field is the sum of contributions from EVERY main lens galaxy at z=0.5, rather than\n", + "from a lone lens galaxy. `Tracer` handles the summation automatically \u2014 no new code is required.\n", + "\n", + "The deflection scaling factor `beta_01` between source_0 and source_1 (and therefore the cosmological\n", + "sensitivity of the system) is unchanged from the imaging case. Group DSPLs are valuable for cosmology because\n", + "they tend to have larger Einstein radii, which improves astrometric precision and therefore the angular\n", + "diameter distance ratios that `beta_01` depends on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/double_einstein_ring/modeling.ipynb b/notebooks/group/features/advanced/double_einstein_ring/modeling.ipynb index be9cf6d92..ecc40290f 100644 --- a/notebooks/group/features/advanced/double_einstein_ring/modeling.ipynb +++ b/notebooks/group/features/advanced/double_einstein_ring/modeling.ipynb @@ -1,467 +1,504 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Group Double Einstein Ring\n", - "=============================================\n", - "\n", - "A group-scale double Einstein ring lens, where two source galaxies at different redshifts are lensed by multiple\n", - "main lens galaxies at the lens-plane redshift.\n", - "\n", - "This script illustrates the PyAutoLens API for modeling such a system in a single non-linear search. The lens\n", - "plane is composed via the group `lens_dict` convention (loaded from a JSON file of main lens galaxy centres),\n", - "and the source plane has two galaxies \u2014 `source_0` at z=1.0 with light + mass, and `source_1` at z=2.0 with\n", - "light only.\n", - "\n", - "__Practical Use: Read This First__\n", - "\n", - "This script is a tutorial. It produces a working fit by \"cheating\" \u2014 every prior is initialised at the true\n", - "simulator value, narrowed by a small Gaussian. On real data this is impossible, and a single Nautilus search on\n", - "this many-parameter group DSPL model would almost certainly converge to a local maximum.\n", - "\n", - "The script you will actually use to fit a group-scale double Einstein ring on real data is\n", - "`autolens_workspace/scripts/group/features/advanced/double_einstein_ring/chaining.py`, which runs two chained\n", - "non-linear searches: the first initialises the main lens galaxies' mass and `source_0` using a smaller mask\n", - "that excludes `source_1`, the second introduces `source_1` and frees `source_0`'s mass.\n", - "\n", - "For production-quality modeling, see `slam.py` in the same directory.\n", - "\n", - "Read this script to understand the model composition API, then jump to `chaining.py`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset, Mask, Over Sampling:** Standard set up.\n", - "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", - "- **Model Composition:** Build `lens_dict` via the group `lens_dict` convention, plus `source_0` and `source_1`.\n", - "- **Cheating:** Override priors with narrow Gaussians around the true simulator values.\n", - "- **Cosmology:** Fixed at Planck18 by default \u2014 a commented-out snippet shows how to make `Om0` free.\n", - "- **Search, Analysis, Run-Time, Result, Wrap Up.**" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the group double Einstein ring `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Centres__\n", - "\n", - "The two main lens galaxy centres are loaded from the JSON file saved by the simulator. They are used to define\n", - "the over-sampling pattern and as centres for each lens galaxy's MGE bulge in the model composition below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "A 4.0\" circular mask that encloses both Einstein rings around the centroid of the two main lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 4.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling at each main lens galaxy centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - Each main lens galaxy has an MGE bulge (20 Gaussians) and an `Isothermal` mass profile. Centres start fixed at\n", - " the simulator values for over-sampling and are then given narrow Gaussian priors in the \"Cheating\" section.\n", - "\n", - " - `source_0` (z=1.0) has an MGE bulge and an `IsothermalSph` mass, since it acts as a secondary deflector for\n", - " `source_1`.\n", - "\n", - " - `source_1` (z=2.0) has an MGE bulge only.\n", - "\n", - "Total non-linear parameter count: ~16 per main lens galaxy mass + bulge centre (~5 each), plus ~6 for source_0\n", - "and ~3 for source_1 \u2014 comparable to the imaging double Einstein ring example.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook:\n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies via lens_dict:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " centre_prior_is_uniform=True,\n", - " centre=(centre[0], centre[1]),\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = (centre[0], centre[1])\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " )\n", - "\n", - "# Source 0:\n", - "\n", - "bulge_source_0 = af.Model(al.lp_linear.ExponentialCoreSph)\n", - "mass_source_0 = af.Model(al.mp.IsothermalSph)\n", - "\n", - "source_0 = af.Model(\n", - " al.Galaxy, redshift=1.0, bulge=bulge_source_0, mass=mass_source_0, centre=(0.0, 0.0)\n", - ")\n", - "\n", - "# Source 1:\n", - "\n", - "bulge_source_1 = af.Model(al.lp_linear.ExponentialCoreSph)\n", - "\n", - "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=bulge_source_1, centre=(0.0, 0.0))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cheating__\n", - "\n", - "Initializing a group double Einstein ring lens model is even harder than the single-lens case, due to the larger\n", - "parameter space. To make the single-search example tractable, we override key priors with narrow Gaussians\n", - "centred at the simulator values. On real data this is impossible.\n", - "\n", - "The main lens galaxy centres are kept fixed at their loaded positions (this is the standard group convention \u2014\n", - "centres are determined from independent photometry). We \"cheat\" by narrowing `source_0`'s mass priors around\n", - "its true position." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_0.mass.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.2)\n", - "source_0.mass.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.2)\n", - "source_0.mass.einstein_radius = af.GaussianPrior(mean=0.25, sigma=0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cosmology__\n", - "\n", - "Double Einstein rings constrain cosmology via the angular diameter distance ratios between the lens, source_0,\n", - "and source_1. This script uses a fixed Planck18 cosmology to keep the parameter count manageable for the\n", - "single-search \"cheating\" workflow. A realistic cosmological constraint requires the chained-search workflow in\n", - "`chaining.py` and significantly more data than one simulated system.\n", - "\n", - "To make `Om0` (Omega_m) a free parameter, uncomment the three lines below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# cosmology = af.Model(al.cosmo.FlatLambdaCDM)\n", - "# cosmology.Om0 = af.GaussianPrior(mean=0.3, sigma=0.1)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source_0=source_0, source_1=source_1),\n", - " # cosmology=cosmology,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "We use the nested sampling algorithm Nautilus." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"double_einstein_ring\",\n", - " unique_tag=dataset_name,\n", - " n_live=200,\n", - " n_batch=50,\n", - " iterations_per_quick_update=2000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "Group-scale double Einstein ring lenses are VRAM-intensive on GPU: multi-plane ray-tracing, multiple main lens\n", - "galaxies and batched non-linear search samples all multiply VRAM use." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "Run times for this model are quite long, because (a) multi-plane ray-tracing is expensive, (b) the model has\n", - "many more parameters than a single-plane lens, and (c) the \"cheating\" workflow is only viable when the priors\n", - "are narrow enough to keep the search away from local maxima.\n", - "\n", - "For real data, use `chaining.py` instead \u2014 the chained-search approach is significantly more efficient AND more\n", - "robust.\n", - "\n", - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Group-scale double Einstein ring systems can be fit in PyAutoLens, but this script \"cheats\" by initialising\n", - "priors at their true values. For real data, use `chaining.py` (two chained searches) or `slam.py` (the full\n", - "SLaM pipeline), and consult the imaging double Einstein ring example for an introduction to the multi-plane\n", - "ray-tracing API." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Group Double Einstein Ring\n", + "=============================================\n", + "\n", + "A group-scale double Einstein ring lens, where two source galaxies at different redshifts are lensed by multiple\n", + "main lens galaxies at the lens-plane redshift.\n", + "\n", + "This script illustrates the PyAutoLens API for modeling such a system in a single non-linear search. The lens\n", + "plane is composed via the group `lens_dict` convention (loaded from a JSON file of main lens galaxy centres),\n", + "and the source plane has two galaxies \u2014 `source_0` at z=1.0 with light + mass, and `source_1` at z=2.0 with\n", + "light only.\n", + "\n", + "__Practical Use: Read This First__\n", + "\n", + "This script is a tutorial. It produces a working fit by \"cheating\" \u2014 every prior is initialised at the true\n", + "simulator value, narrowed by a small Gaussian. On real data this is impossible, and a single Nautilus search on\n", + "this many-parameter group DSPL model would almost certainly converge to a local maximum.\n", + "\n", + "The script you will actually use to fit a group-scale double Einstein ring on real data is\n", + "`autolens_workspace/scripts/group/features/advanced/double_einstein_ring/chaining.py`, which runs two chained\n", + "non-linear searches: the first initialises the main lens galaxies' mass and `source_0` using a smaller mask\n", + "that excludes `source_1`, the second introduces `source_1` and frees `source_0`'s mass.\n", + "\n", + "For production-quality modeling, see `slam.py` in the same directory.\n", + "\n", + "Read this script to understand the model composition API, then jump to `chaining.py`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset, Mask, Over Sampling:** Standard set up.\n", + "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", + "- **Model Composition:** Build `lens_dict` via the group `lens_dict` convention, plus `source_0` and `source_1`.\n", + "- **Cheating:** Override priors with narrow Gaussians around the true simulator values.\n", + "- **Cosmology:** Fixed at Planck18 by default \u2014 a commented-out snippet shows how to make `Om0` free.\n", + "- **Search, Analysis, Run-Time, Result, Wrap Up.**" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the group double Einstein ring `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Centres__\n", + "\n", + "The two main lens galaxy centres are loaded from the JSON file saved by the simulator. They are used to define\n", + "the over-sampling pattern and as centres for each lens galaxy's MGE bulge in the model composition below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "A 4.0\" circular mask that encloses both Einstein rings around the centroid of the two main lens galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 4.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling at each main lens galaxy centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - Each main lens galaxy has an MGE bulge (20 Gaussians) and an `Isothermal` mass profile. Centres start fixed at\n", + " the simulator values for over-sampling and are then given narrow Gaussian priors in the \"Cheating\" section.\n", + "\n", + " - `source_0` (z=1.0) has an MGE bulge and an `IsothermalSph` mass, since it acts as a secondary deflector for\n", + " `source_1`.\n", + "\n", + " - `source_1` (z=2.0) has an MGE bulge only.\n", + "\n", + "Total non-linear parameter count: ~16 per main lens galaxy mass + bulge centre (~5 each), plus ~6 for source_0\n", + "and ~3 for source_1 \u2014 comparable to the imaging double Einstein ring example.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook:\n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies via lens_dict:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " centre_prior_is_uniform=True,\n", + " centre=(centre[0], centre[1]),\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = (centre[0], centre[1])\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " )\n", + "\n", + "# Source 0:\n", + "\n", + "bulge_source_0 = af.Model(al.lp_linear.ExponentialCoreSph)\n", + "mass_source_0 = af.Model(al.mp.IsothermalSph)\n", + "\n", + "source_0 = af.Model(\n", + " al.Galaxy, redshift=1.0, bulge=bulge_source_0, mass=mass_source_0, centre=(0.0, 0.0)\n", + ")\n", + "\n", + "# Source 1:\n", + "\n", + "bulge_source_1 = af.Model(al.lp_linear.ExponentialCoreSph)\n", + "\n", + "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=bulge_source_1, centre=(0.0, 0.0))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cheating__\n", + "\n", + "Initializing a group double Einstein ring lens model is even harder than the single-lens case, due to the larger\n", + "parameter space. To make the single-search example tractable, we override key priors with narrow Gaussians\n", + "centred at the simulator values. On real data this is impossible.\n", + "\n", + "The main lens galaxy centres are kept fixed at their loaded positions (this is the standard group convention \u2014\n", + "centres are determined from independent photometry). We \"cheat\" by narrowing `source_0`'s mass priors around\n", + "its true position." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_0.mass.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.2)\n", + "source_0.mass.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.2)\n", + "source_0.mass.einstein_radius = af.GaussianPrior(mean=0.25, sigma=0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cosmology__\n", + "\n", + "Double Einstein rings constrain cosmology via the angular diameter distance ratios between the lens, source_0,\n", + "and source_1. This script uses a fixed Planck18 cosmology to keep the parameter count manageable for the\n", + "single-search \"cheating\" workflow. A realistic cosmological constraint requires the chained-search workflow in\n", + "`chaining.py` and significantly more data than one simulated system.\n", + "\n", + "To make `Om0` (Omega_m) a free parameter, uncomment the three lines below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# cosmology = af.Model(al.cosmo.FlatLambdaCDM)\n", + "# cosmology.Om0 = af.GaussianPrior(mean=0.3, sigma=0.1)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source_0=source_0, source_1=source_1),\n", + " # cosmology=cosmology,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "We use the nested sampling algorithm Nautilus." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"double_einstein_ring\",\n", + " unique_tag=dataset_name,\n", + " n_live=200,\n", + " n_batch=50,\n", + " iterations_per_quick_update=2000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "Group-scale double Einstein ring lenses are VRAM-intensive on GPU: multi-plane ray-tracing, multiple main lens\n", + "galaxies and batched non-linear search samples all multiply VRAM use." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "Run times for this model are quite long, because (a) multi-plane ray-tracing is expensive, (b) the model has\n", + "many more parameters than a single-plane lens, and (c) the \"cheating\" workflow is only viable when the priors\n", + "are narrow enough to keep the search away from local maxima.\n", + "\n", + "For real data, use `chaining.py` instead \u2014 the chained-search approach is significantly more efficient AND more\n", + "robust.\n", + "\n", + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Group-scale double Einstein ring systems can be fit in PyAutoLens, but this script \"cheats\" by initialising\n", + "priors at their true values. For real data, use `chaining.py` (two chained searches) or `slam.py` (the full\n", + "SLaM pipeline), and consult the imaging double Einstein ring example for an introduction to the multi-plane\n", + "ray-tracing API." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/double_einstein_ring/simulator.ipynb b/notebooks/group/features/advanced/double_einstein_ring/simulator.ipynb index b125e2c24..43c81d6d9 100644 --- a/notebooks/group/features/advanced/double_einstein_ring/simulator.ipynb +++ b/notebooks/group/features/advanced/double_einstein_ring/simulator.ipynb @@ -1,414 +1,451 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Group Double Einstein Ring\n", - "=====================================\n", - "\n", - "A double Einstein ring lens in a group-scale context, where two source galaxies at different redshifts are lensed\n", - "by multiple main lens galaxies at the lens-plane redshift.\n", - "\n", - "This script simulates an `Imaging` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - There are TWO main lens galaxies at z=0.5, each with a `SersicSph` light profile and an `IsothermalSph` mass\n", - " profile.\n", - " - The first source galaxy at z=1.0 has a `SersicSph` light profile AND an `IsothermalSph` mass profile. This\n", - " source acts as both a light source AND a deflector for the second source.\n", - " - The second source galaxy at z=2.0 has a `SersicSph` light profile only.\n", - "\n", - "The multi-plane ray-tracing accounts for all main lens galaxy masses at the lens redshift, plus the intermediate\n", - "source galaxy's mass acting as a secondary lens for the more distant source.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** The dataset name and output folder.\n", - "- **Grid:** Define the 2D image-plane grid.\n", - "- **Galaxy Centres:** Centres of the two main lens galaxies, saved as JSON for the modeling scripts to load.\n", - "- **Over Sampling:** Adaptive over-sampling at the main lens galaxy centres.\n", - "- **Main Lens Galaxies:** Two galaxies at z=0.5.\n", - "- **Source Galaxies:** source_0 at z=1.0 (light + mass), source_1 at z=2.0 (light only).\n", - "- **Ray Tracing:** Build the multi-plane Tracer.\n", - "- **Dataset:** Simulate the imaging dataset and write .fits.\n", - "- **Tracer JSON:** Save the simulator Tracer for provenance.\n", - "- **Centres JSON:** Save the main lens centres for the modeling scripts.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to:\n", - "\n", - " - `autolens_workspace/scripts/group/simulator.py` \u2014 the canonical group-scale simulator.\n", - " - `autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/simulator.py` \u2014 the single-lens\n", - " double Einstein ring simulator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The dataset is written to `dataset/group/double_einstein_ring/`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "A 200x200 grid at 0.1\"/px gives a 20\" field of view, large enough to contain a group-scale double Einstein ring\n", - "where the two main lens galaxies are separated by ~1.5\"." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(200, 200),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "The two main lens galaxies are separated by ~1.5\" along the x-axis. These centres are saved as JSON so the\n", - "modeling and fit scripts can load them via `al.from_json(...)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [(0.0, -0.75), (0.0, 0.75)]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling is applied at each main lens galaxy's centre and at the origin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=main_lens_centres,\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Imaging simulator: exposure time, background sky, noise, PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "Two `SersicSph` + `IsothermalSph` galaxies at z=0.5. Their Einstein radii are chosen so the combined deflection\n", - "produces a clearly visible primary Einstein ring around the group's centroid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=main_lens_centres[0],\n", - " intensity=0.7,\n", - " effective_radius=1.0,\n", - " sersic_index=4.0,\n", - " ),\n", - " mass=al.mp.IsothermalSph(\n", - " centre=main_lens_centres[0],\n", - " einstein_radius=1.2,\n", - " ),\n", - ")\n", - "\n", - "lens_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=main_lens_centres[1],\n", - " intensity=0.7,\n", - " effective_radius=1.0,\n", - " sersic_index=4.0,\n", - " ),\n", - " mass=al.mp.IsothermalSph(\n", - " centre=main_lens_centres[1],\n", - " einstein_radius=1.2,\n", - " ),\n", - ")\n", - "\n", - "main_lens_galaxies = [lens_0, lens_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxies__\n", - "\n", - "`source_0` at z=1.0 has BOTH a light profile (its light forms the primary Einstein ring) AND a mass profile\n", - "(its mass deflects the light from `source_1` and contributes to the second Einstein ring).\n", - "\n", - "`source_1` at z=2.0 has a light profile only, and is offset so its lensed images form a distinct second ring." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.ExponentialCoreSph(\n", - " centre=(0.0, 0.0),\n", - " intensity=1.2,\n", - " effective_radius=0.1,\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.25),\n", - ")\n", - "\n", - "source_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " bulge=al.lp.ExponentialCoreSph(\n", - " centre=(-0.3, 0.3),\n", - " intensity=0.6,\n", - " effective_radius=0.07,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The tracer is composed of the main lens galaxies followed by the two source galaxies. PyAutoLens orders galaxies\n", - "internally by redshift, so the multi-plane deflection chain runs:\n", - "\n", - " image-plane \u2192 source_0-plane (deflected by both main lens galaxies)\n", - " \u2192 source_1-plane (deflected by both main lens galaxies AND source_0's mass)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=main_lens_galaxies + [source_0, source_1])\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Group DSPL Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Pass the simulator a tracer to produce the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Write the simulated dataset to .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer JSON__\n", - "\n", - "Save the simulator `Tracer` so the true profiles can be inspected later." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres JSON__\n", - "\n", - "Save the main lens centres so the modeling scripts can load them via `al.from_json(...)`. This mirrors the\n", - "canonical `group/simulator.py` convention." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Group Double Einstein Ring\n", + "=====================================\n", + "\n", + "A double Einstein ring lens in a group-scale context, where two source galaxies at different redshifts are lensed\n", + "by multiple main lens galaxies at the lens-plane redshift.\n", + "\n", + "This script simulates an `Imaging` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - There are TWO main lens galaxies at z=0.5, each with a `SersicSph` light profile and an `IsothermalSph` mass\n", + " profile.\n", + " - The first source galaxy at z=1.0 has a `SersicSph` light profile AND an `IsothermalSph` mass profile. This\n", + " source acts as both a light source AND a deflector for the second source.\n", + " - The second source galaxy at z=2.0 has a `SersicSph` light profile only.\n", + "\n", + "The multi-plane ray-tracing accounts for all main lens galaxy masses at the lens redshift, plus the intermediate\n", + "source galaxy's mass acting as a secondary lens for the more distant source.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** The dataset name and output folder.\n", + "- **Grid:** Define the 2D image-plane grid.\n", + "- **Galaxy Centres:** Centres of the two main lens galaxies, saved as JSON for the modeling scripts to load.\n", + "- **Over Sampling:** Adaptive over-sampling at the main lens galaxy centres.\n", + "- **Main Lens Galaxies:** Two galaxies at z=0.5.\n", + "- **Source Galaxies:** source_0 at z=1.0 (light + mass), source_1 at z=2.0 (light only).\n", + "- **Ray Tracing:** Build the multi-plane Tracer.\n", + "- **Dataset:** Simulate the imaging dataset and write .fits.\n", + "- **Tracer JSON:** Save the simulator Tracer for provenance.\n", + "- **Centres JSON:** Save the main lens centres for the modeling scripts.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to:\n", + "\n", + " - `autolens_workspace/scripts/group/simulator.py` \u2014 the canonical group-scale simulator.\n", + " - `autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/simulator.py` \u2014 the single-lens\n", + " double Einstein ring simulator." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The dataset is written to `dataset/group/double_einstein_ring/`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "A 200x200 grid at 0.1\"/px gives a 20\" field of view, large enough to contain a group-scale double Einstein ring\n", + "where the two main lens galaxies are separated by ~1.5\"." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(200, 200),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "The two main lens galaxies are separated by ~1.5\" along the x-axis. These centres are saved as JSON so the\n", + "modeling and fit scripts can load them via `al.from_json(...)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [(0.0, -0.75), (0.0, 0.75)]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling is applied at each main lens galaxy's centre and at the origin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=main_lens_centres,\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Imaging simulator: exposure time, background sky, noise, PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "Two `SersicSph` + `IsothermalSph` galaxies at z=0.5. Their Einstein radii are chosen so the combined deflection\n", + "produces a clearly visible primary Einstein ring around the group's centroid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=main_lens_centres[0],\n", + " intensity=0.7,\n", + " effective_radius=1.0,\n", + " sersic_index=4.0,\n", + " ),\n", + " mass=al.mp.IsothermalSph(\n", + " centre=main_lens_centres[0],\n", + " einstein_radius=1.2,\n", + " ),\n", + ")\n", + "\n", + "lens_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=main_lens_centres[1],\n", + " intensity=0.7,\n", + " effective_radius=1.0,\n", + " sersic_index=4.0,\n", + " ),\n", + " mass=al.mp.IsothermalSph(\n", + " centre=main_lens_centres[1],\n", + " einstein_radius=1.2,\n", + " ),\n", + ")\n", + "\n", + "main_lens_galaxies = [lens_0, lens_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxies__\n", + "\n", + "`source_0` at z=1.0 has BOTH a light profile (its light forms the primary Einstein ring) AND a mass profile\n", + "(its mass deflects the light from `source_1` and contributes to the second Einstein ring).\n", + "\n", + "`source_1` at z=2.0 has a light profile only, and is offset so its lensed images form a distinct second ring." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.ExponentialCoreSph(\n", + " centre=(0.0, 0.0),\n", + " intensity=1.2,\n", + " effective_radius=0.1,\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.25),\n", + ")\n", + "\n", + "source_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " bulge=al.lp.ExponentialCoreSph(\n", + " centre=(-0.3, 0.3),\n", + " intensity=0.6,\n", + " effective_radius=0.07,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The tracer is composed of the main lens galaxies followed by the two source galaxies. PyAutoLens orders galaxies\n", + "internally by redshift, so the multi-plane deflection chain runs:\n", + "\n", + " image-plane \u2192 source_0-plane (deflected by both main lens galaxies)\n", + " \u2192 source_1-plane (deflected by both main lens galaxies AND source_0's mass)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=main_lens_galaxies + [source_0, source_1])\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Group DSPL Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Pass the simulator a tracer to produce the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Write the simulated dataset to .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer JSON__\n", + "\n", + "Save the simulator `Tracer` so the true profiles can be inspected later." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres JSON__\n", + "\n", + "Save the main lens centres so the modeling scripts can load them via `al.from_json(...)`. This mirrors the\n", + "canonical `group/simulator.py` convention." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/double_einstein_ring/slam.ipynb b/notebooks/group/features/advanced/double_einstein_ring/slam.ipynb index 6f4484f13..466dcabc7 100644 --- a/notebooks/group/features/advanced/double_einstein_ring/slam.ipynb +++ b/notebooks/group/features/advanced/double_einstein_ring/slam.ipynb @@ -1,878 +1,915 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM (Source, Light and Mass): Group Double Einstein Ring\n", - "=========================================================\n", - "\n", - "This script adapts the SLaM (Source, Light and Mass) pipelines to a group-scale Double Source Plane Lens (DSPL)\n", - "system, where two source galaxies at different redshifts behind multiple main lens galaxies form a double\n", - "Einstein ring.\n", - "\n", - "This script is the group + DSPL analogue of `guides/modeling/slam_start_here.py` and\n", - "`scripts/imaging/features/advanced/double_einstein_ring/slam.py`. Each pipeline stage is a plain inline Python\n", - "function, priors are chained via `al.util.chaining.mass_from`, image positions are derived automatically via\n", - "`positions_likelihood_from`, and MGE light profiles are constructed via `al.model_util.mge_model_from`.\n", - "\n", - "__Group + DSPL-Specific Differences From Standard SLaM__\n", - "\n", - " - The lens-plane (z=0.5) is composed via the group `lens_dict` convention: one `af.Model(al.Galaxy)` entry per\n", - " main lens galaxy centre, with the `ExternalShear` attached only to `lens_0`.\n", - " - There are two source galaxies (`source_0` at redshift 1.0, `source_1` at redshift 2.0). `source_0` is a light\n", - " source AND a mass deflector for `source_1`.\n", - " - The SOURCE LP PIPELINE is split into two searches: the first fits all main lens galaxies + `source_0` only;\n", - " the second frees `source_0`'s mass and adds `source_1`'s light.\n", - " - The SOURCE PIX PIPELINE has an extra search: one pixelizes `source_0` while `source_1` is a bare ray-tracing\n", - " galaxy, and the next pixelizes `source_1` with `source_0`'s mass fixed from the previous search.\n", - " - Two `PositionsLH` likelihoods are used once both sources are active, one per source-plane redshift.\n", - " - Adapt images are stitched across pipeline stages and across multiple main lens galaxies.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE and SOURCE PIX PIPELINE, this DSPL group SLaM modeling script fits an `Imaging`\n", - "dataset of a group-scale double Einstein ring system where in the final model:\n", - "\n", - " - Each main lens galaxy's light is a bulge with an MGE light profile.\n", - " - Each main lens galaxy's total mass distribution is an `Isothermal`. `lens_0` carries an `ExternalShear`.\n", - " - The first source galaxy's light is a `Pixelization` and its mass is an `Isothermal`.\n", - " - The second source galaxy's light is a `Pixelization`.\n", - "\n", - "Optional LIGHT LP and MASS TOTAL stages are left as a follow-up exercise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Helpers__\n", - "\n", - "`build_lens_dict_model` constructs a dictionary of `af.Model(al.Galaxy)` entries \u2014 one per main lens galaxy\n", - "centre \u2014 with each galaxy's bulge built from `mge_model_from` and a free `Isothermal` mass profile. The shear\n", - "is added to `lens_0` only, mirroring the canonical group convention.\n", - "\n", - "`lens_dict_instance` and `lens_dict_model` extract per-lens result objects from a search result for forward\n", - "passing to the next stage." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def build_lens_dict_model(\n", - " main_lens_centres,\n", - " redshift_lens: float,\n", - " mask_radius: float,\n", - " total_gaussians: int = 20,\n", - " gaussian_per_basis: int = 2,\n", - "):\n", - " lens_dict = {}\n", - " for i, centre in enumerate(main_lens_centres):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=total_gaussians,\n", - " gaussian_per_basis=gaussian_per_basis,\n", - " centre_prior_is_uniform=True,\n", - " centre=(centre[0], centre[1]),\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = (centre[0], centre[1])\n", - "\n", - " kwargs = dict(\n", - " redshift=redshift_lens,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = af.Model(al.mp.ExternalShear)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - " return lens_dict\n", - "\n", - "\n", - "def lens_dict_from_result(result, n_lens: int, *, attr: str = \"instance\"):\n", - " \"\"\"Return a dict of lens-galaxy result objects (instance or model) keyed by `lens_{i}`.\"\"\"\n", - " galaxies = getattr(result, attr).galaxies\n", - " return {f\"lens_{i}\": getattr(galaxies, f\"lens_{i}\") for i in range(n_lens)}\n", - "\n", - "\n", - "def galaxy_image_dict_for_all_lenses(result, n_lens: int):\n", - " \"\"\"Stitch per-lens galaxy-name image-dict entries from `galaxy_name_image_dict_via_result_from`.\"\"\"\n", - " full = al.galaxy_name_image_dict_via_result_from(result=result)\n", - " out = {}\n", - " for i in range(n_lens):\n", - " key = f\"('galaxies', 'lens_{i}')\"\n", - " out[key] = full[key]\n", - " return out\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 1__\n", - "\n", - "The first SOURCE LP PIPELINE search initializes a model where `source_1` is ignored and only the main lens\n", - "galaxies and `source_0` are fit. This single-plane fit provides robust initial priors for each main lens\n", - "galaxy's light, mass, the external shear on `lens_0`, and `source_0`'s light before the more complex DSPL model\n", - "is introduced.\n", - "\n", - "Model:\n", - " - Per main lens galaxy: MGE bulge (2 x 20 Gaussians), `Isothermal` mass. `lens_0` also has an `ExternalShear`.\n", - " - `source_0` light: MGE with 1 x 20 Gaussians.\n", - " - `source_1`: absent." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " main_lens_centres,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source_0: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_dict = build_lens_dict_model(\n", - " main_lens_centres=main_lens_centres,\n", - " redshift_lens=redshift_lens,\n", - " mask_radius=mask_radius,\n", - " )\n", - "\n", - " source_0_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " **lens_dict,\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source_0,\n", - " bulge=source_0_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 2__\n", - "\n", - "The second SOURCE LP PIPELINE search introduces `source_1`. Each main lens galaxy's bulge, mass and (for\n", - "`lens_0`) shear are fixed to the search-1 instance values, and `source_0`'s light is also fixed. New free\n", - "parameters:\n", - "\n", - " - `source_0`'s mass: `Isothermal` with a narrow prior centred near the origin (the first source typically sits\n", - " close to the lens centroid).\n", - " - `source_1`'s light: MGE with 1 x 20 Gaussians.\n", - "\n", - "A `PositionsLH` for `source_0` is attached to the analysis to prevent unphysical mass models during the DSPL\n", - "fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result_1: af.Result,\n", - " mask_radius: float,\n", - " redshift_source_1: float,\n", - " n_batch: int = 30,\n", - ") -> af.Result:\n", - " n_lens = len(\n", - " [\n", - " name\n", - " for name in vars(source_lp_result_1.instance.galaxies)\n", - " if name.startswith(\"lens_\")\n", - " ]\n", - " )\n", - "\n", - " positions_likelihood_source_0 = source_lp_result_1.positions_likelihood_from(\n", - " factor=3.0,\n", - " minimum_threshold=0.3,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[positions_likelihood_source_0],\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lens):\n", - " lens_inst = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", - " kwargs = dict(\n", - " redshift=lens_inst.redshift,\n", - " bulge=lens_inst.bulge,\n", - " mass=lens_inst.mass,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = lens_inst.shear\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - "\n", - " source_0_mass = af.Model(al.mp.Isothermal)\n", - " source_0_mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - " source_0_mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - " source_0_mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - " source_1_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " **lens_dict,\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.source_0.redshift,\n", - " bulge=source_lp_result_1.instance.galaxies.source_0.bulge,\n", - " mass=source_0_mass,\n", - " ),\n", - " source_1=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source_1,\n", - " bulge=source_1_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1 \u2014 source_0__\n", - "\n", - "Pixelizes `source_0` while `source_1` is present only as a bare galaxy for ray-tracing purposes. Each main\n", - "lens galaxy's mass is freed with priors initialized from the SOURCE LP PIPELINE result 2, and `source_1` is not\n", - "fit (no light, no mass).\n", - "\n", - "Adapt images for each main lens galaxy come from the SOURCE LP PIPELINE result 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1_source_0(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result_1: af.Result,\n", - " source_lp_result_2: af.Result,\n", - " redshift_source_1: float,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " n_lens = len(\n", - " [\n", - " name\n", - " for name in vars(source_lp_result_2.instance.galaxies)\n", - " if name.startswith(\"lens_\")\n", - " ]\n", - " )\n", - "\n", - " galaxy_name_image_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result_1\n", - " )\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", - "\n", - " positions_likelihood_source_0 = source_lp_result_1.positions_likelihood_from(\n", - " factor=3.0,\n", - " minimum_threshold=0.2,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[positions_likelihood_source_0],\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lens):\n", - " lens_inst = getattr(source_lp_result_2.instance.galaxies, f\"lens_{i}\")\n", - " lens_model_mass = getattr(source_lp_result_2.model.galaxies, f\"lens_{i}\").mass\n", - " mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.Isothermal),\n", - " mass_result=lens_model_mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " kwargs = dict(\n", - " redshift=lens_inst.redshift,\n", - " bulge=lens_inst.bulge,\n", - " mass=mass,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = source_lp_result_2.model.galaxies.lens_0.shear\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " **lens_dict,\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " source_1=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source_1,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]_source_0\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1 \u2014 source_1__\n", - "\n", - "Pixelizes `source_1`. `source_0`'s mass is freed with priors initialized from the previous pixelized search.\n", - "Each main lens galaxy's mass is fixed from `source_pix_result_1_source_0`.\n", - "\n", - "Two `PositionsLH` are attached \u2014 one per source plane \u2014 to prevent unphysical reconstructions.\n", - "\n", - "Adapt images are stitched: per-lens adapt images come from the LP pipeline result 2; `source_0`'s adapt image\n", - "comes from the pixelized search above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1_source_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result_2: af.Result,\n", - " source_pix_result_1_source_0: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " n_lens = len(\n", - " [\n", - " name\n", - " for name in vars(source_lp_result_2.instance.galaxies)\n", - " if name.startswith(\"lens_\")\n", - " ]\n", - " )\n", - "\n", - " lp2_dict = al.galaxy_name_image_dict_via_result_from(result=source_lp_result_2)\n", - " pix0_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1_source_0\n", - " )\n", - "\n", - " galaxy_name_image_dict = {}\n", - " for i in range(n_lens):\n", - " key = f\"('galaxies', 'lens_{i}')\"\n", - " galaxy_name_image_dict[key] = lp2_dict[key]\n", - " galaxy_name_image_dict[\"('galaxies', 'source_0')\"] = pix0_dict[\n", - " \"('galaxies', 'source_0')\"\n", - " ]\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", - "\n", - " positions_likelihood_source_0 = (\n", - " source_pix_result_1_source_0.positions_likelihood_from(\n", - " factor=3.0,\n", - " minimum_threshold=0.2,\n", - " plane_redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", - " )\n", - " )\n", - " positions_likelihood_source_1 = source_lp_result_2.positions_likelihood_from(\n", - " factor=3.0,\n", - " minimum_threshold=0.2,\n", - " plane_redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " positions_likelihood_source_0,\n", - " positions_likelihood_source_1,\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " source_0_mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.Isothermal),\n", - " mass_result=source_lp_result_2.model.galaxies.source_0.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lens):\n", - " lens_inst_lp = getattr(source_lp_result_2.instance.galaxies, f\"lens_{i}\")\n", - " lens_inst_pix = getattr(\n", - " source_pix_result_1_source_0.instance.galaxies, f\"lens_{i}\"\n", - " )\n", - "\n", - " kwargs = dict(\n", - " redshift=lens_inst_lp.redshift,\n", - " bulge=lens_inst_lp.bulge,\n", - " mass=lens_inst_pix.mass,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = lens_inst_pix.shear\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " **lens_dict,\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", - " mass=source_0_mass,\n", - " ),\n", - " source_1=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]_source_1\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "The final SOURCE PIX PIPELINE search fits both source galaxies simultaneously with adaptive pixelizations.\n", - "Per-lens mass, shear (on `lens_0`) and `source_0`'s mass are all fixed to the maximum-likelihood instances of\n", - "the previous pixelized searches; only the pixelization regularization parameters are free." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result_2: af.Result,\n", - " source_pix_result_1_source_0: af.Result,\n", - " source_pix_result_1_source_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " n_lens = len(\n", - " [\n", - " name\n", - " for name in vars(source_lp_result_2.instance.galaxies)\n", - " if name.startswith(\"lens_\")\n", - " ]\n", - " )\n", - "\n", - " lp2_dict = al.galaxy_name_image_dict_via_result_from(result=source_lp_result_2)\n", - " pix0_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1_source_0\n", - " )\n", - " pix1_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1_source_1\n", - " )\n", - "\n", - " galaxy_name_image_dict = {}\n", - " for i in range(n_lens):\n", - " key = f\"('galaxies', 'lens_{i}')\"\n", - " galaxy_name_image_dict[key] = lp2_dict[key]\n", - " galaxy_name_image_dict[\"('galaxies', 'source_0')\"] = pix0_dict[\n", - " \"('galaxies', 'source_0')\"\n", - " ]\n", - " galaxy_name_image_dict[\"('galaxies', 'source_1')\"] = pix1_dict[\n", - " \"('galaxies', 'source_1')\"\n", - " ]\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lens):\n", - " lens_inst_lp = getattr(source_lp_result_2.instance.galaxies, f\"lens_{i}\")\n", - " lens_inst_pix = getattr(\n", - " source_pix_result_1_source_1.instance.galaxies, f\"lens_{i}\"\n", - " )\n", - "\n", - " kwargs = dict(\n", - " redshift=lens_inst_lp.redshift,\n", - " bulge=lens_inst_lp.bulge,\n", - " mass=lens_inst_pix.mass,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = lens_inst_pix.shear\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " **lens_dict,\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", - " mass=source_pix_result_1_source_1.instance.galaxies.source_0.mass,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " source_1=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the group double Einstein ring `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "\n", - "mask_radius = 4.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"group\") / \"slam_dspl\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens-plane and the two source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source_0 = 1.0\n", - "redshift_source_1 = 2.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The pixelization mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below runs the full DSPL group SLaM pipeline. See the docstring above each function for a description\n", - "of each stage." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result_1 = source_lp_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " main_lens_centres=main_lens_centres,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source_0=redshift_source_0,\n", - ")\n", - "\n", - "source_lp_result_2 = source_lp_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result_1=source_lp_result_1,\n", - " mask_radius=mask_radius,\n", - " redshift_source_1=redshift_source_1,\n", - ")\n", - "\n", - "source_pix_result_1_source_0 = source_pix_1_source_0(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result_1=source_lp_result_1,\n", - " source_lp_result_2=source_lp_result_2,\n", - " redshift_source_1=redshift_source_1,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "source_pix_result_1_source_1 = source_pix_1_source_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result_2=source_lp_result_2,\n", - " source_pix_result_1_source_0=source_pix_result_1_source_0,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result_2=source_lp_result_2,\n", - " source_pix_result_1_source_0=source_pix_result_1_source_0,\n", - " source_pix_result_1_source_1=source_pix_result_1_source_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM (Source, Light and Mass): Group Double Einstein Ring\n", + "=========================================================\n", + "\n", + "This script adapts the SLaM (Source, Light and Mass) pipelines to a group-scale Double Source Plane Lens (DSPL)\n", + "system, where two source galaxies at different redshifts behind multiple main lens galaxies form a double\n", + "Einstein ring.\n", + "\n", + "This script is the group + DSPL analogue of `guides/modeling/slam_start_here.py` and\n", + "`scripts/imaging/features/advanced/double_einstein_ring/slam.py`. Each pipeline stage is a plain inline Python\n", + "function, priors are chained via `al.util.chaining.mass_from`, image positions are derived automatically via\n", + "`positions_likelihood_from`, and MGE light profiles are constructed via `al.model_util.mge_model_from`.\n", + "\n", + "__Group + DSPL-Specific Differences From Standard SLaM__\n", + "\n", + " - The lens-plane (z=0.5) is composed via the group `lens_dict` convention: one `af.Model(al.Galaxy)` entry per\n", + " main lens galaxy centre, with the `ExternalShear` attached only to `lens_0`.\n", + " - There are two source galaxies (`source_0` at redshift 1.0, `source_1` at redshift 2.0). `source_0` is a light\n", + " source AND a mass deflector for `source_1`.\n", + " - The SOURCE LP PIPELINE is split into two searches: the first fits all main lens galaxies + `source_0` only;\n", + " the second frees `source_0`'s mass and adds `source_1`'s light.\n", + " - The SOURCE PIX PIPELINE has an extra search: one pixelizes `source_0` while `source_1` is a bare ray-tracing\n", + " galaxy, and the next pixelizes `source_1` with `source_0`'s mass fixed from the previous search.\n", + " - Two `PositionsLH` likelihoods are used once both sources are active, one per source-plane redshift.\n", + " - Adapt images are stitched across pipeline stages and across multiple main lens galaxies.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE and SOURCE PIX PIPELINE, this DSPL group SLaM modeling script fits an `Imaging`\n", + "dataset of a group-scale double Einstein ring system where in the final model:\n", + "\n", + " - Each main lens galaxy's light is a bulge with an MGE light profile.\n", + " - Each main lens galaxy's total mass distribution is an `Isothermal`. `lens_0` carries an `ExternalShear`.\n", + " - The first source galaxy's light is a `Pixelization` and its mass is an `Isothermal`.\n", + " - The second source galaxy's light is a `Pixelization`.\n", + "\n", + "Optional LIGHT LP and MASS TOTAL stages are left as a follow-up exercise." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Helpers__\n", + "\n", + "`build_lens_dict_model` constructs a dictionary of `af.Model(al.Galaxy)` entries \u2014 one per main lens galaxy\n", + "centre \u2014 with each galaxy's bulge built from `mge_model_from` and a free `Isothermal` mass profile. The shear\n", + "is added to `lens_0` only, mirroring the canonical group convention.\n", + "\n", + "`lens_dict_instance` and `lens_dict_model` extract per-lens result objects from a search result for forward\n", + "passing to the next stage." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def build_lens_dict_model(\n", + " main_lens_centres,\n", + " redshift_lens: float,\n", + " mask_radius: float,\n", + " total_gaussians: int = 20,\n", + " gaussian_per_basis: int = 2,\n", + "):\n", + " lens_dict = {}\n", + " for i, centre in enumerate(main_lens_centres):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=total_gaussians,\n", + " gaussian_per_basis=gaussian_per_basis,\n", + " centre_prior_is_uniform=True,\n", + " centre=(centre[0], centre[1]),\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = (centre[0], centre[1])\n", + "\n", + " kwargs = dict(\n", + " redshift=redshift_lens,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = af.Model(al.mp.ExternalShear)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + " return lens_dict\n", + "\n", + "\n", + "def lens_dict_from_result(result, n_lens: int, *, attr: str = \"instance\"):\n", + " \"\"\"Return a dict of lens-galaxy result objects (instance or model) keyed by `lens_{i}`.\"\"\"\n", + " galaxies = getattr(result, attr).galaxies\n", + " return {f\"lens_{i}\": getattr(galaxies, f\"lens_{i}\") for i in range(n_lens)}\n", + "\n", + "\n", + "def galaxy_image_dict_for_all_lenses(result, n_lens: int):\n", + " \"\"\"Stitch per-lens galaxy-name image-dict entries from `galaxy_name_image_dict_via_result_from`.\"\"\"\n", + " full = al.galaxy_name_image_dict_via_result_from(result=result)\n", + " out = {}\n", + " for i in range(n_lens):\n", + " key = f\"('galaxies', 'lens_{i}')\"\n", + " out[key] = full[key]\n", + " return out\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 1__\n", + "\n", + "The first SOURCE LP PIPELINE search initializes a model where `source_1` is ignored and only the main lens\n", + "galaxies and `source_0` are fit. This single-plane fit provides robust initial priors for each main lens\n", + "galaxy's light, mass, the external shear on `lens_0`, and `source_0`'s light before the more complex DSPL model\n", + "is introduced.\n", + "\n", + "Model:\n", + " - Per main lens galaxy: MGE bulge (2 x 20 Gaussians), `Isothermal` mass. `lens_0` also has an `ExternalShear`.\n", + " - `source_0` light: MGE with 1 x 20 Gaussians.\n", + " - `source_1`: absent." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " main_lens_centres,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source_0: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_dict = build_lens_dict_model(\n", + " main_lens_centres=main_lens_centres,\n", + " redshift_lens=redshift_lens,\n", + " mask_radius=mask_radius,\n", + " )\n", + "\n", + " source_0_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " **lens_dict,\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source_0,\n", + " bulge=source_0_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 2__\n", + "\n", + "The second SOURCE LP PIPELINE search introduces `source_1`. Each main lens galaxy's bulge, mass and (for\n", + "`lens_0`) shear are fixed to the search-1 instance values, and `source_0`'s light is also fixed. New free\n", + "parameters:\n", + "\n", + " - `source_0`'s mass: `Isothermal` with a narrow prior centred near the origin (the first source typically sits\n", + " close to the lens centroid).\n", + " - `source_1`'s light: MGE with 1 x 20 Gaussians.\n", + "\n", + "A `PositionsLH` for `source_0` is attached to the analysis to prevent unphysical mass models during the DSPL\n", + "fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result_1: af.Result,\n", + " mask_radius: float,\n", + " redshift_source_1: float,\n", + " n_batch: int = 30,\n", + ") -> af.Result:\n", + " n_lens = len(\n", + " [\n", + " name\n", + " for name in vars(source_lp_result_1.instance.galaxies)\n", + " if name.startswith(\"lens_\")\n", + " ]\n", + " )\n", + "\n", + " positions_likelihood_source_0 = source_lp_result_1.positions_likelihood_from(\n", + " factor=3.0,\n", + " minimum_threshold=0.3,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[positions_likelihood_source_0],\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lens):\n", + " lens_inst = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", + " kwargs = dict(\n", + " redshift=lens_inst.redshift,\n", + " bulge=lens_inst.bulge,\n", + " mass=lens_inst.mass,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = lens_inst.shear\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + "\n", + " source_0_mass = af.Model(al.mp.Isothermal)\n", + " source_0_mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + " source_0_mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + " source_0_mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + " source_1_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " **lens_dict,\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.source_0.redshift,\n", + " bulge=source_lp_result_1.instance.galaxies.source_0.bulge,\n", + " mass=source_0_mass,\n", + " ),\n", + " source_1=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source_1,\n", + " bulge=source_1_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1 \u2014 source_0__\n", + "\n", + "Pixelizes `source_0` while `source_1` is present only as a bare galaxy for ray-tracing purposes. Each main\n", + "lens galaxy's mass is freed with priors initialized from the SOURCE LP PIPELINE result 2, and `source_1` is not\n", + "fit (no light, no mass).\n", + "\n", + "Adapt images for each main lens galaxy come from the SOURCE LP PIPELINE result 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1_source_0(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result_1: af.Result,\n", + " source_lp_result_2: af.Result,\n", + " redshift_source_1: float,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " n_lens = len(\n", + " [\n", + " name\n", + " for name in vars(source_lp_result_2.instance.galaxies)\n", + " if name.startswith(\"lens_\")\n", + " ]\n", + " )\n", + "\n", + " galaxy_name_image_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result_1\n", + " )\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", + "\n", + " positions_likelihood_source_0 = source_lp_result_1.positions_likelihood_from(\n", + " factor=3.0,\n", + " minimum_threshold=0.2,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[positions_likelihood_source_0],\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lens):\n", + " lens_inst = getattr(source_lp_result_2.instance.galaxies, f\"lens_{i}\")\n", + " lens_model_mass = getattr(source_lp_result_2.model.galaxies, f\"lens_{i}\").mass\n", + " mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.Isothermal),\n", + " mass_result=lens_model_mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " kwargs = dict(\n", + " redshift=lens_inst.redshift,\n", + " bulge=lens_inst.bulge,\n", + " mass=mass,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = source_lp_result_2.model.galaxies.lens_0.shear\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " **lens_dict,\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " source_1=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source_1,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]_source_0\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1 \u2014 source_1__\n", + "\n", + "Pixelizes `source_1`. `source_0`'s mass is freed with priors initialized from the previous pixelized search.\n", + "Each main lens galaxy's mass is fixed from `source_pix_result_1_source_0`.\n", + "\n", + "Two `PositionsLH` are attached \u2014 one per source plane \u2014 to prevent unphysical reconstructions.\n", + "\n", + "Adapt images are stitched: per-lens adapt images come from the LP pipeline result 2; `source_0`'s adapt image\n", + "comes from the pixelized search above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1_source_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result_2: af.Result,\n", + " source_pix_result_1_source_0: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " n_lens = len(\n", + " [\n", + " name\n", + " for name in vars(source_lp_result_2.instance.galaxies)\n", + " if name.startswith(\"lens_\")\n", + " ]\n", + " )\n", + "\n", + " lp2_dict = al.galaxy_name_image_dict_via_result_from(result=source_lp_result_2)\n", + " pix0_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1_source_0\n", + " )\n", + "\n", + " galaxy_name_image_dict = {}\n", + " for i in range(n_lens):\n", + " key = f\"('galaxies', 'lens_{i}')\"\n", + " galaxy_name_image_dict[key] = lp2_dict[key]\n", + " galaxy_name_image_dict[\"('galaxies', 'source_0')\"] = pix0_dict[\n", + " \"('galaxies', 'source_0')\"\n", + " ]\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", + "\n", + " positions_likelihood_source_0 = (\n", + " source_pix_result_1_source_0.positions_likelihood_from(\n", + " factor=3.0,\n", + " minimum_threshold=0.2,\n", + " plane_redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", + " )\n", + " )\n", + " positions_likelihood_source_1 = source_lp_result_2.positions_likelihood_from(\n", + " factor=3.0,\n", + " minimum_threshold=0.2,\n", + " plane_redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " positions_likelihood_source_0,\n", + " positions_likelihood_source_1,\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " source_0_mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.Isothermal),\n", + " mass_result=source_lp_result_2.model.galaxies.source_0.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lens):\n", + " lens_inst_lp = getattr(source_lp_result_2.instance.galaxies, f\"lens_{i}\")\n", + " lens_inst_pix = getattr(\n", + " source_pix_result_1_source_0.instance.galaxies, f\"lens_{i}\"\n", + " )\n", + "\n", + " kwargs = dict(\n", + " redshift=lens_inst_lp.redshift,\n", + " bulge=lens_inst_lp.bulge,\n", + " mass=lens_inst_pix.mass,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = lens_inst_pix.shear\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " **lens_dict,\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", + " mass=source_0_mass,\n", + " ),\n", + " source_1=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]_source_1\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "The final SOURCE PIX PIPELINE search fits both source galaxies simultaneously with adaptive pixelizations.\n", + "Per-lens mass, shear (on `lens_0`) and `source_0`'s mass are all fixed to the maximum-likelihood instances of\n", + "the previous pixelized searches; only the pixelization regularization parameters are free." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result_2: af.Result,\n", + " source_pix_result_1_source_0: af.Result,\n", + " source_pix_result_1_source_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " n_lens = len(\n", + " [\n", + " name\n", + " for name in vars(source_lp_result_2.instance.galaxies)\n", + " if name.startswith(\"lens_\")\n", + " ]\n", + " )\n", + "\n", + " lp2_dict = al.galaxy_name_image_dict_via_result_from(result=source_lp_result_2)\n", + " pix0_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1_source_0\n", + " )\n", + " pix1_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1_source_1\n", + " )\n", + "\n", + " galaxy_name_image_dict = {}\n", + " for i in range(n_lens):\n", + " key = f\"('galaxies', 'lens_{i}')\"\n", + " galaxy_name_image_dict[key] = lp2_dict[key]\n", + " galaxy_name_image_dict[\"('galaxies', 'source_0')\"] = pix0_dict[\n", + " \"('galaxies', 'source_0')\"\n", + " ]\n", + " galaxy_name_image_dict[\"('galaxies', 'source_1')\"] = pix1_dict[\n", + " \"('galaxies', 'source_1')\"\n", + " ]\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lens):\n", + " lens_inst_lp = getattr(source_lp_result_2.instance.galaxies, f\"lens_{i}\")\n", + " lens_inst_pix = getattr(\n", + " source_pix_result_1_source_1.instance.galaxies, f\"lens_{i}\"\n", + " )\n", + "\n", + " kwargs = dict(\n", + " redshift=lens_inst_lp.redshift,\n", + " bulge=lens_inst_lp.bulge,\n", + " mass=lens_inst_pix.mass,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = lens_inst_pix.shear\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " **lens_dict,\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", + " mass=source_pix_result_1_source_1.instance.galaxies.source_0.mass,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " source_1=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the group double Einstein ring `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "\n", + "mask_radius = 4.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"group\") / \"slam_dspl\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens-plane and the two source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source_0 = 1.0\n", + "redshift_source_1 = 2.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The pixelization mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below runs the full DSPL group SLaM pipeline. See the docstring above each function for a description\n", + "of each stage." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result_1 = source_lp_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " main_lens_centres=main_lens_centres,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source_0=redshift_source_0,\n", + ")\n", + "\n", + "source_lp_result_2 = source_lp_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result_1=source_lp_result_1,\n", + " mask_radius=mask_radius,\n", + " redshift_source_1=redshift_source_1,\n", + ")\n", + "\n", + "source_pix_result_1_source_0 = source_pix_1_source_0(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result_1=source_lp_result_1,\n", + " source_lp_result_2=source_lp_result_2,\n", + " redshift_source_1=redshift_source_1,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "source_pix_result_1_source_1 = source_pix_1_source_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result_2=source_lp_result_2,\n", + " source_pix_result_1_source_0=source_pix_result_1_source_0,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result_2=source_lp_result_2,\n", + " source_pix_result_1_source_0=source_pix_result_1_source_0,\n", + " source_pix_result_1_source_1=source_pix_result_1_source_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/mass_stellar_dark/chaining.ipynb b/notebooks/group/features/advanced/mass_stellar_dark/chaining.ipynb index 22d7a941c..d1d7980e3 100644 --- a/notebooks/group/features/advanced/mass_stellar_dark/chaining.ipynb +++ b/notebooks/group/features/advanced/mass_stellar_dark/chaining.ipynb @@ -1,374 +1,411 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Chaining: Group Mass Stellar Dark\n", - "=================================\n", - "\n", - "This script chains two searches to fit `Imaging` data of a 'group-scale' strong lens where each main lens\n", - "galaxy is decomposed into stellar + dark matter components.\n", - "\n", - "The two searches break down as follows:\n", - "\n", - " 1) Models each main lens galaxy's light using a `lp.Sersic` bulge (a pure light profile, no stellar mass\n", - " coupling) and the source as an MGE. No dark matter component is included. This fits the geometric and\n", - " photometric properties of each galaxy's bulge with a relatively small parameter space.\n", - "\n", - " 2) Reintroduces the stellar-mass coupling by swapping each galaxy's bulge to a `lmp.Sersic` (a light+mass\n", - " profile), adds an `NFWSph` dark matter halo per galaxy, and adds an `ExternalShear` on `lens_0`. Priors on\n", - " the bulge geometry (centre, ell_comps, intensity, effective_radius, sersic_index) are passed from search\n", - " 1, leaving only the `mass_to_light_ratio` and the dark halo parameters as new free parameters in search 2.\n", - "\n", - "__Why Chain?__\n", - "\n", - "A group-scale decomposed-mass model is hard to fit in a single Nautilus run. With two main lens galaxies plus\n", - "shear, the search 2 model carries:\n", - "\n", - " - 2 x 6 = 12 bulge parameters (linear intensity is solved separately).\n", - " - 2 x 2 = 4 dark NFW parameters (`kappa_s`, `scale_radius` per galaxy).\n", - " - 2 x 1 = 2 mass-to-light ratio parameters.\n", - " - 2 external shear parameters.\n", - " - source MGE parameters.\n", - "\n", - "That parameter space has too many degeneracies for a single search starting from broad priors. Chaining lets\n", - "search 1 lock down the bulge geometry of every main lens galaxy first, leaving search 2 free to focus on the\n", - "mass-to-light ratios, dark NFW parameters, and the shear, with bulge priors tight from search 1.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset + Masking:** Load, plot and mask the `Imaging` data.\n", - "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", - "- **Paths:** The output path for both chained searches.\n", - "- **Model (Search 1):** Per-lens `lp.Sersic` bulge + MGE source. No mass.\n", - "- **Search 1:** Nautilus fit, returns `result_1`.\n", - "- **Model (Search 2):** Per-lens `lmp.Sersic` bulge (priors from `result_1`) + `NFWSph` dark + shear on\n", - " `lens_0` + MGE source.\n", - "- **Search 2:** Nautilus fit, returns `result_2`.\n", - "- **Wrap Up:** Summary and pointer to `slam.py`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to:\n", - "\n", - " - `autolens_workspace/scripts/group/start_here.py` \u2014 the canonical group walkthrough.\n", - " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/chaining.py` \u2014 the single-galaxy\n", - " decomposed-mass chaining walkthrough this script generalises across multiple main lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Centres__\n", - "\n", - "Load the two main lens galaxy centres from JSON. These are fixed on each galaxy's bulge centre in both\n", - "searches." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "\n", - "mask_radius = 3.7\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__\n", - "\n", - "The path where the results of both chained searches are output:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"group\") / \"chaining\" / \"mass_stellar_dark\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Search 1 fits each main lens galaxy's bulge as a pure `lp.Sersic` light profile (no stellar mass coupling,\n", - "no dark NFW). The source is modelled as an MGE. Each bulge's centre is fixed to the main lens centre loaded\n", - "from JSON.\n", - "\n", - "For two main lens galaxies, the lens-plane carries 2 x 6 = 12 bulge parameters (linear intensities solved\n", - "separately). The source MGE adds its own parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict_1 = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " bulge = af.Model(al.lp.Sersic)\n", - " bulge.centre = (centre[0], centre[1])\n", - " lens_dict_1[f\"lens_{i}\"] = af.Model(al.Galaxy, redshift=0.5, bulge=bulge)\n", - "\n", - "source_bulge_1 = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source_1 = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge_1)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(**lens_dict_1, source=source_1))\n", - "\n", - "print(model_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__lens_light\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)\n", - "\n", - "print(result_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "Search 2 reintroduces the stellar-mass coupling. For each main lens galaxy:\n", - "\n", - " - Replace `lp.Sersic` bulge with `lmp.Sersic` (a light AND mass profile). Pass bulge geometric / photometric\n", - " priors from `result_1.model` via `take_attributes` \u2014 `centre`, `ell_comps`, `intensity`,\n", - " `effective_radius`, `sersic_index` all transfer because they share the same names between `lp.Sersic` and\n", - " `lmp.Sersic`. The new parameter introduced by the swap is `mass_to_light_ratio`.\n", - " - Add an `NFWSph` dark matter halo with `centre` fixed to the bulge centre.\n", - " - Add an `ExternalShear` on `lens_0` only.\n", - "\n", - "The source MGE bulge is fixed to its `result_1.instance` value \u2014 search 2 does not re-optimise the source\n", - "geometry, only the lens-plane mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict_2 = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " bulge = af.Model(al.lmp.Sersic)\n", - " bulge.take_attributes(source=getattr(result_1.model.galaxies, f\"lens_{i}\"))\n", - " bulge.centre = (centre[0], centre[1])\n", - "\n", - " dark = af.Model(al.mp.NFWSph)\n", - " dark.centre = (centre[0], centre[1])\n", - "\n", - " galaxy_kwargs = dict(redshift=0.5, bulge=bulge, dark=dark)\n", - "\n", - " if i == 0:\n", - " galaxy_kwargs[\"shear\"] = af.Model(al.mp.ExternalShear)\n", - "\n", - " lens_dict_2[f\"lens_{i}\"] = af.Model(al.Galaxy, **galaxy_kwargs)\n", - "\n", - "source_2 = af.Model(\n", - " al.Galaxy, redshift=1.0, bulge=result_1.instance.galaxies.source.bulge\n", - ")\n", - "\n", - "model_2 = af.Collection(galaxies=af.Collection(**lens_dict_2, source=source_2))\n", - "\n", - "print(model_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 2)__\n", - "\n", - "You may wish to inspect the `model.info` file of the search 2 model-fit to ensure the priors were passed\n", - "correctly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__mass_stellar_dark\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)\n", - "\n", - "print(result_2.info)\n", - "\n", - "aplt.subplot_tracer(tracer=result_2.max_log_likelihood_tracer, grid=result_2.grids.lp)\n", - "aplt.subplot_fit_imaging(fit=result_2.max_log_likelihood_fit)\n", - "aplt.corner_anesthetic(samples=result_2.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "In this example, we passed each main lens galaxy's bulge light model to a per-galaxy decomposed stellar +\n", - "dark matter mass model. Search 1 constrained the bulge geometry of every main lens galaxy from the lens\n", - "light alone; search 2 then used those tight priors as the starting point for the harder mass + shear fit,\n", - "introducing only the `mass_to_light_ratio` and NFW parameters as new free directions.\n", - "\n", - "__SLaM (Source, Light and Mass)__\n", - "\n", - "An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens\n", - "modeling processing into a series of fits that first perfect the source model, then the lens light model and\n", - "finally the lens mass model. See `slam.py` in this directory." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Chaining: Group Mass Stellar Dark\n", + "=================================\n", + "\n", + "This script chains two searches to fit `Imaging` data of a 'group-scale' strong lens where each main lens\n", + "galaxy is decomposed into stellar + dark matter components.\n", + "\n", + "The two searches break down as follows:\n", + "\n", + " 1) Models each main lens galaxy's light using a `lp.Sersic` bulge (a pure light profile, no stellar mass\n", + " coupling) and the source as an MGE. No dark matter component is included. This fits the geometric and\n", + " photometric properties of each galaxy's bulge with a relatively small parameter space.\n", + "\n", + " 2) Reintroduces the stellar-mass coupling by swapping each galaxy's bulge to a `lmp.Sersic` (a light+mass\n", + " profile), adds an `NFWSph` dark matter halo per galaxy, and adds an `ExternalShear` on `lens_0`. Priors on\n", + " the bulge geometry (centre, ell_comps, intensity, effective_radius, sersic_index) are passed from search\n", + " 1, leaving only the `mass_to_light_ratio` and the dark halo parameters as new free parameters in search 2.\n", + "\n", + "__Why Chain?__\n", + "\n", + "A group-scale decomposed-mass model is hard to fit in a single Nautilus run. With two main lens galaxies plus\n", + "shear, the search 2 model carries:\n", + "\n", + " - 2 x 6 = 12 bulge parameters (linear intensity is solved separately).\n", + " - 2 x 2 = 4 dark NFW parameters (`kappa_s`, `scale_radius` per galaxy).\n", + " - 2 x 1 = 2 mass-to-light ratio parameters.\n", + " - 2 external shear parameters.\n", + " - source MGE parameters.\n", + "\n", + "That parameter space has too many degeneracies for a single search starting from broad priors. Chaining lets\n", + "search 1 lock down the bulge geometry of every main lens galaxy first, leaving search 2 free to focus on the\n", + "mass-to-light ratios, dark NFW parameters, and the shear, with bulge priors tight from search 1.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset + Masking:** Load, plot and mask the `Imaging` data.\n", + "- **Main Lens Centres:** Load the two main lens galaxy centres from JSON.\n", + "- **Paths:** The output path for both chained searches.\n", + "- **Model (Search 1):** Per-lens `lp.Sersic` bulge + MGE source. No mass.\n", + "- **Search 1:** Nautilus fit, returns `result_1`.\n", + "- **Model (Search 2):** Per-lens `lmp.Sersic` bulge (priors from `result_1`) + `NFWSph` dark + shear on\n", + " `lens_0` + MGE source.\n", + "- **Search 2:** Nautilus fit, returns `result_2`.\n", + "- **Wrap Up:** Summary and pointer to `slam.py`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to:\n", + "\n", + " - `autolens_workspace/scripts/group/start_here.py` \u2014 the canonical group walkthrough.\n", + " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/chaining.py` \u2014 the single-galaxy\n", + " decomposed-mass chaining walkthrough this script generalises across multiple main lens galaxies." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Centres__\n", + "\n", + "Load the two main lens galaxy centres from JSON. These are fixed on each galaxy's bulge centre in both\n", + "searches." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "\n", + "mask_radius = 3.7\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__\n", + "\n", + "The path where the results of both chained searches are output:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"group\") / \"chaining\" / \"mass_stellar_dark\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Search 1 fits each main lens galaxy's bulge as a pure `lp.Sersic` light profile (no stellar mass coupling,\n", + "no dark NFW). The source is modelled as an MGE. Each bulge's centre is fixed to the main lens centre loaded\n", + "from JSON.\n", + "\n", + "For two main lens galaxies, the lens-plane carries 2 x 6 = 12 bulge parameters (linear intensities solved\n", + "separately). The source MGE adds its own parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict_1 = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " bulge = af.Model(al.lp.Sersic)\n", + " bulge.centre = (centre[0], centre[1])\n", + " lens_dict_1[f\"lens_{i}\"] = af.Model(al.Galaxy, redshift=0.5, bulge=bulge)\n", + "\n", + "source_bulge_1 = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source_1 = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge_1)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(**lens_dict_1, source=source_1))\n", + "\n", + "print(model_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__lens_light\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)\n", + "\n", + "print(result_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "Search 2 reintroduces the stellar-mass coupling. For each main lens galaxy:\n", + "\n", + " - Replace `lp.Sersic` bulge with `lmp.Sersic` (a light AND mass profile). Pass bulge geometric / photometric\n", + " priors from `result_1.model` via `take_attributes` \u2014 `centre`, `ell_comps`, `intensity`,\n", + " `effective_radius`, `sersic_index` all transfer because they share the same names between `lp.Sersic` and\n", + " `lmp.Sersic`. The new parameter introduced by the swap is `mass_to_light_ratio`.\n", + " - Add an `NFWSph` dark matter halo with `centre` fixed to the bulge centre.\n", + " - Add an `ExternalShear` on `lens_0` only.\n", + "\n", + "The source MGE bulge is fixed to its `result_1.instance` value \u2014 search 2 does not re-optimise the source\n", + "geometry, only the lens-plane mass." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict_2 = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " bulge = af.Model(al.lmp.Sersic)\n", + " bulge.take_attributes(source=getattr(result_1.model.galaxies, f\"lens_{i}\"))\n", + " bulge.centre = (centre[0], centre[1])\n", + "\n", + " dark = af.Model(al.mp.NFWSph)\n", + " dark.centre = (centre[0], centre[1])\n", + "\n", + " galaxy_kwargs = dict(redshift=0.5, bulge=bulge, dark=dark)\n", + "\n", + " if i == 0:\n", + " galaxy_kwargs[\"shear\"] = af.Model(al.mp.ExternalShear)\n", + "\n", + " lens_dict_2[f\"lens_{i}\"] = af.Model(al.Galaxy, **galaxy_kwargs)\n", + "\n", + "source_2 = af.Model(\n", + " al.Galaxy, redshift=1.0, bulge=result_1.instance.galaxies.source.bulge\n", + ")\n", + "\n", + "model_2 = af.Collection(galaxies=af.Collection(**lens_dict_2, source=source_2))\n", + "\n", + "print(model_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 2)__\n", + "\n", + "You may wish to inspect the `model.info` file of the search 2 model-fit to ensure the priors were passed\n", + "correctly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__mass_stellar_dark\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)\n", + "\n", + "print(result_2.info)\n", + "\n", + "aplt.subplot_tracer(tracer=result_2.max_log_likelihood_tracer, grid=result_2.grids.lp)\n", + "aplt.subplot_fit_imaging(fit=result_2.max_log_likelihood_fit)\n", + "aplt.corner_anesthetic(samples=result_2.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "In this example, we passed each main lens galaxy's bulge light model to a per-galaxy decomposed stellar +\n", + "dark matter mass model. Search 1 constrained the bulge geometry of every main lens galaxy from the lens\n", + "light alone; search 2 then used those tight priors as the starting point for the harder mass + shear fit,\n", + "introducing only the `mass_to_light_ratio` and NFW parameters as new free directions.\n", + "\n", + "__SLaM (Source, Light and Mass)__\n", + "\n", + "An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens\n", + "modeling processing into a series of fits that first perfect the source model, then the lens light model and\n", + "finally the lens mass model. See `slam.py` in this directory." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/mass_stellar_dark/fit.ipynb b/notebooks/group/features/advanced/mass_stellar_dark/fit.ipynb index e9f4efc5c..00ec944cc 100644 --- a/notebooks/group/features/advanced/mass_stellar_dark/fit.ipynb +++ b/notebooks/group/features/advanced/mass_stellar_dark/fit.ipynb @@ -1,527 +1,564 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Group Mass Stellar Dark Fit\n", - "=====================================\n", - "\n", - "A group-scale strong lens where each main lens galaxy carries a decomposed mass model \u2014 a stellar component\n", - "tied to its observed light via a mass-to-light ratio, plus a separately-parameterized dark matter halo. The\n", - "total deflection at every image-plane coordinate is the sum over all main lens galaxies of the per-galaxy\n", - "stellar + dark contributions, plus a single external shear.\n", - "\n", - "This script illustrates the API for performing a fit to a group-scale decomposed-mass lens via the standard\n", - "`Tracer` and `FitImaging` objects, without invoking a non-linear search. It is intended to make the per-galaxy\n", - "deflection decomposition concrete in the group context before the reader moves on to `modeling.py`\n", - "(search-based) or `chaining.py` / `slam.py` (realistic, robust modeling).\n", - "\n", - "The source galaxy is modelled with a Multi Gaussian Expansion (MGE), the same source parameterization used in\n", - "`chaining.py` and `slam.py`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Adaptive over-sampling at the main lens galaxy centres.\n", - "- **MGE Basis:** Build a `Basis` of linear Gaussians for the source.\n", - "- **Main Lens Centres:** Load the centres of the two main lens galaxies from JSON.\n", - "- **Galaxies:** Compose the lens galaxies (via a `lens_dict` loop) plus the MGE source.\n", - "- **Tracer:** Build the two-plane `Tracer` that performs the ray-tracing.\n", - "- **Fit:** Create a `FitImaging` and inspect the fit.\n", - "- **Decomposed Deflection (Multi-Galaxy):** A short tour of how the total lens-plane deflection is the sum,\n", - " over every main lens galaxy, of stellar + dark contributions, plus the single external shear.\n", - "- **Intensities:** The solved-for linear light profile `intensity` values for each MGE Gaussian.\n", - "- **Wrap Up:** Summary and next steps.\n", - "\n", - "__Prerequisites__\n", - "\n", - "This script focuses on the API specific to a group-scale decomposed-mass fit. For background on the underlying\n", - "single-galaxy decomposition and the group `lens_dict` API, you should read first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/fit.py` \u2014 the single-galaxy\n", - " decomposition tour. The Decomposed Deflection section below generalises that walkthrough across multiple\n", - " main lens galaxies.\n", - " - `autolens_workspace/scripts/group/start_here.py` \u2014 the group-scale `lens_dict` API, including how\n", - " `main_lens_centres.json` is loaded and used to drive a per-galaxy loop.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py` \u2014 the MGE `Basis` API for\n", - " the source.\n", - "\n", - "The galaxy redshifts (`lenses=0.5`, `source=1.0`) and per-galaxy mass parameters match those used by the\n", - "simulator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `mass_stellar_dark` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Centres__\n", - "\n", - "Load the centres of the two main lens galaxies from JSON. The same file is used by every other script in this\n", - "directory (modeling.py, chaining.py, slam.py)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.7\" circular mask, which includes both main lens galaxies and the lensed source emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.7\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling at each main lens galaxy centre, so the stellar mass-to-light coupling is\n", - "evaluated accurately at the peak of each bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE Basis__\n", - "\n", - "We build a `Basis` of 30 linear Gaussians as the source-galaxy light model, centred on the simulator's source\n", - "position." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "source_bulge = build_source_basis(centre=(0.0, 0.0))\n", - "\n", - "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "We compose each main lens galaxy via a `lens_dict` loop over `main_lens_centres`. Each galaxy gets:\n", - "\n", - " - a `lmp.Sersic` bulge (acts as light AND stellar mass via `mass_to_light_ratio`),\n", - " - a `NFWSph` dark matter halo aligned with its bulge,\n", - " - and (for the first lens galaxy only) an `ExternalShear` representing the group-wide shear field.\n", - "\n", - "All non-linear parameters are set to the simulator's true values, so the fit visibly recovers the lensing\n", - "configuration without a search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge_params = [\n", - " dict(axis_ratio=0.9, angle=45.0, intensity=1.0, effective_radius=0.8, m_to_l=0.20),\n", - " dict(axis_ratio=0.8, angle=120.0, intensity=0.8, effective_radius=0.7, m_to_l=0.25),\n", - "]\n", - "\n", - "dark_params = [\n", - " dict(kappa_s=0.10, scale_radius=20.0),\n", - " dict(kappa_s=0.08, scale_radius=20.0),\n", - "]\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " galaxy_kwargs = dict(\n", - " redshift=0.5,\n", - " bulge=al.lmp.Sersic(\n", - " centre=(centre[0], centre[1]),\n", - " ell_comps=al.convert.ell_comps_from(\n", - " axis_ratio=bulge_params[i][\"axis_ratio\"],\n", - " angle=bulge_params[i][\"angle\"],\n", - " ),\n", - " intensity=bulge_params[i][\"intensity\"],\n", - " effective_radius=bulge_params[i][\"effective_radius\"],\n", - " sersic_index=4.0,\n", - " mass_to_light_ratio=bulge_params[i][\"m_to_l\"],\n", - " ),\n", - " dark=al.mp.NFWSph(\n", - " centre=(centre[0], centre[1]),\n", - " kappa_s=dark_params[i][\"kappa_s\"],\n", - " scale_radius=dark_params[i][\"scale_radius\"],\n", - " ),\n", - " )\n", - "\n", - " if i == 0:\n", - " galaxy_kwargs[\"shear\"] = al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = al.Galaxy(**galaxy_kwargs)\n", - "\n", - "source = al.Galaxy(redshift=1.0, bulge=source_bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "The `Tracer` performs the ray-tracing. Internally it queries every mass profile attached to every galaxy in\n", - "the lens plane and sums their deflections. For our group lens, this means each `lens_i`'s `bulge` contributes\n", - "a stellar mass deflection (`(M/L)_i * alpha_light_i`), each `dark` halo contributes an `NFWSph` deflection,\n", - "and `lens_0`'s `shear` contributes the external shear \u2014 all summed before mapping image-plane coordinates\n", - "onto the source-plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=list(lens_dict.values()) + [source])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We pass the `Tracer` to a `FitImaging` to fit the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Decomposed Deflection (Multi-Galaxy)__\n", - "\n", - "This is the section that makes the group-scale decomposed-mass fit conceptually distinct. The lens-plane\n", - "deflection map is the SUM, over every main lens galaxy AND every mass component in that galaxy, of independent\n", - "deflection contributions:\n", - "\n", - " alpha_lens(theta) = sum_i [ alpha_stellar_i(theta) + alpha_dark_i(theta) ] + alpha_shear(theta)\n", - " = sum_i [ (M/L)_i * alpha_light_i + alpha_NFW_i ] + alpha_shear\n", - "\n", - "Every individual deflection is a public method on the corresponding profile.\n", - "\n", - "We verify this by computing each contribution explicitly and confirming the sum equals what the `Tracer`\n", - "returns." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = dataset.grid\n", - "\n", - "alpha_stellar_list = [\n", - " lens.bulge.deflections_yx_2d_from(grid=grid) for lens in lens_dict.values()\n", - "]\n", - "alpha_dark_list = [\n", - " lens.dark.deflections_yx_2d_from(grid=grid) for lens in lens_dict.values()\n", - "]\n", - "alpha_shear = lens_dict[\"lens_0\"].shear.deflections_yx_2d_from(grid=grid)\n", - "\n", - "print(f\"alpha_stellar[lens_0] (first coord): {alpha_stellar_list[0][0]}\")\n", - "print(f\"alpha_dark [lens_0] (first coord): {alpha_dark_list[0][0]}\")\n", - "print(f\"alpha_stellar[lens_1] (first coord): {alpha_stellar_list[1][0]}\")\n", - "print(f\"alpha_dark [lens_1] (first coord): {alpha_dark_list[1][0]}\")\n", - "print(f\"alpha_shear (first coord): {alpha_shear[0]}\")\n", - "\n", - "alpha_total_summed = sum(alpha_stellar_list) + sum(alpha_dark_list) + alpha_shear" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The tracer-produced source-plane grid is just `grid - alpha_total_internal`. Recovering the internal total\n", - "deflection from `traced_grid_2d_list_from` and comparing to our hand-summed `alpha_total` is the cleanest\n", - "end-to-end check that we have accounted for every per-galaxy contribution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grids = tracer.traced_grid_2d_list_from(grid=grid)\n", - "alpha_total_tracer = grid - traced_grids[1]\n", - "\n", - "print(f\"alpha_total (summed by hand, first 3): {alpha_total_summed[:3]}\")\n", - "print(f\"alpha_total (from tracer, first 3): {alpha_total_tracer[:3]}\")\n", - "\n", - "assert np.allclose(np.asarray(alpha_total_summed), np.asarray(alpha_total_tracer))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The same component-wise decomposition shows up in the convergence (kappa) map. We sum each component's\n", - "contribution across all main lens galaxies and plot the result, which highlights that the stellar component is\n", - "peaked at the two galaxy centres while the dark halos extend more diffusely." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "kappa_stellar_total = sum(\n", - " lens.bulge.convergence_2d_from(grid=plot_grid) for lens in lens_dict.values()\n", - ")\n", - "kappa_dark_total = sum(\n", - " lens.dark.convergence_2d_from(grid=plot_grid) for lens in lens_dict.values()\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=kappa_stellar_total, title=\"Stellar convergence (sum over galaxies)\"\n", - ")\n", - "aplt.plot_array(\n", - " array=kappa_dark_total, title=\"Dark matter convergence (sum over galaxies)\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "After the fit, every linear Gaussian in the source MGE basis has been assigned an `intensity` via linear\n", - "algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"\\nFirst Gaussian intensity, source = \"\n", - " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", - ")\n", - "\n", - "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer_fitted.galaxies[-1].bulge, grid=plot_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated the group-scale decomposed-mass API and the per-galaxy deflection decomposition,\n", - "without invoking a non-linear search. Each main lens galaxy's `bulge` simultaneously acts as a light profile\n", - "and a stellar mass profile (coupled by its own `mass_to_light_ratio`), and each separately-parameterized\n", - "`dark` NFW halo adds an independent dark mass contribution. A single `ExternalShear` is attached to `lens_0`\n", - "representing the group-wide shear field.\n", - "\n", - "In a real modeling workflow:\n", - "\n", - " - `modeling.py` shows how to fit the same system using `Nautilus`, but \"cheats\" by initialising priors at\n", - " the true values. It is therefore only useful as a tutorial.\n", - " - `chaining.py` is the practical workflow \u2014 two chained searches that fit the lens light first (treating each\n", - " bulge as a pure light profile), then reintroduce the stellar-mass coupling and add the dark NFW halos.\n", - " - `slam.py` is the most robust pipeline for production-quality group-scale decomposed-mass modeling, chaining\n", - " through SOURCE LP, SOURCE PIX, LIGHT LP, and MASS_LIGHT_DARK pipelines and ending in a pixelized source\n", - " reconstruction.\n", - "\n", - "The key takeaway from this script is that a group-scale decomposed-mass lens is fit with the same `Tracer` +\n", - "`FitImaging` objects as any other lens; the only difference is that the lens plane carries MULTIPLE galaxies,\n", - "each with multiple independent mass components, all of whose deflections sum into the total lens-plane\n", - "deflection. The `lens_dict` API scales naturally to any number of main lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Group Mass Stellar Dark Fit\n", + "=====================================\n", + "\n", + "A group-scale strong lens where each main lens galaxy carries a decomposed mass model \u2014 a stellar component\n", + "tied to its observed light via a mass-to-light ratio, plus a separately-parameterized dark matter halo. The\n", + "total deflection at every image-plane coordinate is the sum over all main lens galaxies of the per-galaxy\n", + "stellar + dark contributions, plus a single external shear.\n", + "\n", + "This script illustrates the API for performing a fit to a group-scale decomposed-mass lens via the standard\n", + "`Tracer` and `FitImaging` objects, without invoking a non-linear search. It is intended to make the per-galaxy\n", + "deflection decomposition concrete in the group context before the reader moves on to `modeling.py`\n", + "(search-based) or `chaining.py` / `slam.py` (realistic, robust modeling).\n", + "\n", + "The source galaxy is modelled with a Multi Gaussian Expansion (MGE), the same source parameterization used in\n", + "`chaining.py` and `slam.py`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Adaptive over-sampling at the main lens galaxy centres.\n", + "- **MGE Basis:** Build a `Basis` of linear Gaussians for the source.\n", + "- **Main Lens Centres:** Load the centres of the two main lens galaxies from JSON.\n", + "- **Galaxies:** Compose the lens galaxies (via a `lens_dict` loop) plus the MGE source.\n", + "- **Tracer:** Build the two-plane `Tracer` that performs the ray-tracing.\n", + "- **Fit:** Create a `FitImaging` and inspect the fit.\n", + "- **Decomposed Deflection (Multi-Galaxy):** A short tour of how the total lens-plane deflection is the sum,\n", + " over every main lens galaxy, of stellar + dark contributions, plus the single external shear.\n", + "- **Intensities:** The solved-for linear light profile `intensity` values for each MGE Gaussian.\n", + "- **Wrap Up:** Summary and next steps.\n", + "\n", + "__Prerequisites__\n", + "\n", + "This script focuses on the API specific to a group-scale decomposed-mass fit. For background on the underlying\n", + "single-galaxy decomposition and the group `lens_dict` API, you should read first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/fit.py` \u2014 the single-galaxy\n", + " decomposition tour. The Decomposed Deflection section below generalises that walkthrough across multiple\n", + " main lens galaxies.\n", + " - `autolens_workspace/scripts/group/start_here.py` \u2014 the group-scale `lens_dict` API, including how\n", + " `main_lens_centres.json` is loaded and used to drive a per-galaxy loop.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py` \u2014 the MGE `Basis` API for\n", + " the source.\n", + "\n", + "The galaxy redshifts (`lenses=0.5`, `source=1.0`) and per-galaxy mass parameters match those used by the\n", + "simulator." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `mass_stellar_dark` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Centres__\n", + "\n", + "Load the centres of the two main lens galaxies from JSON. The same file is used by every other script in this\n", + "directory (modeling.py, chaining.py, slam.py)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.7\" circular mask, which includes both main lens galaxies and the lensed source emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.7\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling at each main lens galaxy centre, so the stellar mass-to-light coupling is\n", + "evaluated accurately at the peak of each bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE Basis__\n", + "\n", + "We build a `Basis` of 30 linear Gaussians as the source-galaxy light model, centred on the simulator's source\n", + "position." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "source_bulge = build_source_basis(centre=(0.0, 0.0))\n", + "\n", + "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "We compose each main lens galaxy via a `lens_dict` loop over `main_lens_centres`. Each galaxy gets:\n", + "\n", + " - a `lmp.Sersic` bulge (acts as light AND stellar mass via `mass_to_light_ratio`),\n", + " - a `NFWSph` dark matter halo aligned with its bulge,\n", + " - and (for the first lens galaxy only) an `ExternalShear` representing the group-wide shear field.\n", + "\n", + "All non-linear parameters are set to the simulator's true values, so the fit visibly recovers the lensing\n", + "configuration without a search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge_params = [\n", + " dict(axis_ratio=0.9, angle=45.0, intensity=1.0, effective_radius=0.8, m_to_l=0.20),\n", + " dict(axis_ratio=0.8, angle=120.0, intensity=0.8, effective_radius=0.7, m_to_l=0.25),\n", + "]\n", + "\n", + "dark_params = [\n", + " dict(kappa_s=0.10, scale_radius=20.0),\n", + " dict(kappa_s=0.08, scale_radius=20.0),\n", + "]\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " galaxy_kwargs = dict(\n", + " redshift=0.5,\n", + " bulge=al.lmp.Sersic(\n", + " centre=(centre[0], centre[1]),\n", + " ell_comps=al.convert.ell_comps_from(\n", + " axis_ratio=bulge_params[i][\"axis_ratio\"],\n", + " angle=bulge_params[i][\"angle\"],\n", + " ),\n", + " intensity=bulge_params[i][\"intensity\"],\n", + " effective_radius=bulge_params[i][\"effective_radius\"],\n", + " sersic_index=4.0,\n", + " mass_to_light_ratio=bulge_params[i][\"m_to_l\"],\n", + " ),\n", + " dark=al.mp.NFWSph(\n", + " centre=(centre[0], centre[1]),\n", + " kappa_s=dark_params[i][\"kappa_s\"],\n", + " scale_radius=dark_params[i][\"scale_radius\"],\n", + " ),\n", + " )\n", + "\n", + " if i == 0:\n", + " galaxy_kwargs[\"shear\"] = al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = al.Galaxy(**galaxy_kwargs)\n", + "\n", + "source = al.Galaxy(redshift=1.0, bulge=source_bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "The `Tracer` performs the ray-tracing. Internally it queries every mass profile attached to every galaxy in\n", + "the lens plane and sums their deflections. For our group lens, this means each `lens_i`'s `bulge` contributes\n", + "a stellar mass deflection (`(M/L)_i * alpha_light_i`), each `dark` halo contributes an `NFWSph` deflection,\n", + "and `lens_0`'s `shear` contributes the external shear \u2014 all summed before mapping image-plane coordinates\n", + "onto the source-plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=list(lens_dict.values()) + [source])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We pass the `Tracer` to a `FitImaging` to fit the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Decomposed Deflection (Multi-Galaxy)__\n", + "\n", + "This is the section that makes the group-scale decomposed-mass fit conceptually distinct. The lens-plane\n", + "deflection map is the SUM, over every main lens galaxy AND every mass component in that galaxy, of independent\n", + "deflection contributions:\n", + "\n", + " alpha_lens(theta) = sum_i [ alpha_stellar_i(theta) + alpha_dark_i(theta) ] + alpha_shear(theta)\n", + " = sum_i [ (M/L)_i * alpha_light_i + alpha_NFW_i ] + alpha_shear\n", + "\n", + "Every individual deflection is a public method on the corresponding profile.\n", + "\n", + "We verify this by computing each contribution explicitly and confirming the sum equals what the `Tracer`\n", + "returns." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = dataset.grid\n", + "\n", + "alpha_stellar_list = [\n", + " lens.bulge.deflections_yx_2d_from(grid=grid) for lens in lens_dict.values()\n", + "]\n", + "alpha_dark_list = [\n", + " lens.dark.deflections_yx_2d_from(grid=grid) for lens in lens_dict.values()\n", + "]\n", + "alpha_shear = lens_dict[\"lens_0\"].shear.deflections_yx_2d_from(grid=grid)\n", + "\n", + "print(f\"alpha_stellar[lens_0] (first coord): {alpha_stellar_list[0][0]}\")\n", + "print(f\"alpha_dark [lens_0] (first coord): {alpha_dark_list[0][0]}\")\n", + "print(f\"alpha_stellar[lens_1] (first coord): {alpha_stellar_list[1][0]}\")\n", + "print(f\"alpha_dark [lens_1] (first coord): {alpha_dark_list[1][0]}\")\n", + "print(f\"alpha_shear (first coord): {alpha_shear[0]}\")\n", + "\n", + "alpha_total_summed = sum(alpha_stellar_list) + sum(alpha_dark_list) + alpha_shear" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The tracer-produced source-plane grid is just `grid - alpha_total_internal`. Recovering the internal total\n", + "deflection from `traced_grid_2d_list_from` and comparing to our hand-summed `alpha_total` is the cleanest\n", + "end-to-end check that we have accounted for every per-galaxy contribution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grids = tracer.traced_grid_2d_list_from(grid=grid)\n", + "alpha_total_tracer = grid - traced_grids[1]\n", + "\n", + "print(f\"alpha_total (summed by hand, first 3): {alpha_total_summed[:3]}\")\n", + "print(f\"alpha_total (from tracer, first 3): {alpha_total_tracer[:3]}\")\n", + "\n", + "assert np.allclose(np.asarray(alpha_total_summed), np.asarray(alpha_total_tracer))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The same component-wise decomposition shows up in the convergence (kappa) map. We sum each component's\n", + "contribution across all main lens galaxies and plot the result, which highlights that the stellar component is\n", + "peaked at the two galaxy centres while the dark halos extend more diffusely." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "kappa_stellar_total = sum(\n", + " lens.bulge.convergence_2d_from(grid=plot_grid) for lens in lens_dict.values()\n", + ")\n", + "kappa_dark_total = sum(\n", + " lens.dark.convergence_2d_from(grid=plot_grid) for lens in lens_dict.values()\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=kappa_stellar_total, title=\"Stellar convergence (sum over galaxies)\"\n", + ")\n", + "aplt.plot_array(\n", + " array=kappa_dark_total, title=\"Dark matter convergence (sum over galaxies)\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "After the fit, every linear Gaussian in the source MGE basis has been assigned an `intensity` via linear\n", + "algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"\\nFirst Gaussian intensity, source = \"\n", + " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", + ")\n", + "\n", + "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer_fitted.galaxies[-1].bulge, grid=plot_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated the group-scale decomposed-mass API and the per-galaxy deflection decomposition,\n", + "without invoking a non-linear search. Each main lens galaxy's `bulge` simultaneously acts as a light profile\n", + "and a stellar mass profile (coupled by its own `mass_to_light_ratio`), and each separately-parameterized\n", + "`dark` NFW halo adds an independent dark mass contribution. A single `ExternalShear` is attached to `lens_0`\n", + "representing the group-wide shear field.\n", + "\n", + "In a real modeling workflow:\n", + "\n", + " - `modeling.py` shows how to fit the same system using `Nautilus`, but \"cheats\" by initialising priors at\n", + " the true values. It is therefore only useful as a tutorial.\n", + " - `chaining.py` is the practical workflow \u2014 two chained searches that fit the lens light first (treating each\n", + " bulge as a pure light profile), then reintroduce the stellar-mass coupling and add the dark NFW halos.\n", + " - `slam.py` is the most robust pipeline for production-quality group-scale decomposed-mass modeling, chaining\n", + " through SOURCE LP, SOURCE PIX, LIGHT LP, and MASS_LIGHT_DARK pipelines and ending in a pixelized source\n", + " reconstruction.\n", + "\n", + "The key takeaway from this script is that a group-scale decomposed-mass lens is fit with the same `Tracer` +\n", + "`FitImaging` objects as any other lens; the only difference is that the lens plane carries MULTIPLE galaxies,\n", + "each with multiple independent mass components, all of whose deflections sum into the total lens-plane\n", + "deflection. The `lens_dict` API scales naturally to any number of main lens galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/mass_stellar_dark/likelihood_function.ipynb b/notebooks/group/features/advanced/mass_stellar_dark/likelihood_function.ipynb index 4656d0474..d63fffcb7 100644 --- a/notebooks/group/features/advanced/mass_stellar_dark/likelihood_function.ipynb +++ b/notebooks/group/features/advanced/mass_stellar_dark/likelihood_function.ipynb @@ -1,435 +1,472 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Group Mass Stellar Dark__\n", - "\n", - "This script describes the additional steps required to compute the `log_likelihood` for a group-scale strong\n", - "lens whose mass model decomposes each main lens galaxy's mass into a stellar component (tied to its light via\n", - "a mass-to-light ratio) and a separately-parameterized dark matter halo.\n", - "\n", - "This script does NOT repeat the steps shared with single-galaxy lensing (mask, image-plane grid, PSF\n", - "convolution, chi-squared, noise normalization, linear-algebra solver for MGE source intensities). It documents\n", - "only the parts of the likelihood function which are specific to a group-scale decomposed-mass lens: the\n", - "multi-galaxy deflection composition.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function below builds directly on the single-galaxy decomposed-mass likelihood and the standard\n", - "imaging / MGE likelihood functions. You should read these notebooks first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/likelihood_function.py` \u2014 the\n", - " single-galaxy decomposed-mass walkthrough, covering the lens-plane deflection composition\n", - " `(M/L) * alpha_light + alpha_NFW + alpha_shear` for one galaxy. This script generalises that walkthrough\n", - " across multiple main lens galaxies.\n", - " - `autolens_workspace/scripts/group/start_here.py` \u2014 the group-scale `lens_dict` API, including how\n", - " `main_lens_centres.json` is loaded.\n", - " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 the canonical single-plane log likelihood,\n", - " covering image-plane grids, ray-tracing, source-plane evaluation, PSF convolution, chi-squared and the\n", - " noise normalization term.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a\n", - " `Basis` of linear Gaussians is solved for via linear algebra.\n", - "\n", - "Sections covered in those scripts (e.g. \"Chi Squared\", \"Noise Normalization Term\", \"Mapping Matrix\") are not\n", - "repeated here.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script (see above).\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Main Lens Centres:** Load the centres of the two main lens galaxies from JSON.\n", - "- **Galaxies:** Build the `lens_dict` and the MGE source.\n", - "- **Decomposed Deflection (Multi-Galaxy):** Sum of per-galaxy stellar + dark + shear deflections.\n", - "- **Manual Ray-Tracing:** Hand-compute the source-plane grid and confirm it matches `Tracer`.\n", - "- **Source-Plane Image:** Source MGE evaluated at the ray-traced grid.\n", - "- **Model Image:** Reference up to the canonical chi-squared / noise normalization.\n", - "- **Fit Check:** `FitImaging.log_likelihood`.\n", - "- **Wrap Up.**\n", - "\n", - "__What Changes For A Group-Scale Decomposed Mass Model__\n", - "\n", - "For a single-galaxy lens with a single total-mass profile (e.g. `Isothermal`), the lens-plane deflection at\n", - "every image-plane coordinate is produced by ONE mass profile:\n", - "\n", - " alpha_lens(theta) = alpha_total(theta ; Isothermal parameters)\n", - "\n", - "For a single-galaxy DECOMPOSED-mass lens, the same galaxy carries multiple independent mass components and\n", - "the lens-plane deflection is their sum (see the single-galaxy prerequisite above):\n", - "\n", - " alpha_lens(theta) = (M/L) * alpha_light(theta) + alpha_NFW(theta) + alpha_shear(theta)\n", - "\n", - "For a GROUP-SCALE decomposed-mass lens, every main lens galaxy carries its own decomposition, and the\n", - "lens-plane deflection is the SUM of every galaxy's contribution:\n", - "\n", - " alpha_lens(theta) = sum_i [ (M/L)_i * alpha_light_i(theta) + alpha_NFW_i(theta) ] + alpha_shear(theta)\n", - "\n", - "A single `ExternalShear` is attached to `lens_0` representing the group-wide shear field. Every other step of\n", - "the likelihood (PSF convolution, chi-squared, noise normalization, MGE linear-algebra solver) is unchanged." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the group mass_stellar_dark dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.7\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Centres__\n", - "\n", - "Load the two main lens galaxy centres from JSON, the same file used by every other script in this directory." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "The main lens galaxies + the source that participate in the ray-tracing:\n", - "\n", - " - Each `lens_i` (z=0.5): an `lmp.Sersic` bulge (acting as light + stellar mass via a single\n", - " `mass_to_light_ratio`), an `NFWSph` dark matter halo aligned with the bulge. Only `lens_0` carries an\n", - " `ExternalShear`.\n", - " - `source` (z=1.0): an MGE light component (a simple basis of 10 linear Gaussians).\n", - "\n", - "The mass-profile parameters are set to the simulator's true values so the manual likelihood computation below\n", - "produces a sensible-looking model image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 10\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "bulge_params = [\n", - " dict(axis_ratio=0.9, angle=45.0, intensity=1.0, effective_radius=0.8, m_to_l=0.20),\n", - " dict(axis_ratio=0.8, angle=120.0, intensity=0.8, effective_radius=0.7, m_to_l=0.25),\n", - "]\n", - "\n", - "dark_params = [\n", - " dict(kappa_s=0.10, scale_radius=20.0),\n", - " dict(kappa_s=0.08, scale_radius=20.0),\n", - "]\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " galaxy_kwargs = dict(\n", - " redshift=0.5,\n", - " bulge=al.lmp.Sersic(\n", - " centre=(centre[0], centre[1]),\n", - " ell_comps=al.convert.ell_comps_from(\n", - " axis_ratio=bulge_params[i][\"axis_ratio\"],\n", - " angle=bulge_params[i][\"angle\"],\n", - " ),\n", - " intensity=bulge_params[i][\"intensity\"],\n", - " effective_radius=bulge_params[i][\"effective_radius\"],\n", - " sersic_index=4.0,\n", - " mass_to_light_ratio=bulge_params[i][\"m_to_l\"],\n", - " ),\n", - " dark=al.mp.NFWSph(\n", - " centre=(centre[0], centre[1]),\n", - " kappa_s=dark_params[i][\"kappa_s\"],\n", - " scale_radius=dark_params[i][\"scale_radius\"],\n", - " ),\n", - " )\n", - "\n", - " if i == 0:\n", - " galaxy_kwargs[\"shear\"] = al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = al.Galaxy(**galaxy_kwargs)\n", - "\n", - "source = al.Galaxy(redshift=1.0, bulge=build_source_basis(centre=(0.0, 0.0)))\n", - "\n", - "tracer = al.Tracer(galaxies=list(lens_dict.values()) + [source])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Decomposed Deflection (Multi-Galaxy)__\n", - "\n", - "The single `Tracer.traced_grid_2d_list_from(...)` call performs the standard image-plane \u2192 source-plane\n", - "ray-tracing. Internally it queries every mass profile on every galaxy in the lens plane and sums their\n", - "deflections.\n", - "\n", - "To make the decomposition concrete we re-compute the same source-plane grid by hand. Each profile exposes its\n", - "own `deflections_yx_2d_from`; the SUM of all per-galaxy stellar + dark contributions, plus the single external\n", - "shear, is what the tracer applies internally:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_grid = dataset.grid\n", - "\n", - "alpha_stellar_list = [\n", - " lens.bulge.deflections_yx_2d_from(grid=masked_grid) for lens in lens_dict.values()\n", - "]\n", - "alpha_dark_list = [\n", - " lens.dark.deflections_yx_2d_from(grid=masked_grid) for lens in lens_dict.values()\n", - "]\n", - "alpha_shear = lens_dict[\"lens_0\"].shear.deflections_yx_2d_from(grid=masked_grid)\n", - "\n", - "alpha_total = sum(alpha_stellar_list) + sum(alpha_dark_list) + alpha_shear\n", - "\n", - "print(f\"alpha_stellar[lens_0] (first coord): {alpha_stellar_list[0][0]}\")\n", - "print(f\"alpha_dark [lens_0] (first coord): {alpha_dark_list[0][0]}\")\n", - "print(f\"alpha_stellar[lens_1] (first coord): {alpha_stellar_list[1][0]}\")\n", - "print(f\"alpha_dark [lens_1] (first coord): {alpha_dark_list[1][0]}\")\n", - "print(f\"alpha_shear (first coord): {alpha_shear[0]}\")\n", - "print(f\"alpha_total (first coord): {alpha_total[0]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Manual Ray-Tracing__\n", - "\n", - "The source-plane grid is the image-plane grid minus the total deflection. We compute it by hand and compare to\n", - "the `Tracer`-produced grid; they should be identical to within floating-point precision." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_source_manual = masked_grid - alpha_total\n", - "\n", - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=masked_grid)\n", - "grid_source_tracer = traced_grid_list[1]\n", - "\n", - "print(f\"\\nsource-plane grid (first coord, manual): {grid_source_manual[0]}\")\n", - "print(f\"source-plane grid (first coord, tracer): {grid_source_tracer[0]}\")\n", - "\n", - "assert np.allclose(np.asarray(grid_source_manual), np.asarray(grid_source_tracer))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Image__\n", - "\n", - "The source galaxy's MGE basis is evaluated at the ray-traced source-plane grid, producing image-plane pixel\n", - "values per Gaussian with a placeholder `intensity=1.0`. The true `intensity` of each Gaussian is solved for at\n", - "the linear-algebra step (see the MGE likelihood prerequisite).\n", - "\n", - "For this manual walkthrough we use the convenience method `image_2d_from` on the `Tracer`, which evaluates the\n", - "source MGE at the correct (ray-traced) plane and projects it into the image plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image_unconvolved = tracer.image_2d_from(grid=masked_grid)\n", - "\n", - "aplt.plot_array(\n", - " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "What `image_2d_from` does internally for our group-scale decomposed-mass lens:\n", - "\n", - " 1. Computes `alpha_lens(theta) = sum_i [ alpha_stellar_i + alpha_dark_i ] + alpha_shear` (the decomposition\n", - " above).\n", - " 2. Ray-traces the image-plane grid to obtain `grid_source = grid - alpha_lens`.\n", - " 3. Evaluates the source MGE at `grid_source`, producing its image-plane contribution.\n", - "\n", - "For a single-galaxy lens there is just one galaxy contributing to step 1; for a group with N main lens\n", - "galaxies there are 2N + 1 mass contributions (stellar + dark per galaxy, plus one shear).\n", - "\n", - "__Model Image__\n", - "\n", - "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model\n", - "image above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", - "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`. The MGE source's per-Gaussian\n", - "intensities are solved for via the linear-algebra step documented in\n", - "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py`.\n", - "\n", - "We delegate the remaining steps to `FitImaging`, which handles the linear-algebra step that solves for each\n", - "Gaussian's `intensity` and assembles the full `log_likelihood`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(\n", - " f\"\\nLog likelihood of the manual group mass-stellar-dark fit: {fit.log_likelihood}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood__\n", - "\n", - "The final `log_likelihood` combines:\n", - "\n", - " - The chi-squared term, computed from the residuals between the PSF-convolved model image and the data,\n", - " weighted by the noise map.\n", - " - The noise normalization term, the standard Gaussian normalization over all unmasked pixels.\n", - " - The linear algebra terms (regularization and curvature determinants) introduced by the MGE\n", - " `Basis` of linear Gaussians.\n", - "\n", - "The first two are documented in `imaging/likelihood_function.py`; the third in\n", - "`imaging/features/multi_gaussian_expansion/likelihood_function.py`. No new terms are introduced by the\n", - "group-scale decomposed mass model \u2014 the only change is the lens-plane deflection composition described above.\n", - "\n", - "__Wrap Up__\n", - "\n", - "The group-scale decomposed-mass `log_likelihood` differs from the single-galaxy decomposed-mass case in exactly\n", - "one place: the lens-plane deflection is a sum over MULTIPLE galaxies, each contributing its own stellar +\n", - "dark decomposition, plus a single external shear. Every other step (ray-tracing, source-plane evaluation, PSF\n", - "convolution, chi-squared, noise normalization, linear algebra) is shared with the standard imaging likelihood\n", - "and documented in the prerequisite scripts.\n", - "\n", - "This per-galaxy decomposition is the standard tool for studying mass-to-light variation across a group\n", - "environment: each `mass_to_light_ratio` is constrained independently by the data, so the relative stellar vs\n", - "dark contribution of each main lens galaxy can be measured separately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Group Mass Stellar Dark__\n", + "\n", + "This script describes the additional steps required to compute the `log_likelihood` for a group-scale strong\n", + "lens whose mass model decomposes each main lens galaxy's mass into a stellar component (tied to its light via\n", + "a mass-to-light ratio) and a separately-parameterized dark matter halo.\n", + "\n", + "This script does NOT repeat the steps shared with single-galaxy lensing (mask, image-plane grid, PSF\n", + "convolution, chi-squared, noise normalization, linear-algebra solver for MGE source intensities). It documents\n", + "only the parts of the likelihood function which are specific to a group-scale decomposed-mass lens: the\n", + "multi-galaxy deflection composition.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function below builds directly on the single-galaxy decomposed-mass likelihood and the standard\n", + "imaging / MGE likelihood functions. You should read these notebooks first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/likelihood_function.py` \u2014 the\n", + " single-galaxy decomposed-mass walkthrough, covering the lens-plane deflection composition\n", + " `(M/L) * alpha_light + alpha_NFW + alpha_shear` for one galaxy. This script generalises that walkthrough\n", + " across multiple main lens galaxies.\n", + " - `autolens_workspace/scripts/group/start_here.py` \u2014 the group-scale `lens_dict` API, including how\n", + " `main_lens_centres.json` is loaded.\n", + " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 the canonical single-plane log likelihood,\n", + " covering image-plane grids, ray-tracing, source-plane evaluation, PSF convolution, chi-squared and the\n", + " noise normalization term.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a\n", + " `Basis` of linear Gaussians is solved for via linear algebra.\n", + "\n", + "Sections covered in those scripts (e.g. \"Chi Squared\", \"Noise Normalization Term\", \"Mapping Matrix\") are not\n", + "repeated here.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script (see above).\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Main Lens Centres:** Load the centres of the two main lens galaxies from JSON.\n", + "- **Galaxies:** Build the `lens_dict` and the MGE source.\n", + "- **Decomposed Deflection (Multi-Galaxy):** Sum of per-galaxy stellar + dark + shear deflections.\n", + "- **Manual Ray-Tracing:** Hand-compute the source-plane grid and confirm it matches `Tracer`.\n", + "- **Source-Plane Image:** Source MGE evaluated at the ray-traced grid.\n", + "- **Model Image:** Reference up to the canonical chi-squared / noise normalization.\n", + "- **Fit Check:** `FitImaging.log_likelihood`.\n", + "- **Wrap Up.**\n", + "\n", + "__What Changes For A Group-Scale Decomposed Mass Model__\n", + "\n", + "For a single-galaxy lens with a single total-mass profile (e.g. `Isothermal`), the lens-plane deflection at\n", + "every image-plane coordinate is produced by ONE mass profile:\n", + "\n", + " alpha_lens(theta) = alpha_total(theta ; Isothermal parameters)\n", + "\n", + "For a single-galaxy DECOMPOSED-mass lens, the same galaxy carries multiple independent mass components and\n", + "the lens-plane deflection is their sum (see the single-galaxy prerequisite above):\n", + "\n", + " alpha_lens(theta) = (M/L) * alpha_light(theta) + alpha_NFW(theta) + alpha_shear(theta)\n", + "\n", + "For a GROUP-SCALE decomposed-mass lens, every main lens galaxy carries its own decomposition, and the\n", + "lens-plane deflection is the SUM of every galaxy's contribution:\n", + "\n", + " alpha_lens(theta) = sum_i [ (M/L)_i * alpha_light_i(theta) + alpha_NFW_i(theta) ] + alpha_shear(theta)\n", + "\n", + "A single `ExternalShear` is attached to `lens_0` representing the group-wide shear field. Every other step of\n", + "the likelihood (PSF convolution, chi-squared, noise normalization, MGE linear-algebra solver) is unchanged." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the group mass_stellar_dark dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.7\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Centres__\n", + "\n", + "Load the two main lens galaxy centres from JSON, the same file used by every other script in this directory." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "The main lens galaxies + the source that participate in the ray-tracing:\n", + "\n", + " - Each `lens_i` (z=0.5): an `lmp.Sersic` bulge (acting as light + stellar mass via a single\n", + " `mass_to_light_ratio`), an `NFWSph` dark matter halo aligned with the bulge. Only `lens_0` carries an\n", + " `ExternalShear`.\n", + " - `source` (z=1.0): an MGE light component (a simple basis of 10 linear Gaussians).\n", + "\n", + "The mass-profile parameters are set to the simulator's true values so the manual likelihood computation below\n", + "produces a sensible-looking model image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 10\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "bulge_params = [\n", + " dict(axis_ratio=0.9, angle=45.0, intensity=1.0, effective_radius=0.8, m_to_l=0.20),\n", + " dict(axis_ratio=0.8, angle=120.0, intensity=0.8, effective_radius=0.7, m_to_l=0.25),\n", + "]\n", + "\n", + "dark_params = [\n", + " dict(kappa_s=0.10, scale_radius=20.0),\n", + " dict(kappa_s=0.08, scale_radius=20.0),\n", + "]\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " galaxy_kwargs = dict(\n", + " redshift=0.5,\n", + " bulge=al.lmp.Sersic(\n", + " centre=(centre[0], centre[1]),\n", + " ell_comps=al.convert.ell_comps_from(\n", + " axis_ratio=bulge_params[i][\"axis_ratio\"],\n", + " angle=bulge_params[i][\"angle\"],\n", + " ),\n", + " intensity=bulge_params[i][\"intensity\"],\n", + " effective_radius=bulge_params[i][\"effective_radius\"],\n", + " sersic_index=4.0,\n", + " mass_to_light_ratio=bulge_params[i][\"m_to_l\"],\n", + " ),\n", + " dark=al.mp.NFWSph(\n", + " centre=(centre[0], centre[1]),\n", + " kappa_s=dark_params[i][\"kappa_s\"],\n", + " scale_radius=dark_params[i][\"scale_radius\"],\n", + " ),\n", + " )\n", + "\n", + " if i == 0:\n", + " galaxy_kwargs[\"shear\"] = al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = al.Galaxy(**galaxy_kwargs)\n", + "\n", + "source = al.Galaxy(redshift=1.0, bulge=build_source_basis(centre=(0.0, 0.0)))\n", + "\n", + "tracer = al.Tracer(galaxies=list(lens_dict.values()) + [source])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Decomposed Deflection (Multi-Galaxy)__\n", + "\n", + "The single `Tracer.traced_grid_2d_list_from(...)` call performs the standard image-plane \u2192 source-plane\n", + "ray-tracing. Internally it queries every mass profile on every galaxy in the lens plane and sums their\n", + "deflections.\n", + "\n", + "To make the decomposition concrete we re-compute the same source-plane grid by hand. Each profile exposes its\n", + "own `deflections_yx_2d_from`; the SUM of all per-galaxy stellar + dark contributions, plus the single external\n", + "shear, is what the tracer applies internally:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_grid = dataset.grid\n", + "\n", + "alpha_stellar_list = [\n", + " lens.bulge.deflections_yx_2d_from(grid=masked_grid) for lens in lens_dict.values()\n", + "]\n", + "alpha_dark_list = [\n", + " lens.dark.deflections_yx_2d_from(grid=masked_grid) for lens in lens_dict.values()\n", + "]\n", + "alpha_shear = lens_dict[\"lens_0\"].shear.deflections_yx_2d_from(grid=masked_grid)\n", + "\n", + "alpha_total = sum(alpha_stellar_list) + sum(alpha_dark_list) + alpha_shear\n", + "\n", + "print(f\"alpha_stellar[lens_0] (first coord): {alpha_stellar_list[0][0]}\")\n", + "print(f\"alpha_dark [lens_0] (first coord): {alpha_dark_list[0][0]}\")\n", + "print(f\"alpha_stellar[lens_1] (first coord): {alpha_stellar_list[1][0]}\")\n", + "print(f\"alpha_dark [lens_1] (first coord): {alpha_dark_list[1][0]}\")\n", + "print(f\"alpha_shear (first coord): {alpha_shear[0]}\")\n", + "print(f\"alpha_total (first coord): {alpha_total[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Manual Ray-Tracing__\n", + "\n", + "The source-plane grid is the image-plane grid minus the total deflection. We compute it by hand and compare to\n", + "the `Tracer`-produced grid; they should be identical to within floating-point precision." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_source_manual = masked_grid - alpha_total\n", + "\n", + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=masked_grid)\n", + "grid_source_tracer = traced_grid_list[1]\n", + "\n", + "print(f\"\\nsource-plane grid (first coord, manual): {grid_source_manual[0]}\")\n", + "print(f\"source-plane grid (first coord, tracer): {grid_source_tracer[0]}\")\n", + "\n", + "assert np.allclose(np.asarray(grid_source_manual), np.asarray(grid_source_tracer))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Image__\n", + "\n", + "The source galaxy's MGE basis is evaluated at the ray-traced source-plane grid, producing image-plane pixel\n", + "values per Gaussian with a placeholder `intensity=1.0`. The true `intensity` of each Gaussian is solved for at\n", + "the linear-algebra step (see the MGE likelihood prerequisite).\n", + "\n", + "For this manual walkthrough we use the convenience method `image_2d_from` on the `Tracer`, which evaluates the\n", + "source MGE at the correct (ray-traced) plane and projects it into the image plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image_unconvolved = tracer.image_2d_from(grid=masked_grid)\n", + "\n", + "aplt.plot_array(\n", + " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "What `image_2d_from` does internally for our group-scale decomposed-mass lens:\n", + "\n", + " 1. Computes `alpha_lens(theta) = sum_i [ alpha_stellar_i + alpha_dark_i ] + alpha_shear` (the decomposition\n", + " above).\n", + " 2. Ray-traces the image-plane grid to obtain `grid_source = grid - alpha_lens`.\n", + " 3. Evaluates the source MGE at `grid_source`, producing its image-plane contribution.\n", + "\n", + "For a single-galaxy lens there is just one galaxy contributing to step 1; for a group with N main lens\n", + "galaxies there are 2N + 1 mass contributions (stellar + dark per galaxy, plus one shear).\n", + "\n", + "__Model Image__\n", + "\n", + "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model\n", + "image above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", + "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`. The MGE source's per-Gaussian\n", + "intensities are solved for via the linear-algebra step documented in\n", + "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py`.\n", + "\n", + "We delegate the remaining steps to `FitImaging`, which handles the linear-algebra step that solves for each\n", + "Gaussian's `intensity` and assembles the full `log_likelihood`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(\n", + " f\"\\nLog likelihood of the manual group mass-stellar-dark fit: {fit.log_likelihood}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood__\n", + "\n", + "The final `log_likelihood` combines:\n", + "\n", + " - The chi-squared term, computed from the residuals between the PSF-convolved model image and the data,\n", + " weighted by the noise map.\n", + " - The noise normalization term, the standard Gaussian normalization over all unmasked pixels.\n", + " - The linear algebra terms (regularization and curvature determinants) introduced by the MGE\n", + " `Basis` of linear Gaussians.\n", + "\n", + "The first two are documented in `imaging/likelihood_function.py`; the third in\n", + "`imaging/features/multi_gaussian_expansion/likelihood_function.py`. No new terms are introduced by the\n", + "group-scale decomposed mass model \u2014 the only change is the lens-plane deflection composition described above.\n", + "\n", + "__Wrap Up__\n", + "\n", + "The group-scale decomposed-mass `log_likelihood` differs from the single-galaxy decomposed-mass case in exactly\n", + "one place: the lens-plane deflection is a sum over MULTIPLE galaxies, each contributing its own stellar +\n", + "dark decomposition, plus a single external shear. Every other step (ray-tracing, source-plane evaluation, PSF\n", + "convolution, chi-squared, noise normalization, linear algebra) is shared with the standard imaging likelihood\n", + "and documented in the prerequisite scripts.\n", + "\n", + "This per-galaxy decomposition is the standard tool for studying mass-to-light variation across a group\n", + "environment: each `mass_to_light_ratio` is constrained independently by the data, so the relative stellar vs\n", + "dark contribution of each main lens galaxy can be measured separately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/mass_stellar_dark/modeling.ipynb b/notebooks/group/features/advanced/mass_stellar_dark/modeling.ipynb index b12bdc3b3..56b1b6abc 100644 --- a/notebooks/group/features/advanced/mass_stellar_dark/modeling.ipynb +++ b/notebooks/group/features/advanced/mass_stellar_dark/modeling.ipynb @@ -1,509 +1,546 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Group Mass Stellar Dark\n", - "==========================================\n", - "\n", - "A group-scale strong lens where each main lens galaxy carries a decomposed mass model: a stellar component\n", - "tied to the galaxy's own light via a mass-to-light ratio, plus a separately-parameterized dark matter halo.\n", - "The total lens-plane deflection is the sum, over every main lens galaxy, of stellar + dark contributions,\n", - "plus a single external shear attached to `lens_0` representing the group-wide shear field.\n", - "\n", - "This script fits a group lens model where each main lens galaxy is decomposed into stellar + dark components.\n", - "Per-galaxy decomposition is the standard tool for studying mass-to-light variation across a group environment.\n", - "\n", - "__Practical Use: Read This First__\n", - "\n", - "This script is a tutorial. It produces a working fit by \"cheating\" \u2014 every prior is initialised at the true\n", - "simulator value, narrowed by a small Gaussian. On real data this is impossible, and a single Nautilus search\n", - "on a group-scale decomposed mass model would almost certainly converge to a local maximum. Two effects compound:\n", - "\n", - " - the per-galaxy `mass_to_light_ratio` couples each bulge's light and stellar mass, creating tight parameter\n", - " degeneracies between the light parameters and the mass deflection;\n", - " - every main lens galaxy adds its own stellar + dark contribution, so the deflection field is the sum of\n", - " many independent components which a single search struggles to disentangle without good starting points.\n", - "\n", - "The script you will actually use to fit a group decomposed-mass model on real data is\n", - "`autolens_workspace/scripts/group/features/advanced/mass_stellar_dark/chaining.py`, which runs two chained\n", - "non-linear searches: the first fits each lens galaxy's bulge as a pure light profile (no stellar mass\n", - "coupling, no dark NFW), the second reintroduces the stellar-mass coupling and adds the dark NFW per galaxy\n", - "with priors carried over from search 1.\n", - "\n", - "For production-quality modeling, see `slam.py` in the same directory, which uses the `MASS_LIGHT_DARK` SLaM\n", - "pipeline.\n", - "\n", - "Read this script to understand the model composition API, then jump to `chaining.py`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Main Lens Centres:** Load the centres of the two main lens galaxies from JSON.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'group-scale' strong lens with a model where:\n", - "\n", - " - Each main lens galaxy's light and stellar mass is a linear `Sersic` [7 parameters per galaxy].\n", - " - Each main lens galaxy's dark matter mass distribution is a `NFWSph` aligned with that galaxy's bulge centre.\n", - " - The first main lens galaxy additionally carries an `ExternalShear` [2 parameters].\n", - " - The source galaxy's light is a Multi Gaussian Expansion.\n", - "\n", - "For two main lens galaxies, the lens-plane carries (7 + 4) * 2 = 22 free parameters, plus 2 for the shear,\n", - "plus the source MGE parameters.\n", - "\n", - "Note that for each main lens galaxy's stellar light and mass, we use a \"light and mass profile\" via the `.lmp`\n", - "package. This profile simultaneously acts like a light and mass profile.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook:\n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to:\n", - "\n", - " - `autolens_workspace/scripts/group/start_here.py` \u2014 the canonical group modeling walkthrough.\n", - " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/modeling.py` \u2014 the single-galaxy\n", - " decomposed-mass walkthrough." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `mass_stellar_dark` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Centres__\n", - "\n", - "Load the centres of the two main lens galaxies from JSON. These centres are fixed on each galaxy's bulge and\n", - "dark profile in the model below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.7\" circular mask, which includes both main lens galaxies and the lensed source emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.7\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling at each main lens galaxy centre, so the stellar mass-to-light coupling is\n", - "evaluated accurately at the peak of each bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - Each main lens galaxy's light and stellar mass is a linear `Sersic` [7 parameters per galaxy].\n", - " - Each main lens galaxy's dark matter mass distribution is a `NFWSph` whose centre is fixed to the bulge\n", - " centre (i.e. that galaxy's `main_lens_centres` entry) [3 parameters per galaxy].\n", - " - The first main lens galaxy additionally carries an `ExternalShear` [2 parameters].\n", - " - The source galaxy's light is a Multi Gaussian Expansion.\n", - "\n", - "The bulge and dark centres are fixed (no priors) because the centres are determined externally (e.g. by the\n", - "GUI used in `group/start_here.py`).\n", - "\n", - "__List-Based Model Composition__\n", - "\n", - "For group-scale lenses, we compose the lens-plane model via a `for i, centre in enumerate(main_lens_centres)`\n", - "loop. Each main lens galaxy is created in a loop and stored in a dictionary as `lens_0`, `lens_1`, etc. This\n", - "API scales naturally to groups with any number of main lens galaxies.\n", - "\n", - "Only the first lens galaxy (`lens_0`) carries an `ExternalShear`, as the group system has one overall external\n", - "shear." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " bulge = af.Model(al.lmp.Sersic)\n", - " bulge.centre = (centre[0], centre[1])\n", - "\n", - " dark = af.Model(al.mp.NFWSph)\n", - " dark.centre = (centre[0], centre[1])\n", - "\n", - " galaxy_kwargs = dict(redshift=0.5, bulge=bulge, dark=dark)\n", - "\n", - " if i == 0:\n", - " galaxy_kwargs[\"shear\"] = af.Model(al.mp.ExternalShear)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **galaxy_kwargs)\n", - "\n", - "# Source:\n", - "\n", - "source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", - "\n", - "# Overall Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(**lens_dict, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", - "refer to `start_here.ipynb` for a description of how to fix this)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a full\n", - "description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"mass_stellar_dark\",\n", - " unique_tag=dataset_name,\n", - " n_live=200,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=2000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", - "required by a model.\n", - "\n", - "Deflection angle calculations of stellar mass models and dark matter mass models can use techniques which\n", - "store more data in VRAM than other methods.\n", - "\n", - "Given VRAM use is an important consideration for group-scale lenses (which carry many mass components), we\n", - "print out the estimated VRAM required for this model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The likelihood evaluation time for analysing decomposed stellar and dark matter mass models is longer than for\n", - "total mass models like the isothermal or power-law. This is because the deflection angles of these mass\n", - "profiles are more expensive to compute, requiring a Gaussian expansion or numerical calculation.\n", - "\n", - "For a group lens with multiple main galaxies, each galaxy adds its own pair of (stellar, dark) deflection\n", - "evaluations, scaling the per-likelihood cost roughly linearly with the number of main lens galaxies.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output\n", - "folder for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", - "refer to `start_here.ipynb` for a description of how to fix this)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "These plots show that the per-galaxy decomposed stars + dark matter model recovers the ray-traced source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "These examples include a results API with specific tools for visualizing and analysing decomposed mass models,\n", - "for example 1D plots which separately show the density of stars and dark matter as a function of radius for\n", - "each main lens galaxy.\n", - "\n", - "__Wrap Up__\n", - "\n", - "Group-scale decomposed mass models have advantages and disadvantages compared to total mass models.\n", - "\n", - "The model which is best suited to your needs depends on the science you are hoping to undertake and the\n", - "quality of the data you are fitting.\n", - "\n", - "In general, it is recommended that you first get fits going using total mass models, because they are simpler\n", - "and make fewer assumptions regarding how light is tied to mass. Once you have robust results, decomposed mass\n", - "models can then be fitted and compared in order to gain deeper insight into the per-galaxy stellar and dark\n", - "contributions across the group." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Group Mass Stellar Dark\n", + "==========================================\n", + "\n", + "A group-scale strong lens where each main lens galaxy carries a decomposed mass model: a stellar component\n", + "tied to the galaxy's own light via a mass-to-light ratio, plus a separately-parameterized dark matter halo.\n", + "The total lens-plane deflection is the sum, over every main lens galaxy, of stellar + dark contributions,\n", + "plus a single external shear attached to `lens_0` representing the group-wide shear field.\n", + "\n", + "This script fits a group lens model where each main lens galaxy is decomposed into stellar + dark components.\n", + "Per-galaxy decomposition is the standard tool for studying mass-to-light variation across a group environment.\n", + "\n", + "__Practical Use: Read This First__\n", + "\n", + "This script is a tutorial. It produces a working fit by \"cheating\" \u2014 every prior is initialised at the true\n", + "simulator value, narrowed by a small Gaussian. On real data this is impossible, and a single Nautilus search\n", + "on a group-scale decomposed mass model would almost certainly converge to a local maximum. Two effects compound:\n", + "\n", + " - the per-galaxy `mass_to_light_ratio` couples each bulge's light and stellar mass, creating tight parameter\n", + " degeneracies between the light parameters and the mass deflection;\n", + " - every main lens galaxy adds its own stellar + dark contribution, so the deflection field is the sum of\n", + " many independent components which a single search struggles to disentangle without good starting points.\n", + "\n", + "The script you will actually use to fit a group decomposed-mass model on real data is\n", + "`autolens_workspace/scripts/group/features/advanced/mass_stellar_dark/chaining.py`, which runs two chained\n", + "non-linear searches: the first fits each lens galaxy's bulge as a pure light profile (no stellar mass\n", + "coupling, no dark NFW), the second reintroduces the stellar-mass coupling and adds the dark NFW per galaxy\n", + "with priors carried over from search 1.\n", + "\n", + "For production-quality modeling, see `slam.py` in the same directory, which uses the `MASS_LIGHT_DARK` SLaM\n", + "pipeline.\n", + "\n", + "Read this script to understand the model composition API, then jump to `chaining.py`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Main Lens Centres:** Load the centres of the two main lens galaxies from JSON.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'group-scale' strong lens with a model where:\n", + "\n", + " - Each main lens galaxy's light and stellar mass is a linear `Sersic` [7 parameters per galaxy].\n", + " - Each main lens galaxy's dark matter mass distribution is a `NFWSph` aligned with that galaxy's bulge centre.\n", + " - The first main lens galaxy additionally carries an `ExternalShear` [2 parameters].\n", + " - The source galaxy's light is a Multi Gaussian Expansion.\n", + "\n", + "For two main lens galaxies, the lens-plane carries (7 + 4) * 2 = 22 free parameters, plus 2 for the shear,\n", + "plus the source MGE parameters.\n", + "\n", + "Note that for each main lens galaxy's stellar light and mass, we use a \"light and mass profile\" via the `.lmp`\n", + "package. This profile simultaneously acts like a light and mass profile.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook:\n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to:\n", + "\n", + " - `autolens_workspace/scripts/group/start_here.py` \u2014 the canonical group modeling walkthrough.\n", + " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/modeling.py` \u2014 the single-galaxy\n", + " decomposed-mass walkthrough." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `mass_stellar_dark` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Centres__\n", + "\n", + "Load the centres of the two main lens galaxies from JSON. These centres are fixed on each galaxy's bulge and\n", + "dark profile in the model below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.7\" circular mask, which includes both main lens galaxies and the lensed source emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.7\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling at each main lens galaxy centre, so the stellar mass-to-light coupling is\n", + "evaluated accurately at the peak of each bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - Each main lens galaxy's light and stellar mass is a linear `Sersic` [7 parameters per galaxy].\n", + " - Each main lens galaxy's dark matter mass distribution is a `NFWSph` whose centre is fixed to the bulge\n", + " centre (i.e. that galaxy's `main_lens_centres` entry) [3 parameters per galaxy].\n", + " - The first main lens galaxy additionally carries an `ExternalShear` [2 parameters].\n", + " - The source galaxy's light is a Multi Gaussian Expansion.\n", + "\n", + "The bulge and dark centres are fixed (no priors) because the centres are determined externally (e.g. by the\n", + "GUI used in `group/start_here.py`).\n", + "\n", + "__List-Based Model Composition__\n", + "\n", + "For group-scale lenses, we compose the lens-plane model via a `for i, centre in enumerate(main_lens_centres)`\n", + "loop. Each main lens galaxy is created in a loop and stored in a dictionary as `lens_0`, `lens_1`, etc. This\n", + "API scales naturally to groups with any number of main lens galaxies.\n", + "\n", + "Only the first lens galaxy (`lens_0`) carries an `ExternalShear`, as the group system has one overall external\n", + "shear." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " bulge = af.Model(al.lmp.Sersic)\n", + " bulge.centre = (centre[0], centre[1])\n", + "\n", + " dark = af.Model(al.mp.NFWSph)\n", + " dark.centre = (centre[0], centre[1])\n", + "\n", + " galaxy_kwargs = dict(redshift=0.5, bulge=bulge, dark=dark)\n", + "\n", + " if i == 0:\n", + " galaxy_kwargs[\"shear\"] = af.Model(al.mp.ExternalShear)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **galaxy_kwargs)\n", + "\n", + "# Source:\n", + "\n", + "source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", + "\n", + "# Overall Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(**lens_dict, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", + "refer to `start_here.ipynb` for a description of how to fix this)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a full\n", + "description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"mass_stellar_dark\",\n", + " unique_tag=dataset_name,\n", + " n_live=200,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=2000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", + "required by a model.\n", + "\n", + "Deflection angle calculations of stellar mass models and dark matter mass models can use techniques which\n", + "store more data in VRAM than other methods.\n", + "\n", + "Given VRAM use is an important consideration for group-scale lenses (which carry many mass components), we\n", + "print out the estimated VRAM required for this model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The likelihood evaluation time for analysing decomposed stellar and dark matter mass models is longer than for\n", + "total mass models like the isothermal or power-law. This is because the deflection angles of these mass\n", + "profiles are more expensive to compute, requiring a Gaussian expansion or numerical calculation.\n", + "\n", + "For a group lens with multiple main galaxies, each galaxy adds its own pair of (stellar, dark) deflection\n", + "evaluations, scaling the per-likelihood cost roughly linearly with the number of main lens galaxies.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output\n", + "folder for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", + "refer to `start_here.ipynb` for a description of how to fix this)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "These plots show that the per-galaxy decomposed stars + dark matter model recovers the ray-traced source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "These examples include a results API with specific tools for visualizing and analysing decomposed mass models,\n", + "for example 1D plots which separately show the density of stars and dark matter as a function of radius for\n", + "each main lens galaxy.\n", + "\n", + "__Wrap Up__\n", + "\n", + "Group-scale decomposed mass models have advantages and disadvantages compared to total mass models.\n", + "\n", + "The model which is best suited to your needs depends on the science you are hoping to undertake and the\n", + "quality of the data you are fitting.\n", + "\n", + "In general, it is recommended that you first get fits going using total mass models, because they are simpler\n", + "and make fewer assumptions regarding how light is tied to mass. Once you have robust results, decomposed mass\n", + "models can then be fitted and compared in order to gain deeper insight into the per-galaxy stellar and dark\n", + "contributions across the group." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/mass_stellar_dark/simulator.ipynb b/notebooks/group/features/advanced/mass_stellar_dark/simulator.ipynb index 982d1dd38..2db0f57ad 100644 --- a/notebooks/group/features/advanced/mass_stellar_dark/simulator.ipynb +++ b/notebooks/group/features/advanced/mass_stellar_dark/simulator.ipynb @@ -1,413 +1,450 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Group Mass Stellar Dark\n", - "==================================\n", - "\n", - "A group-scale strong lens where each main lens galaxy carries a decomposed mass model: a stellar component tied\n", - "to the galaxy's own light via a mass-to-light ratio, plus a separately-parameterized dark matter halo.\n", - "\n", - "This script simulates an `Imaging` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - There are TWO main lens galaxies at z=0.5. Each carries a `lmp.Sersic` bulge (acting as light AND stellar\n", - " mass via `mass_to_light_ratio`) and a spherical `NFWSph` dark matter halo aligned with the bulge centre.\n", - " - The first main lens galaxy additionally carries an `ExternalShear` representing the group-scale shear from\n", - " the wider environment.\n", - " - The source galaxy at z=1.0 has a `SersicCore` light profile.\n", - "\n", - "The total deflection at every image-plane coordinate is the SUM over all main lens galaxies of the per-galaxy\n", - "stellar + dark deflections, plus the external shear contribution.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** The dataset name and output folder.\n", - "- **Grid:** Define the 2D image-plane grid.\n", - "- **Galaxy Centres:** Centres of the two main lens galaxies, saved as JSON for the modeling scripts to load.\n", - "- **Over Sampling:** Adaptive over-sampling at the main lens galaxy centres.\n", - "- **Main Lens Galaxies:** Two galaxies at z=0.5, each decomposed into stellar (lmp.Sersic) and dark (NFWSph).\n", - "- **Source Galaxy:** A single source at z=1.0 with a SersicCore light profile.\n", - "- **Ray Tracing:** Build the Tracer.\n", - "- **Dataset:** Simulate the imaging dataset and write .fits.\n", - "- **Tracer JSON:** Save the simulator Tracer for provenance.\n", - "- **Centres JSON:** Save the main lens centres for the modeling scripts.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to:\n", - "\n", - " - `autolens_workspace/scripts/group/simulator.py` \u2014 the canonical group-scale simulator.\n", - " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/simulator.py` \u2014 the single-galaxy\n", - " decomposed-mass simulator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The dataset is written to `dataset/group/mass_stellar_dark/`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "A 200x200 grid at 0.1\"/px gives a 20\" field of view, large enough to contain a group-scale lens where the two\n", - "main lens galaxies are separated by ~1.5\"." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(200, 200),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "The two main lens galaxies are separated by ~1.5\" along the y-axis. These centres are saved as JSON so the\n", - "modeling and fit scripts can load them via `al.from_json(...)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [(-0.75, 0.0), (0.75, 0.0)]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling is applied at each main lens galaxy's centre, so the stellar mass-to-light coupling is\n", - "evaluated accurately at the peak of each bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=main_lens_centres,\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Imaging simulator: exposure time, background sky, noise, PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "Two galaxies at z=0.5, each with a `lmp.Sersic` bulge (coupled to its own stellar mass via `mass_to_light_ratio`)\n", - "and an `NFWSph` dark matter halo aligned with the bulge.\n", - "\n", - "The first galaxy additionally carries an `ExternalShear` \u2014 a single shear field representing the wider group\n", - "environment, conventionally attached to `lens_0` (matches the group SLaM convention in\n", - "`scripts/group/features/advanced/double_einstein_ring/slam.py`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lmp.Sersic(\n", - " centre=main_lens_centres[0],\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " mass_to_light_ratio=0.2,\n", - " ),\n", - " dark=al.mp.NFWSph(centre=main_lens_centres[0], kappa_s=0.1, scale_radius=20.0),\n", - " shear=al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005),\n", - ")\n", - "\n", - "lens_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lmp.Sersic(\n", - " centre=main_lens_centres[1],\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=120.0),\n", - " intensity=0.8,\n", - " effective_radius=0.7,\n", - " sersic_index=4.0,\n", - " mass_to_light_ratio=0.25,\n", - " ),\n", - " dark=al.mp.NFWSph(centre=main_lens_centres[1], kappa_s=0.08, scale_radius=20.0),\n", - ")\n", - "\n", - "main_lens_galaxies = [lens_0, lens_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy__\n", - "\n", - "A single compact source at z=1.0 with a `SersicCore` light profile, positioned near the group centre so its\n", - "lensed image forms a clearly visible Einstein-ring-like configuration around the two main lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The tracer is composed of the two main lens galaxies followed by the source. PyAutoLens orders galaxies\n", - "internally by redshift, so the deflection chain runs:\n", - "\n", - " image-plane \u2192 source-plane (deflected by both main lens galaxies' stellar + dark + shear contributions)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=main_lens_galaxies + [source])\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid), title=\"Group Mass Stellar Dark Image\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Pass the simulator a tracer to produce the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Write the simulated dataset to .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer JSON__\n", - "\n", - "Save the simulator `Tracer` so the true profiles can be inspected later." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres JSON__\n", - "\n", - "Save the main lens centres so the modeling scripts can load them via `al.from_json(...)`. This mirrors the\n", - "canonical `group/simulator.py` and `group/features/advanced/double_einstein_ring/simulator.py` conventions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Group Mass Stellar Dark\n", + "==================================\n", + "\n", + "A group-scale strong lens where each main lens galaxy carries a decomposed mass model: a stellar component tied\n", + "to the galaxy's own light via a mass-to-light ratio, plus a separately-parameterized dark matter halo.\n", + "\n", + "This script simulates an `Imaging` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - There are TWO main lens galaxies at z=0.5. Each carries a `lmp.Sersic` bulge (acting as light AND stellar\n", + " mass via `mass_to_light_ratio`) and a spherical `NFWSph` dark matter halo aligned with the bulge centre.\n", + " - The first main lens galaxy additionally carries an `ExternalShear` representing the group-scale shear from\n", + " the wider environment.\n", + " - The source galaxy at z=1.0 has a `SersicCore` light profile.\n", + "\n", + "The total deflection at every image-plane coordinate is the SUM over all main lens galaxies of the per-galaxy\n", + "stellar + dark deflections, plus the external shear contribution.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** The dataset name and output folder.\n", + "- **Grid:** Define the 2D image-plane grid.\n", + "- **Galaxy Centres:** Centres of the two main lens galaxies, saved as JSON for the modeling scripts to load.\n", + "- **Over Sampling:** Adaptive over-sampling at the main lens galaxy centres.\n", + "- **Main Lens Galaxies:** Two galaxies at z=0.5, each decomposed into stellar (lmp.Sersic) and dark (NFWSph).\n", + "- **Source Galaxy:** A single source at z=1.0 with a SersicCore light profile.\n", + "- **Ray Tracing:** Build the Tracer.\n", + "- **Dataset:** Simulate the imaging dataset and write .fits.\n", + "- **Tracer JSON:** Save the simulator Tracer for provenance.\n", + "- **Centres JSON:** Save the main lens centres for the modeling scripts.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to:\n", + "\n", + " - `autolens_workspace/scripts/group/simulator.py` \u2014 the canonical group-scale simulator.\n", + " - `autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/simulator.py` \u2014 the single-galaxy\n", + " decomposed-mass simulator." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The dataset is written to `dataset/group/mass_stellar_dark/`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "A 200x200 grid at 0.1\"/px gives a 20\" field of view, large enough to contain a group-scale lens where the two\n", + "main lens galaxies are separated by ~1.5\"." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(200, 200),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "The two main lens galaxies are separated by ~1.5\" along the y-axis. These centres are saved as JSON so the\n", + "modeling and fit scripts can load them via `al.from_json(...)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [(-0.75, 0.0), (0.75, 0.0)]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling is applied at each main lens galaxy's centre, so the stellar mass-to-light coupling is\n", + "evaluated accurately at the peak of each bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=main_lens_centres,\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Imaging simulator: exposure time, background sky, noise, PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "Two galaxies at z=0.5, each with a `lmp.Sersic` bulge (coupled to its own stellar mass via `mass_to_light_ratio`)\n", + "and an `NFWSph` dark matter halo aligned with the bulge.\n", + "\n", + "The first galaxy additionally carries an `ExternalShear` \u2014 a single shear field representing the wider group\n", + "environment, conventionally attached to `lens_0` (matches the group SLaM convention in\n", + "`scripts/group/features/advanced/double_einstein_ring/slam.py`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lmp.Sersic(\n", + " centre=main_lens_centres[0],\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " mass_to_light_ratio=0.2,\n", + " ),\n", + " dark=al.mp.NFWSph(centre=main_lens_centres[0], kappa_s=0.1, scale_radius=20.0),\n", + " shear=al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005),\n", + ")\n", + "\n", + "lens_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lmp.Sersic(\n", + " centre=main_lens_centres[1],\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=120.0),\n", + " intensity=0.8,\n", + " effective_radius=0.7,\n", + " sersic_index=4.0,\n", + " mass_to_light_ratio=0.25,\n", + " ),\n", + " dark=al.mp.NFWSph(centre=main_lens_centres[1], kappa_s=0.08, scale_radius=20.0),\n", + ")\n", + "\n", + "main_lens_galaxies = [lens_0, lens_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy__\n", + "\n", + "A single compact source at z=1.0 with a `SersicCore` light profile, positioned near the group centre so its\n", + "lensed image forms a clearly visible Einstein-ring-like configuration around the two main lens galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The tracer is composed of the two main lens galaxies followed by the source. PyAutoLens orders galaxies\n", + "internally by redshift, so the deflection chain runs:\n", + "\n", + " image-plane \u2192 source-plane (deflected by both main lens galaxies' stellar + dark + shear contributions)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=main_lens_galaxies + [source])\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid), title=\"Group Mass Stellar Dark Image\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Pass the simulator a tracer to produce the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Write the simulated dataset to .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer JSON__\n", + "\n", + "Save the simulator `Tracer` so the true profiles can be inspected later." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres JSON__\n", + "\n", + "Save the main lens centres so the modeling scripts can load them via `al.from_json(...)`. This mirrors the\n", + "canonical `group/simulator.py` and `group/features/advanced/double_einstein_ring/simulator.py` conventions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/mass_stellar_dark/slam.ipynb b/notebooks/group/features/advanced/mass_stellar_dark/slam.ipynb index 2ce3b2592..f72e39558 100644 --- a/notebooks/group/features/advanced/mass_stellar_dark/slam.ipynb +++ b/notebooks/group/features/advanced/mass_stellar_dark/slam.ipynb @@ -1,742 +1,779 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM (Source, Light and Mass): Group Mass Stellar Dark\n", - "======================================================\n", - "\n", - "This script adapts the SLaM (Source, Light and Mass) pipelines to a group-scale strong lens where each main\n", - "lens galaxy is decomposed into a stellar component (tied to its light via a mass-to-light ratio) and a dark\n", - "matter halo.\n", - "\n", - "This script is the group analogue of:\n", - "\n", - " - `guides/modeling/slam_start_here.py` \u2014 the canonical single-galaxy SLaM walkthrough.\n", - " - `scripts/imaging/features/advanced/mass_stellar_dark/slam.py` \u2014 the single-galaxy decomposed-mass SLaM.\n", - " - `scripts/group/features/advanced/double_einstein_ring/slam.py` \u2014 the group SLaM with two source planes (the\n", - " `lens_dict` plumbing in that script is the structural template used here, simplified for a single source\n", - " plane).\n", - "\n", - "Each pipeline stage is a plain inline Python function. Per-lens priors are chained via\n", - "`al.util.chaining.mass_from`, image positions are derived automatically via `positions_likelihood_from`, and\n", - "MGE light profiles are constructed via `al.model_util.mge_model_from`.\n", - "\n", - "__Group-Specific Differences From Standard SLaM__\n", - "\n", - " - The lens-plane (z=0.5) is composed via the group `lens_dict` convention: one `af.Model(al.Galaxy)` entry per\n", - " main lens galaxy centre, with the `ExternalShear` attached only to `lens_0`.\n", - " - Each pipeline iterates over the main lens galaxies via `lens_{i}` keys rather than referencing a single\n", - " `lens` attribute.\n", - " - The MASS LIGHT DARK pipeline constructs the per-galaxy `lmp.Sersic + NFWSph` manually rather than calling\n", - " `al.util.chaining.mass_light_dark_from`, which only supports the single-`lens` layout.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP, SOURCE PIX 1, SOURCE PIX 2, LIGHT LP and a MASS LIGHT DARK pipeline, this group SLaM script\n", - "fits an `Imaging` dataset where in the final model:\n", - "\n", - " - Each main lens galaxy's light is a `Sersic` linear light profile.\n", - " - Each main lens galaxy's stellar mass distribution is a `Sersic` tied to its OWN light via a\n", - " `mass_to_light_ratio` (one per galaxy, free parameters).\n", - " - Each main lens galaxy's dark matter mass distribution is an `NFWSph` aligned with the bulge centre.\n", - " - The first main lens galaxy carries an `ExternalShear`.\n", - " - The source galaxy's light is a `Pixelization`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Helpers__\n", - "\n", - "`build_lens_dict_source_lp` constructs an `af.Model` lens_dict for the SOURCE LP pipeline: each main lens\n", - "galaxy has an MGE bulge plus a free `Isothermal` mass profile, with `ExternalShear` on `lens_0`.\n", - "\n", - "`build_lens_dict_light_lp` constructs an `af.Model` lens_dict for the LIGHT LP pipeline: each main lens galaxy\n", - "has a `lp_linear.Sersic` bulge (chosen because the MASS LIGHT DARK pipeline requires a `LightMassProfile`\n", - "sharing the same profile type), with mass and shear fixed from the SOURCE PIX 1 result.\n", - "\n", - "`build_lens_dict_mass_light_dark` constructs the final MASS LIGHT DARK lens_dict: each lens galaxy's bulge is\n", - "swapped to `lmp.Sersic` with priors carried over from LIGHT LP, plus a new `NFWSph` dark halo." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def n_lens_from(result) -> int:\n", - " return len(\n", - " [name for name in vars(result.instance.galaxies) if name.startswith(\"lens_\")]\n", - " )\n", - "\n", - "\n", - "def lens_galaxy_image_dict(result, n_lens: int) -> dict:\n", - " full = al.galaxy_name_image_dict_via_result_from(result=result)\n", - " return {\n", - " f\"('galaxies', 'lens_{i}')\": full[f\"('galaxies', 'lens_{i}')\"]\n", - " for i in range(n_lens)\n", - " }\n", - "\n", - "\n", - "def build_lens_dict_source_lp(\n", - " main_lens_centres,\n", - " redshift_lens: float,\n", - " mask_radius: float,\n", - " total_gaussians: int = 20,\n", - " gaussian_per_basis: int = 2,\n", - "):\n", - " lens_dict = {}\n", - " for i, centre in enumerate(main_lens_centres):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=total_gaussians,\n", - " gaussian_per_basis=gaussian_per_basis,\n", - " centre_prior_is_uniform=True,\n", - " centre=(centre[0], centre[1]),\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = (centre[0], centre[1])\n", - "\n", - " kwargs = dict(redshift=redshift_lens, bulge=bulge, mass=mass)\n", - " if i == 0:\n", - " kwargs[\"shear\"] = af.Model(al.mp.ExternalShear)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - " return lens_dict\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Initial fit using MGE bulges + `Isothermal` mass for each main lens galaxy, with an MGE source. This\n", - "constrains the source-plane geometry and the per-galaxy total-mass before subsequent pipelines decompose it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " main_lens_centres,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_dict = build_lens_dict_source_lp(\n", - " main_lens_centres=main_lens_centres,\n", - " redshift_lens=redshift_lens,\n", - " mask_radius=mask_radius,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " **lens_dict,\n", - " source=af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Pixelize the source using an initial mesh / regularization. Each main lens galaxy's mass is freed with priors\n", - "chained from the SOURCE LP pipeline via `al.util.chaining.mass_from`. Adapt images are stitched per-lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " n_lens = n_lens_from(source_lp_result)\n", - "\n", - " galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", - "\n", - " positions_likelihood = source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lens):\n", - " lens_inst = getattr(source_lp_result.instance.galaxies, f\"lens_{i}\")\n", - " lens_model_mass = getattr(source_lp_result.model.galaxies, f\"lens_{i}\").mass\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.Isothermal),\n", - " mass_result=lens_model_mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " kwargs = dict(\n", - " redshift=lens_inst.redshift,\n", - " bulge=lens_inst.bulge,\n", - " mass=mass,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = source_lp_result.model.galaxies.lens_0.shear\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " **lens_dict,\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Refines the source pixelization with an adapt mesh derived from the SOURCE PIX 1 source reconstruction. Each\n", - "main lens galaxy's bulge, mass and (for `lens_0`) shear are fixed to the SOURCE PIX 1 instance." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " n_lens = n_lens_from(source_lp_result)\n", - "\n", - " galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lens):\n", - " lens_inst_lp = getattr(source_lp_result.instance.galaxies, f\"lens_{i}\")\n", - " lens_inst_pix = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " kwargs = dict(\n", - " redshift=lens_inst_lp.redshift,\n", - " bulge=lens_inst_lp.bulge,\n", - " mass=lens_inst_pix.mass,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = lens_inst_pix.shear\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " **lens_dict,\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Refits each main lens galaxy's light using a `lp_linear.Sersic` bulge (replacing the SOURCE LP MGE bulge). The\n", - "linear `Sersic` is the light profile type the MASS LIGHT DARK pipeline pairs with `lmp.Sersic` for the\n", - "stellar-mass coupling. Mass + shear + source pixelization are fixed from SOURCE PIX 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " main_lens_centres,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " redshift_lens: float,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " n_lens = n_lens_from(source_result_for_lens)\n", - "\n", - " galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset, adapt_images=adapt_images)\n", - "\n", - " lens_dict = {}\n", - " for i, centre in enumerate(main_lens_centres):\n", - " bulge = af.Model(al.lp_linear.Sersic)\n", - " bulge.centre = (centre[0], centre[1])\n", - "\n", - " lens_inst = getattr(source_result_for_lens.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " kwargs = dict(\n", - " redshift=redshift_lens,\n", - " bulge=bulge,\n", - " mass=lens_inst.mass,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = source_result_for_lens.instance.galaxies.lens_0.shear\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(galaxies=af.Collection(**lens_dict, source=source))\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS LIGHT DARK PIPELINE__\n", - "\n", - "The terminal pipeline. Each main lens galaxy's bulge is swapped from the LIGHT LP `lp_linear.Sersic` to a\n", - "`lmp.Sersic` (light + stellar mass coupled via `mass_to_light_ratio`); priors on the bulge geometry / intensity\n", - "are carried over via `take_attributes`. A separate `NFWSph` dark halo is added per galaxy with its centre fixed\n", - "to the bulge centre.\n", - "\n", - "This pipeline is the per-galaxy generalisation of the imaging\n", - "`scripts/imaging/features/advanced/mass_stellar_dark/slam.py` MASS LIGHT DARK stage. Because\n", - "`al.util.chaining.mass_light_dark_from` accesses `light_result.instance.galaxies.lens.` directly (a\n", - "hardcoded single-lens path), the per-galaxy decomposition is constructed manually here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_light_dark(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " main_lens_centres,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " n_lens = n_lens_from(light_result)\n", - "\n", - " galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i, centre in enumerate(main_lens_centres):\n", - " bulge = af.Model(al.lmp.Sersic)\n", - " bulge.take_attributes(source=getattr(light_result.model.galaxies, f\"lens_{i}\"))\n", - " bulge.centre = (centre[0], centre[1])\n", - " bulge.mass_to_light_ratio = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - " dark = af.Model(al.mp.NFWSph)\n", - " dark.centre = (centre[0], centre[1])\n", - " dark.kappa_s = af.UniformPrior(lower_limit=0.0, upper_limit=1.0)\n", - " dark.scale_radius = af.UniformPrior(lower_limit=5.0, upper_limit=50.0)\n", - "\n", - " kwargs = dict(\n", - " redshift=light_result.instance.galaxies.lens_0.redshift,\n", - " bulge=bulge,\n", - " dark=dark,\n", - " )\n", - " if i == 0:\n", - " kwargs[\"shear\"] = source_result_for_lens.model.galaxies.lens_0.shear\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(galaxies=af.Collection(**lens_dict, source=source))\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_light_dark[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=250,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "\n", - "mask_radius = 3.7\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=list(main_lens_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"group\") / \"slam\" / \"mass_stellar_dark\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below runs the full group decomposed-mass SLaM pipeline. See the docstring above each function for a\n", - "description of each stage." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " main_lens_centres=main_lens_centres,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - ")\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", - "over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", - " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", - " noise_map=dataset.noise_map,\n", - ")\n", - "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=over_sampling)\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " main_lens_centres=main_lens_centres,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " redshift_lens=redshift_lens,\n", - ")\n", - "\n", - "mass_result = mass_light_dark(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " main_lens_centres=main_lens_centres,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM (Source, Light and Mass): Group Mass Stellar Dark\n", + "======================================================\n", + "\n", + "This script adapts the SLaM (Source, Light and Mass) pipelines to a group-scale strong lens where each main\n", + "lens galaxy is decomposed into a stellar component (tied to its light via a mass-to-light ratio) and a dark\n", + "matter halo.\n", + "\n", + "This script is the group analogue of:\n", + "\n", + " - `guides/modeling/slam_start_here.py` \u2014 the canonical single-galaxy SLaM walkthrough.\n", + " - `scripts/imaging/features/advanced/mass_stellar_dark/slam.py` \u2014 the single-galaxy decomposed-mass SLaM.\n", + " - `scripts/group/features/advanced/double_einstein_ring/slam.py` \u2014 the group SLaM with two source planes (the\n", + " `lens_dict` plumbing in that script is the structural template used here, simplified for a single source\n", + " plane).\n", + "\n", + "Each pipeline stage is a plain inline Python function. Per-lens priors are chained via\n", + "`al.util.chaining.mass_from`, image positions are derived automatically via `positions_likelihood_from`, and\n", + "MGE light profiles are constructed via `al.model_util.mge_model_from`.\n", + "\n", + "__Group-Specific Differences From Standard SLaM__\n", + "\n", + " - The lens-plane (z=0.5) is composed via the group `lens_dict` convention: one `af.Model(al.Galaxy)` entry per\n", + " main lens galaxy centre, with the `ExternalShear` attached only to `lens_0`.\n", + " - Each pipeline iterates over the main lens galaxies via `lens_{i}` keys rather than referencing a single\n", + " `lens` attribute.\n", + " - The MASS LIGHT DARK pipeline constructs the per-galaxy `lmp.Sersic + NFWSph` manually rather than calling\n", + " `al.util.chaining.mass_light_dark_from`, which only supports the single-`lens` layout.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP, SOURCE PIX 1, SOURCE PIX 2, LIGHT LP and a MASS LIGHT DARK pipeline, this group SLaM script\n", + "fits an `Imaging` dataset where in the final model:\n", + "\n", + " - Each main lens galaxy's light is a `Sersic` linear light profile.\n", + " - Each main lens galaxy's stellar mass distribution is a `Sersic` tied to its OWN light via a\n", + " `mass_to_light_ratio` (one per galaxy, free parameters).\n", + " - Each main lens galaxy's dark matter mass distribution is an `NFWSph` aligned with the bulge centre.\n", + " - The first main lens galaxy carries an `ExternalShear`.\n", + " - The source galaxy's light is a `Pixelization`." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Helpers__\n", + "\n", + "`build_lens_dict_source_lp` constructs an `af.Model` lens_dict for the SOURCE LP pipeline: each main lens\n", + "galaxy has an MGE bulge plus a free `Isothermal` mass profile, with `ExternalShear` on `lens_0`.\n", + "\n", + "`build_lens_dict_light_lp` constructs an `af.Model` lens_dict for the LIGHT LP pipeline: each main lens galaxy\n", + "has a `lp_linear.Sersic` bulge (chosen because the MASS LIGHT DARK pipeline requires a `LightMassProfile`\n", + "sharing the same profile type), with mass and shear fixed from the SOURCE PIX 1 result.\n", + "\n", + "`build_lens_dict_mass_light_dark` constructs the final MASS LIGHT DARK lens_dict: each lens galaxy's bulge is\n", + "swapped to `lmp.Sersic` with priors carried over from LIGHT LP, plus a new `NFWSph` dark halo." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def n_lens_from(result) -> int:\n", + " return len(\n", + " [name for name in vars(result.instance.galaxies) if name.startswith(\"lens_\")]\n", + " )\n", + "\n", + "\n", + "def lens_galaxy_image_dict(result, n_lens: int) -> dict:\n", + " full = al.galaxy_name_image_dict_via_result_from(result=result)\n", + " return {\n", + " f\"('galaxies', 'lens_{i}')\": full[f\"('galaxies', 'lens_{i}')\"]\n", + " for i in range(n_lens)\n", + " }\n", + "\n", + "\n", + "def build_lens_dict_source_lp(\n", + " main_lens_centres,\n", + " redshift_lens: float,\n", + " mask_radius: float,\n", + " total_gaussians: int = 20,\n", + " gaussian_per_basis: int = 2,\n", + "):\n", + " lens_dict = {}\n", + " for i, centre in enumerate(main_lens_centres):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=total_gaussians,\n", + " gaussian_per_basis=gaussian_per_basis,\n", + " centre_prior_is_uniform=True,\n", + " centre=(centre[0], centre[1]),\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = (centre[0], centre[1])\n", + "\n", + " kwargs = dict(redshift=redshift_lens, bulge=bulge, mass=mass)\n", + " if i == 0:\n", + " kwargs[\"shear\"] = af.Model(al.mp.ExternalShear)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + " return lens_dict\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Initial fit using MGE bulges + `Isothermal` mass for each main lens galaxy, with an MGE source. This\n", + "constrains the source-plane geometry and the per-galaxy total-mass before subsequent pipelines decompose it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " main_lens_centres,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_dict = build_lens_dict_source_lp(\n", + " main_lens_centres=main_lens_centres,\n", + " redshift_lens=redshift_lens,\n", + " mask_radius=mask_radius,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " **lens_dict,\n", + " source=af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Pixelize the source using an initial mesh / regularization. Each main lens galaxy's mass is freed with priors\n", + "chained from the SOURCE LP pipeline via `al.util.chaining.mass_from`. Adapt images are stitched per-lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " n_lens = n_lens_from(source_lp_result)\n", + "\n", + " galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", + "\n", + " positions_likelihood = source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lens):\n", + " lens_inst = getattr(source_lp_result.instance.galaxies, f\"lens_{i}\")\n", + " lens_model_mass = getattr(source_lp_result.model.galaxies, f\"lens_{i}\").mass\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.Isothermal),\n", + " mass_result=lens_model_mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " kwargs = dict(\n", + " redshift=lens_inst.redshift,\n", + " bulge=lens_inst.bulge,\n", + " mass=mass,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = source_lp_result.model.galaxies.lens_0.shear\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " **lens_dict,\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Refines the source pixelization with an adapt mesh derived from the SOURCE PIX 1 source reconstruction. Each\n", + "main lens galaxy's bulge, mass and (for `lens_0`) shear are fixed to the SOURCE PIX 1 instance." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " n_lens = n_lens_from(source_lp_result)\n", + "\n", + " galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lens):\n", + " lens_inst_lp = getattr(source_lp_result.instance.galaxies, f\"lens_{i}\")\n", + " lens_inst_pix = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " kwargs = dict(\n", + " redshift=lens_inst_lp.redshift,\n", + " bulge=lens_inst_lp.bulge,\n", + " mass=lens_inst_pix.mass,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = lens_inst_pix.shear\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " **lens_dict,\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Refits each main lens galaxy's light using a `lp_linear.Sersic` bulge (replacing the SOURCE LP MGE bulge). The\n", + "linear `Sersic` is the light profile type the MASS LIGHT DARK pipeline pairs with `lmp.Sersic` for the\n", + "stellar-mass coupling. Mass + shear + source pixelization are fixed from SOURCE PIX 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " main_lens_centres,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " redshift_lens: float,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " n_lens = n_lens_from(source_result_for_lens)\n", + "\n", + " galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset, adapt_images=adapt_images)\n", + "\n", + " lens_dict = {}\n", + " for i, centre in enumerate(main_lens_centres):\n", + " bulge = af.Model(al.lp_linear.Sersic)\n", + " bulge.centre = (centre[0], centre[1])\n", + "\n", + " lens_inst = getattr(source_result_for_lens.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " kwargs = dict(\n", + " redshift=redshift_lens,\n", + " bulge=bulge,\n", + " mass=lens_inst.mass,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = source_result_for_lens.instance.galaxies.lens_0.shear\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(galaxies=af.Collection(**lens_dict, source=source))\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS LIGHT DARK PIPELINE__\n", + "\n", + "The terminal pipeline. Each main lens galaxy's bulge is swapped from the LIGHT LP `lp_linear.Sersic` to a\n", + "`lmp.Sersic` (light + stellar mass coupled via `mass_to_light_ratio`); priors on the bulge geometry / intensity\n", + "are carried over via `take_attributes`. A separate `NFWSph` dark halo is added per galaxy with its centre fixed\n", + "to the bulge centre.\n", + "\n", + "This pipeline is the per-galaxy generalisation of the imaging\n", + "`scripts/imaging/features/advanced/mass_stellar_dark/slam.py` MASS LIGHT DARK stage. Because\n", + "`al.util.chaining.mass_light_dark_from` accesses `light_result.instance.galaxies.lens.` directly (a\n", + "hardcoded single-lens path), the per-galaxy decomposition is constructed manually here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_light_dark(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " main_lens_centres,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " n_lens = n_lens_from(light_result)\n", + "\n", + " galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i, centre in enumerate(main_lens_centres):\n", + " bulge = af.Model(al.lmp.Sersic)\n", + " bulge.take_attributes(source=getattr(light_result.model.galaxies, f\"lens_{i}\"))\n", + " bulge.centre = (centre[0], centre[1])\n", + " bulge.mass_to_light_ratio = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + " dark = af.Model(al.mp.NFWSph)\n", + " dark.centre = (centre[0], centre[1])\n", + " dark.kappa_s = af.UniformPrior(lower_limit=0.0, upper_limit=1.0)\n", + " dark.scale_radius = af.UniformPrior(lower_limit=5.0, upper_limit=50.0)\n", + "\n", + " kwargs = dict(\n", + " redshift=light_result.instance.galaxies.lens_0.redshift,\n", + " bulge=bulge,\n", + " dark=dark,\n", + " )\n", + " if i == 0:\n", + " kwargs[\"shear\"] = source_result_for_lens.model.galaxies.lens_0.shear\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, **kwargs)\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(galaxies=af.Collection(**lens_dict, source=source))\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_light_dark[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=250,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "\n", + "mask_radius = 3.7\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=list(main_lens_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"group\") / \"slam\" / \"mass_stellar_dark\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below runs the full group decomposed-mass SLaM pipeline. See the docstring above each function for a\n", + "description of each stage." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " main_lens_centres=main_lens_centres,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "galaxy_image_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + ")\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_dict)\n", + "over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", + " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", + " noise_map=dataset.noise_map,\n", + ")\n", + "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=over_sampling)\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " main_lens_centres=main_lens_centres,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " redshift_lens=redshift_lens,\n", + ")\n", + "\n", + "mass_result = mass_light_dark(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " main_lens_centres=main_lens_centres,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/operated_light_profile/modeling.ipynb b/notebooks/group/features/advanced/operated_light_profile/modeling.ipynb index 11410213a..a65905088 100644 --- a/notebooks/group/features/advanced/operated_light_profile/modeling.ipynb +++ b/notebooks/group/features/advanced/operated_light_profile/modeling.ipynb @@ -1,538 +1,575 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Operated Light Profiles (Group)\n", - "==================================================\n", - "\n", - "Operated light profiles are light profiles which are assumed to have already been convolved with the PSF. This means\n", - "that during the model-fitting process, these profiles are NOT convolved again with the PSF, unlike standard light\n", - "profiles.\n", - "\n", - "This is useful when modeling data where the lens light subtraction was performed using PSF-convolved models, or\n", - "when a galaxy has compact point-source emission (e.g. an AGN) that has already been blurred by the telescope optics.\n", - "The operated profile is fitted directly to this already-convolved emission.\n", - "\n", - "For a group-scale lens, this feature can be applied to both the main lens galaxies and the extra galaxies. Each\n", - "galaxy may have a compact nuclear component that is best represented as an operated light profile.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres:** The centres of the main lens galaxies and extra galaxies are loaded from JSON files.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - Each main lens galaxy's light is a linear ``Sersic`` bulge plus an operated linear ``Gaussian`` PSF component.\n", - " - The first main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear``.\n", - " - There are two extra lens galaxies with linear operated ``Sersic`` light and ``IsothermalSph`` total mass\n", - " distributions, with centres fixed to the observed centres of light.\n", - " - The source galaxy's light is a linear ``SersicCore`` (which IS convolved with the PSF as normal).\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the ``group/modeling`` and\n", - "``imaging/features/advanced/operated_light_profile/modeling`` notebooks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `operated`, which includes lens light that has already been convolved with\n", - "the PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"operated\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/operated_light_profile/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask for group-scale lenses, which is larger than galaxy-scale because the\n", - "group has emission spread over a wider area." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) is performed to ensure the lens\n", - "calculations are accurate across the full field of the group." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a group lens model where:\n", - "\n", - " - The main lens galaxy's light is a linear ``Sersic`` bulge plus an operated linear ``Gaussian`` PSF component.\n", - " The operated ``Gaussian`` represents compact point-source emission (e.g. AGN) that has already been convolved\n", - " with the telescope PSF. It is NOT convolved again during fitting.\n", - "\n", - " - The main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear``.\n", - "\n", - " - The extra galaxies use linear operated ``Sersic`` light profiles. Because the data was simulated with\n", - " operated profiles for these galaxies, we model them with the same type. Their centres are fixed to the\n", - " observed centres of light.\n", - "\n", - " - The source galaxy's light is a linear ``SersicCore``, which IS convolved with the PSF as it represents\n", - " genuine unconvolved source emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = af.Model(al.lp_linear.Sersic)\n", - " psf = af.Model(al.lp_linear_operated.Gaussian)\n", - " bulge.centre = psf.centre\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " psf=psf,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light (operated -- not convolved again with PSF)\n", - "\n", - " bulge = af.Model(al.lp_linear_operated.Gaussian)\n", - " bulge.centre = centre\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ``info`` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Improved Lens Model__\n", - "\n", - "The model above uses simple parametric profiles. For better performance, we replace the main lens galaxy's light\n", - "with an MGE model and the extra galaxies with MGE models. However, here the operated Gaussian PSF component\n", - "remains to demonstrate the operated light profile feature." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " psf = af.Model(al.lp_linear_operated.Gaussian)\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " psf=psf,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light (MGE)\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"operated_light_profiles\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the ``AnalysisImaging`` object defining how the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The likelihood evaluation time for operated light profiles is faster than standard light profiles because the\n", - "PSF convolution step is omitted for those components. The overall run-time may be slightly slower due to the\n", - "additional ``psf`` parameter on the main lens galaxy.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result contains entries for each main lens galaxy (with its operated ``psf`` component), the source galaxy\n", - "and the extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script shows how to fit a group-scale lens model where compact emission is modeled using operated light\n", - "profiles. The operated profiles bypass PSF convolution, making them ideal for point-source emission or\n", - "pre-convolved light components." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Operated Light Profiles (Group)\n", + "==================================================\n", + "\n", + "Operated light profiles are light profiles which are assumed to have already been convolved with the PSF. This means\n", + "that during the model-fitting process, these profiles are NOT convolved again with the PSF, unlike standard light\n", + "profiles.\n", + "\n", + "This is useful when modeling data where the lens light subtraction was performed using PSF-convolved models, or\n", + "when a galaxy has compact point-source emission (e.g. an AGN) that has already been blurred by the telescope optics.\n", + "The operated profile is fitted directly to this already-convolved emission.\n", + "\n", + "For a group-scale lens, this feature can be applied to both the main lens galaxies and the extra galaxies. Each\n", + "galaxy may have a compact nuclear component that is best represented as an operated light profile.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres:** The centres of the main lens galaxies and extra galaxies are loaded from JSON files.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - Each main lens galaxy's light is a linear ``Sersic`` bulge plus an operated linear ``Gaussian`` PSF component.\n", + " - The first main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear``.\n", + " - There are two extra lens galaxies with linear operated ``Sersic`` light and ``IsothermalSph`` total mass\n", + " distributions, with centres fixed to the observed centres of light.\n", + " - The source galaxy's light is a linear ``SersicCore`` (which IS convolved with the PSF as normal).\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the ``group/modeling`` and\n", + "``imaging/features/advanced/operated_light_profile/modeling`` notebooks." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `operated`, which includes lens light that has already been convolved with\n", + "the PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"operated\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/operated_light_profile/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask for group-scale lenses, which is larger than galaxy-scale because the\n", + "group has emission spread over a wider area." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) is performed to ensure the lens\n", + "calculations are accurate across the full field of the group." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a group lens model where:\n", + "\n", + " - The main lens galaxy's light is a linear ``Sersic`` bulge plus an operated linear ``Gaussian`` PSF component.\n", + " The operated ``Gaussian`` represents compact point-source emission (e.g. AGN) that has already been convolved\n", + " with the telescope PSF. It is NOT convolved again during fitting.\n", + "\n", + " - The main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear``.\n", + "\n", + " - The extra galaxies use linear operated ``Sersic`` light profiles. Because the data was simulated with\n", + " operated profiles for these galaxies, we model them with the same type. Their centres are fixed to the\n", + " observed centres of light.\n", + "\n", + " - The source galaxy's light is a linear ``SersicCore``, which IS convolved with the PSF as it represents\n", + " genuine unconvolved source emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = af.Model(al.lp_linear.Sersic)\n", + " psf = af.Model(al.lp_linear_operated.Gaussian)\n", + " bulge.centre = psf.centre\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " psf=psf,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light (operated -- not convolved again with PSF)\n", + "\n", + " bulge = af.Model(al.lp_linear_operated.Gaussian)\n", + " bulge.centre = centre\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The ``info`` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Improved Lens Model__\n", + "\n", + "The model above uses simple parametric profiles. For better performance, we replace the main lens galaxy's light\n", + "with an MGE model and the extra galaxies with MGE models. However, here the operated Gaussian PSF component\n", + "remains to demonstrate the operated light profile feature." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " psf = af.Model(al.lp_linear_operated.Gaussian)\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " psf=psf,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light (MGE)\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"operated_light_profiles\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the ``AnalysisImaging`` object defining how the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The likelihood evaluation time for operated light profiles is faster than standard light profiles because the\n", + "PSF convolution step is omitted for those components. The overall run-time may be slightly slower due to the\n", + "additional ``psf`` parameter on the main lens galaxy.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result contains entries for each main lens galaxy (with its operated ``psf`` component), the source galaxy\n", + "and the extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script shows how to fit a group-scale lens model where compact emission is modeled using operated light\n", + "profiles. The operated profiles bypass PSF convolution, making them ideal for point-source emission or\n", + "pre-convolved light components." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/operated_light_profile/simulator.ipynb b/notebooks/group/features/advanced/operated_light_profile/simulator.ipynb index d74e0d0c9..f7c5c07ca 100644 --- a/notebooks/group/features/advanced/operated_light_profile/simulator.ipynb +++ b/notebooks/group/features/advanced/operated_light_profile/simulator.ipynb @@ -1,448 +1,485 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Operated Light Profiles (Group)\n", - "==========================================\n", - "\n", - "This script simulates a group-scale strong lens where the galaxy light uses operated light profiles. Operated\n", - "light profiles represent emission that has already been convolved with the PSF, meaning they are NOT convolved\n", - "again during the simulation.\n", - "\n", - "This creates a dataset where the lens light appears more smoothed, as would be the case if lens light subtraction\n", - "was performed using PSF-convolved models and the residual light retains that convolution.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** The ``dataset_type`` describes the type of data being simulated.\n", - "- **Grid:** Define the 2d grid of (y,x) coordinates for the simulation.\n", - "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Main Lens Galaxies:** The main lens galaxy with an operated ``Sersic`` light profile.\n", - "- **Extra Galaxies:** The extra galaxies with operated ``Sersic`` light profiles.\n", - "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", - "- **Ray Tracing:** Use all galaxies to set up a tracer.\n", - "- **Dataset:** Simulate and output the dataset.\n", - "- **Centre JSON Files:** Save the centres as JSON files.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates ``Imaging`` of a 'group-scale' strong lens where:\n", - "\n", - " - The main lens galaxy's light is an operated ``Sersic`` (already PSF-convolved).\n", - " - The main lens galaxy's total mass distribution is an ``IsothermalSph``.\n", - " - The extra galaxies' light profiles are operated ``Sersic`` profiles.\n", - " - The source galaxy's light is a ``SersicCore``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The dataset is output to ``dataset/group/operated``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"operated\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(250, 250),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Define the centres of the main lens galaxies and extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [(0.0, 0.0)]\n", - "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive oversampling at all galaxy centres for accurate light profile evaluation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=main_lens_centres + extra_galaxies_centres,\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "The main lens galaxy uses an operated ``Sersic`` light profile. This means the light is assumed to have\n", - "already been convolved with the PSF and will not be convolved again during simulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_operated.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=0.7,\n", - " effective_radius=2.0,\n", - " sersic_index=4.0,\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "main_lens_galaxies = [lens_0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "The extra galaxies also use operated ``Sersic`` light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_operated.Sersic(\n", - " centre=(3.5, 2.5),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=0.9,\n", - " effective_radius=0.8,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_operated.Sersic(\n", - " centre=(-4.4, -5.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=0.9,\n", - " effective_radius=0.8,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy__\n", - "\n", - "The source galaxy uses a standard (non-operated) ``SersicCore`` profile, as the source light has not been\n", - "pre-convolved." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Use all galaxies to set up a tracer, which will generate the image for the simulated ``Imaging`` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the ``Tracer`` in the dataset folder as a .json file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centre JSON Files__\n", - "\n", - "Save the centres of the main lens galaxies and extra galaxies as JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder ``autolens_workspace/dataset/group/operated``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Operated Light Profiles (Group)\n", + "==========================================\n", + "\n", + "This script simulates a group-scale strong lens where the galaxy light uses operated light profiles. Operated\n", + "light profiles represent emission that has already been convolved with the PSF, meaning they are NOT convolved\n", + "again during the simulation.\n", + "\n", + "This creates a dataset where the lens light appears more smoothed, as would be the case if lens light subtraction\n", + "was performed using PSF-convolved models and the residual light retains that convolution.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** The ``dataset_type`` describes the type of data being simulated.\n", + "- **Grid:** Define the 2d grid of (y,x) coordinates for the simulation.\n", + "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Main Lens Galaxies:** The main lens galaxy with an operated ``Sersic`` light profile.\n", + "- **Extra Galaxies:** The extra galaxies with operated ``Sersic`` light profiles.\n", + "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", + "- **Ray Tracing:** Use all galaxies to set up a tracer.\n", + "- **Dataset:** Simulate and output the dataset.\n", + "- **Centre JSON Files:** Save the centres as JSON files.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates ``Imaging`` of a 'group-scale' strong lens where:\n", + "\n", + " - The main lens galaxy's light is an operated ``Sersic`` (already PSF-convolved).\n", + " - The main lens galaxy's total mass distribution is an ``IsothermalSph``.\n", + " - The extra galaxies' light profiles are operated ``Sersic`` profiles.\n", + " - The source galaxy's light is a ``SersicCore``." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The dataset is output to ``dataset/group/operated``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"operated\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(250, 250),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Define the centres of the main lens galaxies and extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [(0.0, 0.0)]\n", + "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive oversampling at all galaxy centres for accurate light profile evaluation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=main_lens_centres + extra_galaxies_centres,\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "The main lens galaxy uses an operated ``Sersic`` light profile. This means the light is assumed to have\n", + "already been convolved with the PSF and will not be convolved again during simulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_operated.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=0.7,\n", + " effective_radius=2.0,\n", + " sersic_index=4.0,\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "main_lens_galaxies = [lens_0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "The extra galaxies also use operated ``Sersic`` light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_operated.Sersic(\n", + " centre=(3.5, 2.5),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=0.9,\n", + " effective_radius=0.8,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_operated.Sersic(\n", + " centre=(-4.4, -5.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=0.9,\n", + " effective_radius=0.8,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy__\n", + "\n", + "The source galaxy uses a standard (non-operated) ``SersicCore`` profile, as the source light has not been\n", + "pre-convolved." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Use all galaxies to set up a tracer, which will generate the image for the simulated ``Imaging`` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the ``Tracer`` in the dataset folder as a .json file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centre JSON Files__\n", + "\n", + "Save the centres of the main lens galaxies and extra galaxies as JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder ``autolens_workspace/dataset/group/operated``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/shapelets/fit.ipynb b/notebooks/group/features/advanced/shapelets/fit.ipynb index 2e7265fd5..67a203750 100644 --- a/notebooks/group/features/advanced/shapelets/fit.ipynb +++ b/notebooks/group/features/advanced/shapelets/fit.ipynb @@ -1,413 +1,450 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fit Features: Shapelets (Group)\n", - "===============================\n", - "\n", - "A shapelet is a basis function that is appropriate for capturing the exponential / disk-like features of a galaxy.\n", - "This script demonstrates how to create a fit using shapelet light profiles for a group-scale strong lens, without\n", - "performing a non-linear search.\n", - "\n", - "This is useful for understanding the shapelet API and visualizing how shapelets decompose the source galaxy's light\n", - "in a group lens context. The lens mass model and galaxy parameters are specified manually using the true values\n", - "from the simulation.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres:** Load galaxy centres from JSON files.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Basis:** Build a ``Basis`` of multiple linear shapelet light profiles.\n", - "- **Fit:** Fit the lens model to the dataset using the shapelet basis.\n", - "- **Intensities:** Extract the solved-for intensity values from the fit.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - The main lens galaxy and extra galaxies use the true simulation parameters.\n", - " - The source galaxy's light is a superposition of ~20 linear ``ShapeletPolar`` profiles.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the ``group/fit`` and\n", - "``imaging/features/advanced/shapelets/fit`` notebooks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset ``simple``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask for group-scale lenses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Basis__\n", - "\n", - "We build a ``Basis`` of ~20 linear polar shapelet light profiles for the source galaxy. These shapelets:\n", - "\n", - " - All share the same centre and elliptical components.\n", - " - The size of the shapelet basis is controlled by a ``beta`` parameter.\n", - " - The ``intensity`` of each shapelet is solved via linear algebra.\n", - "\n", - "We use the true source centre (0.0, 0.1) and a reasonable guess for ``beta``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_n = 5\n", - "total_m = sum(range(2, total_n + 1)) + 1\n", - "\n", - "n_count = 1\n", - "m_count = -1\n", - "\n", - "shapelets_bulge_list = []\n", - "\n", - "shapelet_0 = al.lp_linear.ShapeletPolar(\n", - " n=0,\n", - " m=0,\n", - " centre=(0.0, 0.1),\n", - " ell_comps=(0.0, 0.0),\n", - " beta=1.0,\n", - ")\n", - "\n", - "shapelets_bulge_list.append(shapelet_0)\n", - "\n", - "for i in range(total_n + total_m):\n", - " shapelet = al.lp_linear.ShapeletPolar(\n", - " n=n_count,\n", - " m=m_count,\n", - " centre=(0.0, 0.1),\n", - " ell_comps=(0.0, 0.0),\n", - " beta=1.0,\n", - " )\n", - "\n", - " shapelets_bulge_list.append(shapelet)\n", - "\n", - " m_count += 2\n", - "\n", - " if m_count > n_count:\n", - " n_count += 1\n", - " m_count = -n_count\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We now create a fit using the true group lens galaxy parameters and the shapelet source model.\n", - "\n", - "The main lens galaxy and extra galaxies use the true simulation parameters, while the source galaxy uses\n", - "the shapelet basis whose intensities will be solved via linear algebra.\n", - "\n", - "We set ``use_positive_only_solver=False`` because shapelets require negative intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=bulge,\n", - ")\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_positive_only_solver=False),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the shapelets do a reasonable job at capturing the appearance of the source galaxy\n", - "in the group lens context, with faint residuals where the lensed source is located." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The fit contains the solved-for intensity values for each shapelet. These can be extracted from the\n", - "``linear_light_profile_intensity_dict``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_bulge = fit.tracer.galaxies[-1].bulge\n", - "\n", - "print(\n", - " f\"\\nIntensity of source galaxy's first shapelet = \"\n", - " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A ``Tracer`` where all linear light profiles are replaced with standard light profiles using the solved-for\n", - "intensities is accessible from the fit. This can be used for visualization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "subplot_basis_image(basis=tracer.galaxies[-1].bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script shows how to fit a group-scale lens using shapelets for the source galaxy light, demonstrating\n", - "the ``Basis`` API and how shapelet intensities are solved via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fit Features: Shapelets (Group)\n", + "===============================\n", + "\n", + "A shapelet is a basis function that is appropriate for capturing the exponential / disk-like features of a galaxy.\n", + "This script demonstrates how to create a fit using shapelet light profiles for a group-scale strong lens, without\n", + "performing a non-linear search.\n", + "\n", + "This is useful for understanding the shapelet API and visualizing how shapelets decompose the source galaxy's light\n", + "in a group lens context. The lens mass model and galaxy parameters are specified manually using the true values\n", + "from the simulation.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres:** Load galaxy centres from JSON files.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Basis:** Build a ``Basis`` of multiple linear shapelet light profiles.\n", + "- **Fit:** Fit the lens model to the dataset using the shapelet basis.\n", + "- **Intensities:** Extract the solved-for intensity values from the fit.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - The main lens galaxy and extra galaxies use the true simulation parameters.\n", + " - The source galaxy's light is a superposition of ~20 linear ``ShapeletPolar`` profiles.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the ``group/fit`` and\n", + "``imaging/features/advanced/shapelets/fit`` notebooks." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset ``simple``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask for group-scale lenses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Basis__\n", + "\n", + "We build a ``Basis`` of ~20 linear polar shapelet light profiles for the source galaxy. These shapelets:\n", + "\n", + " - All share the same centre and elliptical components.\n", + " - The size of the shapelet basis is controlled by a ``beta`` parameter.\n", + " - The ``intensity`` of each shapelet is solved via linear algebra.\n", + "\n", + "We use the true source centre (0.0, 0.1) and a reasonable guess for ``beta``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_n = 5\n", + "total_m = sum(range(2, total_n + 1)) + 1\n", + "\n", + "n_count = 1\n", + "m_count = -1\n", + "\n", + "shapelets_bulge_list = []\n", + "\n", + "shapelet_0 = al.lp_linear.ShapeletPolar(\n", + " n=0,\n", + " m=0,\n", + " centre=(0.0, 0.1),\n", + " ell_comps=(0.0, 0.0),\n", + " beta=1.0,\n", + ")\n", + "\n", + "shapelets_bulge_list.append(shapelet_0)\n", + "\n", + "for i in range(total_n + total_m):\n", + " shapelet = al.lp_linear.ShapeletPolar(\n", + " n=n_count,\n", + " m=m_count,\n", + " centre=(0.0, 0.1),\n", + " ell_comps=(0.0, 0.0),\n", + " beta=1.0,\n", + " )\n", + "\n", + " shapelets_bulge_list.append(shapelet)\n", + "\n", + " m_count += 2\n", + "\n", + " if m_count > n_count:\n", + " n_count += 1\n", + " m_count = -n_count\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We now create a fit using the true group lens galaxy parameters and the shapelet source model.\n", + "\n", + "The main lens galaxy and extra galaxies use the true simulation parameters, while the source galaxy uses\n", + "the shapelet basis whose intensities will be solved via linear algebra.\n", + "\n", + "We set ``use_positive_only_solver=False`` because shapelets require negative intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=bulge,\n", + ")\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_positive_only_solver=False),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the shapelets do a reasonable job at capturing the appearance of the source galaxy\n", + "in the group lens context, with faint residuals where the lensed source is located." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The fit contains the solved-for intensity values for each shapelet. These can be extracted from the\n", + "``linear_light_profile_intensity_dict``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_bulge = fit.tracer.galaxies[-1].bulge\n", + "\n", + "print(\n", + " f\"\\nIntensity of source galaxy's first shapelet = \"\n", + " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A ``Tracer`` where all linear light profiles are replaced with standard light profiles using the solved-for\n", + "intensities is accessible from the fit. This can be used for visualization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "subplot_basis_image(basis=tracer.galaxies[-1].bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script shows how to fit a group-scale lens using shapelets for the source galaxy light, demonstrating\n", + "the ``Basis`` API and how shapelet intensities are solved via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/shapelets/modeling.ipynb b/notebooks/group/features/advanced/shapelets/modeling.ipynb index 0bcf3a6c3..2a42fc99c 100644 --- a/notebooks/group/features/advanced/shapelets/modeling.ipynb +++ b/notebooks/group/features/advanced/shapelets/modeling.ipynb @@ -1,481 +1,518 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Shapelets (Group)\n", - "====================================\n", - "\n", - "A shapelet is a basis function that is appropriate for capturing the exponential / disk-like features of a galaxy.\n", - "It has been employed in many strong lensing studies to model the light of lensed source galaxies, because it can\n", - "represent features of disky star forming galaxies that a single Sersic function cannot.\n", - "\n", - "Shapelets are described in full in the following paper:\n", - "\n", - " https://arxiv.org/abs/astro-ph/0105178\n", - "\n", - "This script performs a group-scale model-fit using shapelets to decompose the source galaxy's light into ~20\n", - "shapelet basis functions. The ``intensity`` of every shapelet is solved for via linear algebra.\n", - "\n", - "For group-scale lenses, the main lens galaxies and extra galaxies are modeled with MGE light profiles (which are\n", - "more efficient than shapelets for smooth elliptical galaxies), while the source galaxy benefits from the\n", - "flexibility of shapelets to capture complex morphology.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres:** The centres of the main lens galaxies and extra galaxies are loaded from JSON files.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - Each main lens galaxy's light is an MGE bulge.\n", - " - The first main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear``.\n", - " - There are two extra lens galaxies with MGE light and ``IsothermalSph`` total mass distributions.\n", - " - The source galaxy's light is a superposition of ~20 linear ``ShapeletPolar`` profiles.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the ``group/modeling`` and\n", - "``imaging/features/advanced/shapelets/modeling`` notebooks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset ``simple``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask for group-scale lenses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a group lens model where:\n", - "\n", - " - Main lens galaxies use MGE light profiles (efficient for smooth elliptical galaxies).\n", - " - Extra galaxies use MGE light profiles with fixed centres.\n", - " - The source galaxy uses a shapelet basis decomposition, which captures complex morphology using fewer\n", - " non-linear parameters than a Sersic profile.\n", - "\n", - "The shapelets are composed as a ``Basis`` of ~20 ``ShapeletPolar`` profiles with linked centres, elliptical\n", - "components and beta parameters. Only the centre, ellipticity and beta size parameter are non-linear; the\n", - "intensity of each shapelet is solved via linear algebra.\n", - "\n", - "__Positive Negative Solver__\n", - "\n", - "Shapelets require the ability to use negative intensities in the linear algebra solution (unlike MGE which\n", - "uses positive-only). We therefore use ``use_positive_only_solver=False`` in the analysis settings." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source (Shapelet Basis):\n", - "\n", - "total_n = 10\n", - "total_m = sum(range(2, total_n + 1)) + 1\n", - "\n", - "shapelets_bulge_list = af.Collection(\n", - " af.Model(al.lp_linear.ShapeletPolar) for _ in range(total_n + total_m + 1)\n", - ")\n", - "\n", - "n_count = 1\n", - "m_count = -1\n", - "\n", - "for i, shapelet in enumerate(shapelets_bulge_list):\n", - " if i == 0:\n", - " shapelet.n = 0\n", - " shapelet.m = 0\n", - "\n", - " else:\n", - " shapelet.n = n_count\n", - " shapelet.m = m_count\n", - "\n", - " m_count += 2\n", - "\n", - " if m_count > n_count:\n", - " n_count += 1\n", - " m_count = -n_count\n", - "\n", - " shapelet.centre = shapelets_bulge_list[0].centre\n", - " shapelet.ell_comps = shapelets_bulge_list[0].ell_comps\n", - " shapelet.beta = shapelets_bulge_list[0].beta\n", - "\n", - "bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=shapelets_bulge_list,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ``info`` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"shapelets\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the ``AnalysisImaging`` object. We set ``use_positive_only_solver=False`` because shapelets require\n", - "negative intensities in the linear solution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " settings=al.Settings(use_positive_only_solver=False),\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The likelihood evaluation time for shapelets is significantly slower than standard light profiles because\n", - "every shapelet image must be computed and convolved with the PSF. However, gains are made from the reduced\n", - "number of non-linear parameters (the source has only ~3 free parameters: centre, ellipticity, beta).\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result contains entries for each main lens galaxy, the source galaxy (with shapelet basis) and the\n", - "extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script shows how to fit a group-scale lens model where the source galaxy light is decomposed into a\n", - "shapelet basis. Shapelets capture complex morphology and are useful for irregular star-forming source galaxies,\n", - "while the group lens galaxies are efficiently modeled with MGE profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Shapelets (Group)\n", + "====================================\n", + "\n", + "A shapelet is a basis function that is appropriate for capturing the exponential / disk-like features of a galaxy.\n", + "It has been employed in many strong lensing studies to model the light of lensed source galaxies, because it can\n", + "represent features of disky star forming galaxies that a single Sersic function cannot.\n", + "\n", + "Shapelets are described in full in the following paper:\n", + "\n", + " https://arxiv.org/abs/astro-ph/0105178\n", + "\n", + "This script performs a group-scale model-fit using shapelets to decompose the source galaxy's light into ~20\n", + "shapelet basis functions. The ``intensity`` of every shapelet is solved for via linear algebra.\n", + "\n", + "For group-scale lenses, the main lens galaxies and extra galaxies are modeled with MGE light profiles (which are\n", + "more efficient than shapelets for smooth elliptical galaxies), while the source galaxy benefits from the\n", + "flexibility of shapelets to capture complex morphology.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres:** The centres of the main lens galaxies and extra galaxies are loaded from JSON files.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - Each main lens galaxy's light is an MGE bulge.\n", + " - The first main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear``.\n", + " - There are two extra lens galaxies with MGE light and ``IsothermalSph`` total mass distributions.\n", + " - The source galaxy's light is a superposition of ~20 linear ``ShapeletPolar`` profiles.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the ``group/modeling`` and\n", + "``imaging/features/advanced/shapelets/modeling`` notebooks." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset ``simple``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask for group-scale lenses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a group lens model where:\n", + "\n", + " - Main lens galaxies use MGE light profiles (efficient for smooth elliptical galaxies).\n", + " - Extra galaxies use MGE light profiles with fixed centres.\n", + " - The source galaxy uses a shapelet basis decomposition, which captures complex morphology using fewer\n", + " non-linear parameters than a Sersic profile.\n", + "\n", + "The shapelets are composed as a ``Basis`` of ~20 ``ShapeletPolar`` profiles with linked centres, elliptical\n", + "components and beta parameters. Only the centre, ellipticity and beta size parameter are non-linear; the\n", + "intensity of each shapelet is solved via linear algebra.\n", + "\n", + "__Positive Negative Solver__\n", + "\n", + "Shapelets require the ability to use negative intensities in the linear algebra solution (unlike MGE which\n", + "uses positive-only). We therefore use ``use_positive_only_solver=False`` in the analysis settings." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source (Shapelet Basis):\n", + "\n", + "total_n = 10\n", + "total_m = sum(range(2, total_n + 1)) + 1\n", + "\n", + "shapelets_bulge_list = af.Collection(\n", + " af.Model(al.lp_linear.ShapeletPolar) for _ in range(total_n + total_m + 1)\n", + ")\n", + "\n", + "n_count = 1\n", + "m_count = -1\n", + "\n", + "for i, shapelet in enumerate(shapelets_bulge_list):\n", + " if i == 0:\n", + " shapelet.n = 0\n", + " shapelet.m = 0\n", + "\n", + " else:\n", + " shapelet.n = n_count\n", + " shapelet.m = m_count\n", + "\n", + " m_count += 2\n", + "\n", + " if m_count > n_count:\n", + " n_count += 1\n", + " m_count = -n_count\n", + "\n", + " shapelet.centre = shapelets_bulge_list[0].centre\n", + " shapelet.ell_comps = shapelets_bulge_list[0].ell_comps\n", + " shapelet.beta = shapelets_bulge_list[0].beta\n", + "\n", + "bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=shapelets_bulge_list,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The ``info`` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"shapelets\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the ``AnalysisImaging`` object. We set ``use_positive_only_solver=False`` because shapelets require\n", + "negative intensities in the linear solution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " settings=al.Settings(use_positive_only_solver=False),\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The likelihood evaluation time for shapelets is significantly slower than standard light profiles because\n", + "every shapelet image must be computed and convolved with the PSF. However, gains are made from the reduced\n", + "number of non-linear parameters (the source has only ~3 free parameters: centre, ellipticity, beta).\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result contains entries for each main lens galaxy, the source galaxy (with shapelet basis) and the\n", + "extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script shows how to fit a group-scale lens model where the source galaxy light is decomposed into a\n", + "shapelet basis. Shapelets capture complex morphology and are useful for irregular star-forming source galaxies,\n", + "while the group lens galaxies are efficiently modeled with MGE profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/sky_background/fit.ipynb b/notebooks/group/features/advanced/sky_background/fit.ipynb index 1a0298fef..b0e5b4ab6 100644 --- a/notebooks/group/features/advanced/sky_background/fit.ipynb +++ b/notebooks/group/features/advanced/sky_background/fit.ipynb @@ -1,310 +1,347 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fit Features: Sky Background (Group)\n", - "=====================================\n", - "\n", - "The background of an image is the light that is not associated with the strong lens we are interested in. This\n", - "script demonstrates how to include the sky background in a fit for a group-scale strong lens, without performing\n", - "a non-linear search.\n", - "\n", - "This illustrates the ``DatasetModel`` API for sky background subtraction using standard objects like a ``Galaxy``,\n", - "``Tracer`` and ``FitImaging``.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres:** Load galaxy centres from JSON files.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Fit:** Demonstrate fitting with a ``DatasetModel`` that includes sky background.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - The main lens galaxy and extra galaxies use the true simulation parameters.\n", - " - The sky background is modeled using a ``DatasetModel`` with the true sky level.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the ``group/fit`` and\n", - "``imaging/features/advanced/sky_background/fit`` notebooks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset ``sky_background``, which has not had the sky background subtracted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"sky_background\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/sky_background/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask for group-scale lenses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We create a fit using the true group lens galaxy parameters and a ``DatasetModel`` that includes the true\n", - "sky background level of 5.0 electrons per second.\n", - "\n", - "For the galaxies, we use the true parameters from the simulation. The key addition is the ``DatasetModel``\n", - "with ``background_sky_level=5.0``, which subtracts the sky from the data during the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "dataset_model = al.DatasetModel(background_sky_level=5.0)\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer, dataset_model=dataset_model)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the sky is subtracted from the data such that the outskirts are zero.\n", - "The group galaxies and lensed source are well fitted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script shows how to include the sky background in a group-scale lens fit using a ``DatasetModel`` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fit Features: Sky Background (Group)\n", + "=====================================\n", + "\n", + "The background of an image is the light that is not associated with the strong lens we are interested in. This\n", + "script demonstrates how to include the sky background in a fit for a group-scale strong lens, without performing\n", + "a non-linear search.\n", + "\n", + "This illustrates the ``DatasetModel`` API for sky background subtraction using standard objects like a ``Galaxy``,\n", + "``Tracer`` and ``FitImaging``.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres:** Load galaxy centres from JSON files.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Fit:** Demonstrate fitting with a ``DatasetModel`` that includes sky background.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - The main lens galaxy and extra galaxies use the true simulation parameters.\n", + " - The sky background is modeled using a ``DatasetModel`` with the true sky level.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the ``group/fit`` and\n", + "``imaging/features/advanced/sky_background/fit`` notebooks." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset ``sky_background``, which has not had the sky background subtracted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"sky_background\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/sky_background/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask for group-scale lenses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We create a fit using the true group lens galaxy parameters and a ``DatasetModel`` that includes the true\n", + "sky background level of 5.0 electrons per second.\n", + "\n", + "For the galaxies, we use the true parameters from the simulation. The key addition is the ``DatasetModel``\n", + "with ``background_sky_level=5.0``, which subtracts the sky from the data during the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "dataset_model = al.DatasetModel(background_sky_level=5.0)\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer, dataset_model=dataset_model)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the sky is subtracted from the data such that the outskirts are zero.\n", + "The group galaxies and lensed source are well fitted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script shows how to include the sky background in a group-scale lens fit using a ``DatasetModel`` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/sky_background/modeling.ipynb b/notebooks/group/features/advanced/sky_background/modeling.ipynb index 4b5073702..d882ce9c5 100644 --- a/notebooks/group/features/advanced/sky_background/modeling.ipynb +++ b/notebooks/group/features/advanced/sky_background/modeling.ipynb @@ -1,459 +1,496 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Sky Background (Group)\n", - "=========================================\n", - "\n", - "The background of an image is the light that is not associated with the strong lens we are interested in. This is\n", - "due to light from the sky, zodiacal light, and light from other galaxies in the field of view.\n", - "\n", - "The background sky is often subtracted from image data during data reduction. If this subtraction is perfect,\n", - "there is no need to include the sky in the model. However, achieving a perfect subtraction is difficult, and\n", - "residuals can leave a signal degenerate with the lens galaxy light, especially for low surface brightness features.\n", - "\n", - "For group-scale lenses, this is particularly important because the larger field of view means more area is\n", - "affected by sky background uncertainties, and the faint outskirts of multiple group galaxies can be affected.\n", - "\n", - "This example illustrates how to include the sky background in the model-fitting of a group-scale ``Imaging``\n", - "dataset as a non-linear free parameter.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres:** The centres of the main lens galaxies and extra galaxies are loaded from JSON files.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - The sky background is included as a ``DatasetModel`` with a free ``background_sky_level`` parameter.\n", - " - Each main lens galaxy's light is an MGE bulge.\n", - " - The first main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear``.\n", - " - There are two extra lens galaxies with MGE light and ``IsothermalSph`` total mass distributions.\n", - " - The source galaxy's light is an MGE.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the ``group/modeling`` and\n", - "``imaging/features/advanced/sky_background/modeling`` notebooks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset ``sky_background``, which has not had the sky background subtracted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"sky_background\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/sky_background/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask for group-scale lenses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a group lens model that includes a sky background component:\n", - "\n", - " - The sky background is modeled as a ``DatasetModel`` with a free ``background_sky_level`` parameter.\n", - " This is not part of the ``galaxies`` collection but is a separate model component.\n", - "\n", - " - The main lens galaxies use MGE light profiles and isothermal mass profiles. Only the first main lens\n", - " galaxy carries an ``ExternalShear``.\n", - "\n", - " - The extra galaxies use MGE light profiles with fixed centres and isothermal mass profiles.\n", - "\n", - " - The source galaxy uses an MGE light profile.\n", - "\n", - "The prior on ``background_sky_level`` must be set manually based on the expected sky level in the data.\n", - "In this example, the true sky level is 5.0 electrons per second." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Sky Background:\n", - "\n", - "dataset_model = af.Model(al.DatasetModel)\n", - "dataset_model.background_sky_level = af.UniformPrior(lower_limit=0.0, upper_limit=10.0)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " dataset_model=dataset_model,\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ``info`` attribute shows the model in a readable format. This confirms that the sky is a model component\n", - "that is not part of the ``galaxies`` collection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"sky_background\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the ``AnalysisImaging`` object defining how the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "Adding the background sky model to the analysis has a negligible impact on the run time, as it simply adds a\n", - "constant value to the data. The run time is dominated by the group galaxy light and mass model evaluation.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result contains the inferred ``background_sky_level`` alongside the group galaxy model parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "print(result.instance.dataset_model.background_sky_level)\n", - "\n", - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script shows how to include the sky background as part of a group-scale lens model using a ``DatasetModel``\n", - "object. This ensures uncertainties on galaxy light profile parameters fully account for sky background\n", - "subtraction errors, which is especially important for the extended faint emission of group-scale lenses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Sky Background (Group)\n", + "=========================================\n", + "\n", + "The background of an image is the light that is not associated with the strong lens we are interested in. This is\n", + "due to light from the sky, zodiacal light, and light from other galaxies in the field of view.\n", + "\n", + "The background sky is often subtracted from image data during data reduction. If this subtraction is perfect,\n", + "there is no need to include the sky in the model. However, achieving a perfect subtraction is difficult, and\n", + "residuals can leave a signal degenerate with the lens galaxy light, especially for low surface brightness features.\n", + "\n", + "For group-scale lenses, this is particularly important because the larger field of view means more area is\n", + "affected by sky background uncertainties, and the faint outskirts of multiple group galaxies can be affected.\n", + "\n", + "This example illustrates how to include the sky background in the model-fitting of a group-scale ``Imaging``\n", + "dataset as a non-linear free parameter.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres:** The centres of the main lens galaxies and extra galaxies are loaded from JSON files.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - The sky background is included as a ``DatasetModel`` with a free ``background_sky_level`` parameter.\n", + " - Each main lens galaxy's light is an MGE bulge.\n", + " - The first main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear``.\n", + " - There are two extra lens galaxies with MGE light and ``IsothermalSph`` total mass distributions.\n", + " - The source galaxy's light is an MGE.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the ``group/modeling`` and\n", + "``imaging/features/advanced/sky_background/modeling`` notebooks." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset ``sky_background``, which has not had the sky background subtracted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"sky_background\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/sky_background/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask for group-scale lenses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a group lens model that includes a sky background component:\n", + "\n", + " - The sky background is modeled as a ``DatasetModel`` with a free ``background_sky_level`` parameter.\n", + " This is not part of the ``galaxies`` collection but is a separate model component.\n", + "\n", + " - The main lens galaxies use MGE light profiles and isothermal mass profiles. Only the first main lens\n", + " galaxy carries an ``ExternalShear``.\n", + "\n", + " - The extra galaxies use MGE light profiles with fixed centres and isothermal mass profiles.\n", + "\n", + " - The source galaxy uses an MGE light profile.\n", + "\n", + "The prior on ``background_sky_level`` must be set manually based on the expected sky level in the data.\n", + "In this example, the true sky level is 5.0 electrons per second." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Sky Background:\n", + "\n", + "dataset_model = af.Model(al.DatasetModel)\n", + "dataset_model.background_sky_level = af.UniformPrior(lower_limit=0.0, upper_limit=10.0)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " dataset_model=dataset_model,\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The ``info`` attribute shows the model in a readable format. This confirms that the sky is a model component\n", + "that is not part of the ``galaxies`` collection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"sky_background\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the ``AnalysisImaging`` object defining how the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "Adding the background sky model to the analysis has a negligible impact on the run time, as it simply adds a\n", + "constant value to the data. The run time is dominated by the group galaxy light and mass model evaluation.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result contains the inferred ``background_sky_level`` alongside the group galaxy model parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "print(result.instance.dataset_model.background_sky_level)\n", + "\n", + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script shows how to include the sky background as part of a group-scale lens model using a ``DatasetModel``\n", + "object. This ensures uncertainties on galaxy light profile parameters fully account for sky background\n", + "subtraction errors, which is especially important for the extended faint emission of group-scale lenses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/sky_background/simulator.ipynb b/notebooks/group/features/advanced/sky_background/simulator.ipynb index 304370bd1..6795d5c67 100644 --- a/notebooks/group/features/advanced/sky_background/simulator.ipynb +++ b/notebooks/group/features/advanced/sky_background/simulator.ipynb @@ -1,437 +1,474 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Sky Background (Group)\n", - "=================================\n", - "\n", - "This script simulates a group-scale strong lens dataset where the sky background is NOT subtracted from the image.\n", - "The sky background is therefore present in the data and must be accounted for during model-fitting.\n", - "\n", - "This is used to demonstrate sky background modeling in the\n", - "``group/features/advanced/sky_background/modeling.py`` example.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** The ``dataset_type`` describes the type of data being simulated.\n", - "- **Grid:** Define the 2d grid of (y,x) coordinates for the simulation.\n", - "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid.\n", - "- **Main Lens Galaxies:** The main lens galaxy at the origin.\n", - "- **Extra Galaxies:** Two companion galaxies near the lens system.\n", - "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", - "- **Ray Tracing:** Use all galaxies to set up a tracer.\n", - "- **Dataset:** Simulate and output the dataset.\n", - "- **Centre JSON Files:** Save the centres as JSON files.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates ``Imaging`` of a 'group-scale' strong lens where:\n", - "\n", - " - The group consists of one main lens galaxy and two extra galaxies with ``SersicSph`` light and\n", - " ``IsothermalSph`` mass profiles.\n", - " - A single source galaxy with ``SersicCore`` light.\n", - " - The sky background level is 5.0 electrons per second and is NOT subtracted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The dataset is output to ``dataset/group/sky_background``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"sky_background\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2d grid of (y,x) coordinates for the simulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(250, 250),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Define the centres of the main lens galaxies and extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [(0.0, 0.0)]\n", - "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive oversampling at all galaxy centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=main_lens_centres + extra_galaxies_centres,\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data.\n", - "\n", - "The ``background_sky_level`` is set to 5.0 electrons per second, much higher than the standard 0.1. The\n", - "``subtract_background_sky=False`` flag ensures the sky is NOT subtracted, so it remains in the output data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=5.0,\n", - " add_poisson_noise_to_data=True,\n", - " subtract_background_sky=False,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "The main lens galaxy at the origin with spherical Sersic light and isothermal mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "main_lens_galaxies = [lens_0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "Two companion galaxies near the lens system." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy__\n", - "\n", - "The source galaxy whose lensed images we simulate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Use all galaxies to set up a tracer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the ``Tracer`` in the dataset folder as a .json file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centre JSON Files__\n", - "\n", - "Save the centres of the main lens galaxies and extra galaxies as JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder ``autolens_workspace/dataset/group/sky_background``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Sky Background (Group)\n", + "=================================\n", + "\n", + "This script simulates a group-scale strong lens dataset where the sky background is NOT subtracted from the image.\n", + "The sky background is therefore present in the data and must be accounted for during model-fitting.\n", + "\n", + "This is used to demonstrate sky background modeling in the\n", + "``group/features/advanced/sky_background/modeling.py`` example.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** The ``dataset_type`` describes the type of data being simulated.\n", + "- **Grid:** Define the 2d grid of (y,x) coordinates for the simulation.\n", + "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid.\n", + "- **Main Lens Galaxies:** The main lens galaxy at the origin.\n", + "- **Extra Galaxies:** Two companion galaxies near the lens system.\n", + "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", + "- **Ray Tracing:** Use all galaxies to set up a tracer.\n", + "- **Dataset:** Simulate and output the dataset.\n", + "- **Centre JSON Files:** Save the centres as JSON files.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates ``Imaging`` of a 'group-scale' strong lens where:\n", + "\n", + " - The group consists of one main lens galaxy and two extra galaxies with ``SersicSph`` light and\n", + " ``IsothermalSph`` mass profiles.\n", + " - A single source galaxy with ``SersicCore`` light.\n", + " - The sky background level is 5.0 electrons per second and is NOT subtracted." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The dataset is output to ``dataset/group/sky_background``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"sky_background\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2d grid of (y,x) coordinates for the simulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(250, 250),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Define the centres of the main lens galaxies and extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [(0.0, 0.0)]\n", + "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive oversampling at all galaxy centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=main_lens_centres + extra_galaxies_centres,\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data.\n", + "\n", + "The ``background_sky_level`` is set to 5.0 electrons per second, much higher than the standard 0.1. The\n", + "``subtract_background_sky=False`` flag ensures the sky is NOT subtracted, so it remains in the output data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=5.0,\n", + " add_poisson_noise_to_data=True,\n", + " subtract_background_sky=False,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "The main lens galaxy at the origin with spherical Sersic light and isothermal mass." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "main_lens_galaxies = [lens_0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "Two companion galaxies near the lens system." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy__\n", + "\n", + "The source galaxy whose lensed images we simulate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Use all galaxies to set up a tracer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the ``Tracer`` in the dataset folder as a .json file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centre JSON Files__\n", + "\n", + "Save the centres of the main lens galaxies and extra galaxies as JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder ``autolens_workspace/dataset/group/sky_background``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/subhalo/detect/start_here.ipynb b/notebooks/group/features/advanced/subhalo/detect/start_here.ipynb index 563362ced..8e0a20339 100644 --- a/notebooks/group/features/advanced/subhalo/detect/start_here.ipynb +++ b/notebooks/group/features/advanced/subhalo/detect/start_here.ipynb @@ -1,1152 +1,1189 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Subhalo Detection: Group\n", - "=========================\n", - "\n", - "Strong gravitational lenses can be used to detect the presence of small-scale dark matter (DM) subhalos. This occurs\n", - "when the DM subhalo overlaps the lensed source emission, and therefore gravitationally perturbs the observed image of\n", - "the lensed source galaxy.\n", - "\n", - "This script extends the standard DM subhalo detection pipeline to group-scale lenses. The key adaptation is that the\n", - "lens model includes ALL group galaxies (main lens galaxies + extra galaxies) with their mass profiles, alongside\n", - "the dark matter subhalo.\n", - "\n", - "__Contents__\n", - "\n", - "- **SLaM Pipelines:** The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines.\n", - "- **Grid Search:** The second stage uses a grid-search of non-linear searches.\n", - "- **Group Adaptation:** The lens model includes all group galaxies.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres:** Load galaxy centres from JSON files.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid.\n", - "- **SLaM Pipeline Functions:** The pipeline functions adapted for group-scale lenses.\n", - "- **Bayesian Evidence:** Determine if a DM subhalo was detected.\n", - "- **Grid Search Result:** Inspect the grid search results.\n", - "\n", - "__SLaM Pipelines__\n", - "\n", - "The Source, (lens) Light and Mass (SLaM) pipelines are used for all DM subhalo detection analyses. The SLaM\n", - "pipelines are extended with a SUBHALO PIPELINE which performs three chained non-linear searches:\n", - "\n", - " 1) Fits the lens model without a DM subhalo to establish a Bayesian evidence baseline.\n", - " 2) Performs a grid-search where each cell includes a DM subhalo confined to a small 2D region.\n", - " 3) Refines the best-fit subhalo model initialized from the highest evidence grid cell.\n", - "\n", - "__Group Adaptation__\n", - "\n", - "For group-scale lenses, the key differences from galaxy-scale subhalo detection are:\n", - "\n", - " - The lens model includes all main lens galaxies and extra galaxies with their mass profiles.\n", - " - The subhalo is added as an additional galaxy in the model alongside the group galaxies.\n", - " - The larger field of view and more complex mass distribution mean the subhalo must be searched\n", - " across a wider area.\n", - "\n", - "__Model__\n", - "\n", - "Using SOURCE LP, SOURCE PIX, LIGHT LP, MASS TOTAL and SUBHALO PIPELINES this script fits ``Imaging`` of a\n", - "group-scale strong lens where in the final model:\n", - "\n", - " - The main lens galaxy's light is an MGE bulge.\n", - " - The main lens galaxy's total mass distribution is a ``PowerLaw``.\n", - " - Extra galaxies have MGE light and ``IsothermalSph`` mass with fixed centres.\n", - " - A dark matter subhalo near the lens galaxy mass is included as an ``NFWMCRLudlowSph``.\n", - " - The source galaxy is an ``Inversion``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Fits the group lens with parametric light profiles for the source, using MGE for all galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " main_lens_centres,\n", - " extra_galaxies_centres,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " # Main Lens Galaxies:\n", - "\n", - " lens_dict = {}\n", - "\n", - " for i, centre in enumerate(main_lens_centres):\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - " # Extra Galaxies:\n", - "\n", - " extra_galaxies_list = []\n", - "\n", - " for centre in extra_galaxies_centres:\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_galaxy = af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=mass\n", - " )\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - " extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - " # Source:\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", - "\n", - " # Overall Model:\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Fits a pixelized source reconstruction, using the lens model from the SOURCE LP PIPELINE." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " # Use the lens model from the SOURCE LP result, fixing light and freeing mass.\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens_0.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens_0.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens_0.shear\n", - "\n", - " lens_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens_0.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens_0.bulge,\n", - " mass=mass,\n", - " shear=shear,\n", - " )\n", - "\n", - " lens_dict = {\"lens_0\": lens_0}\n", - "\n", - " # Fix extra galaxies to their best-fit values.\n", - "\n", - " extra_galaxies = source_lp_result.instance.extra_galaxies\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Refines the pixelized source reconstruction with adaptive mesh." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens_0.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens_0.bulge,\n", - " mass=source_pix_result_1.instance.galaxies.lens_0.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens_0.shear,\n", - " )\n", - "\n", - " lens_dict = {\"lens_0\": lens_0}\n", - "\n", - " extra_galaxies = source_pix_result_1.instance.extra_galaxies\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Refits the lens light using MGE, with the source pixelization fixed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " main_lens_centres,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " lens_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens_0.redshift,\n", - " bulge=lens_bulge,\n", - " mass=source_result_for_lens.instance.galaxies.lens_0.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens_0.shear,\n", - " )\n", - "\n", - " lens_dict = {\"lens_0\": lens_0}\n", - "\n", - " extra_galaxies = source_result_for_lens.instance.extra_galaxies\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Fits the total mass distribution as a ``PowerLaw``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens_0.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " bulge = light_result.instance.galaxies.lens_0.bulge\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " lens_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens_0.redshift,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens_0.shear,\n", - " )\n", - "\n", - " lens_dict = {\"lens_0\": lens_0}\n", - "\n", - " extra_galaxies = light_result.instance.extra_galaxies\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (no subhalo)__\n", - "\n", - "Refits the lens model without a DM subhalo to establish a Bayesian evidence baseline." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def subhalo_no_subhalo(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " mass_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=mass_result)\n", - " lens_0 = mass_result.model.galaxies.lens_0\n", - "\n", - " lens_dict = {\"lens_0\": lens_0}\n", - "\n", - " extra_galaxies = mass_result.instance.extra_galaxies\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (grid search)__\n", - "\n", - "Performs a grid search where each cell includes a DM subhalo confined to a small 2D region of the image plane.\n", - "\n", - "For group-scale lenses, the grid search area may need to be larger than galaxy-scale lenses because the\n", - "lensed source images can span a wider area due to the more complex mass distribution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def subhalo_grid_search(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " mass_result: af.Result,\n", - " subhalo_no_subhalo_result: af.Result,\n", - " subhalo_mass: af.Model,\n", - " grid_dimension_arcsec: float = 5.0,\n", - " number_of_steps: int = 2,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " )\n", - "\n", - " subhalo = af.Model(al.Galaxy, mass=subhalo_mass)\n", - "\n", - " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", - " subhalo.mass.centre_0 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - " )\n", - " subhalo.mass.centre_1 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - " )\n", - "\n", - " subhalo.redshift = subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift\n", - " subhalo.mass.redshift_object = (\n", - " subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift\n", - " )\n", - " subhalo.mass.redshift_source = (\n", - " subhalo_no_subhalo_result.instance.galaxies.source.redshift\n", - " )\n", - "\n", - " lens_0 = mass_result.model.galaxies.lens_0\n", - " source = al.util.chaining.source_from(result=mass_result)\n", - "\n", - " lens_dict = {\"lens_0\": lens_0, \"subhalo\": subhalo}\n", - "\n", - " extra_galaxies = mass_result.instance.extra_galaxies\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[2]_[search_lens_plane]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " subhalo_grid_search = af.SearchGridSearch(\n", - " search=search,\n", - " number_of_steps=number_of_steps,\n", - " )\n", - "\n", - " return subhalo_grid_search.fit(\n", - " model=model,\n", - " analysis=analysis,\n", - " grid_priors=[\n", - " model.galaxies.subhalo.mass.centre_1,\n", - " model.galaxies.subhalo.mass.centre_0,\n", - " ],\n", - " info=settings_search.info,\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (refine)__\n", - "\n", - "Refines the best-fit subhalo model, initializing the subhalo centre from the highest evidence grid cell." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def subhalo_refine(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " mass_result: af.Result,\n", - " subhalo_no_subhalo_result: af.Result,\n", - " subhalo_grid_search_result: af.Result,\n", - " subhalo_mass: af.Model,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " )\n", - "\n", - " subhalo = af.Model(\n", - " al.Galaxy,\n", - " redshift=subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift,\n", - " mass=subhalo_mass,\n", - " )\n", - "\n", - " subhalo.redshift = subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift\n", - " subhalo.mass.redshift_object = (\n", - " subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift\n", - " )\n", - " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", - " subhalo.mass.centre = subhalo_grid_search_result.model_centred_absolute(\n", - " a=1.0\n", - " ).galaxies.subhalo.mass.centre\n", - "\n", - " subhalo.redshift = subhalo_grid_search_result.model.galaxies.subhalo.redshift\n", - " subhalo.mass.redshift_object = subhalo.redshift\n", - "\n", - " source = subhalo_grid_search_result.model.galaxies.source\n", - "\n", - " lens_dict = {\n", - " \"lens_0\": subhalo_grid_search_result.model.galaxies.lens_0,\n", - " \"subhalo\": subhalo,\n", - " }\n", - "\n", - " extra_galaxies = mass_result.instance.extra_galaxies\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[3]_[single_plane_refine]\",\n", - " **settings_search.search_dict,\n", - " n_live=600,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load, plot and mask the ``Imaging`` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"dark_matter_subhalo\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/advanced/subhalo/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at all galaxy centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"subhalo_detect_group\"),\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The mesh shape for the pixelized source reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE adapted for group-scale lenses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " main_lens_centres=main_lens_centres,\n", - " extra_galaxies_centres=extra_galaxies_centres,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", - " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", - " noise_map=dataset.noise_map,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=over_sampling)\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " main_lens_centres=main_lens_centres,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")\n", - "\n", - "result_no_subhalo = subhalo_no_subhalo(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mass_result=mass_result,\n", - ")\n", - "\n", - "result_subhalo_grid_search = subhalo_grid_search(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mass_result=mass_result,\n", - " subhalo_no_subhalo_result=result_no_subhalo,\n", - " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", - " grid_dimension_arcsec=5.0,\n", - " number_of_steps=2,\n", - ")\n", - "\n", - "result_with_subhalo = subhalo_refine(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mass_result=mass_result,\n", - " subhalo_no_subhalo_result=result_no_subhalo,\n", - " subhalo_grid_search_result=result_subhalo_grid_search,\n", - " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Bayesian Evidence__\n", - "\n", - "To determine if a DM subhalo was detected by the pipeline, we compare the log of the Bayesian evidences of\n", - "the model-fits performed with and without a subhalo.\n", - "\n", - "The following scale describes how different log evidence increases correspond to detection significances:\n", - "\n", - " - Negative log evidence increase: No detection.\n", - " - Log evidence increase between 0 and 3: No detection.\n", - " - Log evidence increase between 3 and 5: Weak evidence, should consider it a non-detection.\n", - " - Log evidence increase between 5 and 10: Medium evidence, but still inconclusive.\n", - " - Log evidence increase between 10 and 20: Strong evidence, consider it a detection.\n", - " - Log evidence increase > 20: Very strong evidence, definitive detection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "evidence_no_subhalo = result_no_subhalo.samples.log_evidence\n", - "evidence_with_subhalo = result_with_subhalo.samples.log_evidence\n", - "\n", - "log_evidence_increase = evidence_with_subhalo - evidence_no_subhalo\n", - "\n", - "print(\"Evidence Increase: \", log_evidence_increase)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood__\n", - "\n", - "The log likelihood increase provides a simpler metric for how well the subhalo model fits the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_likelihood_no_subhalo = result_no_subhalo.samples.log_likelihood\n", - "log_likelihood_with_subhalo = result_with_subhalo.samples.log_likelihood\n", - "\n", - "log_likelihood_increase = log_likelihood_with_subhalo - log_likelihood_no_subhalo\n", - "\n", - "print(\"Log Likelihood Increase: \", log_likelihood_increase)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid Search Result__\n", - "\n", - "The grid search results can be used to inspect where in the image plane a subhalo provides the best fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "subhalo_grid_search_result = al.subhalo.SubhaloGridSearchResult(\n", - " result=result_subhalo_grid_search\n", - ")\n", - "\n", - "log_evidence_array = subhalo_grid_search_result.figure_of_merit_array(\n", - " use_log_evidences=True,\n", - " relative_to_value=result_no_subhalo.samples.log_evidence,\n", - ")\n", - "\n", - "print(\"Log Evidence Array: \\n\")\n", - "print(log_evidence_array)\n", - "\n", - "aplt.plot_array(array=log_evidence_array, title=\"\")\n", - "\n", - "mass_array = subhalo_grid_search_result.subhalo_mass_array\n", - "\n", - "print(\"Mass Array: \\n\")\n", - "print(mass_array)\n", - "\n", - "subhalo_centres_grid = subhalo_grid_search_result.subhalo_centres_grid\n", - "\n", - "print(\"Subhalo Centres Grid: \\n\")\n", - "print(subhalo_centres_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Subhalo Detection: Group\n", + "=========================\n", + "\n", + "Strong gravitational lenses can be used to detect the presence of small-scale dark matter (DM) subhalos. This occurs\n", + "when the DM subhalo overlaps the lensed source emission, and therefore gravitationally perturbs the observed image of\n", + "the lensed source galaxy.\n", + "\n", + "This script extends the standard DM subhalo detection pipeline to group-scale lenses. The key adaptation is that the\n", + "lens model includes ALL group galaxies (main lens galaxies + extra galaxies) with their mass profiles, alongside\n", + "the dark matter subhalo.\n", + "\n", + "__Contents__\n", + "\n", + "- **SLaM Pipelines:** The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines.\n", + "- **Grid Search:** The second stage uses a grid-search of non-linear searches.\n", + "- **Group Adaptation:** The lens model includes all group galaxies.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres:** Load galaxy centres from JSON files.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid.\n", + "- **SLaM Pipeline Functions:** The pipeline functions adapted for group-scale lenses.\n", + "- **Bayesian Evidence:** Determine if a DM subhalo was detected.\n", + "- **Grid Search Result:** Inspect the grid search results.\n", + "\n", + "__SLaM Pipelines__\n", + "\n", + "The Source, (lens) Light and Mass (SLaM) pipelines are used for all DM subhalo detection analyses. The SLaM\n", + "pipelines are extended with a SUBHALO PIPELINE which performs three chained non-linear searches:\n", + "\n", + " 1) Fits the lens model without a DM subhalo to establish a Bayesian evidence baseline.\n", + " 2) Performs a grid-search where each cell includes a DM subhalo confined to a small 2D region.\n", + " 3) Refines the best-fit subhalo model initialized from the highest evidence grid cell.\n", + "\n", + "__Group Adaptation__\n", + "\n", + "For group-scale lenses, the key differences from galaxy-scale subhalo detection are:\n", + "\n", + " - The lens model includes all main lens galaxies and extra galaxies with their mass profiles.\n", + " - The subhalo is added as an additional galaxy in the model alongside the group galaxies.\n", + " - The larger field of view and more complex mass distribution mean the subhalo must be searched\n", + " across a wider area.\n", + "\n", + "__Model__\n", + "\n", + "Using SOURCE LP, SOURCE PIX, LIGHT LP, MASS TOTAL and SUBHALO PIPELINES this script fits ``Imaging`` of a\n", + "group-scale strong lens where in the final model:\n", + "\n", + " - The main lens galaxy's light is an MGE bulge.\n", + " - The main lens galaxy's total mass distribution is a ``PowerLaw``.\n", + " - Extra galaxies have MGE light and ``IsothermalSph`` mass with fixed centres.\n", + " - A dark matter subhalo near the lens galaxy mass is included as an ``NFWMCRLudlowSph``.\n", + " - The source galaxy is an ``Inversion``." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Fits the group lens with parametric light profiles for the source, using MGE for all galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " main_lens_centres,\n", + " extra_galaxies_centres,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " # Main Lens Galaxies:\n", + "\n", + " lens_dict = {}\n", + "\n", + " for i, centre in enumerate(main_lens_centres):\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + " # Extra Galaxies:\n", + "\n", + " extra_galaxies_list = []\n", + "\n", + " for centre in extra_galaxies_centres:\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_galaxy = af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=mass\n", + " )\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + " extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + " # Source:\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", + "\n", + " # Overall Model:\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Fits a pixelized source reconstruction, using the lens model from the SOURCE LP PIPELINE." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " # Use the lens model from the SOURCE LP result, fixing light and freeing mass.\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens_0.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens_0.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens_0.shear\n", + "\n", + " lens_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens_0.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens_0.bulge,\n", + " mass=mass,\n", + " shear=shear,\n", + " )\n", + "\n", + " lens_dict = {\"lens_0\": lens_0}\n", + "\n", + " # Fix extra galaxies to their best-fit values.\n", + "\n", + " extra_galaxies = source_lp_result.instance.extra_galaxies\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Refines the pixelized source reconstruction with adaptive mesh." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens_0.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens_0.bulge,\n", + " mass=source_pix_result_1.instance.galaxies.lens_0.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens_0.shear,\n", + " )\n", + "\n", + " lens_dict = {\"lens_0\": lens_0}\n", + "\n", + " extra_galaxies = source_pix_result_1.instance.extra_galaxies\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Refits the lens light using MGE, with the source pixelization fixed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " main_lens_centres,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " lens_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens_0.redshift,\n", + " bulge=lens_bulge,\n", + " mass=source_result_for_lens.instance.galaxies.lens_0.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens_0.shear,\n", + " )\n", + "\n", + " lens_dict = {\"lens_0\": lens_0}\n", + "\n", + " extra_galaxies = source_result_for_lens.instance.extra_galaxies\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Fits the total mass distribution as a ``PowerLaw``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens_0.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " bulge = light_result.instance.galaxies.lens_0.bulge\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " lens_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens_0.redshift,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens_0.shear,\n", + " )\n", + "\n", + " lens_dict = {\"lens_0\": lens_0}\n", + "\n", + " extra_galaxies = light_result.instance.extra_galaxies\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (no subhalo)__\n", + "\n", + "Refits the lens model without a DM subhalo to establish a Bayesian evidence baseline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def subhalo_no_subhalo(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " mass_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=mass_result)\n", + " lens_0 = mass_result.model.galaxies.lens_0\n", + "\n", + " lens_dict = {\"lens_0\": lens_0}\n", + "\n", + " extra_galaxies = mass_result.instance.extra_galaxies\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (grid search)__\n", + "\n", + "Performs a grid search where each cell includes a DM subhalo confined to a small 2D region of the image plane.\n", + "\n", + "For group-scale lenses, the grid search area may need to be larger than galaxy-scale lenses because the\n", + "lensed source images can span a wider area due to the more complex mass distribution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def subhalo_grid_search(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " mass_result: af.Result,\n", + " subhalo_no_subhalo_result: af.Result,\n", + " subhalo_mass: af.Model,\n", + " grid_dimension_arcsec: float = 5.0,\n", + " number_of_steps: int = 2,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " )\n", + "\n", + " subhalo = af.Model(al.Galaxy, mass=subhalo_mass)\n", + "\n", + " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", + " subhalo.mass.centre_0 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + " )\n", + " subhalo.mass.centre_1 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + " )\n", + "\n", + " subhalo.redshift = subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift\n", + " subhalo.mass.redshift_object = (\n", + " subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift\n", + " )\n", + " subhalo.mass.redshift_source = (\n", + " subhalo_no_subhalo_result.instance.galaxies.source.redshift\n", + " )\n", + "\n", + " lens_0 = mass_result.model.galaxies.lens_0\n", + " source = al.util.chaining.source_from(result=mass_result)\n", + "\n", + " lens_dict = {\"lens_0\": lens_0, \"subhalo\": subhalo}\n", + "\n", + " extra_galaxies = mass_result.instance.extra_galaxies\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[2]_[search_lens_plane]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " subhalo_grid_search = af.SearchGridSearch(\n", + " search=search,\n", + " number_of_steps=number_of_steps,\n", + " )\n", + "\n", + " return subhalo_grid_search.fit(\n", + " model=model,\n", + " analysis=analysis,\n", + " grid_priors=[\n", + " model.galaxies.subhalo.mass.centre_1,\n", + " model.galaxies.subhalo.mass.centre_0,\n", + " ],\n", + " info=settings_search.info,\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (refine)__\n", + "\n", + "Refines the best-fit subhalo model, initializing the subhalo centre from the highest evidence grid cell." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def subhalo_refine(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " mass_result: af.Result,\n", + " subhalo_no_subhalo_result: af.Result,\n", + " subhalo_grid_search_result: af.Result,\n", + " subhalo_mass: af.Model,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " )\n", + "\n", + " subhalo = af.Model(\n", + " al.Galaxy,\n", + " redshift=subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift,\n", + " mass=subhalo_mass,\n", + " )\n", + "\n", + " subhalo.redshift = subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift\n", + " subhalo.mass.redshift_object = (\n", + " subhalo_no_subhalo_result.instance.galaxies.lens_0.redshift\n", + " )\n", + " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", + " subhalo.mass.centre = subhalo_grid_search_result.model_centred_absolute(\n", + " a=1.0\n", + " ).galaxies.subhalo.mass.centre\n", + "\n", + " subhalo.redshift = subhalo_grid_search_result.model.galaxies.subhalo.redshift\n", + " subhalo.mass.redshift_object = subhalo.redshift\n", + "\n", + " source = subhalo_grid_search_result.model.galaxies.source\n", + "\n", + " lens_dict = {\n", + " \"lens_0\": subhalo_grid_search_result.model.galaxies.lens_0,\n", + " \"subhalo\": subhalo,\n", + " }\n", + "\n", + " extra_galaxies = mass_result.instance.extra_galaxies\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[3]_[single_plane_refine]\",\n", + " **settings_search.search_dict,\n", + " n_live=600,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load, plot and mask the ``Imaging`` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"dark_matter_subhalo\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/advanced/subhalo/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at all galaxy centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"subhalo_detect_group\"),\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The mesh shape for the pixelized source reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE adapted for group-scale lenses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " main_lens_centres=main_lens_centres,\n", + " extra_galaxies_centres=extra_galaxies_centres,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", + " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", + " noise_map=dataset.noise_map,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=over_sampling)\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " main_lens_centres=main_lens_centres,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")\n", + "\n", + "result_no_subhalo = subhalo_no_subhalo(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mass_result=mass_result,\n", + ")\n", + "\n", + "result_subhalo_grid_search = subhalo_grid_search(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mass_result=mass_result,\n", + " subhalo_no_subhalo_result=result_no_subhalo,\n", + " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", + " grid_dimension_arcsec=5.0,\n", + " number_of_steps=2,\n", + ")\n", + "\n", + "result_with_subhalo = subhalo_refine(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mass_result=mass_result,\n", + " subhalo_no_subhalo_result=result_no_subhalo,\n", + " subhalo_grid_search_result=result_subhalo_grid_search,\n", + " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Bayesian Evidence__\n", + "\n", + "To determine if a DM subhalo was detected by the pipeline, we compare the log of the Bayesian evidences of\n", + "the model-fits performed with and without a subhalo.\n", + "\n", + "The following scale describes how different log evidence increases correspond to detection significances:\n", + "\n", + " - Negative log evidence increase: No detection.\n", + " - Log evidence increase between 0 and 3: No detection.\n", + " - Log evidence increase between 3 and 5: Weak evidence, should consider it a non-detection.\n", + " - Log evidence increase between 5 and 10: Medium evidence, but still inconclusive.\n", + " - Log evidence increase between 10 and 20: Strong evidence, consider it a detection.\n", + " - Log evidence increase > 20: Very strong evidence, definitive detection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "evidence_no_subhalo = result_no_subhalo.samples.log_evidence\n", + "evidence_with_subhalo = result_with_subhalo.samples.log_evidence\n", + "\n", + "log_evidence_increase = evidence_with_subhalo - evidence_no_subhalo\n", + "\n", + "print(\"Evidence Increase: \", log_evidence_increase)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood__\n", + "\n", + "The log likelihood increase provides a simpler metric for how well the subhalo model fits the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_likelihood_no_subhalo = result_no_subhalo.samples.log_likelihood\n", + "log_likelihood_with_subhalo = result_with_subhalo.samples.log_likelihood\n", + "\n", + "log_likelihood_increase = log_likelihood_with_subhalo - log_likelihood_no_subhalo\n", + "\n", + "print(\"Log Likelihood Increase: \", log_likelihood_increase)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid Search Result__\n", + "\n", + "The grid search results can be used to inspect where in the image plane a subhalo provides the best fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "subhalo_grid_search_result = al.subhalo.SubhaloGridSearchResult(\n", + " result=result_subhalo_grid_search\n", + ")\n", + "\n", + "log_evidence_array = subhalo_grid_search_result.figure_of_merit_array(\n", + " use_log_evidences=True,\n", + " relative_to_value=result_no_subhalo.samples.log_evidence,\n", + ")\n", + "\n", + "print(\"Log Evidence Array: \\n\")\n", + "print(log_evidence_array)\n", + "\n", + "aplt.plot_array(array=log_evidence_array, title=\"\")\n", + "\n", + "mass_array = subhalo_grid_search_result.subhalo_mass_array\n", + "\n", + "print(\"Mass Array: \\n\")\n", + "print(mass_array)\n", + "\n", + "subhalo_centres_grid = subhalo_grid_search_result.subhalo_centres_grid\n", + "\n", + "print(\"Subhalo Centres Grid: \\n\")\n", + "print(subhalo_centres_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/advanced/subhalo/simulator.ipynb b/notebooks/group/features/advanced/subhalo/simulator.ipynb index cb44abca0..225a5cab2 100644 --- a/notebooks/group/features/advanced/subhalo/simulator.ipynb +++ b/notebooks/group/features/advanced/subhalo/simulator.ipynb @@ -1,482 +1,519 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Subhalo (Group)\n", - "==========================\n", - "\n", - "If a low mass dark matter halo overlaps the lensed source emission, it perturbs it in a unique and observable way.\n", - "\n", - "This script simulates a group-scale strong lens dataset that includes a dark matter subhalo. The subhalo is a\n", - "small mass perturber near the lensed images, and its effect on the lensed source can be detected through\n", - "careful lens modeling.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** The ``dataset_type`` describes the type of data being simulated.\n", - "- **Grid:** Define the 2d grid of (y,x) coordinates for the simulation.\n", - "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid.\n", - "- **Main Lens Galaxies:** The main lens galaxy, which includes a dark matter subhalo.\n", - "- **Extra Galaxies:** Two companion galaxies near the lens system.\n", - "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", - "- **Ray Tracing:** Use all galaxies to set up a tracer.\n", - "- **Dataset:** Simulate and output the dataset.\n", - "- **Centre JSON Files:** Save the centres as JSON files.\n", - "- **Subhalo Difference Image:** Visualize the effect of the subhalo.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates ``Imaging`` of a 'group-scale' strong lens where:\n", - "\n", - " - The main lens galaxy's light is a ``SersicSph``, total mass is an ``IsothermalSph``, and it includes\n", - " a dark matter subhalo modeled as an ``NFWTruncatedMCRLudlowSph``.\n", - " - The extra galaxies have ``SersicSph`` light and ``IsothermalSph`` mass profiles.\n", - " - A single source galaxy with ``SersicCore`` light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The dataset is output to ``dataset/group/dark_matter_subhalo``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"dark_matter_subhalo\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2d grid of (y,x) coordinates for the simulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(250, 250),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Define the centres of the main lens galaxies and extra galaxies. We also include the subhalo centre\n", - "in the over-sampling to ensure accurate evaluation near the subhalo." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [(0.0, 0.0)]\n", - "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]\n", - "subhalo_centre = (1.6, 0.0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive oversampling at all galaxy centres including the subhalo position." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=main_lens_centres + extra_galaxies_centres + [subhalo_centre],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "The main lens galaxy includes a dark matter subhalo as an ``NFWTruncatedMCRLudlowSph`` mass component. The\n", - "subhalo is positioned near the lensed images at (1.6, 0.0) with a mass of 1e10 solar masses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - " subhalo=al.mp.NFWTruncatedMCRLudlowSph(centre=(1.6, 0.0), mass_at_200=1.0e10),\n", - ")\n", - "\n", - "main_lens_galaxies = [lens_0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "Two companion galaxies near the lens system." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy__\n", - "\n", - "The source galaxy whose lensed images we simulate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Use all galaxies to set up a tracer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the ``Tracer`` in the dataset folder as a .json file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centre JSON Files__\n", - "\n", - "Save the centres of the main lens galaxies and extra galaxies as JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Subhalo Difference Image__\n", - "\n", - "An informative way to visualize the effect of a subhalo on a strong lens is to subtract the image-plane image of\n", - "the tracer with and without the subhalo.\n", - "\n", - "This creates a subhalo residual-map showing the regions of the image plane where the subhalo's effects are\n", - "located." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0_no_subhalo = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "tracer_no_subhalo = al.Tracer(\n", - " galaxies=[lens_0_no_subhalo] + extra_galaxies + [source_galaxy]\n", - ")\n", - "\n", - "image = tracer.image_2d_from(grid=grid)\n", - "image_no_subhalo = tracer_no_subhalo.image_2d_from(grid=grid)\n", - "\n", - "subhalo_residual_image = image - image_no_subhalo\n", - "\n", - "aplt.plot_array(\n", - " array=subhalo_residual_image,\n", - " title=\"Subhalo Residual Image\",\n", - " output_path=dataset_path,\n", - " output_format=\"png\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder ``autolens_workspace/dataset/group/dark_matter_subhalo``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Subhalo (Group)\n", + "==========================\n", + "\n", + "If a low mass dark matter halo overlaps the lensed source emission, it perturbs it in a unique and observable way.\n", + "\n", + "This script simulates a group-scale strong lens dataset that includes a dark matter subhalo. The subhalo is a\n", + "small mass perturber near the lensed images, and its effect on the lensed source can be detected through\n", + "careful lens modeling.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** The ``dataset_type`` describes the type of data being simulated.\n", + "- **Grid:** Define the 2d grid of (y,x) coordinates for the simulation.\n", + "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid.\n", + "- **Main Lens Galaxies:** The main lens galaxy, which includes a dark matter subhalo.\n", + "- **Extra Galaxies:** Two companion galaxies near the lens system.\n", + "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", + "- **Ray Tracing:** Use all galaxies to set up a tracer.\n", + "- **Dataset:** Simulate and output the dataset.\n", + "- **Centre JSON Files:** Save the centres as JSON files.\n", + "- **Subhalo Difference Image:** Visualize the effect of the subhalo.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates ``Imaging`` of a 'group-scale' strong lens where:\n", + "\n", + " - The main lens galaxy's light is a ``SersicSph``, total mass is an ``IsothermalSph``, and it includes\n", + " a dark matter subhalo modeled as an ``NFWTruncatedMCRLudlowSph``.\n", + " - The extra galaxies have ``SersicSph`` light and ``IsothermalSph`` mass profiles.\n", + " - A single source galaxy with ``SersicCore`` light." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The dataset is output to ``dataset/group/dark_matter_subhalo``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"dark_matter_subhalo\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2d grid of (y,x) coordinates for the simulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(250, 250),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Define the centres of the main lens galaxies and extra galaxies. We also include the subhalo centre\n", + "in the over-sampling to ensure accurate evaluation near the subhalo." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [(0.0, 0.0)]\n", + "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]\n", + "subhalo_centre = (1.6, 0.0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive oversampling at all galaxy centres including the subhalo position." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=main_lens_centres + extra_galaxies_centres + [subhalo_centre],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "The main lens galaxy includes a dark matter subhalo as an ``NFWTruncatedMCRLudlowSph`` mass component. The\n", + "subhalo is positioned near the lensed images at (1.6, 0.0) with a mass of 1e10 solar masses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + " subhalo=al.mp.NFWTruncatedMCRLudlowSph(centre=(1.6, 0.0), mass_at_200=1.0e10),\n", + ")\n", + "\n", + "main_lens_galaxies = [lens_0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "Two companion galaxies near the lens system." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy__\n", + "\n", + "The source galaxy whose lensed images we simulate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Use all galaxies to set up a tracer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the ``Tracer`` in the dataset folder as a .json file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centre JSON Files__\n", + "\n", + "Save the centres of the main lens galaxies and extra galaxies as JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Subhalo Difference Image__\n", + "\n", + "An informative way to visualize the effect of a subhalo on a strong lens is to subtract the image-plane image of\n", + "the tracer with and without the subhalo.\n", + "\n", + "This creates a subhalo residual-map showing the regions of the image plane where the subhalo's effects are\n", + "located." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0_no_subhalo = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "tracer_no_subhalo = al.Tracer(\n", + " galaxies=[lens_0_no_subhalo] + extra_galaxies + [source_galaxy]\n", + ")\n", + "\n", + "image = tracer.image_2d_from(grid=grid)\n", + "image_no_subhalo = tracer_no_subhalo.image_2d_from(grid=grid)\n", + "\n", + "subhalo_residual_image = image - image_no_subhalo\n", + "\n", + "aplt.plot_array(\n", + " array=subhalo_residual_image,\n", + " title=\"Subhalo Residual Image\",\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder ``autolens_workspace/dataset/group/dark_matter_subhalo``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/linear_light_profiles/fit.ipynb b/notebooks/group/features/linear_light_profiles/fit.ipynb index 758f146ae..b3010b926 100644 --- a/notebooks/group/features/linear_light_profiles/fit.ipynb +++ b/notebooks/group/features/linear_light_profiles/fit.ipynb @@ -1,363 +1,400 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fit Features: Linear Light Profiles (Group)\n", - "===========================================\n", - "\n", - "This script shows how to fit data using the ``FitImaging`` object for group-scale strong lenses when using\n", - "**linear light profiles**, where the ``intensity`` of every light profile is solved via linear algebra rather\n", - "than being specified as a parameter.\n", - "\n", - "A group-scale lens differs from a galaxy-scale lens in that there are multiple lens galaxies contributing to the\n", - "lensing. In this example, there is a single main lens galaxy and two extra galaxies nearby whose mass contributes\n", - "significantly to the ray-tracing.\n", - "\n", - "The key difference from the standard group fit script is that all light profiles use the ``lp_linear`` module\n", - "and therefore do not have an ``intensity`` parameter -- it is solved automatically when the ``FitImaging``\n", - "object is created.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Galaxy Centres:** Load the centres of the main lens galaxies and extra galaxies from JSON.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Fitting:** Fit the lens model to the dataset using linear light profiles.\n", - "- **Intensities:** Extract the solved-for intensity values.\n", - "- **Visualization:** Convert linear light profiles to ordinary profiles for visualization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset ``simple`` from .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.plot_array(array=dataset.data, title=\"Image Data With Mask Applied\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "print(f\"Main lens centres: {main_lens_centres}\")\n", - "print(f\"Extra galaxies centres: {extra_galaxies_centres}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre ensures that light profiles are accurately evaluated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fitting__\n", - "\n", - "We now create galaxies using linear light profiles (via the ``lp_linear`` module). Note that no ``intensity``\n", - "parameter is specified for any light profile -- it will be solved via linear algebra when the ``FitImaging``\n", - "object is created.\n", - "\n", - "All other parameters (centre, effective_radius, sersic_index, etc.) are specified with known values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.SersicSph(\n", - " centre=(0.0, 0.0), effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.SersicSph(\n", - " centre=(3.5, 2.5), effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.SersicSph(\n", - " centre=(-4.4, -5.0), effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now use a ``FitImaging`` object to fit this tracer to the dataset. The fit automatically detects that\n", - "the tracer contains linear light profiles and solves for their ``intensity`` values via a linear inversion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The fit contains the solved-for intensity values. These are accessible via the\n", - "``linear_light_profile_intensity_dict``, which maps each linear light profile to its solved ``intensity``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_bulge = tracer.galaxies[0].bulge\n", - "extra_0_bulge = tracer.galaxies[1].bulge\n", - "extra_1_bulge = tracer.galaxies[2].bulge\n", - "source_bulge = tracer.galaxies[3].bulge\n", - "\n", - "print(fit.linear_light_profile_intensity_dict)\n", - "\n", - "print(\n", - " f\"\\n Intensity of main lens bulge = {fit.linear_light_profile_intensity_dict[lens_bulge]}\"\n", - ")\n", - "print(\n", - " f\"\\n Intensity of extra galaxy 0 bulge = {fit.linear_light_profile_intensity_dict[extra_0_bulge]}\"\n", - ")\n", - "print(\n", - " f\"\\n Intensity of extra galaxy 1 bulge = {fit.linear_light_profile_intensity_dict[extra_1_bulge]}\"\n", - ")\n", - "print(\n", - " f\"\\n Intensity of source bulge = {fit.linear_light_profile_intensity_dict[source_bulge]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization__\n", - "\n", - "Linear light profiles cannot be plotted directly because they do not have an ``intensity`` value until the\n", - "inversion is performed. The ``model_obj_linear_light_profiles_to_light_profiles`` property returns a ``Tracer``\n", - "where all linear light profiles are replaced with ordinary light profiles using the solved-for ``intensity``\n", - "values. This can be used for visualization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "print(f\"Main lens bulge intensity: {tracer.galaxies[0].bulge.intensity}\")\n", - "print(f\"Extra galaxy 0 bulge intensity: {tracer.galaxies[1].bulge.intensity}\")\n", - "print(f\"Extra galaxy 1 bulge intensity: {tracer.galaxies[2].bulge.intensity}\")\n", - "print(f\"Source bulge intensity: {tracer.galaxies[3].bulge.intensity}\")\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fit Features: Linear Light Profiles (Group)\n", + "===========================================\n", + "\n", + "This script shows how to fit data using the ``FitImaging`` object for group-scale strong lenses when using\n", + "**linear light profiles**, where the ``intensity`` of every light profile is solved via linear algebra rather\n", + "than being specified as a parameter.\n", + "\n", + "A group-scale lens differs from a galaxy-scale lens in that there are multiple lens galaxies contributing to the\n", + "lensing. In this example, there is a single main lens galaxy and two extra galaxies nearby whose mass contributes\n", + "significantly to the ray-tracing.\n", + "\n", + "The key difference from the standard group fit script is that all light profiles use the ``lp_linear`` module\n", + "and therefore do not have an ``intensity`` parameter -- it is solved automatically when the ``FitImaging``\n", + "object is created.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Galaxy Centres:** Load the centres of the main lens galaxies and extra galaxies from JSON.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Fitting:** Fit the lens model to the dataset using linear light profiles.\n", + "- **Intensities:** Extract the solved-for intensity values.\n", + "- **Visualization:** Convert linear light profiles to ordinary profiles for visualization." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset ``simple`` from .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.plot_array(array=dataset.data, title=\"Image Data With Mask Applied\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "print(f\"Main lens centres: {main_lens_centres}\")\n", + "print(f\"Extra galaxies centres: {extra_galaxies_centres}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre ensures that light profiles are accurately evaluated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fitting__\n", + "\n", + "We now create galaxies using linear light profiles (via the ``lp_linear`` module). Note that no ``intensity``\n", + "parameter is specified for any light profile -- it will be solved via linear algebra when the ``FitImaging``\n", + "object is created.\n", + "\n", + "All other parameters (centre, effective_radius, sersic_index, etc.) are specified with known values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.SersicSph(\n", + " centre=(0.0, 0.0), effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.SersicSph(\n", + " centre=(3.5, 2.5), effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.SersicSph(\n", + " centre=(-4.4, -5.0), effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now use a ``FitImaging`` object to fit this tracer to the dataset. The fit automatically detects that\n", + "the tracer contains linear light profiles and solves for their ``intensity`` values via a linear inversion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The fit contains the solved-for intensity values. These are accessible via the\n", + "``linear_light_profile_intensity_dict``, which maps each linear light profile to its solved ``intensity``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_bulge = tracer.galaxies[0].bulge\n", + "extra_0_bulge = tracer.galaxies[1].bulge\n", + "extra_1_bulge = tracer.galaxies[2].bulge\n", + "source_bulge = tracer.galaxies[3].bulge\n", + "\n", + "print(fit.linear_light_profile_intensity_dict)\n", + "\n", + "print(\n", + " f\"\\n Intensity of main lens bulge = {fit.linear_light_profile_intensity_dict[lens_bulge]}\"\n", + ")\n", + "print(\n", + " f\"\\n Intensity of extra galaxy 0 bulge = {fit.linear_light_profile_intensity_dict[extra_0_bulge]}\"\n", + ")\n", + "print(\n", + " f\"\\n Intensity of extra galaxy 1 bulge = {fit.linear_light_profile_intensity_dict[extra_1_bulge]}\"\n", + ")\n", + "print(\n", + " f\"\\n Intensity of source bulge = {fit.linear_light_profile_intensity_dict[source_bulge]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization__\n", + "\n", + "Linear light profiles cannot be plotted directly because they do not have an ``intensity`` value until the\n", + "inversion is performed. The ``model_obj_linear_light_profiles_to_light_profiles`` property returns a ``Tracer``\n", + "where all linear light profiles are replaced with ordinary light profiles using the solved-for ``intensity``\n", + "values. This can be used for visualization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "print(f\"Main lens bulge intensity: {tracer.galaxies[0].bulge.intensity}\")\n", + "print(f\"Extra galaxy 0 bulge intensity: {tracer.galaxies[1].bulge.intensity}\")\n", + "print(f\"Extra galaxy 1 bulge intensity: {tracer.galaxies[2].bulge.intensity}\")\n", + "print(f\"Source bulge intensity: {tracer.galaxies[3].bulge.intensity}\")\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/linear_light_profiles/likelihood_function.ipynb b/notebooks/group/features/linear_light_profiles/likelihood_function.ipynb index 81ca67d22..ffdb08431 100644 --- a/notebooks/group/features/linear_light_profiles/likelihood_function.ipynb +++ b/notebooks/group/features/linear_light_profiles/likelihood_function.ipynb @@ -1,517 +1,554 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Linear Light Profiles (Group)__\n", - "\n", - "This script provides a step-by-step guide of the ``log_likelihood_function`` which is used to fit ``Imaging``\n", - "data of a group-scale strong lens using **linear light profiles**.\n", - "\n", - "The key difference from the standard group likelihood function is that the ``intensity`` of each light\n", - "profile is not a free parameter. Instead, intensities are solved via linear algebra (an \"inversion\") every\n", - "time the model is evaluated, always finding the intensity values that maximize the likelihood given all\n", - "other parameters.\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers, so that readers can understand the likelihood\n", - " function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - " - To illustrate how the linear inversion works for group-scale lenses with multiple galaxies.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Linear Light Profiles:** Define the lens, extra galaxies and source using ``lp_linear`` profiles.\n", - "- **Lens Galaxy Mass:** Compute deflection angles from all mass profiles.\n", - "- **Ray Tracing:** Trace image-plane coordinates to the source plane.\n", - "- **Linear Inversion:** Explain how intensities are solved via the linear algebra system.\n", - "- **Fit:** Use the ``FitImaging`` object which handles all steps automatically.\n", - "- **Likelihood Function:** Compute the log likelihood step by step.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function of a linear light profile builds on that used for standard light profiles,\n", - "therefore you should read the following before this script:\n", - "\n", - "- ``group/likelihood_function.py``\n", - "- ``imaging/features/linear_light_profiles/likelihood_function.py``" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autoarray as aa\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the group-scale strong lens dataset with 0.1 arcsecond-per-pixel resolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"group\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define a 7.5\" circular mask. This is larger than a typical galaxy-scale mask because group-scale systems\n", - "have lensed images that extend over a wider area." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=masked_dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "For simplicity, we disable over sampling in this guide by setting ``sub_size=1``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masked Image Grid__\n", - "\n", - "The masked dataset provides a 2D grid of (y,x) coordinates used for all calculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=masked_dataset.grids.lp, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Light Profiles__\n", - "\n", - "We define all galaxies using linear light profiles from the ``lp_linear`` module. These profiles have no\n", - "``intensity`` parameter -- internally they use an intensity of 1.0, and the true intensity is solved via\n", - "linear algebra.\n", - "\n", - "__Main Lens Galaxy__\n", - "\n", - "The main lens galaxy has a spherical linear Sersic light profile and an isothermal mass profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.SersicSph(\n", - " centre=(0.0, 0.0), effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "The two extra galaxies are companion galaxies near the main lens. They each have linear SersicSph light\n", - "profiles and isothermal mass profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.SersicSph(\n", - " centre=(3.5, 2.5), effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.SersicSph(\n", - " centre=(-4.4, -5.0), effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Light Profile__\n", - "\n", - "The source galaxy uses a linear cored elliptical Sersic -- again with no intensity parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Internally, linear light profiles have an ``intensity`` parameter set to 1.0. This is because each profile's\n", - "image (with unit intensity) forms a column in the \"mapping matrix\" used in the linear inversion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Lens bulge internal intensity: {lens_galaxy.bulge.intensity}\")\n", - "print(f\"Extra galaxy 0 internal intensity: {extra_galaxy_0.bulge.intensity}\")\n", - "print(f\"Source internal intensity: {source_galaxy.bulge.intensity}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy Mass__\n", - "\n", - "We compute the deflection angles from ALL galaxies in the group. For group-scale lensing, the total\n", - "deflection is the sum of deflections from every galaxy's mass profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_lens = lens_galaxy.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_extra_0 = extra_galaxy_0.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_extra_1 = extra_galaxy_1.deflections_yx_2d_from(grid=masked_dataset.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The total deflection is the sum of deflections from all galaxies. We ray-trace every image-plane\n", - "coordinate to the source plane using these combined deflections." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", - "\n", - "aplt.plot_grid(grid=traced_grid, title=\"Source Plane Grid (Traced)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Inversion__\n", - "\n", - "For standard light profiles, we would now evaluate each profile's image (with known intensity) and sum them\n", - "to get the model image. With linear light profiles, the process is different:\n", - "\n", - "1. Each linear light profile's image is computed with unit intensity (intensity = 1.0). These images form\n", - " the columns of a \"mapping matrix\" (or \"blurred mapping matrix\" after PSF convolution).\n", - "\n", - "2. The mapping matrices from ALL planes (image-plane for lens galaxies, source-plane for the source) are\n", - " combined into a single system.\n", - "\n", - "3. A linear algebra solver finds the intensity values that minimize chi-squared. This is done using a\n", - " positive-only solver to ensure all intensities are physical (non-negative).\n", - "\n", - "4. The solved intensities are multiplied by the unit-intensity images to produce the final model image.\n", - "\n", - "The ``FitImaging`` object handles all of this automatically. When it detects linear light profiles in the\n", - "tracer, it sets up and solves the linear system." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Fit__\n", - "\n", - "The ``FitImaging`` object performs the full likelihood evaluation, including the linear inversion. We can\n", - "verify that the figure of merit matches what we would compute manually.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)\n", - "fit_figure_of_merit = fit.figure_of_merit\n", - "print(fit_figure_of_merit)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "The likelihood function for linear light profiles includes the same terms as the standard parametric case:\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "The difference is that the model image used to compute the chi-squared is itself the output of the linear\n", - "inversion, rather than being computed from fixed intensity values.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The chi-squared measures how well the model (with solved intensities) fits the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image = fit.model_data\n", - "residual_map = masked_dataset.data - model_image\n", - "normalized_residual_map = residual_map / masked_dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(f\"Chi-squared: {chi_squared}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "The noise normalization term is the same as for standard light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "We compute the log likelihood by combining the chi-squared and noise normalization terms." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(f\"Log likelihood (manual): {figure_of_merit}\")\n", - "print(f\"Log likelihood (fit): {fit_figure_of_merit}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm ``Nautilus``. The linear inversion for intensities\n", - "happens automatically at every likelihood evaluation, meaning the non-linear search only needs to explore\n", - "the structural parameters (centres, effective radii, sersic indices, mass parameters, etc.).\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the group-scale linear light profile likelihood function.\n", - "\n", - "The key differences from the standard group likelihood function are:\n", - "\n", - " - Intensities are not free parameters but are solved via linear algebra at each likelihood evaluation.\n", - " - The mapping matrix encodes each profile's unit-intensity image; the inversion finds the best intensities.\n", - " - For group-scale lenses with many galaxies, this dramatically reduces the non-linear parameter space.\n", - " - The ``FitImaging`` and ``Tracer`` objects handle all of this automatically." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Linear Light Profiles (Group)__\n", + "\n", + "This script provides a step-by-step guide of the ``log_likelihood_function`` which is used to fit ``Imaging``\n", + "data of a group-scale strong lens using **linear light profiles**.\n", + "\n", + "The key difference from the standard group likelihood function is that the ``intensity`` of each light\n", + "profile is not a free parameter. Instead, intensities are solved via linear algebra (an \"inversion\") every\n", + "time the model is evaluated, always finding the intensity values that maximize the likelihood given all\n", + "other parameters.\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers, so that readers can understand the likelihood\n", + " function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + " - To illustrate how the linear inversion works for group-scale lenses with multiple galaxies.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Linear Light Profiles:** Define the lens, extra galaxies and source using ``lp_linear`` profiles.\n", + "- **Lens Galaxy Mass:** Compute deflection angles from all mass profiles.\n", + "- **Ray Tracing:** Trace image-plane coordinates to the source plane.\n", + "- **Linear Inversion:** Explain how intensities are solved via the linear algebra system.\n", + "- **Fit:** Use the ``FitImaging`` object which handles all steps automatically.\n", + "- **Likelihood Function:** Compute the log likelihood step by step.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function of a linear light profile builds on that used for standard light profiles,\n", + "therefore you should read the following before this script:\n", + "\n", + "- ``group/likelihood_function.py``\n", + "- ``imaging/features/linear_light_profiles/likelihood_function.py``" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autoarray as aa\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the group-scale strong lens dataset with 0.1 arcsecond-per-pixel resolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"group\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define a 7.5\" circular mask. This is larger than a typical galaxy-scale mask because group-scale systems\n", + "have lensed images that extend over a wider area." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=masked_dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "For simplicity, we disable over sampling in this guide by setting ``sub_size=1``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masked Image Grid__\n", + "\n", + "The masked dataset provides a 2D grid of (y,x) coordinates used for all calculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=masked_dataset.grids.lp, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Light Profiles__\n", + "\n", + "We define all galaxies using linear light profiles from the ``lp_linear`` module. These profiles have no\n", + "``intensity`` parameter -- internally they use an intensity of 1.0, and the true intensity is solved via\n", + "linear algebra.\n", + "\n", + "__Main Lens Galaxy__\n", + "\n", + "The main lens galaxy has a spherical linear Sersic light profile and an isothermal mass profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.SersicSph(\n", + " centre=(0.0, 0.0), effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "The two extra galaxies are companion galaxies near the main lens. They each have linear SersicSph light\n", + "profiles and isothermal mass profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.SersicSph(\n", + " centre=(3.5, 2.5), effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.SersicSph(\n", + " centre=(-4.4, -5.0), effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Light Profile__\n", + "\n", + "The source galaxy uses a linear cored elliptical Sersic -- again with no intensity parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Internally, linear light profiles have an ``intensity`` parameter set to 1.0. This is because each profile's\n", + "image (with unit intensity) forms a column in the \"mapping matrix\" used in the linear inversion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Lens bulge internal intensity: {lens_galaxy.bulge.intensity}\")\n", + "print(f\"Extra galaxy 0 internal intensity: {extra_galaxy_0.bulge.intensity}\")\n", + "print(f\"Source internal intensity: {source_galaxy.bulge.intensity}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy Mass__\n", + "\n", + "We compute the deflection angles from ALL galaxies in the group. For group-scale lensing, the total\n", + "deflection is the sum of deflections from every galaxy's mass profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_lens = lens_galaxy.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_extra_0 = extra_galaxy_0.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_extra_1 = extra_galaxy_1.deflections_yx_2d_from(grid=masked_dataset.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The total deflection is the sum of deflections from all galaxies. We ray-trace every image-plane\n", + "coordinate to the source plane using these combined deflections." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", + "\n", + "aplt.plot_grid(grid=traced_grid, title=\"Source Plane Grid (Traced)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Inversion__\n", + "\n", + "For standard light profiles, we would now evaluate each profile's image (with known intensity) and sum them\n", + "to get the model image. With linear light profiles, the process is different:\n", + "\n", + "1. Each linear light profile's image is computed with unit intensity (intensity = 1.0). These images form\n", + " the columns of a \"mapping matrix\" (or \"blurred mapping matrix\" after PSF convolution).\n", + "\n", + "2. The mapping matrices from ALL planes (image-plane for lens galaxies, source-plane for the source) are\n", + " combined into a single system.\n", + "\n", + "3. A linear algebra solver finds the intensity values that minimize chi-squared. This is done using a\n", + " positive-only solver to ensure all intensities are physical (non-negative).\n", + "\n", + "4. The solved intensities are multiplied by the unit-intensity images to produce the final model image.\n", + "\n", + "The ``FitImaging`` object handles all of this automatically. When it detects linear light profiles in the\n", + "tracer, it sets up and solves the linear system." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Fit__\n", + "\n", + "The ``FitImaging`` object performs the full likelihood evaluation, including the linear inversion. We can\n", + "verify that the figure of merit matches what we would compute manually.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)\n", + "fit_figure_of_merit = fit.figure_of_merit\n", + "print(fit_figure_of_merit)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "The likelihood function for linear light profiles includes the same terms as the standard parametric case:\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "The difference is that the model image used to compute the chi-squared is itself the output of the linear\n", + "inversion, rather than being computed from fixed intensity values.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The chi-squared measures how well the model (with solved intensities) fits the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image = fit.model_data\n", + "residual_map = masked_dataset.data - model_image\n", + "normalized_residual_map = residual_map / masked_dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(f\"Chi-squared: {chi_squared}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "The noise normalization term is the same as for standard light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "We compute the log likelihood by combining the chi-squared and noise normalization terms." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(f\"Log likelihood (manual): {figure_of_merit}\")\n", + "print(f\"Log likelihood (fit): {fit_figure_of_merit}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm ``Nautilus``. The linear inversion for intensities\n", + "happens automatically at every likelihood evaluation, meaning the non-linear search only needs to explore\n", + "the structural parameters (centres, effective radii, sersic indices, mass parameters, etc.).\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the group-scale linear light profile likelihood function.\n", + "\n", + "The key differences from the standard group likelihood function are:\n", + "\n", + " - Intensities are not free parameters but are solved via linear algebra at each likelihood evaluation.\n", + " - The mapping matrix encodes each profile's unit-intensity image; the inversion finds the best intensities.\n", + " - For group-scale lenses with many galaxies, this dramatically reduces the non-linear parameter space.\n", + " - The ``FitImaging`` and ``Tracer`` objects handle all of this automatically." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/linear_light_profiles/modeling.ipynb b/notebooks/group/features/linear_light_profiles/modeling.ipynb index 09983345f..c504bf78e 100644 --- a/notebooks/group/features/linear_light_profiles/modeling.ipynb +++ b/notebooks/group/features/linear_light_profiles/modeling.ipynb @@ -1,495 +1,532 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Linear Light Profiles (Group)\n", - "================================================\n", - "\n", - "This script fits a group-scale strong lens using **linear light profiles**, where the ``intensity`` of every\n", - "light profile is solved analytically via linear algebra rather than being a free parameter in the non-linear\n", - "search.\n", - "\n", - "For a group-scale lens this is especially beneficial: the group model contains many galaxies (main lenses and\n", - "extra galaxies), and each galaxy's light profile would normally add an ``intensity`` parameter. By using linear\n", - "light profiles, none of these contribute to the non-linear parameter space, significantly reducing dimensionality\n", - "and improving sampling efficiency.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages:** Linear light profiles remove `intensity` from the non-linear parameter space.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres:** The centres of the main lens galaxies and extra galaxies are loaded from JSON files.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Intensities:** How to extract the solved-for intensity values from the result.\n", - "\n", - "__Advantages__\n", - "\n", - "Each light profile's ``intensity`` parameter is solved via a linear inversion, reducing the dimensionality of\n", - "non-linear parameter space by the number of light profiles. This also removes the degeneracies between\n", - "``intensity`` and other light profile parameters (e.g. ``effective_radius``, ``sersic_index``), producing\n", - "more reliable lens model results that converge in fewer iterations.\n", - "\n", - "For group-scale lenses with many galaxies, this reduction is particularly impactful: every main lens galaxy\n", - "and every extra galaxy has its intensity removed from the search.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - Each main lens galaxy's light is a linear ``Sersic`` bulge [6 parameters].\n", - " - The first main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear`` [7 parameters].\n", - " - There are two extra lens galaxies with linear ``SersicSph`` light and ``IsothermalSph`` total mass\n", - " distributions, with centres fixed to the observed centres of light [8 parameters].\n", - " - The source galaxy's light is a linear ``SersicCore`` [5 parameters].\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the ``group/modeling`` and\n", - "``imaging/features/linear_light_profiles/modeling`` notebooks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`, which is the dataset we will use to perform lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask, which is larger than a typical galaxy-scale mask because the group-scale\n", - "lens has emission spread over a wider area." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files in the dataset\n", - "directory. This makes the script reusable across different datasets without hardcoding centre values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model using linear light profiles for all galaxies. The key difference from the standard\n", - "group modeling script is that we use ``al.lp_linear.Sersic``, ``al.lp_linear.SersicSph`` and\n", - "``al.lp_linear.SersicCore`` instead of their standard ``al.lp`` counterparts. These linear profiles have no\n", - "``intensity`` parameter -- it is solved via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = af.Model(al.lp_linear.Sersic)\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light (linear -- no intensity parameter)\n", - "\n", - " bulge = af.Model(al.lp_linear.SersicSph)\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ``info`` attribute shows the model in a readable format.\n", - "\n", - "This confirms that light profiles of all galaxies do not include an ``intensity`` parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) is performed to ensure\n", - "the lens calculations are accurate across the full field of the group." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The lens model is fitted to the data using a non-linear search. Because the linear light profiles reduce the\n", - "dimensionality of the parameter space, the fit is more efficient than using standard light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"linear_light_profiles\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We create an ``AnalysisImaging`` object. The linear light profiles are handled automatically by the analysis\n", - "object -- when it detects linear light profiles in the model, it performs the linear inversion to solve for\n", - "their intensities during every likelihood evaluation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result ``info`` confirms that ``intensity`` parameters are not inferred by the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The intensities of linear light profiles are not a part of the model parameterization and therefore are not\n", - "displayed in the ``model.results`` file.\n", - "\n", - "To extract the ``intensity`` values of a specific component in the model, we use the\n", - "``max_log_likelihood_tracer``, which has already performed the inversion and therefore the galaxy light\n", - "profiles have their solved-for ``intensity`` values associated with them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "print(f\"Main lens galaxy bulge intensity: {tracer.galaxies[0].bulge.intensity}\")\n", - "print(f\"Source galaxy bulge intensity: {tracer.galaxies[-1].bulge.intensity}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ``Tracer`` contained in the ``max_log_likelihood_fit`` also has the solved-for ``intensity`` values:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = result.max_log_likelihood_fit\n", - "\n", - "tracer = fit.tracer\n", - "\n", - "print(\n", - " f\"Main lens galaxy bulge intensity (via fit): {tracer.galaxies[0].bulge.intensity}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Linear Light Profiles (Group)\n", + "================================================\n", + "\n", + "This script fits a group-scale strong lens using **linear light profiles**, where the ``intensity`` of every\n", + "light profile is solved analytically via linear algebra rather than being a free parameter in the non-linear\n", + "search.\n", + "\n", + "For a group-scale lens this is especially beneficial: the group model contains many galaxies (main lenses and\n", + "extra galaxies), and each galaxy's light profile would normally add an ``intensity`` parameter. By using linear\n", + "light profiles, none of these contribute to the non-linear parameter space, significantly reducing dimensionality\n", + "and improving sampling efficiency.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages:** Linear light profiles remove `intensity` from the non-linear parameter space.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres:** The centres of the main lens galaxies and extra galaxies are loaded from JSON files.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Intensities:** How to extract the solved-for intensity values from the result.\n", + "\n", + "__Advantages__\n", + "\n", + "Each light profile's ``intensity`` parameter is solved via a linear inversion, reducing the dimensionality of\n", + "non-linear parameter space by the number of light profiles. This also removes the degeneracies between\n", + "``intensity`` and other light profile parameters (e.g. ``effective_radius``, ``sersic_index``), producing\n", + "more reliable lens model results that converge in fewer iterations.\n", + "\n", + "For group-scale lenses with many galaxies, this reduction is particularly impactful: every main lens galaxy\n", + "and every extra galaxy has its intensity removed from the search.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an ``Imaging`` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - Each main lens galaxy's light is a linear ``Sersic`` bulge [6 parameters].\n", + " - The first main lens galaxy's total mass distribution is an ``Isothermal`` and ``ExternalShear`` [7 parameters].\n", + " - There are two extra lens galaxies with linear ``SersicSph`` light and ``IsothermalSph`` total mass\n", + " distributions, with centres fixed to the observed centres of light [8 parameters].\n", + " - The source galaxy's light is a linear ``SersicCore`` [5 parameters].\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the ``group/modeling`` and\n", + "``imaging/features/linear_light_profiles/modeling`` notebooks." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`, which is the dataset we will use to perform lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask, which is larger than a typical galaxy-scale mask because the group-scale\n", + "lens has emission spread over a wider area." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files in the dataset\n", + "directory. This makes the script reusable across different datasets without hardcoding centre values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model using linear light profiles for all galaxies. The key difference from the standard\n", + "group modeling script is that we use ``al.lp_linear.Sersic``, ``al.lp_linear.SersicSph`` and\n", + "``al.lp_linear.SersicCore`` instead of their standard ``al.lp`` counterparts. These linear profiles have no\n", + "``intensity`` parameter -- it is solved via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = af.Model(al.lp_linear.Sersic)\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light (linear -- no intensity parameter)\n", + "\n", + " bulge = af.Model(al.lp_linear.SersicSph)\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The ``info`` attribute shows the model in a readable format.\n", + "\n", + "This confirms that light profiles of all galaxies do not include an ``intensity`` parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) is performed to ensure\n", + "the lens calculations are accurate across the full field of the group." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The lens model is fitted to the data using a non-linear search. Because the linear light profiles reduce the\n", + "dimensionality of the parameter space, the fit is more efficient than using standard light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"linear_light_profiles\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We create an ``AnalysisImaging`` object. The linear light profiles are handled automatically by the analysis\n", + "object -- when it detects linear light profiles in the model, it performs the linear inversion to solve for\n", + "their intensities during every likelihood evaluation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result ``info`` confirms that ``intensity`` parameters are not inferred by the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The intensities of linear light profiles are not a part of the model parameterization and therefore are not\n", + "displayed in the ``model.results`` file.\n", + "\n", + "To extract the ``intensity`` values of a specific component in the model, we use the\n", + "``max_log_likelihood_tracer``, which has already performed the inversion and therefore the galaxy light\n", + "profiles have their solved-for ``intensity`` values associated with them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "print(f\"Main lens galaxy bulge intensity: {tracer.galaxies[0].bulge.intensity}\")\n", + "print(f\"Source galaxy bulge intensity: {tracer.galaxies[-1].bulge.intensity}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The ``Tracer`` contained in the ``max_log_likelihood_fit`` also has the solved-for ``intensity`` values:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = result.max_log_likelihood_fit\n", + "\n", + "tracer = fit.tracer\n", + "\n", + "print(\n", + " f\"Main lens galaxy bulge intensity (via fit): {tracer.galaxies[0].bulge.intensity}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/linear_light_profiles/slam.ipynb b/notebooks/group/features/linear_light_profiles/slam.ipynb index 9ca664a60..9570c74fb 100644 --- a/notebooks/group/features/linear_light_profiles/slam.ipynb +++ b/notebooks/group/features/linear_light_profiles/slam.ipynb @@ -1,1145 +1,1182 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Linear Light Profiles: Group SLaM\n", - "===================================\n", - "\n", - "This script uses the SLaM pipelines to fit a group-scale strong lens using **linear Sersic light profiles**\n", - "instead of Multi-Gaussian Expansion (MGE) profiles for the galaxy light.\n", - "\n", - "The differences from the standard ``group/slam.py`` are:\n", - "\n", - " - The SOURCE LP PIPELINE 0 uses ``al.lp_linear.Sersic`` for main lens galaxies and\n", - " ``al.lp_linear.SersicSph`` for extra galaxies, instead of MGE models.\n", - " - The LIGHT LP PIPELINE uses ``al.lp_linear.Sersic`` for main lens galaxies and\n", - " ``al.lp_linear.SersicSph`` for extra galaxies, instead of MGE models.\n", - "\n", - "Linear light profiles solve for the ``intensity`` analytically via linear algebra, removing it from the\n", - "non-linear parameter space. This provides an alternative to MGE for group-scale modeling: simpler model\n", - "composition with fewer basis functions, while still benefiting from the intensity-free parameter space.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **SOURCE LP PIPELINE 0:** Light-only fit using linear Sersic profiles for all galaxies.\n", - "- **SOURCE LP PIPELINE 1:** Introduces mass and source with light fixed from stage 0.\n", - "- **SOURCE PIX PIPELINE 1:** Pixelized source fitting (identical to group/slam.py).\n", - "- **SOURCE PIX PIPELINE 2:** Refined pixelized source (identical to group/slam.py).\n", - "- **LIGHT LP PIPELINE:** Refits light using linear Sersic profiles.\n", - "- **MASS TOTAL PIPELINE:** Final mass fit with PowerLaw (identical to group/slam.py).\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Galaxy Centres:** Load centres from JSON files.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Settings AutoFit:** The settings of autofit.\n", - "- **SLaM Pipeline:** Run the full pipeline.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (``guides/modeling/slam_start_here``)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines.\n", - "\n", - "- **Group SLaM** (``group/slam``)\n", - " The standard group-scale SLaM pipeline using MGE profiles.\n", - "\n", - "- **Linear Light Profiles** (``features/linear_light_profiles``)\n", - " How linear light profiles work and their advantages." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "\n", - "def _load_centres(path):\n", - " \"\"\"Load a centres JSON file, returning an empty list if the file is absent.\"\"\"\n", - " try:\n", - " return al.Grid2DIrregular(al.from_json(file_path=path))\n", - " except FileNotFoundError:\n", - " return al.Grid2DIrregular([])\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 0__\n", - "\n", - "Light-only fit (no mass, no source) for every galaxy simultaneously, using linear Sersic profiles instead\n", - "of MGE. This gives the next search clean fixed light models to build on.\n", - "\n", - "Main lens galaxies use ``al.lp_linear.Sersic`` and extra galaxies use ``al.lp_linear.SersicSph``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_0(\n", - " dataset,\n", - " settings_search,\n", - " main_lens_centres,\n", - " extra_lens_centres,\n", - " scaling_lens_centres,\n", - " mask_radius,\n", - " redshift_lens,\n", - " n_batch=50,\n", - "):\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " # --- main lens light models (linear Sersic, light only) ---\n", - " lens_dict = {}\n", - " for i, centre in enumerate(main_lens_centres):\n", - " bulge = af.Model(al.lp_linear.Sersic)\n", - " bulge.centre_0 = af.GaussianPrior(mean=centre[0], sigma=0.1)\n", - " bulge.centre_1 = af.GaussianPrior(mean=centre[1], sigma=0.1)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=bulge, disk=None, point=None\n", - " )\n", - "\n", - " # --- extra lens galaxy light models (linear SersicSph) ---\n", - " extra_light_models = []\n", - " for centre in extra_lens_centres:\n", - " bulge = af.Model(al.lp_linear.SersicSph)\n", - " bulge.centre_0 = af.UniformPrior(\n", - " lower_limit=centre[0] - 0.5, upper_limit=centre[0] + 0.5\n", - " )\n", - " bulge.centre_1 = af.UniformPrior(\n", - " lower_limit=centre[1] - 0.5, upper_limit=centre[1] + 0.5\n", - " )\n", - "\n", - " extra_light_models.append(\n", - " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", - "\n", - " # --- scaling galaxy light models (linear SersicSph) ---\n", - " scaling_light_models = []\n", - " for centre in scaling_lens_centres:\n", - " bulge = af.Model(al.lp_linear.SersicSph)\n", - " bulge.centre_0 = af.UniformPrior(\n", - " lower_limit=centre[0] - 0.5, upper_limit=centre[0] + 0.5\n", - " )\n", - " bulge.centre_1 = af.UniformPrior(\n", - " lower_limit=centre[1] - 0.5, upper_limit=centre[1] + 0.5\n", - " )\n", - "\n", - " scaling_light_models.append(\n", - " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", - " )\n", - "\n", - " scaling_galaxies = (\n", - " af.Collection(scaling_light_models) if scaling_light_models else None\n", - " )\n", - "\n", - " n_extra = len(extra_galaxies) if extra_galaxies is not None else 0\n", - " n_scaling = len(scaling_galaxies) if scaling_galaxies is not None else 0\n", - " n_live = 100 + 30 * len(lens_dict) + 30 * n_extra + 30 * n_scaling\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=scaling_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[0]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " n_like_max=1000000,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 1__\n", - "\n", - "Introduces mass and source with light fixed from stage 0. Multiple main-lens galaxies each get an\n", - "``Isothermal`` mass; only ``lens_0`` carries an ``ExternalShear``. Extra-galaxy Einstein radii are\n", - "bounded by a luminosity-derived prior.\n", - "\n", - "This is identical to the standard group/slam.py ``source_lp_1`` -- the light profiles are fixed from\n", - "stage 0 (which used linear Sersic), so no changes are needed here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_1(\n", - " dataset,\n", - " settings_search,\n", - " source_lp_result_0,\n", - " positions,\n", - " pixel_scale,\n", - " redshift_lens,\n", - " redshift_source,\n", - " source_mge_radius,\n", - " n_batch=50,\n", - "):\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(source_lp_result_0.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - " n_extra = (\n", - " len(list(source_lp_result_0.instance.extra_galaxies))\n", - " if source_lp_result_0.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - " n_scaling = (\n", - " len(list(source_lp_result_0.instance.scaling_galaxies))\n", - " if source_lp_result_0.instance.scaling_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " tracer = (\n", - " source_lp_result_0.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", - " )\n", - "\n", - " # Source MGE centred on primary lens bulge from stage 0.\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=source_mge_radius,\n", - " total_gaussians=30,\n", - " centre=source_lp_result_0.instance.galaxies.lens_0.bulge.centre,\n", - " centre_prior_is_uniform=False,\n", - " centre_sigma=0.6,\n", - " )\n", - "\n", - " # --- main lens full models (light fixed from stage 0, mass + shear free) ---\n", - " lens_dict = {}\n", - " for i in range(n_main):\n", - " lp0_lens = getattr(source_lp_result_0.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = lp0_lens.bulge.centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=5.0)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lp0_lens.bulge,\n", - " disk=lp0_lens.disk,\n", - " point=lp0_lens.point,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " # --- extra lens galaxy models (light fixed, mass bounded by luminosity) ---\n", - " extra_mass_models = []\n", - " for i in range(n_extra):\n", - " lp0_extra = source_lp_result_0.instance.extra_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = lp0_extra.bulge.centre\n", - "\n", - " # For linear Sersic profiles, compute luminosity from the solved profile.\n", - " galaxy_with_intensity = tracer.galaxies[n_main + i]\n", - " total_luminosity = (\n", - " abs(galaxy_with_intensity.bulge.luminosity_within_circle_from(radius=10.0))\n", - " / pixel_scale**2\n", - " )\n", - " luminosity_cap = 5 * 0.5 * total_luminosity**0.6\n", - " upper_limit = min(luminosity_cap, 5.0) if luminosity_cap > 0 else 5.0\n", - " mass.einstein_radius = af.UniformPrior(\n", - " lower_limit=0.0,\n", - " upper_limit=upper_limit,\n", - " )\n", - "\n", - " extra_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=lp0_extra.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", - "\n", - " # --- scaling lens galaxy models (light fixed, shared luminosity scaling relation) ---\n", - " scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - " scaling_relation = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - " scaling_mass_models = []\n", - " for i in range(n_scaling):\n", - " lp0_scaling = source_lp_result_0.instance.scaling_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = lp0_scaling.bulge.centre\n", - "\n", - " galaxy_with_intensity = tracer.galaxies[n_main + n_extra + i]\n", - " total_luminosity = (\n", - " abs(galaxy_with_intensity.bulge.luminosity_within_circle_from(radius=10.0))\n", - " / pixel_scale**2\n", - " )\n", - " mass.einstein_radius = scaling_factor * total_luminosity**scaling_relation\n", - "\n", - " scaling_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=lp0_scaling.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " scaling_galaxies = (\n", - " af.Collection(scaling_mass_models) if scaling_mass_models else None\n", - " )\n", - " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=scaling_galaxies,\n", - " )\n", - "\n", - " n_extra_model = len(extra_galaxies) if extra_galaxies is not None else 0\n", - " n_scaling_model = len(scaling_galaxies) if scaling_galaxies is not None else 0\n", - " n_live = 150 + 30 * n_main + 30 * n_extra_model + 30 * n_scaling_model\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " n_like_max=200000,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to ``group/slam.py``. The pixelization pipelines fix the light from previous stages and\n", - "do not change light profile type." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " dataset,\n", - " mask,\n", - " settings_search,\n", - " source_lp_result_1,\n", - " over_sample_size,\n", - " pixel_scale,\n", - " mask_radius,\n", - " positions,\n", - " n_batch=20,\n", - "):\n", - " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", - " edge_pixels_total = 30\n", - " signal_to_noise_threshold = 3.0\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result_1\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Hilbert(\n", - " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", - " )\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=mask,\n", - " adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"],\n", - " )\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " over_sample_size_pixelization = np.where(\n", - " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", - " 4,\n", - " 2,\n", - " )\n", - " over_sample_size_pixelization = al.Array2D(\n", - " values=over_sample_size_pixelization, mask=mask\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=over_sample_size_pixelization,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result_1.positions_likelihood_from(\n", - " factor=2.0, positions=positions, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_lp_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lp_lens_instance = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", - " lp_lens_model = getattr(source_lp_result_1.model.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=lp_lens_model.mass,\n", - " mass_result=lp_lens_model.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lp_lens_instance.redshift,\n", - " bulge=lp_lens_instance.bulge,\n", - " disk=lp_lens_instance.disk,\n", - " point=lp_lens_instance.point,\n", - " mass=mass,\n", - " shear=lp_lens_model.shear,\n", - " )\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=af.Model(al.reg.AdaptSplit),\n", - " ),\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=source_lp_result_1.model.extra_galaxies,\n", - " scaling_galaxies=source_lp_result_1.model.scaling_galaxies,\n", - " )\n", - "\n", - " n_live = 150 + 50 * (n_lenses - 1)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - " return result, dataset, adapt_images\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to ``group/slam.py``. Adapt data for the Hilbert image mesh is capped at a S/N threshold\n", - "of 3.0 to prevent over-concentration of source pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " dataset,\n", - " mask,\n", - " settings_search,\n", - " source_lp_result_1,\n", - " source_pix_result_1,\n", - " over_sample_size,\n", - " pixel_scale,\n", - " mask_radius,\n", - " n_batch=20,\n", - "):\n", - " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", - " edge_pixels_total = 30\n", - " signal_to_noise_threshold = 3.0\n", - " signal_to_noise_threshold_image_mesh = 3.0\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_data_snr_max = galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", - " adapt_data_snr_max[adapt_data_snr_max > signal_to_noise_threshold_image_mesh] = (\n", - " signal_to_noise_threshold_image_mesh\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Hilbert(\n", - " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", - " )\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=mask, adapt_data=adapt_data_snr_max\n", - " )\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " over_sample_size_pixelization = np.where(\n", - " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", - " 4,\n", - " 2,\n", - " )\n", - " over_sample_size_pixelization = al.Array2D(\n", - " values=over_sample_size_pixelization, mask=mask\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=over_sample_size_pixelization,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lp_lens_instance = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", - " pix1_lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lp_lens_instance.redshift,\n", - " bulge=lp_lens_instance.bulge,\n", - " disk=lp_lens_instance.disk,\n", - " point=lp_lens_instance.point,\n", - " mass=pix1_lens_instance.mass,\n", - " shear=pix1_lens_instance.shear,\n", - " )\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=af.Model(al.reg.AdaptSplit),\n", - " ),\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", - " scaling_galaxies=source_pix_result_1.instance.scaling_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - " return result, dataset, adapt_images\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Refits the light using linear Sersic profiles instead of MGE. Main lens galaxies get a fresh\n", - "``al.lp_linear.Sersic`` and extra galaxies get a fresh ``al.lp_linear.SersicSph``, with mass\n", - "fixed from ``source_pix[1]``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " dataset,\n", - " settings_search,\n", - " source_lp_result_0,\n", - " source_pix_result_1,\n", - " source_pix_result_2,\n", - " adapt_images,\n", - " mask_radius,\n", - " redshift_lens,\n", - " n_batch=20,\n", - "):\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - " n_extra = (\n", - " len(list(source_pix_result_1.instance.extra_galaxies))\n", - " if source_pix_result_1.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " # --- main lens light models (fresh linear Sersic) ---\n", - " lens_bulge_list = []\n", - " for i in range(n_lenses):\n", - " bulge = af.Model(al.lp_linear.Sersic)\n", - " lens_bulge_list.append(bulge)\n", - "\n", - " # --- extra lens galaxy light models (fresh linear SersicSph, mass fixed) ---\n", - " extra_light_models = []\n", - " for i in range(n_extra):\n", - " pix1_extra = source_pix_result_1.instance.extra_galaxies[i]\n", - " bulge = af.Model(al.lp_linear.SersicSph)\n", - " bulge.centre = pix1_extra.mass.centre\n", - "\n", - " extra_light_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=pix1_extra.mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_pix_result_2, source_is_model=False\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lens_instance.redshift,\n", - " bulge=lens_bulge_list[i],\n", - " disk=None,\n", - " point=None,\n", - " mass=lens_instance.mass,\n", - " shear=lens_instance.shear,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=source_pix_result_2.instance.scaling_galaxies,\n", - " )\n", - "\n", - " n_live = 300 + 100 * (n_lenses - 1)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to ``group/slam.py``. Extra galaxies receive a new luminosity-bounded ``Isothermal`` mass\n", - "(using ``light[1]`` luminosities) and scaling galaxies receive a new shared luminosity scaling relation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " dataset,\n", - " settings_search,\n", - " source_pix_result_1,\n", - " source_pix_result_2,\n", - " light_result,\n", - " adapt_images,\n", - " positions,\n", - " pixel_scale,\n", - " redshift_lens,\n", - " n_batch=20,\n", - "):\n", - " n_lenses = sum(\n", - " 1 for k in vars(light_result.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - " n_extra = (\n", - " len(list(light_result.instance.extra_galaxies))\n", - " if light_result.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - " n_scaling = (\n", - " len(list(light_result.instance.scaling_galaxies))\n", - " if light_result.instance.scaling_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " # --- extra galaxies: fixed light, free mass ---\n", - " # For linear Sersic profiles, use the tracer with solved intensities.\n", - " tracer = (\n", - " light_result.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", - " )\n", - "\n", - " extra_mass_models = []\n", - " for i in range(n_extra):\n", - " light_extra = light_result.instance.extra_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = light_extra.bulge.centre\n", - "\n", - " galaxy_with_intensity = tracer.galaxies[n_lenses + 1 + i]\n", - " total_luminosity = (\n", - " abs(galaxy_with_intensity.bulge.luminosity_within_circle_from(radius=10.0))\n", - " / pixel_scale**2\n", - " )\n", - " luminosity_cap = 5 * 0.5 * total_luminosity**0.6\n", - " upper_limit = min(luminosity_cap, 5.0) if luminosity_cap > 0 else 5.0\n", - " mass.einstein_radius = af.UniformPrior(\n", - " lower_limit=0.0,\n", - " upper_limit=upper_limit,\n", - " )\n", - "\n", - " extra_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=light_extra.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", - "\n", - " # --- scaling galaxies: fixed light, free shared scaling relation ---\n", - " scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - " scaling_relation = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - " scaling_mass_models = []\n", - " for i in range(n_scaling):\n", - " light_scaling = light_result.instance.scaling_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = light_scaling.bulge.centre\n", - "\n", - " galaxy_with_intensity = tracer.galaxies[n_lenses + 1 + n_extra + i]\n", - " total_luminosity = (\n", - " abs(galaxy_with_intensity.bulge.luminosity_within_circle_from(radius=10.0))\n", - " / pixel_scale**2\n", - " )\n", - " mass.einstein_radius = scaling_factor * total_luminosity**scaling_relation\n", - "\n", - " scaling_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=light_scaling.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " scaling_galaxies = (\n", - " af.Collection(scaling_mass_models) if scaling_mass_models else None\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " light_result.positions_likelihood_from(\n", - " factor=3.0, positions=positions, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lens_model = getattr(source_pix_result_1.model.galaxies, f\"lens_{i}\")\n", - " light_lens_instance = getattr(light_result.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.PowerLaw),\n", - " mass_result=lens_model.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lens_model.redshift,\n", - " bulge=light_lens_instance.bulge,\n", - " disk=light_lens_instance.disk,\n", - " point=light_lens_instance.point,\n", - " mass=mass,\n", - " shear=lens_model.shear,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=scaling_galaxies,\n", - " )\n", - "\n", - " n_live = 200 + 100 * (n_lenses - 1)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the ``Imaging`` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "pixel_scale = 0.1\n", - "mask_radius = 6.0\n", - "mask_centre = (0.0, 0.0)\n", - "redshift_lens = 0.5\n", - "redshift_source = 1.0\n", - "source_mge_radius = 1.0\n", - "n_batch = 20\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=pixel_scale,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "main_lens_centres.json -- required; determines the number of main lenses.\n", - "extra_galaxies_centres.json -- optional; empty list if absent.\n", - "scaling_galaxies_centres.json -- optional; empty list if absent.\n", - "\n", - "All three files contain a list of [y, x] arcsecond coordinates." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = _load_centres(dataset_path / \"main_lens_centres.json\")\n", - "extra_lens_centres = _load_centres(dataset_path / \"extra_galaxies_centres.json\")\n", - "scaling_lens_centres = _load_centres(dataset_path / \"scaling_galaxies_centres.json\")\n", - "\n", - "all_galaxy_centres = al.Grid2DIrregular(\n", - " main_lens_centres.in_list\n", - " + extra_lens_centres.in_list\n", - " + scaling_lens_centres.in_list\n", - ")\n", - "\n", - "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " centre=mask_centre,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=list(all_galaxy_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"group\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result_0 = source_lp_0(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " main_lens_centres=main_lens_centres,\n", - " extra_lens_centres=extra_lens_centres,\n", - " scaling_lens_centres=scaling_lens_centres,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - ")\n", - "\n", - "source_lp_result_1 = source_lp_1(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " source_lp_result_0=source_lp_result_0,\n", - " positions=positions,\n", - " pixel_scale=pixel_scale,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - " source_mge_radius=source_mge_radius,\n", - ")\n", - "\n", - "source_pix_result_1, dataset, adapt_images = source_pix_1(\n", - " dataset=dataset,\n", - " mask=mask,\n", - " settings_search=settings_search,\n", - " source_lp_result_1=source_lp_result_1,\n", - " over_sample_size=over_sample_size,\n", - " pixel_scale=pixel_scale,\n", - " mask_radius=mask_radius,\n", - " positions=positions,\n", - " n_batch=n_batch,\n", - ")\n", - "\n", - "source_pix_result_2, dataset, adapt_images = source_pix_2(\n", - " dataset=dataset,\n", - " mask=mask,\n", - " settings_search=settings_search,\n", - " source_lp_result_1=source_lp_result_1,\n", - " source_pix_result_1=source_pix_result_1,\n", - " over_sample_size=over_sample_size,\n", - " pixel_scale=pixel_scale,\n", - " mask_radius=mask_radius,\n", - " n_batch=n_batch,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " source_lp_result_0=source_lp_result_0,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " adapt_images=adapt_images,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " n_batch=n_batch,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " light_result=light_result,\n", - " adapt_images=adapt_images,\n", - " positions=positions,\n", - " pixel_scale=pixel_scale,\n", - " redshift_lens=redshift_lens,\n", - " n_batch=n_batch,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Linear Light Profiles: Group SLaM\n", + "===================================\n", + "\n", + "This script uses the SLaM pipelines to fit a group-scale strong lens using **linear Sersic light profiles**\n", + "instead of Multi-Gaussian Expansion (MGE) profiles for the galaxy light.\n", + "\n", + "The differences from the standard ``group/slam.py`` are:\n", + "\n", + " - The SOURCE LP PIPELINE 0 uses ``al.lp_linear.Sersic`` for main lens galaxies and\n", + " ``al.lp_linear.SersicSph`` for extra galaxies, instead of MGE models.\n", + " - The LIGHT LP PIPELINE uses ``al.lp_linear.Sersic`` for main lens galaxies and\n", + " ``al.lp_linear.SersicSph`` for extra galaxies, instead of MGE models.\n", + "\n", + "Linear light profiles solve for the ``intensity`` analytically via linear algebra, removing it from the\n", + "non-linear parameter space. This provides an alternative to MGE for group-scale modeling: simpler model\n", + "composition with fewer basis functions, while still benefiting from the intensity-free parameter space.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **SOURCE LP PIPELINE 0:** Light-only fit using linear Sersic profiles for all galaxies.\n", + "- **SOURCE LP PIPELINE 1:** Introduces mass and source with light fixed from stage 0.\n", + "- **SOURCE PIX PIPELINE 1:** Pixelized source fitting (identical to group/slam.py).\n", + "- **SOURCE PIX PIPELINE 2:** Refined pixelized source (identical to group/slam.py).\n", + "- **LIGHT LP PIPELINE:** Refits light using linear Sersic profiles.\n", + "- **MASS TOTAL PIPELINE:** Final mass fit with PowerLaw (identical to group/slam.py).\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Galaxy Centres:** Load centres from JSON files.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Settings AutoFit:** The settings of autofit.\n", + "- **SLaM Pipeline:** Run the full pipeline.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (``guides/modeling/slam_start_here``)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines.\n", + "\n", + "- **Group SLaM** (``group/slam``)\n", + " The standard group-scale SLaM pipeline using MGE profiles.\n", + "\n", + "- **Linear Light Profiles** (``features/linear_light_profiles``)\n", + " How linear light profiles work and their advantages." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "\n", + "def _load_centres(path):\n", + " \"\"\"Load a centres JSON file, returning an empty list if the file is absent.\"\"\"\n", + " try:\n", + " return al.Grid2DIrregular(al.from_json(file_path=path))\n", + " except FileNotFoundError:\n", + " return al.Grid2DIrregular([])\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 0__\n", + "\n", + "Light-only fit (no mass, no source) for every galaxy simultaneously, using linear Sersic profiles instead\n", + "of MGE. This gives the next search clean fixed light models to build on.\n", + "\n", + "Main lens galaxies use ``al.lp_linear.Sersic`` and extra galaxies use ``al.lp_linear.SersicSph``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_0(\n", + " dataset,\n", + " settings_search,\n", + " main_lens_centres,\n", + " extra_lens_centres,\n", + " scaling_lens_centres,\n", + " mask_radius,\n", + " redshift_lens,\n", + " n_batch=50,\n", + "):\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " # --- main lens light models (linear Sersic, light only) ---\n", + " lens_dict = {}\n", + " for i, centre in enumerate(main_lens_centres):\n", + " bulge = af.Model(al.lp_linear.Sersic)\n", + " bulge.centre_0 = af.GaussianPrior(mean=centre[0], sigma=0.1)\n", + " bulge.centre_1 = af.GaussianPrior(mean=centre[1], sigma=0.1)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=bulge, disk=None, point=None\n", + " )\n", + "\n", + " # --- extra lens galaxy light models (linear SersicSph) ---\n", + " extra_light_models = []\n", + " for centre in extra_lens_centres:\n", + " bulge = af.Model(al.lp_linear.SersicSph)\n", + " bulge.centre_0 = af.UniformPrior(\n", + " lower_limit=centre[0] - 0.5, upper_limit=centre[0] + 0.5\n", + " )\n", + " bulge.centre_1 = af.UniformPrior(\n", + " lower_limit=centre[1] - 0.5, upper_limit=centre[1] + 0.5\n", + " )\n", + "\n", + " extra_light_models.append(\n", + " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", + "\n", + " # --- scaling galaxy light models (linear SersicSph) ---\n", + " scaling_light_models = []\n", + " for centre in scaling_lens_centres:\n", + " bulge = af.Model(al.lp_linear.SersicSph)\n", + " bulge.centre_0 = af.UniformPrior(\n", + " lower_limit=centre[0] - 0.5, upper_limit=centre[0] + 0.5\n", + " )\n", + " bulge.centre_1 = af.UniformPrior(\n", + " lower_limit=centre[1] - 0.5, upper_limit=centre[1] + 0.5\n", + " )\n", + "\n", + " scaling_light_models.append(\n", + " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", + " )\n", + "\n", + " scaling_galaxies = (\n", + " af.Collection(scaling_light_models) if scaling_light_models else None\n", + " )\n", + "\n", + " n_extra = len(extra_galaxies) if extra_galaxies is not None else 0\n", + " n_scaling = len(scaling_galaxies) if scaling_galaxies is not None else 0\n", + " n_live = 100 + 30 * len(lens_dict) + 30 * n_extra + 30 * n_scaling\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=scaling_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[0]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " n_like_max=1000000,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 1__\n", + "\n", + "Introduces mass and source with light fixed from stage 0. Multiple main-lens galaxies each get an\n", + "``Isothermal`` mass; only ``lens_0`` carries an ``ExternalShear``. Extra-galaxy Einstein radii are\n", + "bounded by a luminosity-derived prior.\n", + "\n", + "This is identical to the standard group/slam.py ``source_lp_1`` -- the light profiles are fixed from\n", + "stage 0 (which used linear Sersic), so no changes are needed here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_1(\n", + " dataset,\n", + " settings_search,\n", + " source_lp_result_0,\n", + " positions,\n", + " pixel_scale,\n", + " redshift_lens,\n", + " redshift_source,\n", + " source_mge_radius,\n", + " n_batch=50,\n", + "):\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(source_lp_result_0.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + " n_extra = (\n", + " len(list(source_lp_result_0.instance.extra_galaxies))\n", + " if source_lp_result_0.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + " n_scaling = (\n", + " len(list(source_lp_result_0.instance.scaling_galaxies))\n", + " if source_lp_result_0.instance.scaling_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " tracer = (\n", + " source_lp_result_0.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", + " )\n", + "\n", + " # Source MGE centred on primary lens bulge from stage 0.\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=source_mge_radius,\n", + " total_gaussians=30,\n", + " centre=source_lp_result_0.instance.galaxies.lens_0.bulge.centre,\n", + " centre_prior_is_uniform=False,\n", + " centre_sigma=0.6,\n", + " )\n", + "\n", + " # --- main lens full models (light fixed from stage 0, mass + shear free) ---\n", + " lens_dict = {}\n", + " for i in range(n_main):\n", + " lp0_lens = getattr(source_lp_result_0.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = lp0_lens.bulge.centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=5.0)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lp0_lens.bulge,\n", + " disk=lp0_lens.disk,\n", + " point=lp0_lens.point,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " # --- extra lens galaxy models (light fixed, mass bounded by luminosity) ---\n", + " extra_mass_models = []\n", + " for i in range(n_extra):\n", + " lp0_extra = source_lp_result_0.instance.extra_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = lp0_extra.bulge.centre\n", + "\n", + " # For linear Sersic profiles, compute luminosity from the solved profile.\n", + " galaxy_with_intensity = tracer.galaxies[n_main + i]\n", + " total_luminosity = (\n", + " abs(galaxy_with_intensity.bulge.luminosity_within_circle_from(radius=10.0))\n", + " / pixel_scale**2\n", + " )\n", + " luminosity_cap = 5 * 0.5 * total_luminosity**0.6\n", + " upper_limit = min(luminosity_cap, 5.0) if luminosity_cap > 0 else 5.0\n", + " mass.einstein_radius = af.UniformPrior(\n", + " lower_limit=0.0,\n", + " upper_limit=upper_limit,\n", + " )\n", + "\n", + " extra_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=lp0_extra.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", + "\n", + " # --- scaling lens galaxy models (light fixed, shared luminosity scaling relation) ---\n", + " scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + " scaling_relation = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + " scaling_mass_models = []\n", + " for i in range(n_scaling):\n", + " lp0_scaling = source_lp_result_0.instance.scaling_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = lp0_scaling.bulge.centre\n", + "\n", + " galaxy_with_intensity = tracer.galaxies[n_main + n_extra + i]\n", + " total_luminosity = (\n", + " abs(galaxy_with_intensity.bulge.luminosity_within_circle_from(radius=10.0))\n", + " / pixel_scale**2\n", + " )\n", + " mass.einstein_radius = scaling_factor * total_luminosity**scaling_relation\n", + "\n", + " scaling_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=lp0_scaling.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " scaling_galaxies = (\n", + " af.Collection(scaling_mass_models) if scaling_mass_models else None\n", + " )\n", + " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=scaling_galaxies,\n", + " )\n", + "\n", + " n_extra_model = len(extra_galaxies) if extra_galaxies is not None else 0\n", + " n_scaling_model = len(scaling_galaxies) if scaling_galaxies is not None else 0\n", + " n_live = 150 + 30 * n_main + 30 * n_extra_model + 30 * n_scaling_model\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " n_like_max=200000,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to ``group/slam.py``. The pixelization pipelines fix the light from previous stages and\n", + "do not change light profile type." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " dataset,\n", + " mask,\n", + " settings_search,\n", + " source_lp_result_1,\n", + " over_sample_size,\n", + " pixel_scale,\n", + " mask_radius,\n", + " positions,\n", + " n_batch=20,\n", + "):\n", + " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", + " edge_pixels_total = 30\n", + " signal_to_noise_threshold = 3.0\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result_1\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Hilbert(\n", + " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", + " )\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=mask,\n", + " adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"],\n", + " )\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " over_sample_size_pixelization = np.where(\n", + " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", + " 4,\n", + " 2,\n", + " )\n", + " over_sample_size_pixelization = al.Array2D(\n", + " values=over_sample_size_pixelization, mask=mask\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=over_sample_size_pixelization,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result_1.positions_likelihood_from(\n", + " factor=2.0, positions=positions, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_lp_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lp_lens_instance = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", + " lp_lens_model = getattr(source_lp_result_1.model.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=lp_lens_model.mass,\n", + " mass_result=lp_lens_model.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lp_lens_instance.redshift,\n", + " bulge=lp_lens_instance.bulge,\n", + " disk=lp_lens_instance.disk,\n", + " point=lp_lens_instance.point,\n", + " mass=mass,\n", + " shear=lp_lens_model.shear,\n", + " )\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=af.Model(al.reg.AdaptSplit),\n", + " ),\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=source_lp_result_1.model.extra_galaxies,\n", + " scaling_galaxies=source_lp_result_1.model.scaling_galaxies,\n", + " )\n", + "\n", + " n_live = 150 + 50 * (n_lenses - 1)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + " return result, dataset, adapt_images\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to ``group/slam.py``. Adapt data for the Hilbert image mesh is capped at a S/N threshold\n", + "of 3.0 to prevent over-concentration of source pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " dataset,\n", + " mask,\n", + " settings_search,\n", + " source_lp_result_1,\n", + " source_pix_result_1,\n", + " over_sample_size,\n", + " pixel_scale,\n", + " mask_radius,\n", + " n_batch=20,\n", + "):\n", + " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", + " edge_pixels_total = 30\n", + " signal_to_noise_threshold = 3.0\n", + " signal_to_noise_threshold_image_mesh = 3.0\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_data_snr_max = galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", + " adapt_data_snr_max[adapt_data_snr_max > signal_to_noise_threshold_image_mesh] = (\n", + " signal_to_noise_threshold_image_mesh\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Hilbert(\n", + " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", + " )\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=mask, adapt_data=adapt_data_snr_max\n", + " )\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " over_sample_size_pixelization = np.where(\n", + " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", + " 4,\n", + " 2,\n", + " )\n", + " over_sample_size_pixelization = al.Array2D(\n", + " values=over_sample_size_pixelization, mask=mask\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=over_sample_size_pixelization,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lp_lens_instance = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", + " pix1_lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lp_lens_instance.redshift,\n", + " bulge=lp_lens_instance.bulge,\n", + " disk=lp_lens_instance.disk,\n", + " point=lp_lens_instance.point,\n", + " mass=pix1_lens_instance.mass,\n", + " shear=pix1_lens_instance.shear,\n", + " )\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=af.Model(al.reg.AdaptSplit),\n", + " ),\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", + " scaling_galaxies=source_pix_result_1.instance.scaling_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + " return result, dataset, adapt_images\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Refits the light using linear Sersic profiles instead of MGE. Main lens galaxies get a fresh\n", + "``al.lp_linear.Sersic`` and extra galaxies get a fresh ``al.lp_linear.SersicSph``, with mass\n", + "fixed from ``source_pix[1]``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " dataset,\n", + " settings_search,\n", + " source_lp_result_0,\n", + " source_pix_result_1,\n", + " source_pix_result_2,\n", + " adapt_images,\n", + " mask_radius,\n", + " redshift_lens,\n", + " n_batch=20,\n", + "):\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + " n_extra = (\n", + " len(list(source_pix_result_1.instance.extra_galaxies))\n", + " if source_pix_result_1.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " # --- main lens light models (fresh linear Sersic) ---\n", + " lens_bulge_list = []\n", + " for i in range(n_lenses):\n", + " bulge = af.Model(al.lp_linear.Sersic)\n", + " lens_bulge_list.append(bulge)\n", + "\n", + " # --- extra lens galaxy light models (fresh linear SersicSph, mass fixed) ---\n", + " extra_light_models = []\n", + " for i in range(n_extra):\n", + " pix1_extra = source_pix_result_1.instance.extra_galaxies[i]\n", + " bulge = af.Model(al.lp_linear.SersicSph)\n", + " bulge.centre = pix1_extra.mass.centre\n", + "\n", + " extra_light_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=pix1_extra.mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_pix_result_2, source_is_model=False\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lens_instance.redshift,\n", + " bulge=lens_bulge_list[i],\n", + " disk=None,\n", + " point=None,\n", + " mass=lens_instance.mass,\n", + " shear=lens_instance.shear,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=source_pix_result_2.instance.scaling_galaxies,\n", + " )\n", + "\n", + " n_live = 300 + 100 * (n_lenses - 1)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to ``group/slam.py``. Extra galaxies receive a new luminosity-bounded ``Isothermal`` mass\n", + "(using ``light[1]`` luminosities) and scaling galaxies receive a new shared luminosity scaling relation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " dataset,\n", + " settings_search,\n", + " source_pix_result_1,\n", + " source_pix_result_2,\n", + " light_result,\n", + " adapt_images,\n", + " positions,\n", + " pixel_scale,\n", + " redshift_lens,\n", + " n_batch=20,\n", + "):\n", + " n_lenses = sum(\n", + " 1 for k in vars(light_result.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + " n_extra = (\n", + " len(list(light_result.instance.extra_galaxies))\n", + " if light_result.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + " n_scaling = (\n", + " len(list(light_result.instance.scaling_galaxies))\n", + " if light_result.instance.scaling_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " # --- extra galaxies: fixed light, free mass ---\n", + " # For linear Sersic profiles, use the tracer with solved intensities.\n", + " tracer = (\n", + " light_result.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", + " )\n", + "\n", + " extra_mass_models = []\n", + " for i in range(n_extra):\n", + " light_extra = light_result.instance.extra_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = light_extra.bulge.centre\n", + "\n", + " galaxy_with_intensity = tracer.galaxies[n_lenses + 1 + i]\n", + " total_luminosity = (\n", + " abs(galaxy_with_intensity.bulge.luminosity_within_circle_from(radius=10.0))\n", + " / pixel_scale**2\n", + " )\n", + " luminosity_cap = 5 * 0.5 * total_luminosity**0.6\n", + " upper_limit = min(luminosity_cap, 5.0) if luminosity_cap > 0 else 5.0\n", + " mass.einstein_radius = af.UniformPrior(\n", + " lower_limit=0.0,\n", + " upper_limit=upper_limit,\n", + " )\n", + "\n", + " extra_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=light_extra.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", + "\n", + " # --- scaling galaxies: fixed light, free shared scaling relation ---\n", + " scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + " scaling_relation = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + " scaling_mass_models = []\n", + " for i in range(n_scaling):\n", + " light_scaling = light_result.instance.scaling_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = light_scaling.bulge.centre\n", + "\n", + " galaxy_with_intensity = tracer.galaxies[n_lenses + 1 + n_extra + i]\n", + " total_luminosity = (\n", + " abs(galaxy_with_intensity.bulge.luminosity_within_circle_from(radius=10.0))\n", + " / pixel_scale**2\n", + " )\n", + " mass.einstein_radius = scaling_factor * total_luminosity**scaling_relation\n", + "\n", + " scaling_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=light_scaling.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " scaling_galaxies = (\n", + " af.Collection(scaling_mass_models) if scaling_mass_models else None\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " light_result.positions_likelihood_from(\n", + " factor=3.0, positions=positions, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lens_model = getattr(source_pix_result_1.model.galaxies, f\"lens_{i}\")\n", + " light_lens_instance = getattr(light_result.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.PowerLaw),\n", + " mass_result=lens_model.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lens_model.redshift,\n", + " bulge=light_lens_instance.bulge,\n", + " disk=light_lens_instance.disk,\n", + " point=light_lens_instance.point,\n", + " mass=mass,\n", + " shear=lens_model.shear,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=scaling_galaxies,\n", + " )\n", + "\n", + " n_live = 200 + 100 * (n_lenses - 1)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the ``Imaging`` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "pixel_scale = 0.1\n", + "mask_radius = 6.0\n", + "mask_centre = (0.0, 0.0)\n", + "redshift_lens = 0.5\n", + "redshift_source = 1.0\n", + "source_mge_radius = 1.0\n", + "n_batch = 20\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=pixel_scale,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "main_lens_centres.json -- required; determines the number of main lenses.\n", + "extra_galaxies_centres.json -- optional; empty list if absent.\n", + "scaling_galaxies_centres.json -- optional; empty list if absent.\n", + "\n", + "All three files contain a list of [y, x] arcsecond coordinates." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = _load_centres(dataset_path / \"main_lens_centres.json\")\n", + "extra_lens_centres = _load_centres(dataset_path / \"extra_galaxies_centres.json\")\n", + "scaling_lens_centres = _load_centres(dataset_path / \"scaling_galaxies_centres.json\")\n", + "\n", + "all_galaxy_centres = al.Grid2DIrregular(\n", + " main_lens_centres.in_list\n", + " + extra_lens_centres.in_list\n", + " + scaling_lens_centres.in_list\n", + ")\n", + "\n", + "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " centre=mask_centre,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=list(all_galaxy_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"group\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result_0 = source_lp_0(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " main_lens_centres=main_lens_centres,\n", + " extra_lens_centres=extra_lens_centres,\n", + " scaling_lens_centres=scaling_lens_centres,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + ")\n", + "\n", + "source_lp_result_1 = source_lp_1(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " source_lp_result_0=source_lp_result_0,\n", + " positions=positions,\n", + " pixel_scale=pixel_scale,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + " source_mge_radius=source_mge_radius,\n", + ")\n", + "\n", + "source_pix_result_1, dataset, adapt_images = source_pix_1(\n", + " dataset=dataset,\n", + " mask=mask,\n", + " settings_search=settings_search,\n", + " source_lp_result_1=source_lp_result_1,\n", + " over_sample_size=over_sample_size,\n", + " pixel_scale=pixel_scale,\n", + " mask_radius=mask_radius,\n", + " positions=positions,\n", + " n_batch=n_batch,\n", + ")\n", + "\n", + "source_pix_result_2, dataset, adapt_images = source_pix_2(\n", + " dataset=dataset,\n", + " mask=mask,\n", + " settings_search=settings_search,\n", + " source_lp_result_1=source_lp_result_1,\n", + " source_pix_result_1=source_pix_result_1,\n", + " over_sample_size=over_sample_size,\n", + " pixel_scale=pixel_scale,\n", + " mask_radius=mask_radius,\n", + " n_batch=n_batch,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " source_lp_result_0=source_lp_result_0,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " adapt_images=adapt_images,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " n_batch=n_batch,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " light_result=light_result,\n", + " adapt_images=adapt_images,\n", + " positions=positions,\n", + " pixel_scale=pixel_scale,\n", + " redshift_lens=redshift_lens,\n", + " n_batch=n_batch,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/multi_gaussian_expansion/fit.ipynb b/notebooks/group/features/multi_gaussian_expansion/fit.ipynb index 4285e8141..356b5309c 100644 --- a/notebooks/group/features/multi_gaussian_expansion/fit.ipynb +++ b/notebooks/group/features/multi_gaussian_expansion/fit.ipynb @@ -1,470 +1,507 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fit Features: Multi Gaussian Expansion (Group)\n", - "===============================================\n", - "\n", - "This guide shows how to fit data using the `FitImaging` object for group-scale strong lenses, including visualizing\n", - "and interpreting its results.\n", - "\n", - "A Multi Gaussian Expansion (MGE) decomposes each galaxy's light into ~10-30+ Gaussians whose intensities are\n", - "solved via linear algebra. For group-scale lenses, this means that adding extra galaxies does not increase the\n", - "number of non-linear parameters, making MGE the recommended approach for group modeling.\n", - "\n", - "In this example, we use simple `SersicSph` light profiles to create concrete galaxy instances for the fit\n", - "demonstration, because specifying concrete MGE instances requires providing intensity and sigma values for\n", - "every Gaussian. In practice, MGE light profiles would be used via lens modeling (see the `modeling.py` example),\n", - "where the intensities are determined automatically via linear algebra.\n", - "\n", - "__Contents__\n", - "\n", - "- **Loading Data:** Load the group-scale strong lens dataset.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Galaxy Centres:** Load centres of main lens galaxies and extra galaxies from JSON files.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", - "- **Bad Fit:** A bad lens model will show features in the residual-map and chi-squared map.\n", - "- **Fit Quantities:** The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", - "- **Figures of Merit:** There are single valued floats which quantify the goodness of fit.\n", - "- **MGE In Practice:** How MGE would be used in a real fit via modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Data__\n", - "\n", - "We begin by loading the group-scale strong lens dataset `simple` from .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5\" circular mask, which is larger than a typical galaxy-scale lens mask because the group-scale\n", - "lens has emission spread over a wider area due to the multiple lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.plot_array(array=dataset.data, title=\"Image Data With Mask Applied\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "For group-scale lenses we load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "print(f\"Main lens centres: {main_lens_centres}\")\n", - "print(f\"Extra galaxies centres: {extra_galaxies_centres}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre ensures that the light profiles of every galaxy in the group are\n", - "accurately evaluated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fitting__\n", - "\n", - "We create a tracer from a collection of light profiles, mass profiles and galaxies.\n", - "\n", - "For this fit demonstration, we use simple `SersicSph` light profiles for the lens galaxies. In a real\n", - "analysis, these would be replaced with MGE light profiles determined via lens modeling, which would capture\n", - "more complex morphological features.\n", - "\n", - "The combination of light and mass profiles below is the same as those used to generate the simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now use a `FitImaging` object to fit this tracer to the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.plot_array(array=fit.model_data, title=\"Model Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A subplot can be plotted which contains all fit quantities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Bad Fit__\n", - "\n", - "A bad lens model will show features in the residual-map and chi-squared map. We demonstrate this by\n", - "offsetting the main lens galaxy's mass centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.2, 0.2), einstein_radius=4.0),\n", - ")\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit Quantities__\n", - "\n", - "The maximum log likelihood fit contains many 1D and 2D arrays showing the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_data.slim)\n", - "print(fit.residual_map.slim)\n", - "print(fit.normalized_residual_map.slim)\n", - "print(fit.chi_squared_map.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Figures of Merit__\n", - "\n", - "There are single valued floats which quantify the goodness of fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.chi_squared)\n", - "print(fit.noise_normalization)\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Plane Quantities__\n", - "\n", - "The `FitImaging` object has specific quantities which break down each image of each plane. For group-scale\n", - "lenses, all lens galaxies (main and extra) are at the same redshift and therefore in the same plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_images_of_planes_list[0].slim)\n", - "print(fit.model_images_of_planes_list[1].slim)\n", - "\n", - "print(fit.subtracted_images_of_planes_list[0].slim)\n", - "print(fit.subtracted_images_of_planes_list[1].slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Unmasked Quantities__\n", - "\n", - "The `FitImaging` can also compute the unmasked blurred image of each plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.unmasked_blurred_image.native)\n", - "print(fit.unmasked_blurred_image_of_planes_list[0].native)\n", - "print(fit.unmasked_blurred_image_of_planes_list[1].native)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE In Practice__\n", - "\n", - "In a real analysis, the light profiles of each galaxy would be MGE models constructed via\n", - "`al.model_util.mge_model_from`, whose intensities are solved via linear algebra during model fitting.\n", - "\n", - "The `FitImaging` object works identically with MGE light profiles -- the fit subplot, residual maps,\n", - "chi-squared maps, and all other quantities are computed in the same way. The only difference is that the\n", - "light profile images are the sum of many Gaussian components rather than a single analytic profile.\n", - "\n", - "After lens modeling with MGE, you would extract the `max_log_likelihood_fit` from the result object:\n", - "\n", - " fit = result.max_log_likelihood_fit\n", - " aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "This fit object contains the same quantities demonstrated above, but with MGE light profiles providing\n", - "a more accurate representation of each galaxy's light.\n", - "\n", - "See the `modeling.py` example in this folder for the full MGE modeling workflow." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "Fin.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fit Features: Multi Gaussian Expansion (Group)\n", + "===============================================\n", + "\n", + "This guide shows how to fit data using the `FitImaging` object for group-scale strong lenses, including visualizing\n", + "and interpreting its results.\n", + "\n", + "A Multi Gaussian Expansion (MGE) decomposes each galaxy's light into ~10-30+ Gaussians whose intensities are\n", + "solved via linear algebra. For group-scale lenses, this means that adding extra galaxies does not increase the\n", + "number of non-linear parameters, making MGE the recommended approach for group modeling.\n", + "\n", + "In this example, we use simple `SersicSph` light profiles to create concrete galaxy instances for the fit\n", + "demonstration, because specifying concrete MGE instances requires providing intensity and sigma values for\n", + "every Gaussian. In practice, MGE light profiles would be used via lens modeling (see the `modeling.py` example),\n", + "where the intensities are determined automatically via linear algebra.\n", + "\n", + "__Contents__\n", + "\n", + "- **Loading Data:** Load the group-scale strong lens dataset.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Galaxy Centres:** Load centres of main lens galaxies and extra galaxies from JSON files.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", + "- **Bad Fit:** A bad lens model will show features in the residual-map and chi-squared map.\n", + "- **Fit Quantities:** The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", + "- **Figures of Merit:** There are single valued floats which quantify the goodness of fit.\n", + "- **MGE In Practice:** How MGE would be used in a real fit via modeling." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Data__\n", + "\n", + "We begin by loading the group-scale strong lens dataset `simple` from .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5\" circular mask, which is larger than a typical galaxy-scale lens mask because the group-scale\n", + "lens has emission spread over a wider area due to the multiple lens galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.plot_array(array=dataset.data, title=\"Image Data With Mask Applied\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "For group-scale lenses we load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "print(f\"Main lens centres: {main_lens_centres}\")\n", + "print(f\"Extra galaxies centres: {extra_galaxies_centres}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre ensures that the light profiles of every galaxy in the group are\n", + "accurately evaluated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fitting__\n", + "\n", + "We create a tracer from a collection of light profiles, mass profiles and galaxies.\n", + "\n", + "For this fit demonstration, we use simple `SersicSph` light profiles for the lens galaxies. In a real\n", + "analysis, these would be replaced with MGE light profiles determined via lens modeling, which would capture\n", + "more complex morphological features.\n", + "\n", + "The combination of light and mass profiles below is the same as those used to generate the simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now use a `FitImaging` object to fit this tracer to the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.plot_array(array=fit.model_data, title=\"Model Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A subplot can be plotted which contains all fit quantities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Bad Fit__\n", + "\n", + "A bad lens model will show features in the residual-map and chi-squared map. We demonstrate this by\n", + "offsetting the main lens galaxy's mass centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.2, 0.2), einstein_radius=4.0),\n", + ")\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit Quantities__\n", + "\n", + "The maximum log likelihood fit contains many 1D and 2D arrays showing the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_data.slim)\n", + "print(fit.residual_map.slim)\n", + "print(fit.normalized_residual_map.slim)\n", + "print(fit.chi_squared_map.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Figures of Merit__\n", + "\n", + "There are single valued floats which quantify the goodness of fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.chi_squared)\n", + "print(fit.noise_normalization)\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Plane Quantities__\n", + "\n", + "The `FitImaging` object has specific quantities which break down each image of each plane. For group-scale\n", + "lenses, all lens galaxies (main and extra) are at the same redshift and therefore in the same plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_images_of_planes_list[0].slim)\n", + "print(fit.model_images_of_planes_list[1].slim)\n", + "\n", + "print(fit.subtracted_images_of_planes_list[0].slim)\n", + "print(fit.subtracted_images_of_planes_list[1].slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Unmasked Quantities__\n", + "\n", + "The `FitImaging` can also compute the unmasked blurred image of each plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.unmasked_blurred_image.native)\n", + "print(fit.unmasked_blurred_image_of_planes_list[0].native)\n", + "print(fit.unmasked_blurred_image_of_planes_list[1].native)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE In Practice__\n", + "\n", + "In a real analysis, the light profiles of each galaxy would be MGE models constructed via\n", + "`al.model_util.mge_model_from`, whose intensities are solved via linear algebra during model fitting.\n", + "\n", + "The `FitImaging` object works identically with MGE light profiles -- the fit subplot, residual maps,\n", + "chi-squared maps, and all other quantities are computed in the same way. The only difference is that the\n", + "light profile images are the sum of many Gaussian components rather than a single analytic profile.\n", + "\n", + "After lens modeling with MGE, you would extract the `max_log_likelihood_fit` from the result object:\n", + "\n", + " fit = result.max_log_likelihood_fit\n", + " aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "This fit object contains the same quantities demonstrated above, but with MGE light profiles providing\n", + "a more accurate representation of each galaxy's light.\n", + "\n", + "See the `modeling.py` example in this folder for the full MGE modeling workflow." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "Fin.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/multi_gaussian_expansion/likelihood_function.ipynb b/notebooks/group/features/multi_gaussian_expansion/likelihood_function.ipynb index 54a947715..106f084e1 100644 --- a/notebooks/group/features/multi_gaussian_expansion/likelihood_function.ipynb +++ b/notebooks/group/features/multi_gaussian_expansion/likelihood_function.ipynb @@ -1,571 +1,608 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Multi Gaussian Expansion (Group)__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data of\n", - "a group-scale strong lens, with a focus on how Multi Gaussian Expansion (MGE) light profiles fit into the\n", - "likelihood calculation.\n", - "\n", - "An MGE decomposes each galaxy's light into ~10-30+ Gaussians whose intensities are solved via linear algebra.\n", - "For group-scale lenses, this means multiple galaxies can each have their own set of Gaussians, and the total\n", - "model image is the sum of all Gaussian images from all galaxies, convolved with the PSF.\n", - "\n", - "This script uses simple `SersicSph` light profiles for the concrete step-by-step calculation (as specifying\n", - "concrete MGE instances requires many parameters), but explains how the MGE likelihood differs at each step.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Disable over sampling for simplicity.\n", - "- **Main Lens Galaxy:** The main lens galaxy at the centre of the group.\n", - "- **Extra Galaxies:** The two extra galaxies are companion galaxies near the main lens.\n", - "- **Source Galaxy:** The source galaxy light profile.\n", - "- **Lens Light:** Compute a 2D image of each lens galaxy's light and sum them together.\n", - "- **Lens Galaxy Mass:** Compute deflection angles from all mass profiles.\n", - "- **Ray Tracing:** Ray-trace the image-plane grid to the source plane.\n", - "- **Source Image:** Evaluate the source galaxy light on the traced grid.\n", - "- **Convolution:** Convolve the total image with the PSF.\n", - "- **Likelihood Function:** Compute the log likelihood.\n", - "- **MGE Likelihood:** How the likelihood changes when using MGE light profiles.\n", - "- **Fit:** Verify using the FitImaging object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autoarray as aa\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the group-scale strong lens dataset with 0.1 arcsecond-per-pixel resolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"group\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5\" circular mask for the group-scale lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=masked_dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "For simplicity, we disable over sampling in this guide by setting `sub_size=1`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxy__\n", - "\n", - "The main lens galaxy is at the centre of the group. It has a spherical Sersic light profile and a spherical\n", - "isothermal mass profile.\n", - "\n", - "With an MGE, this galaxy's light would instead be decomposed into ~20 Gaussians, each with a different sigma\n", - "value spanning from 0.01\" to the mask radius. The intensities of these Gaussians are solved via linear algebra\n", - "rather than being specified manually." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "The two extra galaxies are companion galaxies near the main lens. With an MGE, each extra galaxy's light\n", - "would be decomposed into ~10 Gaussians with centres fixed to the observed positions. Crucially, this adds\n", - "zero non-linear parameters per extra galaxy, unlike a Sersic profile which adds 5." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Light Profile__\n", - "\n", - "The source galaxy uses a cored elliptical Sersic. With an MGE source model, this would instead be ~20\n", - "Gaussians evaluated on the source-plane grid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Light__\n", - "\n", - "Compute a 2D image of each lens galaxy's light and sum them together.\n", - "\n", - "For an MGE, each galaxy contributes many Gaussian images (e.g. 20 for the main lens, 10 per extra galaxy).\n", - "The total lens light image is the sum of ALL individual Gaussian images from ALL lens galaxies. With linear\n", - "light profiles, the intensity of each Gaussian is optimized to minimize the chi-squared." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", - "extra_0_image_2d = extra_galaxy_0.image_2d_from(grid=masked_dataset.grid)\n", - "extra_1_image_2d = extra_galaxy_1.image_2d_from(grid=masked_dataset.grid)\n", - "\n", - "total_lens_image_2d = lens_image_2d + extra_0_image_2d + extra_1_image_2d\n", - "\n", - "aplt.plot_array(array=total_lens_image_2d, title=\"Total Lens Light (All Galaxies)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Compute blurring images for ALL lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n", - "extra_0_blurring_image_2d = extra_galaxy_0.image_2d_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")\n", - "extra_1_blurring_image_2d = extra_galaxy_1.image_2d_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")\n", - "\n", - "total_lens_blurring_image_2d = (\n", - " lens_blurring_image_2d + extra_0_blurring_image_2d + extra_1_blurring_image_2d\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy Mass__\n", - "\n", - "We compute the deflection angles from all mass profiles in the group." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_lens = lens_galaxy.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_extra_0 = extra_galaxy_0.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_extra_1 = extra_galaxy_1.deflections_yx_2d_from(grid=masked_dataset.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The total deflection angle is the sum of deflection angles from ALL galaxies in the group." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", - "\n", - "aplt.plot_grid(grid=traced_grid, title=\"Source Plane Grid (Traced)\")\n", - "\n", - "traced_blurring_grid = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")[-1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Image__\n", - "\n", - "Evaluate the source galaxy light on the traced grid.\n", - "\n", - "For an MGE source, the source light would be the sum of ~20 Gaussian images evaluated on the traced grid,\n", - "with intensities solved via the same linear algebra inversion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_image_2d = source_galaxy.image_2d_from(grid=traced_grid)\n", - "\n", - "source_blurring_image_2d = source_galaxy.image_2d_from(grid=traced_blurring_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens + Source Light Addition__\n", - "\n", - "Add the total lens light and source image together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = total_lens_image_2d + source_image_2d\n", - "\n", - "aplt.plot_array(array=image, title=\"Total Image (All Galaxies + Source)\")\n", - "\n", - "blurring_image_2d = total_lens_blurring_image_2d + source_blurring_image_2d" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Convolution__\n", - "\n", - "Convolve the 2D image with the PSF in real-space." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", - " image=image, blurring_image=blurring_image_2d\n", - ")\n", - "\n", - "aplt.plot_array(array=convolved_image_2d, title=\"Convolved Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our group-scale lens model.\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "__Chi Squared__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image = convolved_image_2d\n", - "\n", - "residual_map = masked_dataset.data - model_image\n", - "normalized_residual_map = residual_map / masked_dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(chi_squared)\n", - "\n", - "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", - "\n", - "aplt.plot_array(array=chi_squared_map, title=\"Chi-Squared Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(figure_of_merit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE Likelihood__\n", - "\n", - "When using MGE light profiles (linear light profiles), the likelihood calculation differs in an important way.\n", - "Instead of each galaxy having a single light profile with a fixed intensity, each galaxy contributes a set of\n", - "Gaussian basis functions to a `mapping_matrix`.\n", - "\n", - "The mapping matrix has dimensions `(total_image_pixels, total_linear_light_profiles)`, where each column\n", - "is the PSF-convolved image of one Gaussian. For a group with a 20-Gaussian main lens, two 10-Gaussian extras,\n", - "and a 20-Gaussian source, this matrix has 60 columns.\n", - "\n", - "The intensities of all 60 Gaussians are then solved simultaneously via positive-only linear algebra:\n", - "\n", - " s = F^{-1} D\n", - "\n", - "where F is the curvature matrix and D is the data vector (see the `linear_light_profiles` feature for details).\n", - "\n", - "This joint optimization across all galaxies is what makes MGE so powerful for groups: the linear algebra\n", - "automatically determines the optimal intensity decomposition for every galaxy simultaneously, accounting for\n", - "blending between overlapping galaxies.\n", - "\n", - "__Fit__\n", - "\n", - "The `FitImaging` object handles all of the steps above automatically." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)\n", - "fit_figure_of_merit = fit.figure_of_merit\n", - "print(fit_figure_of_merit)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the group-scale parametric likelihood function, with\n", - "explanations of how MGE light profiles modify each step. The key differences when using MGE are:\n", - "\n", - " - Each galaxy's light is a sum of many Gaussians rather than a single analytic profile.\n", - " - The intensities of all Gaussians across all galaxies are solved jointly via linear algebra.\n", - " - This joint linear inversion adds zero non-linear parameters per galaxy, making it ideal for groups\n", - " with many extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Multi Gaussian Expansion (Group)__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data of\n", + "a group-scale strong lens, with a focus on how Multi Gaussian Expansion (MGE) light profiles fit into the\n", + "likelihood calculation.\n", + "\n", + "An MGE decomposes each galaxy's light into ~10-30+ Gaussians whose intensities are solved via linear algebra.\n", + "For group-scale lenses, this means multiple galaxies can each have their own set of Gaussians, and the total\n", + "model image is the sum of all Gaussian images from all galaxies, convolved with the PSF.\n", + "\n", + "This script uses simple `SersicSph` light profiles for the concrete step-by-step calculation (as specifying\n", + "concrete MGE instances requires many parameters), but explains how the MGE likelihood differs at each step.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Disable over sampling for simplicity.\n", + "- **Main Lens Galaxy:** The main lens galaxy at the centre of the group.\n", + "- **Extra Galaxies:** The two extra galaxies are companion galaxies near the main lens.\n", + "- **Source Galaxy:** The source galaxy light profile.\n", + "- **Lens Light:** Compute a 2D image of each lens galaxy's light and sum them together.\n", + "- **Lens Galaxy Mass:** Compute deflection angles from all mass profiles.\n", + "- **Ray Tracing:** Ray-trace the image-plane grid to the source plane.\n", + "- **Source Image:** Evaluate the source galaxy light on the traced grid.\n", + "- **Convolution:** Convolve the total image with the PSF.\n", + "- **Likelihood Function:** Compute the log likelihood.\n", + "- **MGE Likelihood:** How the likelihood changes when using MGE light profiles.\n", + "- **Fit:** Verify using the FitImaging object." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autoarray as aa\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the group-scale strong lens dataset with 0.1 arcsecond-per-pixel resolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"group\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5\" circular mask for the group-scale lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=masked_dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "For simplicity, we disable over sampling in this guide by setting `sub_size=1`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxy__\n", + "\n", + "The main lens galaxy is at the centre of the group. It has a spherical Sersic light profile and a spherical\n", + "isothermal mass profile.\n", + "\n", + "With an MGE, this galaxy's light would instead be decomposed into ~20 Gaussians, each with a different sigma\n", + "value spanning from 0.01\" to the mask radius. The intensities of these Gaussians are solved via linear algebra\n", + "rather than being specified manually." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "The two extra galaxies are companion galaxies near the main lens. With an MGE, each extra galaxy's light\n", + "would be decomposed into ~10 Gaussians with centres fixed to the observed positions. Crucially, this adds\n", + "zero non-linear parameters per extra galaxy, unlike a Sersic profile which adds 5." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Light Profile__\n", + "\n", + "The source galaxy uses a cored elliptical Sersic. With an MGE source model, this would instead be ~20\n", + "Gaussians evaluated on the source-plane grid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Light__\n", + "\n", + "Compute a 2D image of each lens galaxy's light and sum them together.\n", + "\n", + "For an MGE, each galaxy contributes many Gaussian images (e.g. 20 for the main lens, 10 per extra galaxy).\n", + "The total lens light image is the sum of ALL individual Gaussian images from ALL lens galaxies. With linear\n", + "light profiles, the intensity of each Gaussian is optimized to minimize the chi-squared." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", + "extra_0_image_2d = extra_galaxy_0.image_2d_from(grid=masked_dataset.grid)\n", + "extra_1_image_2d = extra_galaxy_1.image_2d_from(grid=masked_dataset.grid)\n", + "\n", + "total_lens_image_2d = lens_image_2d + extra_0_image_2d + extra_1_image_2d\n", + "\n", + "aplt.plot_array(array=total_lens_image_2d, title=\"Total Lens Light (All Galaxies)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compute blurring images for ALL lens galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n", + "extra_0_blurring_image_2d = extra_galaxy_0.image_2d_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")\n", + "extra_1_blurring_image_2d = extra_galaxy_1.image_2d_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")\n", + "\n", + "total_lens_blurring_image_2d = (\n", + " lens_blurring_image_2d + extra_0_blurring_image_2d + extra_1_blurring_image_2d\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy Mass__\n", + "\n", + "We compute the deflection angles from all mass profiles in the group." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_lens = lens_galaxy.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_extra_0 = extra_galaxy_0.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_extra_1 = extra_galaxy_1.deflections_yx_2d_from(grid=masked_dataset.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The total deflection angle is the sum of deflection angles from ALL galaxies in the group." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", + "\n", + "aplt.plot_grid(grid=traced_grid, title=\"Source Plane Grid (Traced)\")\n", + "\n", + "traced_blurring_grid = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")[-1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Image__\n", + "\n", + "Evaluate the source galaxy light on the traced grid.\n", + "\n", + "For an MGE source, the source light would be the sum of ~20 Gaussian images evaluated on the traced grid,\n", + "with intensities solved via the same linear algebra inversion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_image_2d = source_galaxy.image_2d_from(grid=traced_grid)\n", + "\n", + "source_blurring_image_2d = source_galaxy.image_2d_from(grid=traced_blurring_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens + Source Light Addition__\n", + "\n", + "Add the total lens light and source image together." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = total_lens_image_2d + source_image_2d\n", + "\n", + "aplt.plot_array(array=image, title=\"Total Image (All Galaxies + Source)\")\n", + "\n", + "blurring_image_2d = total_lens_blurring_image_2d + source_blurring_image_2d" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Convolution__\n", + "\n", + "Convolve the 2D image with the PSF in real-space." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", + " image=image, blurring_image=blurring_image_2d\n", + ")\n", + "\n", + "aplt.plot_array(array=convolved_image_2d, title=\"Convolved Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our group-scale lens model.\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "__Chi Squared__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image = convolved_image_2d\n", + "\n", + "residual_map = masked_dataset.data - model_image\n", + "normalized_residual_map = residual_map / masked_dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(chi_squared)\n", + "\n", + "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", + "\n", + "aplt.plot_array(array=chi_squared_map, title=\"Chi-Squared Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(figure_of_merit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE Likelihood__\n", + "\n", + "When using MGE light profiles (linear light profiles), the likelihood calculation differs in an important way.\n", + "Instead of each galaxy having a single light profile with a fixed intensity, each galaxy contributes a set of\n", + "Gaussian basis functions to a `mapping_matrix`.\n", + "\n", + "The mapping matrix has dimensions `(total_image_pixels, total_linear_light_profiles)`, where each column\n", + "is the PSF-convolved image of one Gaussian. For a group with a 20-Gaussian main lens, two 10-Gaussian extras,\n", + "and a 20-Gaussian source, this matrix has 60 columns.\n", + "\n", + "The intensities of all 60 Gaussians are then solved simultaneously via positive-only linear algebra:\n", + "\n", + " s = F^{-1} D\n", + "\n", + "where F is the curvature matrix and D is the data vector (see the `linear_light_profiles` feature for details).\n", + "\n", + "This joint optimization across all galaxies is what makes MGE so powerful for groups: the linear algebra\n", + "automatically determines the optimal intensity decomposition for every galaxy simultaneously, accounting for\n", + "blending between overlapping galaxies.\n", + "\n", + "__Fit__\n", + "\n", + "The `FitImaging` object handles all of the steps above automatically." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)\n", + "fit_figure_of_merit = fit.figure_of_merit\n", + "print(fit_figure_of_merit)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the group-scale parametric likelihood function, with\n", + "explanations of how MGE light profiles modify each step. The key differences when using MGE are:\n", + "\n", + " - Each galaxy's light is a sum of many Gaussians rather than a single analytic profile.\n", + " - The intensities of all Gaussians across all galaxies are solved jointly via linear algebra.\n", + " - This joint linear inversion adds zero non-linear parameters per galaxy, making it ideal for groups\n", + " with many extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/multi_gaussian_expansion/modeling.ipynb b/notebooks/group/features/multi_gaussian_expansion/modeling.ipynb index 1ecdfd5d9..f4a94efdb 100644 --- a/notebooks/group/features/multi_gaussian_expansion/modeling.ipynb +++ b/notebooks/group/features/multi_gaussian_expansion/modeling.ipynb @@ -1,523 +1,560 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Multi Gaussian Expansion (Group)\n", - "===================================================\n", - "\n", - "A Multi Gaussian Expansion (MGE) decomposes the light of each galaxy into ~10-30+ Gaussians, where the `intensity`\n", - "of every Gaussian is solved for via linear algebra using a process called an \"inversion\" (see the\n", - "`linear_light_profiles` feature for a full description of this).\n", - "\n", - "This script performs lens modeling of a group-scale strong lens using MGE light profiles for all galaxies:\n", - "the main lens galaxies, the extra galaxies, and the source galaxy. MGE models are constructed using the\n", - "convenience function `al.model_util.mge_model_from`, which handles the setup of Gaussian basis functions\n", - "with appropriate sigma ranges and linked parameters.\n", - "\n", - "__Contents__\n", - "\n", - "- **MGE Advantages for Group Lenses:** The MGE is especially important for group-scale lenses because adding an extra.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Run Times:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__MGE Advantages for Group Lenses__\n", - "\n", - "The MGE is especially important for group-scale lenses because adding an extra galaxy with a traditional Sersic\n", - "light profile introduces 5 non-linear parameters per galaxy. For a group with many extra galaxies, this makes the\n", - "model prohibitively complex and slow to fit.\n", - "\n", - "In contrast, an MGE uses **linear light profiles** whose intensities are solved via linear algebra. This means that\n", - "adding an extra galaxy with an MGE light profile adds **zero** additional non-linear parameters to the model. The\n", - "only non-linear parameters for each galaxy are the centre and elliptical components, which for extra galaxies are\n", - "typically fixed to their observed values.\n", - "\n", - "Furthermore, MGE models capture irregular and asymmetric galaxy morphologies (e.g. isophotal twists, radially\n", - "varying ellipticity) far more effectively than symmetric Sersic profiles, leading to more accurate lens models.\n", - "\n", - "The combination of fewer non-linear parameters and better morphological accuracy makes MGE the recommended\n", - "approach for group-scale lens modeling.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - Each main lens galaxy's light is an MGE with 20 Gaussians [~4 non-linear parameters per galaxy].\n", - " - Each main lens galaxy's total mass distribution is an `Isothermal`, with `ExternalShear` on `lens_0` [7 parameters].\n", - " - Each extra galaxy's light is an MGE with 10 Gaussians with fixed centres [0 non-linear parameters per galaxy].\n", - " - Each extra galaxy's total mass distribution is an `IsothermalSph` with bounded Einstein radius [1 parameter per galaxy].\n", - " - The source galaxy's light is an MGE with 20 Gaussians [~4 non-linear parameters].\n", - "\n", - "__Simulation__\n", - "\n", - "This script fits a simulated `Imaging` dataset of a strong lens, which is produced in the\n", - "script `autolens_workspace/*/group/simulator.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`, which is the dataset we will use to perform lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "The model-fit requires a 2D mask defining the regions of the image we fit the lens model to the data.\n", - "\n", - "We create a 7.5 arcsecond circular mask and apply it to the `Imaging` object that the lens model fits. This is\n", - "larger than a typical galaxy-scale lens mask because the group-scale lens has emission spread over a wider area\n", - "due to the multiple lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "The centres of the main lens galaxies and extra galaxies are loaded from JSON files in the dataset directory.\n", - "These centres are used to set up the MGE models and mass models for each galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where every galaxy uses an MGE for its light profile, constructed via the\n", - "`al.model_util.mge_model_from` convenience function.\n", - "\n", - "For **main lens galaxies**, we use 20 Gaussians with uniform centre priors, allowing the MGE to capture\n", - "the full morphology of the main lens light. Only `lens_0` carries an `ExternalShear`.\n", - "\n", - "For **extra galaxies**, we use 10 Gaussians with centres fixed to the observed positions. This is crucial:\n", - "because the MGE intensities are linear parameters, adding extra galaxies with fixed centres introduces\n", - "**zero** additional non-linear parameters. A Sersic model would add 5 non-linear parameters per galaxy.\n", - "\n", - "For the **source galaxy**, we use 20 Gaussians with a single basis group and Gaussian centre priors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light (MGE with fixed centre -- zero non-linear parameters)\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "This shows the group scale MGE model, with separate entries for each main lens galaxy (e.g. `lens_0`),\n", - "the source galaxy and the extra galaxies collection. Note how the extra galaxies have many Gaussian light\n", - "profiles but no non-linear light parameters -- their intensities are all solved via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) is performed to ensure the lens\n", - "calculations are accurate across the full field of the group." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The lens model is fitted to the data using the nested sampling algorithm Nautilus.\n", - "\n", - "We use 150 live points, which is sufficient for the MGE model despite the group-scale complexity, because\n", - "the MGE parameterization has a much simpler non-linear parameter space than Sersic-based models." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"multi_gaussian_expansion\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We next create an `AnalysisImaging` object with JAX acceleration enabled." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM Use__\n", - "\n", - "For MGE models, VRAM use scales with the total number of Gaussians across all galaxies. With 20 Gaussians\n", - "for the main lens, 10 per extra galaxy, and 20 for the source, the total is moderate but larger than a\n", - "simple Sersic-based model. Check VRAM usage before running on GPU." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "The MGE model has a slower per-evaluation time than Sersic models because the image of every Gaussian must\n", - "be computed. However, the much simpler non-linear parameter space means Nautilus converges in far fewer\n", - "iterations, so the overall run time is typically shorter.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", - "\n", - "The result contains entries for each main lens galaxy (e.g. `lens_0`), the source galaxy and the extra galaxies,\n", - "all with their MGE light profiles and inferred intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script has demonstrated how to use MGE light profiles for group-scale lens modeling. The key advantages are:\n", - "\n", - " - **No additional non-linear parameters per extra galaxy**: MGE intensities are linear, so adding extra galaxies\n", - " does not increase the dimensionality of the non-linear parameter space.\n", - "\n", - " - **Better morphological accuracy**: MGE captures irregular features like isophotal twists and radially varying\n", - " ellipticity that symmetric Sersic profiles cannot.\n", - "\n", - " - **Simpler parameter space**: The MGE parameterization removes degeneracies related to galaxy size parameters,\n", - " making the non-linear search converge faster.\n", - "\n", - "For group-scale lenses with many extra galaxies, MGE is strongly recommended over Sersic-based models.\n", - "\n", - "__Features__\n", - "\n", - "We recommend you also checkout:\n", - "\n", - "- ``scaling_relation``: Model the mass of extra galaxies using a luminosity-to-mass scaling relation.\n", - "- ``pixelization``: Reconstruct the source using an adaptive mesh for even more accurate source modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Multi Gaussian Expansion (Group)\n", + "===================================================\n", + "\n", + "A Multi Gaussian Expansion (MGE) decomposes the light of each galaxy into ~10-30+ Gaussians, where the `intensity`\n", + "of every Gaussian is solved for via linear algebra using a process called an \"inversion\" (see the\n", + "`linear_light_profiles` feature for a full description of this).\n", + "\n", + "This script performs lens modeling of a group-scale strong lens using MGE light profiles for all galaxies:\n", + "the main lens galaxies, the extra galaxies, and the source galaxy. MGE models are constructed using the\n", + "convenience function `al.model_util.mge_model_from`, which handles the setup of Gaussian basis functions\n", + "with appropriate sigma ranges and linked parameters.\n", + "\n", + "__Contents__\n", + "\n", + "- **MGE Advantages for Group Lenses:** The MGE is especially important for group-scale lenses because adding an extra.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Run Times:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__MGE Advantages for Group Lenses__\n", + "\n", + "The MGE is especially important for group-scale lenses because adding an extra galaxy with a traditional Sersic\n", + "light profile introduces 5 non-linear parameters per galaxy. For a group with many extra galaxies, this makes the\n", + "model prohibitively complex and slow to fit.\n", + "\n", + "In contrast, an MGE uses **linear light profiles** whose intensities are solved via linear algebra. This means that\n", + "adding an extra galaxy with an MGE light profile adds **zero** additional non-linear parameters to the model. The\n", + "only non-linear parameters for each galaxy are the centre and elliptical components, which for extra galaxies are\n", + "typically fixed to their observed values.\n", + "\n", + "Furthermore, MGE models capture irregular and asymmetric galaxy morphologies (e.g. isophotal twists, radially\n", + "varying ellipticity) far more effectively than symmetric Sersic profiles, leading to more accurate lens models.\n", + "\n", + "The combination of fewer non-linear parameters and better morphological accuracy makes MGE the recommended\n", + "approach for group-scale lens modeling.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - Each main lens galaxy's light is an MGE with 20 Gaussians [~4 non-linear parameters per galaxy].\n", + " - Each main lens galaxy's total mass distribution is an `Isothermal`, with `ExternalShear` on `lens_0` [7 parameters].\n", + " - Each extra galaxy's light is an MGE with 10 Gaussians with fixed centres [0 non-linear parameters per galaxy].\n", + " - Each extra galaxy's total mass distribution is an `IsothermalSph` with bounded Einstein radius [1 parameter per galaxy].\n", + " - The source galaxy's light is an MGE with 20 Gaussians [~4 non-linear parameters].\n", + "\n", + "__Simulation__\n", + "\n", + "This script fits a simulated `Imaging` dataset of a strong lens, which is produced in the\n", + "script `autolens_workspace/*/group/simulator.py`." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`, which is the dataset we will use to perform lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "The model-fit requires a 2D mask defining the regions of the image we fit the lens model to the data.\n", + "\n", + "We create a 7.5 arcsecond circular mask and apply it to the `Imaging` object that the lens model fits. This is\n", + "larger than a typical galaxy-scale lens mask because the group-scale lens has emission spread over a wider area\n", + "due to the multiple lens galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "The centres of the main lens galaxies and extra galaxies are loaded from JSON files in the dataset directory.\n", + "These centres are used to set up the MGE models and mass models for each galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where every galaxy uses an MGE for its light profile, constructed via the\n", + "`al.model_util.mge_model_from` convenience function.\n", + "\n", + "For **main lens galaxies**, we use 20 Gaussians with uniform centre priors, allowing the MGE to capture\n", + "the full morphology of the main lens light. Only `lens_0` carries an `ExternalShear`.\n", + "\n", + "For **extra galaxies**, we use 10 Gaussians with centres fixed to the observed positions. This is crucial:\n", + "because the MGE intensities are linear parameters, adding extra galaxies with fixed centres introduces\n", + "**zero** additional non-linear parameters. A Sersic model would add 5 non-linear parameters per galaxy.\n", + "\n", + "For the **source galaxy**, we use 20 Gaussians with a single basis group and Gaussian centre priors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light (MGE with fixed centre -- zero non-linear parameters)\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "This shows the group scale MGE model, with separate entries for each main lens galaxy (e.g. `lens_0`),\n", + "the source galaxy and the extra galaxies collection. Note how the extra galaxies have many Gaussian light\n", + "profiles but no non-linear light parameters -- their intensities are all solved via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) is performed to ensure the lens\n", + "calculations are accurate across the full field of the group." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The lens model is fitted to the data using the nested sampling algorithm Nautilus.\n", + "\n", + "We use 150 live points, which is sufficient for the MGE model despite the group-scale complexity, because\n", + "the MGE parameterization has a much simpler non-linear parameter space than Sersic-based models." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"multi_gaussian_expansion\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We next create an `AnalysisImaging` object with JAX acceleration enabled." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM Use__\n", + "\n", + "For MGE models, VRAM use scales with the total number of Gaussians across all galaxies. With 20 Gaussians\n", + "for the main lens, 10 per extra galaxy, and 20 for the source, the total is moderate but larger than a\n", + "simple Sersic-based model. Check VRAM usage before running on GPU." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "The MGE model has a slower per-evaluation time than Sersic models because the image of every Gaussian must\n", + "be computed. However, the much simpler non-linear parameter space means Nautilus converges in far fewer\n", + "iterations, so the overall run time is typically shorter.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", + "\n", + "The result contains entries for each main lens galaxy (e.g. `lens_0`), the source galaxy and the extra galaxies,\n", + "all with their MGE light profiles and inferred intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script has demonstrated how to use MGE light profiles for group-scale lens modeling. The key advantages are:\n", + "\n", + " - **No additional non-linear parameters per extra galaxy**: MGE intensities are linear, so adding extra galaxies\n", + " does not increase the dimensionality of the non-linear parameter space.\n", + "\n", + " - **Better morphological accuracy**: MGE captures irregular features like isophotal twists and radially varying\n", + " ellipticity that symmetric Sersic profiles cannot.\n", + "\n", + " - **Simpler parameter space**: The MGE parameterization removes degeneracies related to galaxy size parameters,\n", + " making the non-linear search converge faster.\n", + "\n", + "For group-scale lenses with many extra galaxies, MGE is strongly recommended over Sersic-based models.\n", + "\n", + "__Features__\n", + "\n", + "We recommend you also checkout:\n", + "\n", + "- ``scaling_relation``: Model the mass of extra galaxies using a luminosity-to-mass scaling relation.\n", + "- ``pixelization``: Reconstruct the source using an adaptive mesh for even more accurate source modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/multi_gaussian_expansion/simulator.ipynb b/notebooks/group/features/multi_gaussian_expansion/simulator.ipynb index 9df0b3567..60e395641 100644 --- a/notebooks/group/features/multi_gaussian_expansion/simulator.ipynb +++ b/notebooks/group/features/multi_gaussian_expansion/simulator.ipynb @@ -1,315 +1,352 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Multi Gaussian Expansion (Group)\n", - "===========================================\n", - "\n", - "This script simulates an example strong lens on the 'group' scale, where there is a single primary lens galaxy\n", - "and two smaller extra galaxies nearby, whose mass contributes significantly to the ray-tracing and is therefore\n", - "included in the strong lens model.\n", - "\n", - "The MGE feature is a **modeling** choice, not a property of the data itself. The simulated dataset is therefore\n", - "identical to the standard group simulator -- the same galaxy light profiles, mass profiles, and source galaxy\n", - "are used. This simulator exists to ensure that the MGE feature scripts can auto-simulate their required dataset.\n", - "\n", - "If the dataset already exists (e.g. because the standard group simulator has already been run), this script\n", - "will not re-simulate it.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** Output path for the simulated dataset.\n", - "- **Grid:** Define the 2d grid of (y,x) coordinates for the simulation.\n", - "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid.\n", - "- **Main Lens Galaxies:** The main lens galaxy at the origin.\n", - "- **Extra Galaxies:** The two extra companion galaxies.\n", - "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", - "- **Ray Tracing:** Use all galaxies to setup a tracer and simulate.\n", - "- **Output:** Output the simulated dataset and metadata." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The dataset is output to `dataset/group/simple`, the same path used by the standard group simulator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"simple\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset already exists, we skip simulation. This prevents overwriting data that may have been\n", - "generated by the standard group simulator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - "\n", - " \"\"\"\n", - " __Grid__\n", - "\n", - " Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and therefore\n", - " simulated on.\n", - " \"\"\"\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=(250, 250),\n", - " pixel_scales=0.1,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Galaxy Centres__\n", - "\n", - " Define the centres of the main lens galaxies and extra galaxies.\n", - " \"\"\"\n", - " main_lens_centres = [(0.0, 0.0)]\n", - " extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]\n", - "\n", - " \"\"\"\n", - " __Over Sampling__\n", - "\n", - " An adaptive oversampling scheme is applied at the centre of every galaxy in the group.\n", - " \"\"\"\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=main_lens_centres + extra_galaxies_centres,\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " \"\"\"\n", - " Simulate a simple Gaussian PSF for the image.\n", - " \"\"\"\n", - " psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - " )\n", - "\n", - " \"\"\"\n", - " To simulate the `Imaging` dataset we first create a simulator.\n", - " \"\"\"\n", - " simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Main Lens Galaxies__\n", - "\n", - " The main lens galaxy is at the origin (0.0, 0.0). It has a spherical Sersic light profile and an isothermal\n", - " mass profile.\n", - " \"\"\"\n", - " lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - " )\n", - "\n", - " main_lens_galaxies = [lens_0]\n", - "\n", - " \"\"\"\n", - " __Extra Galaxies__\n", - "\n", - " The two extra galaxies are companion galaxies near the lens system.\n", - " \"\"\"\n", - " extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - " )\n", - "\n", - " extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - " )\n", - "\n", - " extra_galaxies = [extra_galaxy_0, extra_galaxy_1]\n", - "\n", - " \"\"\"\n", - " __Source Galaxy__\n", - "\n", - " The source galaxy uses a cored Sersic profile so that adaptive over-sampling is not required for the source.\n", - " \"\"\"\n", - " source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - " )\n", - "\n", - " \"\"\"\n", - " __Ray Tracing__\n", - "\n", - " Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset.\n", - " \"\"\"\n", - " tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])\n", - "\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", - "\n", - " \"\"\"\n", - " __Dataset__\n", - "\n", - " Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset.\n", - " \"\"\"\n", - " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " \"\"\"\n", - " Output the simulated dataset to the dataset path as .fits files.\n", - " \"\"\"\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Visualize__\n", - "\n", - " Output a subplot of the simulated dataset.\n", - " \"\"\"\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - " aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - " \"\"\"\n", - " __Tracer json__\n", - "\n", - " Save the `Tracer` in the dataset folder as a .json file.\n", - " \"\"\"\n", - " al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - " )\n", - "\n", - " \"\"\"\n", - " __Centre JSON Files__\n", - "\n", - " Save the centres of the main lens galaxies and extra galaxies as JSON files.\n", - " \"\"\"\n", - " al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - " )\n", - "\n", - " al.output_to_json(\n", - " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - " )\n", - "\n", - " \"\"\"\n", - " __Positions__\n", - "\n", - " Solve for the lensed positions of the source galaxy.\n", - " \"\"\"\n", - " solver = al.PointSolver.for_grid(\n", - " grid=al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1),\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.01,\n", - " )\n", - "\n", - " positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - " )\n", - "\n", - " al.output_to_json(\n", - " obj=positions,\n", - " file_path=dataset_path / \"positions.json\",\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Multi Gaussian Expansion (Group)\n", + "===========================================\n", + "\n", + "This script simulates an example strong lens on the 'group' scale, where there is a single primary lens galaxy\n", + "and two smaller extra galaxies nearby, whose mass contributes significantly to the ray-tracing and is therefore\n", + "included in the strong lens model.\n", + "\n", + "The MGE feature is a **modeling** choice, not a property of the data itself. The simulated dataset is therefore\n", + "identical to the standard group simulator -- the same galaxy light profiles, mass profiles, and source galaxy\n", + "are used. This simulator exists to ensure that the MGE feature scripts can auto-simulate their required dataset.\n", + "\n", + "If the dataset already exists (e.g. because the standard group simulator has already been run), this script\n", + "will not re-simulate it.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** Output path for the simulated dataset.\n", + "- **Grid:** Define the 2d grid of (y,x) coordinates for the simulation.\n", + "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid.\n", + "- **Main Lens Galaxies:** The main lens galaxy at the origin.\n", + "- **Extra Galaxies:** The two extra companion galaxies.\n", + "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", + "- **Ray Tracing:** Use all galaxies to setup a tracer and simulate.\n", + "- **Output:** Output the simulated dataset and metadata." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The dataset is output to `dataset/group/simple`, the same path used by the standard group simulator." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"simple\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset already exists, we skip simulation. This prevents overwriting data that may have been\n", + "generated by the standard group simulator." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + "\n", + " \"\"\"\n", + " __Grid__\n", + "\n", + " Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and therefore\n", + " simulated on.\n", + " \"\"\"\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=(250, 250),\n", + " pixel_scales=0.1,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Galaxy Centres__\n", + "\n", + " Define the centres of the main lens galaxies and extra galaxies.\n", + " \"\"\"\n", + " main_lens_centres = [(0.0, 0.0)]\n", + " extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]\n", + "\n", + " \"\"\"\n", + " __Over Sampling__\n", + "\n", + " An adaptive oversampling scheme is applied at the centre of every galaxy in the group.\n", + " \"\"\"\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=main_lens_centres + extra_galaxies_centres,\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " \"\"\"\n", + " Simulate a simple Gaussian PSF for the image.\n", + " \"\"\"\n", + " psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + " )\n", + "\n", + " \"\"\"\n", + " To simulate the `Imaging` dataset we first create a simulator.\n", + " \"\"\"\n", + " simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Main Lens Galaxies__\n", + "\n", + " The main lens galaxy is at the origin (0.0, 0.0). It has a spherical Sersic light profile and an isothermal\n", + " mass profile.\n", + " \"\"\"\n", + " lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + " )\n", + "\n", + " main_lens_galaxies = [lens_0]\n", + "\n", + " \"\"\"\n", + " __Extra Galaxies__\n", + "\n", + " The two extra galaxies are companion galaxies near the lens system.\n", + " \"\"\"\n", + " extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + " )\n", + "\n", + " extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + " )\n", + "\n", + " extra_galaxies = [extra_galaxy_0, extra_galaxy_1]\n", + "\n", + " \"\"\"\n", + " __Source Galaxy__\n", + "\n", + " The source galaxy uses a cored Sersic profile so that adaptive over-sampling is not required for the source.\n", + " \"\"\"\n", + " source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + " )\n", + "\n", + " \"\"\"\n", + " __Ray Tracing__\n", + "\n", + " Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset.\n", + " \"\"\"\n", + " tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])\n", + "\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", + "\n", + " \"\"\"\n", + " __Dataset__\n", + "\n", + " Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset.\n", + " \"\"\"\n", + " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " \"\"\"\n", + " Output the simulated dataset to the dataset path as .fits files.\n", + " \"\"\"\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Visualize__\n", + "\n", + " Output a subplot of the simulated dataset.\n", + " \"\"\"\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + " aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + " \"\"\"\n", + " __Tracer json__\n", + "\n", + " Save the `Tracer` in the dataset folder as a .json file.\n", + " \"\"\"\n", + " al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + " )\n", + "\n", + " \"\"\"\n", + " __Centre JSON Files__\n", + "\n", + " Save the centres of the main lens galaxies and extra galaxies as JSON files.\n", + " \"\"\"\n", + " al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + " )\n", + "\n", + " al.output_to_json(\n", + " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + " )\n", + "\n", + " \"\"\"\n", + " __Positions__\n", + "\n", + " Solve for the lensed positions of the source galaxy.\n", + " \"\"\"\n", + " solver = al.PointSolver.for_grid(\n", + " grid=al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1),\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.01,\n", + " )\n", + "\n", + " positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + " )\n", + "\n", + " al.output_to_json(\n", + " obj=positions,\n", + " file_path=dataset_path / \"positions.json\",\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/multi_gaussian_expansion/slam.ipynb b/notebooks/group/features/multi_gaussian_expansion/slam.ipynb index bdd92700c..53437bc3f 100644 --- a/notebooks/group/features/multi_gaussian_expansion/slam.ipynb +++ b/notebooks/group/features/multi_gaussian_expansion/slam.ipynb @@ -1,640 +1,677 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Multi Gaussian Expansion: Group SLaM\n", - "=====================================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for fitting a\n", - "group-scale strong lens where all light profiles use Multi Gaussian Expansion (MGE) models.\n", - "\n", - "The standard group SLaM pipeline (``group/slam.py``) already uses MGE by default for all galaxies. This\n", - "feature script serves as a reference implementation documenting the MGE-specific choices and explaining\n", - "why MGE is the default approach for group-scale modeling.\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", - "guide before working through this example.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **This Script:** Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE and MASS TOTAL PIPELINE.\n", - "- **SOURCE LP PIPELINE:** Fits lens light and source light using MGE, with Isothermal mass and ExternalShear.\n", - "- **LIGHT LP PIPELINE:** Refits lens light with a fresh MGE, mass and source fixed from SOURCE LP.\n", - "- **MASS TOTAL PIPELINE:** Refits the mass using a PowerLaw, light and source fixed from previous stages.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Settings AutoFit:** The settings of autofit.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines.\n", - "\n", - "- **Group Modeling** (`group/modeling`)\n", - " How we model group-scale strong lenses, including extra galaxies.\n", - "\n", - "- **MGE Feature** (`imaging/features/multi_gaussian_expansion`)\n", - " The Multi Gaussian Expansion light profile and its advantages.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE and a MASS TOTAL PIPELINE this SLaM modeling script\n", - "fits `Imaging` data of a group-scale strong lens where in the final model:\n", - "\n", - " - Each main lens galaxy has a free MGE bulge and a `PowerLaw` total mass.\n", - " - Each extra galaxy has a free MGE bulge and an `IsothermalSph` mass with bounded Einstein radius.\n", - " - The source galaxy's light is an MGE.\n", - "\n", - "Because the source is parametric (an MGE), the SOURCE PIX PIPELINE is skipped. This is a simpler\n", - "pipeline than the full group SLaM (which transitions to a pixelized source), and is appropriate\n", - "when an MGE source model is sufficient." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Fits the lens light, source light, lens mass and external shear simultaneously using MGE models.\n", - "\n", - "For group-scale lenses:\n", - " - Each main lens galaxy gets a 20-Gaussian MGE with `gaussian_per_basis=1` and uniform centre priors.\n", - " - Each extra galaxy gets a 10-Gaussian MGE with centres fixed to the observed positions and a bounded\n", - " `IsothermalSph` mass.\n", - " - The source galaxy gets a 20-Gaussian MGE with Gaussian centre priors.\n", - " - Only `lens_0` carries an `ExternalShear`.\n", - "\n", - "The MGE source means the SOURCE PIX PIPELINE is not needed, significantly simplifying the overall pipeline." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " main_lens_centres,\n", - " extra_galaxies_centres,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " # Main Lens Galaxies:\n", - "\n", - " lens_dict = {}\n", - "\n", - " for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - " # Extra Galaxies:\n", - "\n", - " extra_galaxies_list = []\n", - "\n", - " for centre in extra_galaxies_centres:\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_galaxy = af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=mass\n", - " )\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - " extra_galaxies = af.Collection(extra_galaxies_list) if extra_galaxies_list else None\n", - "\n", - " # Source:\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", - "\n", - " # Overall Lens Model:\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Refits the lens light with a fresh MGE model, keeping the mass and source fixed from the SOURCE LP result.\n", - "\n", - "For group-scale lenses:\n", - " - Each main lens galaxy gets a fresh 20-Gaussian MGE with uniform centre priors.\n", - " - Each extra galaxy gets a fresh 10-Gaussian MGE with centres fixed to the observed positions.\n", - " - The source is fixed as an instance from the SOURCE LP result.\n", - " - Mass profiles are fixed as instances from the SOURCE LP result." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " main_lens_centres,\n", - " extra_galaxies_centres,\n", - " source_lp_result: af.Result,\n", - " redshift_lens: float,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " # Main Lens Galaxies (fresh MGE, mass fixed):\n", - "\n", - " lens_dict = {}\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(source_lp_result.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " for i in range(n_main):\n", - " lens_instance = getattr(source_lp_result.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=lens_instance.redshift,\n", - " bulge=bulge,\n", - " mass=lens_instance.mass,\n", - " shear=lens_instance.shear,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - " # Extra Galaxies (fresh MGE, mass fixed):\n", - "\n", - " extra_galaxies_list = []\n", - "\n", - " if source_lp_result.instance.extra_galaxies is not None:\n", - " for i, centre in enumerate(extra_galaxies_centres):\n", - " extra_instance = source_lp_result.instance.extra_galaxies[i]\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " extra_galaxy = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=bulge,\n", - " mass=extra_instance.mass,\n", - " )\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - " extra_galaxies = af.Collection(extra_galaxies_list) if extra_galaxies_list else None\n", - "\n", - " # Source (fixed from SOURCE LP):\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " bulge=source_lp_result.instance.galaxies.source.bulge,\n", - " )\n", - "\n", - " # Overall Lens Model:\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Refits the total mass distribution using a `PowerLaw` for the main lens galaxies, with light and source\n", - "fixed from previous stages.\n", - "\n", - "For group-scale lenses:\n", - " - Each main lens galaxy's mass is upgraded to a `PowerLaw`.\n", - " - Each extra galaxy keeps its `IsothermalSph` mass with bounded Einstein radius.\n", - " - Light profiles are fixed as instances from the LIGHT LP result.\n", - " - The source is fixed from the SOURCE LP result." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " light_result: af.Result,\n", - " extra_galaxies_centres,\n", - " redshift_lens: float,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(light_result.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " # Main Lens Galaxies (fixed light, free PowerLaw mass):\n", - "\n", - " lens_dict = {}\n", - "\n", - " for i in range(n_main):\n", - " light_lens_instance = getattr(light_result.instance.galaxies, f\"lens_{i}\")\n", - " source_lp_lens_model = getattr(source_lp_result.model.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.PowerLaw),\n", - " mass_result=source_lp_lens_model.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=light_lens_instance.redshift,\n", - " bulge=light_lens_instance.bulge,\n", - " mass=mass,\n", - " shear=source_lp_result.model.galaxies.lens_0.shear if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - " # Extra Galaxies (fixed light, free bounded mass):\n", - "\n", - " extra_galaxies_list = []\n", - "\n", - " if light_result.instance.extra_galaxies is not None:\n", - " for i, centre in enumerate(extra_galaxies_centres):\n", - " light_extra = light_result.instance.extra_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_galaxy = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=light_extra.bulge,\n", - " mass=mass,\n", - " )\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - " extra_galaxies = af.Collection(extra_galaxies_list) if extra_galaxies_list else None\n", - "\n", - " # Source (fixed from SOURCE LP):\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " bulge=source_lp_result.model.galaxies.source.bulge,\n", - " )\n", - "\n", - " # Overall Lens Model:\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"group\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " main_lens_centres=main_lens_centres,\n", - " extra_galaxies_centres=extra_galaxies_centres,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " main_lens_centres=main_lens_centres,\n", - " extra_galaxies_centres=extra_galaxies_centres,\n", - " source_lp_result=source_lp_result,\n", - " redshift_lens=redshift_lens,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " light_result=light_result,\n", - " extra_galaxies_centres=extra_galaxies_centres,\n", - " redshift_lens=redshift_lens,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Multi Gaussian Expansion: Group SLaM\n", + "=====================================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for fitting a\n", + "group-scale strong lens where all light profiles use Multi Gaussian Expansion (MGE) models.\n", + "\n", + "The standard group SLaM pipeline (``group/slam.py``) already uses MGE by default for all galaxies. This\n", + "feature script serves as a reference implementation documenting the MGE-specific choices and explaining\n", + "why MGE is the default approach for group-scale modeling.\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", + "guide before working through this example.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **This Script:** Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE and MASS TOTAL PIPELINE.\n", + "- **SOURCE LP PIPELINE:** Fits lens light and source light using MGE, with Isothermal mass and ExternalShear.\n", + "- **LIGHT LP PIPELINE:** Refits lens light with a fresh MGE, mass and source fixed from SOURCE LP.\n", + "- **MASS TOTAL PIPELINE:** Refits the mass using a PowerLaw, light and source fixed from previous stages.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Settings AutoFit:** The settings of autofit.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines.\n", + "\n", + "- **Group Modeling** (`group/modeling`)\n", + " How we model group-scale strong lenses, including extra galaxies.\n", + "\n", + "- **MGE Feature** (`imaging/features/multi_gaussian_expansion`)\n", + " The Multi Gaussian Expansion light profile and its advantages.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE and a MASS TOTAL PIPELINE this SLaM modeling script\n", + "fits `Imaging` data of a group-scale strong lens where in the final model:\n", + "\n", + " - Each main lens galaxy has a free MGE bulge and a `PowerLaw` total mass.\n", + " - Each extra galaxy has a free MGE bulge and an `IsothermalSph` mass with bounded Einstein radius.\n", + " - The source galaxy's light is an MGE.\n", + "\n", + "Because the source is parametric (an MGE), the SOURCE PIX PIPELINE is skipped. This is a simpler\n", + "pipeline than the full group SLaM (which transitions to a pixelized source), and is appropriate\n", + "when an MGE source model is sufficient." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Fits the lens light, source light, lens mass and external shear simultaneously using MGE models.\n", + "\n", + "For group-scale lenses:\n", + " - Each main lens galaxy gets a 20-Gaussian MGE with `gaussian_per_basis=1` and uniform centre priors.\n", + " - Each extra galaxy gets a 10-Gaussian MGE with centres fixed to the observed positions and a bounded\n", + " `IsothermalSph` mass.\n", + " - The source galaxy gets a 20-Gaussian MGE with Gaussian centre priors.\n", + " - Only `lens_0` carries an `ExternalShear`.\n", + "\n", + "The MGE source means the SOURCE PIX PIPELINE is not needed, significantly simplifying the overall pipeline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " main_lens_centres,\n", + " extra_galaxies_centres,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " # Main Lens Galaxies:\n", + "\n", + " lens_dict = {}\n", + "\n", + " for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + " # Extra Galaxies:\n", + "\n", + " extra_galaxies_list = []\n", + "\n", + " for centre in extra_galaxies_centres:\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_galaxy = af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=mass\n", + " )\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + " extra_galaxies = af.Collection(extra_galaxies_list) if extra_galaxies_list else None\n", + "\n", + " # Source:\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", + "\n", + " # Overall Lens Model:\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Refits the lens light with a fresh MGE model, keeping the mass and source fixed from the SOURCE LP result.\n", + "\n", + "For group-scale lenses:\n", + " - Each main lens galaxy gets a fresh 20-Gaussian MGE with uniform centre priors.\n", + " - Each extra galaxy gets a fresh 10-Gaussian MGE with centres fixed to the observed positions.\n", + " - The source is fixed as an instance from the SOURCE LP result.\n", + " - Mass profiles are fixed as instances from the SOURCE LP result." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " main_lens_centres,\n", + " extra_galaxies_centres,\n", + " source_lp_result: af.Result,\n", + " redshift_lens: float,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " # Main Lens Galaxies (fresh MGE, mass fixed):\n", + "\n", + " lens_dict = {}\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(source_lp_result.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " for i in range(n_main):\n", + " lens_instance = getattr(source_lp_result.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=lens_instance.redshift,\n", + " bulge=bulge,\n", + " mass=lens_instance.mass,\n", + " shear=lens_instance.shear,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + " # Extra Galaxies (fresh MGE, mass fixed):\n", + "\n", + " extra_galaxies_list = []\n", + "\n", + " if source_lp_result.instance.extra_galaxies is not None:\n", + " for i, centre in enumerate(extra_galaxies_centres):\n", + " extra_instance = source_lp_result.instance.extra_galaxies[i]\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " extra_galaxy = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=bulge,\n", + " mass=extra_instance.mass,\n", + " )\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + " extra_galaxies = af.Collection(extra_galaxies_list) if extra_galaxies_list else None\n", + "\n", + " # Source (fixed from SOURCE LP):\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " bulge=source_lp_result.instance.galaxies.source.bulge,\n", + " )\n", + "\n", + " # Overall Lens Model:\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Refits the total mass distribution using a `PowerLaw` for the main lens galaxies, with light and source\n", + "fixed from previous stages.\n", + "\n", + "For group-scale lenses:\n", + " - Each main lens galaxy's mass is upgraded to a `PowerLaw`.\n", + " - Each extra galaxy keeps its `IsothermalSph` mass with bounded Einstein radius.\n", + " - Light profiles are fixed as instances from the LIGHT LP result.\n", + " - The source is fixed from the SOURCE LP result." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " light_result: af.Result,\n", + " extra_galaxies_centres,\n", + " redshift_lens: float,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(light_result.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " # Main Lens Galaxies (fixed light, free PowerLaw mass):\n", + "\n", + " lens_dict = {}\n", + "\n", + " for i in range(n_main):\n", + " light_lens_instance = getattr(light_result.instance.galaxies, f\"lens_{i}\")\n", + " source_lp_lens_model = getattr(source_lp_result.model.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.PowerLaw),\n", + " mass_result=source_lp_lens_model.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=light_lens_instance.redshift,\n", + " bulge=light_lens_instance.bulge,\n", + " mass=mass,\n", + " shear=source_lp_result.model.galaxies.lens_0.shear if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + " # Extra Galaxies (fixed light, free bounded mass):\n", + "\n", + " extra_galaxies_list = []\n", + "\n", + " if light_result.instance.extra_galaxies is not None:\n", + " for i, centre in enumerate(extra_galaxies_centres):\n", + " light_extra = light_result.instance.extra_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_galaxy = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=light_extra.bulge,\n", + " mass=mass,\n", + " )\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + " extra_galaxies = af.Collection(extra_galaxies_list) if extra_galaxies_list else None\n", + "\n", + " # Source (fixed from SOURCE LP):\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " bulge=source_lp_result.model.galaxies.source.bulge,\n", + " )\n", + "\n", + " # Overall Lens Model:\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"group\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " main_lens_centres=main_lens_centres,\n", + " extra_galaxies_centres=extra_galaxies_centres,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " main_lens_centres=main_lens_centres,\n", + " extra_galaxies_centres=extra_galaxies_centres,\n", + " source_lp_result=source_lp_result,\n", + " redshift_lens=redshift_lens,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " light_result=light_result,\n", + " extra_galaxies_centres=extra_galaxies_centres,\n", + " redshift_lens=redshift_lens,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/multi_gaussian_expansion/source_science.ipynb b/notebooks/group/features/multi_gaussian_expansion/source_science.ipynb index 9702489a6..0bdf36463 100644 --- a/notebooks/group/features/multi_gaussian_expansion/source_science.ipynb +++ b/notebooks/group/features/multi_gaussian_expansion/source_science.ipynb @@ -1,493 +1,530 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Source Science: Multi Gaussian Expansion (Group)\n", - "================================================\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy.\n", - "\n", - "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", - "size of the source.\n", - "\n", - "This example shows how to perform these calculations for a group-scale lens using an MGE source model. The key\n", - "differences from the standard group source science example are:\n", - "\n", - " - The source galaxy is modeled as a Multi Gaussian Expansion (MGE), where the intensities of ~20 Gaussians are\n", - " solved via linear algebra when fitting the data.\n", - " - MGE sources can provide more accurate flux and magnification estimates than single parametric profiles (e.g.\n", - " Sersic) because they capture more complex source morphologies.\n", - " - ALL mass profiles in the group must be included when computing ray-tracing and magnification.\n", - "\n", - "__Contents__\n", - "\n", - "- **Simulated Dataset:** Load the group dataset.\n", - "- **Mask:** Define the 2D mask.\n", - "- **Source Values:** Set up the lens and source galaxies.\n", - "- **MGE Source:** Create an MGE source and fit it to the data to obtain intensities.\n", - "- **Source Flux:** Compute the total source flux.\n", - "- **Source Magnification:** Compute the source magnification using all group galaxies.\n", - "- **Impact of Extra Galaxies:** Show that omitting extra galaxies gives incorrect magnification.\n", - "- **Tracer:** Using the tracer from lens modeling results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulated Dataset__\n", - "\n", - "We load and plot the `simple` group example dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We apply a 7.5 arcsecond circular mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over-sampling at all galaxy centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Values__\n", - "\n", - "For a group-scale lens, we must include ALL galaxies that contribute to the lensing potential when\n", - "performing source science calculations. This includes the main lens galaxy and all extra galaxies.\n", - "\n", - "We first set up the lens and extra galaxies with their mass profiles (no light, since we only need\n", - "the mass for ray-tracing in this context). We also create an MGE source whose intensities will be\n", - "determined by fitting the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE Source__\n", - "\n", - "We set up the source galaxy using an MGE made of 20 Gaussians whose `sigma` values span 0.01\" to the\n", - "mask radius. The intensities are linear light profiles whose values are determined by fitting the data.\n", - "\n", - "The MGE source model offers advantages for source science calculations:\n", - "\n", - " - It captures more complex source morphologies than a single Sersic profile.\n", - " - The intensities are optimized to best reconstruct the lensed source, providing a more accurate\n", - " estimate of the source's total flux.\n", - " - It naturally handles sources with multiple components or asymmetric features." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 20\n", - "\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "bulge_gaussian_list = []\n", - "\n", - "for i in range(total_gaussians):\n", - " gaussian = al.lp_linear.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - "\n", - " bulge_gaussian_list.append(gaussian)\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, bulge=bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We create the tracer using ALL group galaxies and the MGE source, then fit it to the data to\n", - "solve for the Gaussian intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From the fit, extract the tracer with solved-for Gaussian intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Flux__\n", - "\n", - "We compute the total source flux by summing the source galaxy's image over a high-resolution grid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = tracer.galaxies[-1]\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_flux = np.sum(image)\n", - "\n", - "print(f\"Total Source Flux (MGE): {total_flux} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Magnification__\n", - "\n", - "The overall magnification is the ratio of total flux in the image-plane to the source-plane.\n", - "\n", - "For group-scale lenses, the ray-tracing is performed through ALL mass profiles in the group, which\n", - "the tracer handles automatically." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)\n", - "\n", - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_source_plane_flux = np.sum(image)\n", - "\n", - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "source_plane_grid = traced_grid_list[-1]\n", - "\n", - "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(lensed_source_image)\n", - "\n", - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Magnification (all group galaxies, MGE): {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Impact of Extra Galaxies__\n", - "\n", - "For group-scale lenses, the magnification is determined by ALL mass profiles in the group. Omitting\n", - "the extra galaxy masses gives an incorrect magnification estimate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_main_only = al.Tracer(galaxies=[tracer.galaxies[0], source_galaxy])\n", - "\n", - "traced_grid_list_main_only = tracer_main_only.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "source_plane_grid_main_only = traced_grid_list_main_only[-1]\n", - "\n", - "lensed_source_image_main_only = source_galaxy.bulge.image_2d_from(\n", - " grid=source_plane_grid_main_only\n", - ")\n", - "\n", - "total_image_plane_flux_main_only = np.sum(lensed_source_image_main_only)\n", - "\n", - "source_magnification_main_only = (\n", - " total_image_plane_flux_main_only / total_source_plane_flux\n", - ")\n", - "\n", - "print(f\"Source Magnification (main lens only, MGE): {source_magnification_main_only}\")\n", - "print(\n", - " f\"Magnification difference when omitting extra galaxies: \"\n", - " f\"{source_magnification - source_magnification_main_only:.4f}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "Lens modeling returns a `max_log_likelihood_tracer`, which is the object you would use to compute\n", - "source science calculations for real datasets. The code below reproduces the calculations above\n", - "using the tracer's built-in plane functionality." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "image_plane_grid = traced_grid_list[0]\n", - "source_plane_grid = traced_grid_list[-1]\n", - "\n", - "lensed_source_image = tracer.planes[-1].image_2d_from(grid=source_plane_grid)\n", - "source_plane_image = tracer.planes[-1].image_2d_from(grid=image_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(lensed_source_image)\n", - "total_source_plane_flux = np.sum(source_plane_image)\n", - "\n", - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Plane Total Flux via Tracer (MGE): {total_source_plane_flux} e- s^-1\")\n", - "print(f\"Source Magnification via Tracer (MGE): {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parametric Source Models__\n", - "\n", - "If your lens modeling uses a parametric source model (e.g. Sersic or MGE), the only object you need to\n", - "perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", - "\n", - "For group-scale lenses, ensure that the tracer includes all group member galaxies with their mass profiles,\n", - "as each contributes to the ray-tracing and therefore to the magnification of the source.\n", - "\n", - "An MGE source may give different flux and magnification estimates compared to a single Sersic source. If\n", - "precise source science calculations are important for your analysis, we recommend comparing results from\n", - "different source models (e.g. MGE, Sersic, pixelized reconstruction) to estimate systematic uncertainties." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Source Science: Multi Gaussian Expansion (Group)\n", + "================================================\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy.\n", + "\n", + "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", + "size of the source.\n", + "\n", + "This example shows how to perform these calculations for a group-scale lens using an MGE source model. The key\n", + "differences from the standard group source science example are:\n", + "\n", + " - The source galaxy is modeled as a Multi Gaussian Expansion (MGE), where the intensities of ~20 Gaussians are\n", + " solved via linear algebra when fitting the data.\n", + " - MGE sources can provide more accurate flux and magnification estimates than single parametric profiles (e.g.\n", + " Sersic) because they capture more complex source morphologies.\n", + " - ALL mass profiles in the group must be included when computing ray-tracing and magnification.\n", + "\n", + "__Contents__\n", + "\n", + "- **Simulated Dataset:** Load the group dataset.\n", + "- **Mask:** Define the 2D mask.\n", + "- **Source Values:** Set up the lens and source galaxies.\n", + "- **MGE Source:** Create an MGE source and fit it to the data to obtain intensities.\n", + "- **Source Flux:** Compute the total source flux.\n", + "- **Source Magnification:** Compute the source magnification using all group galaxies.\n", + "- **Impact of Extra Galaxies:** Show that omitting extra galaxies gives incorrect magnification.\n", + "- **Tracer:** Using the tracer from lens modeling results." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulated Dataset__\n", + "\n", + "We load and plot the `simple` group example dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We apply a 7.5 arcsecond circular mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over-sampling at all galaxy centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Values__\n", + "\n", + "For a group-scale lens, we must include ALL galaxies that contribute to the lensing potential when\n", + "performing source science calculations. This includes the main lens galaxy and all extra galaxies.\n", + "\n", + "We first set up the lens and extra galaxies with their mass profiles (no light, since we only need\n", + "the mass for ray-tracing in this context). We also create an MGE source whose intensities will be\n", + "determined by fitting the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE Source__\n", + "\n", + "We set up the source galaxy using an MGE made of 20 Gaussians whose `sigma` values span 0.01\" to the\n", + "mask radius. The intensities are linear light profiles whose values are determined by fitting the data.\n", + "\n", + "The MGE source model offers advantages for source science calculations:\n", + "\n", + " - It captures more complex source morphologies than a single Sersic profile.\n", + " - The intensities are optimized to best reconstruct the lensed source, providing a more accurate\n", + " estimate of the source's total flux.\n", + " - It naturally handles sources with multiple components or asymmetric features." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 20\n", + "\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "bulge_gaussian_list = []\n", + "\n", + "for i in range(total_gaussians):\n", + " gaussian = al.lp_linear.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + "\n", + " bulge_gaussian_list.append(gaussian)\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, bulge=bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We create the tracer using ALL group galaxies and the MGE source, then fit it to the data to\n", + "solve for the Gaussian intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the fit, extract the tracer with solved-for Gaussian intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Flux__\n", + "\n", + "We compute the total source flux by summing the source galaxy's image over a high-resolution grid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = tracer.galaxies[-1]\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_flux = np.sum(image)\n", + "\n", + "print(f\"Total Source Flux (MGE): {total_flux} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Magnification__\n", + "\n", + "The overall magnification is the ratio of total flux in the image-plane to the source-plane.\n", + "\n", + "For group-scale lenses, the ray-tracing is performed through ALL mass profiles in the group, which\n", + "the tracer handles automatically." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)\n", + "\n", + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_source_plane_flux = np.sum(image)\n", + "\n", + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "source_plane_grid = traced_grid_list[-1]\n", + "\n", + "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(lensed_source_image)\n", + "\n", + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Magnification (all group galaxies, MGE): {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Impact of Extra Galaxies__\n", + "\n", + "For group-scale lenses, the magnification is determined by ALL mass profiles in the group. Omitting\n", + "the extra galaxy masses gives an incorrect magnification estimate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_main_only = al.Tracer(galaxies=[tracer.galaxies[0], source_galaxy])\n", + "\n", + "traced_grid_list_main_only = tracer_main_only.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "source_plane_grid_main_only = traced_grid_list_main_only[-1]\n", + "\n", + "lensed_source_image_main_only = source_galaxy.bulge.image_2d_from(\n", + " grid=source_plane_grid_main_only\n", + ")\n", + "\n", + "total_image_plane_flux_main_only = np.sum(lensed_source_image_main_only)\n", + "\n", + "source_magnification_main_only = (\n", + " total_image_plane_flux_main_only / total_source_plane_flux\n", + ")\n", + "\n", + "print(f\"Source Magnification (main lens only, MGE): {source_magnification_main_only}\")\n", + "print(\n", + " f\"Magnification difference when omitting extra galaxies: \"\n", + " f\"{source_magnification - source_magnification_main_only:.4f}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "Lens modeling returns a `max_log_likelihood_tracer`, which is the object you would use to compute\n", + "source science calculations for real datasets. The code below reproduces the calculations above\n", + "using the tracer's built-in plane functionality." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "image_plane_grid = traced_grid_list[0]\n", + "source_plane_grid = traced_grid_list[-1]\n", + "\n", + "lensed_source_image = tracer.planes[-1].image_2d_from(grid=source_plane_grid)\n", + "source_plane_image = tracer.planes[-1].image_2d_from(grid=image_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(lensed_source_image)\n", + "total_source_plane_flux = np.sum(source_plane_image)\n", + "\n", + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Plane Total Flux via Tracer (MGE): {total_source_plane_flux} e- s^-1\")\n", + "print(f\"Source Magnification via Tracer (MGE): {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parametric Source Models__\n", + "\n", + "If your lens modeling uses a parametric source model (e.g. Sersic or MGE), the only object you need to\n", + "perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", + "\n", + "For group-scale lenses, ensure that the tracer includes all group member galaxies with their mass profiles,\n", + "as each contributes to the ray-tracing and therefore to the magnification of the source.\n", + "\n", + "An MGE source may give different flux and magnification estimates compared to a single Sersic source. If\n", + "precise source science calculations are important for your analysis, we recommend comparing results from\n", + "different source models (e.g. MGE, Sersic, pixelized reconstruction) to estimate systematic uncertainties." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/no_lens_light/modeling.ipynb b/notebooks/group/features/no_lens_light/modeling.ipynb index 19f7d6e7e..de477493d 100644 --- a/notebooks/group/features/no_lens_light/modeling.ipynb +++ b/notebooks/group/features/no_lens_light/modeling.ipynb @@ -1,474 +1,511 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: No Lens Light (Group)\n", - "========================================\n", - "\n", - "This script models a group-scale strong lens where none of the lens galaxies have visible light emission. In the\n", - "group context, \"no lens light\" means that **all** main lens galaxies **and** all extra galaxies are modeled with\n", - "mass profiles only \u2014 no light profiles at all. Only the source galaxy has light.\n", - "\n", - "This is the group-scale analogue of `imaging/features/no_lens_light/modeling.py`. The key difference is that\n", - "removing light from a group has a much larger impact on the model dimensionality: every main lens galaxy and\n", - "every extra galaxy that would normally require an MGE light model (adding non-linear parameters for centre,\n", - "ellipticity, etc.) now has only mass parameters. For a group with one main lens and two extra galaxies, this\n", - "removes all light-related non-linear parameters from the lens side of the model.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages:** The main advantage of fitting group data without lens light is the dramatic reduction in model.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres:** The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files.\n", - "- **Model:** Compose the lens model fitted to the data \u2014 all galaxies have mass only.\n", - "- **Over Sampling:** Not needed for lens light when there is none.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Advantages__\n", - "\n", - "The main advantage of fitting group data without lens light is the dramatic reduction in model dimensionality.\n", - "In a standard group fit, every main lens galaxy and every extra galaxy requires an MGE light model, each\n", - "contributing non-linear parameters (centre, ellipticity, etc.). By omitting all lens light:\n", - "\n", - " - The number of non-linear parameters drops substantially.\n", - " - The non-linear search converges much faster.\n", - " - The parameter space is simpler and less prone to local maxima.\n", - "\n", - "This is especially powerful for groups with many galaxies, where the light model would otherwise dominate\n", - "the parameter budget.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - There is a main lens galaxy whose total mass distribution is an `Isothermal` and `ExternalShear` \u2014 no light.\n", - " - There are two extra lens galaxies whose total mass distributions are `IsothermalSph` models \u2014 no light.\n", - " - The source galaxy's light is a Multi Gaussian Expansion.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `group/modeling.py` script for the full group modeling API." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple__no_lens_light`, which is the dataset we will use to perform\n", - "lens modeling. This dataset contains only lensed source emission \u2014 no lens galaxy light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "The model-fit requires a 2D mask defining the regions of the image we fit the lens model to the data.\n", - "\n", - "We create a 7.5 arcsecond circular mask and apply it to the `Imaging` object that the lens model fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files in the dataset\n", - "directory. This makes the script reusable across different datasets without hardcoding centre values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where all galaxies have mass only \u2014 no light profiles:\n", - "\n", - " - The main lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - There are two extra lens galaxies with `IsothermalSph` total mass distributions, with centres fixed to\n", - " the observed centres and bounded Einstein radii [2 parameters].\n", - "\n", - " - The source galaxy's light is a Multi Gaussian Expansion [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=15.\n", - "\n", - "Compare this to the standard group model (with MGE light for all galaxies) which would have N=28 or more.\n", - "The removal of all lens light profiles is what makes this so efficient.\n", - "\n", - "__Model Composition (List-Based API)__\n", - "\n", - "The API below uses the same list-based approach as `group/modeling.py`, but every galaxy is created without\n", - "a `bulge` parameter \u2014 mass only." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies (mass only):\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies (mass only):\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source (MGE light):\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "This confirms that no lens galaxy (main or extra) has a light profile \u2014 only mass profiles are present." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "When there is no lens light, we do not need adaptive over-sampling for the lens galaxies. Over-sampling is\n", - "normally applied at the centres of lens galaxies to ensure their light profiles are evaluated accurately on a\n", - "higher-resolution grid. Since no galaxy has a light profile in this model, this step is unnecessary.\n", - "\n", - "The source galaxy uses a cored light profile (`SersicCore`) which changes gradually in its central regions,\n", - "so it does not require over-sampling either." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus.\n", - "\n", - "Given the reduced model complexity (no lens light parameters), we use `n_live=100` which is sufficient for\n", - "this relatively simple model.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"no_lens_light\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format.\n", - "\n", - "This confirms there is no lens galaxy light in the model-fit \u2014 only mass profiles for the main lens and\n", - "extra galaxies, plus the source light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "__Wrap Up__\n", - "\n", - "This script shows how to fit a group-scale lens model to data where no galaxy has visible light.\n", - "\n", - "The key advantage in the group context is the dramatic reduction in model dimensionality. For a group with\n", - "one main lens and two extra galaxies, removing the MGE light model from all three galaxies eliminates all\n", - "light-related non-linear parameters, leaving only the mass and source parameters. This makes the model-fit\n", - "much faster and more robust." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: No Lens Light (Group)\n", + "========================================\n", + "\n", + "This script models a group-scale strong lens where none of the lens galaxies have visible light emission. In the\n", + "group context, \"no lens light\" means that **all** main lens galaxies **and** all extra galaxies are modeled with\n", + "mass profiles only \u2014 no light profiles at all. Only the source galaxy has light.\n", + "\n", + "This is the group-scale analogue of `imaging/features/no_lens_light/modeling.py`. The key difference is that\n", + "removing light from a group has a much larger impact on the model dimensionality: every main lens galaxy and\n", + "every extra galaxy that would normally require an MGE light model (adding non-linear parameters for centre,\n", + "ellipticity, etc.) now has only mass parameters. For a group with one main lens and two extra galaxies, this\n", + "removes all light-related non-linear parameters from the lens side of the model.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages:** The main advantage of fitting group data without lens light is the dramatic reduction in model.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres:** The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files.\n", + "- **Model:** Compose the lens model fitted to the data \u2014 all galaxies have mass only.\n", + "- **Over Sampling:** Not needed for lens light when there is none.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Advantages__\n", + "\n", + "The main advantage of fitting group data without lens light is the dramatic reduction in model dimensionality.\n", + "In a standard group fit, every main lens galaxy and every extra galaxy requires an MGE light model, each\n", + "contributing non-linear parameters (centre, ellipticity, etc.). By omitting all lens light:\n", + "\n", + " - The number of non-linear parameters drops substantially.\n", + " - The non-linear search converges much faster.\n", + " - The parameter space is simpler and less prone to local maxima.\n", + "\n", + "This is especially powerful for groups with many galaxies, where the light model would otherwise dominate\n", + "the parameter budget.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - There is a main lens galaxy whose total mass distribution is an `Isothermal` and `ExternalShear` \u2014 no light.\n", + " - There are two extra lens galaxies whose total mass distributions are `IsothermalSph` models \u2014 no light.\n", + " - The source galaxy's light is a Multi Gaussian Expansion.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `group/modeling.py` script for the full group modeling API." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple__no_lens_light`, which is the dataset we will use to perform\n", + "lens modeling. This dataset contains only lensed source emission \u2014 no lens galaxy light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "The model-fit requires a 2D mask defining the regions of the image we fit the lens model to the data.\n", + "\n", + "We create a 7.5 arcsecond circular mask and apply it to the `Imaging` object that the lens model fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files in the dataset\n", + "directory. This makes the script reusable across different datasets without hardcoding centre values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where all galaxies have mass only \u2014 no light profiles:\n", + "\n", + " - The main lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - There are two extra lens galaxies with `IsothermalSph` total mass distributions, with centres fixed to\n", + " the observed centres and bounded Einstein radii [2 parameters].\n", + "\n", + " - The source galaxy's light is a Multi Gaussian Expansion [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=15.\n", + "\n", + "Compare this to the standard group model (with MGE light for all galaxies) which would have N=28 or more.\n", + "The removal of all lens light profiles is what makes this so efficient.\n", + "\n", + "__Model Composition (List-Based API)__\n", + "\n", + "The API below uses the same list-based approach as `group/modeling.py`, but every galaxy is created without\n", + "a `bulge` parameter \u2014 mass only." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies (mass only):\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies (mass only):\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source (MGE light):\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "This confirms that no lens galaxy (main or extra) has a light profile \u2014 only mass profiles are present." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "When there is no lens light, we do not need adaptive over-sampling for the lens galaxies. Over-sampling is\n", + "normally applied at the centres of lens galaxies to ensure their light profiles are evaluated accurately on a\n", + "higher-resolution grid. Since no galaxy has a light profile in this model, this step is unnecessary.\n", + "\n", + "The source galaxy uses a cored light profile (`SersicCore`) which changes gradually in its central regions,\n", + "so it does not require over-sampling either." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus.\n", + "\n", + "Given the reduced model complexity (no lens light parameters), we use `n_live=100` which is sufficient for\n", + "this relatively simple model.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"no_lens_light\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format.\n", + "\n", + "This confirms there is no lens galaxy light in the model-fit \u2014 only mass profiles for the main lens and\n", + "extra galaxies, plus the source light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "__Wrap Up__\n", + "\n", + "This script shows how to fit a group-scale lens model to data where no galaxy has visible light.\n", + "\n", + "The key advantage in the group context is the dramatic reduction in model dimensionality. For a group with\n", + "one main lens and two extra galaxies, removing the MGE light model from all three galaxies eliminates all\n", + "light-related non-linear parameters, leaving only the mass and source parameters. This makes the model-fit\n", + "much faster and more robust." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/no_lens_light/simulator.ipynb b/notebooks/group/features/no_lens_light/simulator.ipynb index 576091e9a..7b136660e 100644 --- a/notebooks/group/features/no_lens_light/simulator.ipynb +++ b/notebooks/group/features/no_lens_light/simulator.ipynb @@ -1,519 +1,556 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: No Lens Light (Group)\n", - "================================\n", - "\n", - "This script simulates `Imaging` of a 'group-scale' strong lens where none of the lens galaxies have visible\n", - "light emission \u2014 only their mass profiles contribute to the ray-tracing. The source galaxy still has light.\n", - "\n", - "This is the group-scale analogue of `imaging/features/no_lens_light/simulator.py`. In a group context, \"no lens\n", - "light\" means that **all** main lens galaxies **and** all extra galaxies are modeled with mass profiles only.\n", - "\n", - "This script simulates `Imaging` of a 'group-scale' strong lens where:\n", - "\n", - " - The group consists of one main lens galaxy and two extra galaxies, all with `IsothermalSph` mass profiles\n", - " and no light profiles.\n", - " - A single source galaxy is observed whose `LightProfile` is a `SersicCore`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Grid:** Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and.\n", - "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", - "- **Over Sampling:** Set up over-sampling at the source centre for accurate light profile evaluation.\n", - "- **Main Lens Galaxies:** The main lens galaxy at the origin, mass only, no light.\n", - "- **Extra Galaxies:** The two extra galaxies, mass only, no light.\n", - "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", - "- **Ray Tracing:** Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging`.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Visualize:** Output a subplot of the simulated dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file.\n", - "- **Centre JSON Files:** Save the centres of the main lens galaxies and extra galaxies as JSON files.\n", - "- **Positions:** Solve for the lensed positions of the source galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. They\n", - "define the folder the dataset is output to on your hard-disk:\n", - "\n", - " - The image will be output to `/autolens_workspace/dataset/group/simple__no_lens_light/data.fits`.\n", - " - The noise-map will be output to `/autolens_workspace/dataset/group/simple__no_lens_light/noise_map.fits`.\n", - " - The psf will be output to `/autolens_workspace/dataset/group/simple__no_lens_light/psf.fits`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"simple__no_lens_light\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and therefore simulated\n", - "on, via the inputs:\n", - "\n", - " - `shape_native`: The (y_pixels, x_pixels) 2D shape of the grid defining the shape of the data that is simulated.\n", - " - `pixel_scales`: The arc-second to pixel conversion factor of the grid and data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(250, 250),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Define the centres of the main lens galaxies and extra galaxies. These are used for over-sampling and are also\n", - "output to JSON files so that the modeling scripts can load them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [(0.0, 0.0)]\n", - "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Because no galaxy has a light profile, there is no need for adaptive over-sampling at the galaxy centres. However,\n", - "the lensed source light still requires accurate evaluation, so we apply over-sampling at the source centre\n", - "(approximately the primary lens centre where the lensed arcs appear).\n", - "\n", - "An adaptive oversampling grid cannot be defined for the lensed source because its light appears in different regions\n", - "of the image plane for each dataset. For this reason, we use a cored light profile for the source galaxy (`SersicCore`)\n", - "which changes gradually in its central regions, allowing accurate evaluation without requiring heavy oversampling.\n", - "\n", - "We still apply a mild over-sampling at the centre of the image where the lensed arcs are brightest." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the `Imaging` dataset we first create a simulator, which defines the exposure time, background sky,\n", - "noise levels and psf of the dataset that is simulated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "The main lens galaxy is at the origin (0.0, 0.0). It has an isothermal mass profile but **no light profile**.\n", - "\n", - "In the list-based API used by the group modeling scripts, main lens galaxies are stored in a list called\n", - "`main_lens_galaxies`, where each galaxy is referred to as `lens_0`, `lens_1`, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "main_lens_galaxies = [lens_0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "The two extra galaxies are companion galaxies near the lens system. They have isothermal mass profiles\n", - "but **no light profiles**, with centres offset from the origin.\n", - "\n", - "In the list-based API, extra galaxies are stored in a list called `extra_galaxies`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy__\n", - "\n", - "The source galaxy whose lensed images we simulate. It uses a cored Sersic profile so that adaptive over-sampling\n", - "is not required for the source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset.\n", - "\n", - "The tracer combines main lens galaxies, extra galaxies and the source galaxy. Because the lens galaxies have\n", - "no light profiles, the simulated image contains only the lensed source emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets plot the simulated `Imaging` dataset before we output it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future.\n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centre JSON Files__\n", - "\n", - "Save the centres of the main lens galaxies and extra galaxies as JSON files. These are loaded by the group\n", - "modeling scripts to set up the lens model (e.g. fixing centres of extra galaxies, defining scaling relations)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Solve for the lensed positions of the source galaxy, which are used as input for the group\n", - "modeling scripts (e.g. SLaM pipeline) to help the non-linear search converge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1),\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.01,\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=positions,\n", - " file_path=dataset_path / \"positions.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: No Lens Light (Group)\n", + "================================\n", + "\n", + "This script simulates `Imaging` of a 'group-scale' strong lens where none of the lens galaxies have visible\n", + "light emission \u2014 only their mass profiles contribute to the ray-tracing. The source galaxy still has light.\n", + "\n", + "This is the group-scale analogue of `imaging/features/no_lens_light/simulator.py`. In a group context, \"no lens\n", + "light\" means that **all** main lens galaxies **and** all extra galaxies are modeled with mass profiles only.\n", + "\n", + "This script simulates `Imaging` of a 'group-scale' strong lens where:\n", + "\n", + " - The group consists of one main lens galaxy and two extra galaxies, all with `IsothermalSph` mass profiles\n", + " and no light profiles.\n", + " - A single source galaxy is observed whose `LightProfile` is a `SersicCore`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Grid:** Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and.\n", + "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", + "- **Over Sampling:** Set up over-sampling at the source centre for accurate light profile evaluation.\n", + "- **Main Lens Galaxies:** The main lens galaxy at the origin, mass only, no light.\n", + "- **Extra Galaxies:** The two extra galaxies, mass only, no light.\n", + "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", + "- **Ray Tracing:** Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging`.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Visualize:** Output a subplot of the simulated dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file.\n", + "- **Centre JSON Files:** Save the centres of the main lens galaxies and extra galaxies as JSON files.\n", + "- **Positions:** Solve for the lensed positions of the source galaxy." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. They\n", + "define the folder the dataset is output to on your hard-disk:\n", + "\n", + " - The image will be output to `/autolens_workspace/dataset/group/simple__no_lens_light/data.fits`.\n", + " - The noise-map will be output to `/autolens_workspace/dataset/group/simple__no_lens_light/noise_map.fits`.\n", + " - The psf will be output to `/autolens_workspace/dataset/group/simple__no_lens_light/psf.fits`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"simple__no_lens_light\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and therefore simulated\n", + "on, via the inputs:\n", + "\n", + " - `shape_native`: The (y_pixels, x_pixels) 2D shape of the grid defining the shape of the data that is simulated.\n", + " - `pixel_scales`: The arc-second to pixel conversion factor of the grid and data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(250, 250),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Define the centres of the main lens galaxies and extra galaxies. These are used for over-sampling and are also\n", + "output to JSON files so that the modeling scripts can load them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [(0.0, 0.0)]\n", + "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Because no galaxy has a light profile, there is no need for adaptive over-sampling at the galaxy centres. However,\n", + "the lensed source light still requires accurate evaluation, so we apply over-sampling at the source centre\n", + "(approximately the primary lens centre where the lensed arcs appear).\n", + "\n", + "An adaptive oversampling grid cannot be defined for the lensed source because its light appears in different regions\n", + "of the image plane for each dataset. For this reason, we use a cored light profile for the source galaxy (`SersicCore`)\n", + "which changes gradually in its central regions, allowing accurate evaluation without requiring heavy oversampling.\n", + "\n", + "We still apply a mild over-sampling at the centre of the image where the lensed arcs are brightest." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the `Imaging` dataset we first create a simulator, which defines the exposure time, background sky,\n", + "noise levels and psf of the dataset that is simulated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "The main lens galaxy is at the origin (0.0, 0.0). It has an isothermal mass profile but **no light profile**.\n", + "\n", + "In the list-based API used by the group modeling scripts, main lens galaxies are stored in a list called\n", + "`main_lens_galaxies`, where each galaxy is referred to as `lens_0`, `lens_1`, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "main_lens_galaxies = [lens_0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "The two extra galaxies are companion galaxies near the lens system. They have isothermal mass profiles\n", + "but **no light profiles**, with centres offset from the origin.\n", + "\n", + "In the list-based API, extra galaxies are stored in a list called `extra_galaxies`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy__\n", + "\n", + "The source galaxy whose lensed images we simulate. It uses a cored Sersic profile so that adaptive over-sampling\n", + "is not required for the source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset.\n", + "\n", + "The tracer combines main lens galaxies, extra galaxies and the source galaxy. Because the lens galaxies have\n", + "no light profiles, the simulated image contains only the lensed source emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets plot the simulated `Imaging` dataset before we output it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future.\n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centre JSON Files__\n", + "\n", + "Save the centres of the main lens galaxies and extra galaxies as JSON files. These are loaded by the group\n", + "modeling scripts to set up the lens model (e.g. fixing centres of extra galaxies, defining scaling relations)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Solve for the lensed positions of the source galaxy, which are used as input for the group\n", + "modeling scripts (e.g. SLaM pipeline) to help the non-linear search converge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1),\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.01,\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=positions,\n", + " file_path=dataset_path / \"positions.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/no_lens_light/slam.ipynb b/notebooks/group/features/no_lens_light/slam.ipynb index 764bc3fda..e35f18432 100644 --- a/notebooks/group/features/no_lens_light/slam.ipynb +++ b/notebooks/group/features/no_lens_light/slam.ipynb @@ -1,771 +1,808 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "No Lens Light: Group SLaM\n", - "=========================\n", - "\n", - "This script uses the SLaM pipelines to fit a group-scale strong lens where none of the lens galaxies have\n", - "visible light emission. In the group context, \"no lens light\" means that **all** main lens galaxies **and**\n", - "all extra galaxies are modeled with mass profiles only.\n", - "\n", - "This is the group-scale analogue of `imaging/features/no_lens_light/slam.py`. The pipeline is substantially\n", - "simplified compared to the standard group SLaM (`group/slam.py`) because:\n", - "\n", - " - There is no `source_lp_0` stage (no light-only fit needed since there is no lens light).\n", - " - There is no `light_lp` stage (no lens light to refit).\n", - " - The pipeline goes directly: source_lp -> source_pix_1 -> source_pix_2 -> mass_total.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **This Script:** Using a SOURCE LP PIPELINE (one search), SOURCE PIX PIPELINE (two searches), and.\n", - "- **SOURCE LP PIPELINE:** Fits mass + source directly. No lens light for any galaxy.\n", - "- **SOURCE PIX PIPELINE 1:** Pixelized source, mass carried forward. No lens light.\n", - "- **SOURCE PIX PIPELINE 2:** Refined pixelized source. No lens light.\n", - "- **MASS TOTAL PIPELINE:** Final mass fit with PowerLaw. No lens light.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Galaxy Centres:** Load centres from JSON files.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "- **Group** (`group/modeling`):\n", - " How we model group-scale strong lenses in PyAutoLens, including how we include extra galaxies in\n", - " the lens model.\n", - "\n", - "- **No Lens Light** (`imaging/features/no_lens_light`):\n", - " How the SLaM pipeline is adapted when no lens light is present.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE (one search), SOURCE PIX PIPELINE (two searches) and TOTAL MASS PIPELINE this\n", - "SLaM modeling script fits `Imaging` data of a group-scale strong lens where in the final model:\n", - "\n", - " - Each main lens galaxy has a `PowerLaw` total mass \u2014 no light.\n", - " - Each extra galaxy has a bounded `IsothermalSph` mass \u2014 no light.\n", - " - The source galaxy's light is a Delaunay `Pixelization` with `AdaptSplit` regularization.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "\n", - "def _load_centres(path):\n", - " \"\"\"Load a centres JSON file, returning an empty list if the file is absent.\"\"\"\n", - " try:\n", - " return al.Grid2DIrregular(al.from_json(file_path=path))\n", - " except FileNotFoundError:\n", - " return al.Grid2DIrregular([])\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Fits mass + source directly. Because no galaxy has light, there is no need for a light-only stage\n", - "(`source_lp_0` in the standard group SLaM). We go straight to fitting mass and source simultaneously.\n", - "\n", - "Multiple main-lens galaxies each get an `Isothermal` mass; only `lens_0` carries an `ExternalShear`.\n", - "Extra-galaxy Einstein radii are bounded by a uniform prior." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " dataset,\n", - " settings_search,\n", - " main_lens_centres,\n", - " extra_lens_centres,\n", - " mask_radius,\n", - " redshift_lens,\n", - " redshift_source,\n", - " n_batch=50,\n", - "):\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " # Source MGE\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " # --- main lens mass models (no light) ---\n", - " lens_dict = {}\n", - " for i, centre in enumerate(main_lens_centres):\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " # --- extra galaxy mass models (no light) ---\n", - " extra_mass_models = []\n", - " for centre in extra_lens_centres:\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_mass_models.append(af.Model(al.Galaxy, redshift=redshift_lens, mass=mass))\n", - "\n", - " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", - " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " n_extra = len(extra_galaxies) if extra_galaxies is not None else 0\n", - " n_live = 150 + 30 * len(lens_dict) + 30 * n_extra\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Pixelized source with mass carried forward from `source_lp`. No lens light to fix \u2014 the galaxies\n", - "remain mass-only. Extra galaxy models are carried forward as free `model` parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " dataset,\n", - " mask,\n", - " settings_search,\n", - " source_lp_result,\n", - " pixel_scale,\n", - " mask_radius,\n", - " positions,\n", - " n_batch=20,\n", - "):\n", - " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", - " edge_pixels_total = 30\n", - " signal_to_noise_threshold = 3.0\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Hilbert(\n", - " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", - " )\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=mask,\n", - " adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"],\n", - " )\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " over_sample_size_pixelization = np.where(\n", - " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", - " 4,\n", - " 2,\n", - " )\n", - " over_sample_size_pixelization = al.Array2D(\n", - " values=over_sample_size_pixelization, mask=mask\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_pixelization=over_sample_size_pixelization,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=2.0, positions=positions, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_lp_result.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lp_lens_model = getattr(source_lp_result.model.galaxies, f\"lens_{i}\")\n", - " lp_lens_instance = getattr(source_lp_result.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=lp_lens_model.mass,\n", - " mass_result=lp_lens_model.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lp_lens_instance.redshift,\n", - " mass=mass,\n", - " shear=lp_lens_model.shear,\n", - " )\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=af.Model(al.reg.AdaptSplit),\n", - " ),\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=source_lp_result.model.extra_galaxies,\n", - " )\n", - "\n", - " n_live = 150 + 50 * (n_lenses - 1)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - " return result, dataset, adapt_images\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Refined pixelized source. The adapt data for the Hilbert image mesh is capped at a S/N threshold of 3.0\n", - "to prevent over-concentration of source pixels. Extra galaxy models are fixed as instances from `source_pix[1]`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " dataset,\n", - " mask,\n", - " settings_search,\n", - " source_lp_result,\n", - " source_pix_result_1,\n", - " pixel_scale,\n", - " mask_radius,\n", - " n_batch=20,\n", - "):\n", - " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", - " edge_pixels_total = 30\n", - " signal_to_noise_threshold = 3.0\n", - " signal_to_noise_threshold_image_mesh = 3.0\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_data_snr_max = galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", - " adapt_data_snr_max[adapt_data_snr_max > signal_to_noise_threshold_image_mesh] = (\n", - " signal_to_noise_threshold_image_mesh\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Hilbert(\n", - " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", - " )\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=mask, adapt_data=adapt_data_snr_max\n", - " )\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " over_sample_size_pixelization = np.where(\n", - " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", - " 4,\n", - " 2,\n", - " )\n", - " over_sample_size_pixelization = al.Array2D(\n", - " values=over_sample_size_pixelization, mask=mask\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_pixelization=over_sample_size_pixelization,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " pix1_lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=pix1_lens_instance.redshift,\n", - " mass=pix1_lens_instance.mass,\n", - " shear=pix1_lens_instance.shear,\n", - " )\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=af.Model(al.reg.AdaptSplit),\n", - " ),\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - " return result, dataset, adapt_images\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Final mass fit with PowerLaw profiles for the main lens galaxies. No lens light is included.\n", - "Extra galaxies receive new bounded Einstein radii.\n", - "\n", - "Note: there is **no** `light_lp` stage in this pipeline because there is no lens light to refit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " dataset,\n", - " settings_search,\n", - " source_pix_result_1,\n", - " source_pix_result_2,\n", - " adapt_images,\n", - " positions,\n", - " redshift_lens,\n", - " n_batch=20,\n", - "):\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - " n_extra = (\n", - " len(list(source_pix_result_1.instance.extra_galaxies))\n", - " if source_pix_result_1.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " # --- extra galaxies: mass only, new bounded Einstein radii ---\n", - " extra_mass_models = []\n", - " for i in range(n_extra):\n", - " pix1_extra = source_pix_result_1.instance.extra_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = pix1_extra.mass.centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_mass_models.append(af.Model(al.Galaxy, redshift=redshift_lens, mass=mass))\n", - "\n", - " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_1.positions_likelihood_from(\n", - " factor=3.0, positions=positions, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lens_model = getattr(source_pix_result_1.model.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.PowerLaw),\n", - " mass_result=lens_model.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lens_model.redshift,\n", - " mass=mass,\n", - " shear=lens_model.shear,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " n_live = 200 + 100 * (n_lenses - 1)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "pixel_scale = 0.1\n", - "mask_radius = 7.5\n", - "redshift_lens = 0.5\n", - "redshift_source = 1.0\n", - "n_batch = 20\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=pixel_scale,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "main_lens_centres.json \u2014 required; determines the number of main lenses.\n", - "extra_galaxies_centres.json \u2014 optional; empty list if absent.\n", - "\n", - "All files contain a list of [y, x] arcsecond coordinates." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = _load_centres(dataset_path / \"main_lens_centres.json\")\n", - "extra_lens_centres = _load_centres(dataset_path / \"extra_galaxies_centres.json\")\n", - "\n", - "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define the 2D mask applied to the dataset for the model-fit. We use a 7.5 arcsecond circular mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"group\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE for the no-lens-light group case.\n", - "\n", - "The pipeline is:\n", - "\n", - " 1. `source_lp`: Fit mass + MGE source directly \u2014 no lens light for any galaxy.\n", - " 2. `source_pix_1`: Pixelized source, mass carried forward.\n", - " 3. `source_pix_2`: Refined pixelized source.\n", - " 4. `mass_total`: Final PowerLaw mass fit.\n", - "\n", - "There is **no** `source_lp_0` (light-only) stage and **no** `light_lp` stage because there is no lens light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " main_lens_centres=main_lens_centres,\n", - " extra_lens_centres=extra_lens_centres,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1, dataset, adapt_images = source_pix_1(\n", - " dataset=dataset,\n", - " mask=mask,\n", - " settings_search=settings_search,\n", - " source_lp_result=source_lp_result,\n", - " pixel_scale=pixel_scale,\n", - " mask_radius=mask_radius,\n", - " positions=positions,\n", - " n_batch=n_batch,\n", - ")\n", - "\n", - "source_pix_result_2, dataset, adapt_images = source_pix_2(\n", - " dataset=dataset,\n", - " mask=mask,\n", - " settings_search=settings_search,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " pixel_scale=pixel_scale,\n", - " mask_radius=mask_radius,\n", - " n_batch=n_batch,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " adapt_images=adapt_images,\n", - " positions=positions,\n", - " redshift_lens=redshift_lens,\n", - " n_batch=n_batch,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "No Lens Light: Group SLaM\n", + "=========================\n", + "\n", + "This script uses the SLaM pipelines to fit a group-scale strong lens where none of the lens galaxies have\n", + "visible light emission. In the group context, \"no lens light\" means that **all** main lens galaxies **and**\n", + "all extra galaxies are modeled with mass profiles only.\n", + "\n", + "This is the group-scale analogue of `imaging/features/no_lens_light/slam.py`. The pipeline is substantially\n", + "simplified compared to the standard group SLaM (`group/slam.py`) because:\n", + "\n", + " - There is no `source_lp_0` stage (no light-only fit needed since there is no lens light).\n", + " - There is no `light_lp` stage (no lens light to refit).\n", + " - The pipeline goes directly: source_lp -> source_pix_1 -> source_pix_2 -> mass_total.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **This Script:** Using a SOURCE LP PIPELINE (one search), SOURCE PIX PIPELINE (two searches), and.\n", + "- **SOURCE LP PIPELINE:** Fits mass + source directly. No lens light for any galaxy.\n", + "- **SOURCE PIX PIPELINE 1:** Pixelized source, mass carried forward. No lens light.\n", + "- **SOURCE PIX PIPELINE 2:** Refined pixelized source. No lens light.\n", + "- **MASS TOTAL PIPELINE:** Final mass fit with PowerLaw. No lens light.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Galaxy Centres:** Load centres from JSON files.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "- **Group** (`group/modeling`):\n", + " How we model group-scale strong lenses in PyAutoLens, including how we include extra galaxies in\n", + " the lens model.\n", + "\n", + "- **No Lens Light** (`imaging/features/no_lens_light`):\n", + " How the SLaM pipeline is adapted when no lens light is present.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE (one search), SOURCE PIX PIPELINE (two searches) and TOTAL MASS PIPELINE this\n", + "SLaM modeling script fits `Imaging` data of a group-scale strong lens where in the final model:\n", + "\n", + " - Each main lens galaxy has a `PowerLaw` total mass \u2014 no light.\n", + " - Each extra galaxy has a bounded `IsothermalSph` mass \u2014 no light.\n", + " - The source galaxy's light is a Delaunay `Pixelization` with `AdaptSplit` regularization.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "\n", + "def _load_centres(path):\n", + " \"\"\"Load a centres JSON file, returning an empty list if the file is absent.\"\"\"\n", + " try:\n", + " return al.Grid2DIrregular(al.from_json(file_path=path))\n", + " except FileNotFoundError:\n", + " return al.Grid2DIrregular([])\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Fits mass + source directly. Because no galaxy has light, there is no need for a light-only stage\n", + "(`source_lp_0` in the standard group SLaM). We go straight to fitting mass and source simultaneously.\n", + "\n", + "Multiple main-lens galaxies each get an `Isothermal` mass; only `lens_0` carries an `ExternalShear`.\n", + "Extra-galaxy Einstein radii are bounded by a uniform prior." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " dataset,\n", + " settings_search,\n", + " main_lens_centres,\n", + " extra_lens_centres,\n", + " mask_radius,\n", + " redshift_lens,\n", + " redshift_source,\n", + " n_batch=50,\n", + "):\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " # Source MGE\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " # --- main lens mass models (no light) ---\n", + " lens_dict = {}\n", + " for i, centre in enumerate(main_lens_centres):\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " # --- extra galaxy mass models (no light) ---\n", + " extra_mass_models = []\n", + " for centre in extra_lens_centres:\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_mass_models.append(af.Model(al.Galaxy, redshift=redshift_lens, mass=mass))\n", + "\n", + " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", + " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " n_extra = len(extra_galaxies) if extra_galaxies is not None else 0\n", + " n_live = 150 + 30 * len(lens_dict) + 30 * n_extra\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Pixelized source with mass carried forward from `source_lp`. No lens light to fix \u2014 the galaxies\n", + "remain mass-only. Extra galaxy models are carried forward as free `model` parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " dataset,\n", + " mask,\n", + " settings_search,\n", + " source_lp_result,\n", + " pixel_scale,\n", + " mask_radius,\n", + " positions,\n", + " n_batch=20,\n", + "):\n", + " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", + " edge_pixels_total = 30\n", + " signal_to_noise_threshold = 3.0\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Hilbert(\n", + " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", + " )\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=mask,\n", + " adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"],\n", + " )\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " over_sample_size_pixelization = np.where(\n", + " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", + " 4,\n", + " 2,\n", + " )\n", + " over_sample_size_pixelization = al.Array2D(\n", + " values=over_sample_size_pixelization, mask=mask\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_pixelization=over_sample_size_pixelization,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=2.0, positions=positions, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_lp_result.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lp_lens_model = getattr(source_lp_result.model.galaxies, f\"lens_{i}\")\n", + " lp_lens_instance = getattr(source_lp_result.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=lp_lens_model.mass,\n", + " mass_result=lp_lens_model.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lp_lens_instance.redshift,\n", + " mass=mass,\n", + " shear=lp_lens_model.shear,\n", + " )\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=af.Model(al.reg.AdaptSplit),\n", + " ),\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=source_lp_result.model.extra_galaxies,\n", + " )\n", + "\n", + " n_live = 150 + 50 * (n_lenses - 1)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + " return result, dataset, adapt_images\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Refined pixelized source. The adapt data for the Hilbert image mesh is capped at a S/N threshold of 3.0\n", + "to prevent over-concentration of source pixels. Extra galaxy models are fixed as instances from `source_pix[1]`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " dataset,\n", + " mask,\n", + " settings_search,\n", + " source_lp_result,\n", + " source_pix_result_1,\n", + " pixel_scale,\n", + " mask_radius,\n", + " n_batch=20,\n", + "):\n", + " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", + " edge_pixels_total = 30\n", + " signal_to_noise_threshold = 3.0\n", + " signal_to_noise_threshold_image_mesh = 3.0\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_data_snr_max = galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", + " adapt_data_snr_max[adapt_data_snr_max > signal_to_noise_threshold_image_mesh] = (\n", + " signal_to_noise_threshold_image_mesh\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Hilbert(\n", + " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", + " )\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=mask, adapt_data=adapt_data_snr_max\n", + " )\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " over_sample_size_pixelization = np.where(\n", + " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", + " 4,\n", + " 2,\n", + " )\n", + " over_sample_size_pixelization = al.Array2D(\n", + " values=over_sample_size_pixelization, mask=mask\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_pixelization=over_sample_size_pixelization,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " pix1_lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=pix1_lens_instance.redshift,\n", + " mass=pix1_lens_instance.mass,\n", + " shear=pix1_lens_instance.shear,\n", + " )\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=af.Model(al.reg.AdaptSplit),\n", + " ),\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + " return result, dataset, adapt_images\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Final mass fit with PowerLaw profiles for the main lens galaxies. No lens light is included.\n", + "Extra galaxies receive new bounded Einstein radii.\n", + "\n", + "Note: there is **no** `light_lp` stage in this pipeline because there is no lens light to refit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " dataset,\n", + " settings_search,\n", + " source_pix_result_1,\n", + " source_pix_result_2,\n", + " adapt_images,\n", + " positions,\n", + " redshift_lens,\n", + " n_batch=20,\n", + "):\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + " n_extra = (\n", + " len(list(source_pix_result_1.instance.extra_galaxies))\n", + " if source_pix_result_1.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " # --- extra galaxies: mass only, new bounded Einstein radii ---\n", + " extra_mass_models = []\n", + " for i in range(n_extra):\n", + " pix1_extra = source_pix_result_1.instance.extra_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = pix1_extra.mass.centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_mass_models.append(af.Model(al.Galaxy, redshift=redshift_lens, mass=mass))\n", + "\n", + " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_1.positions_likelihood_from(\n", + " factor=3.0, positions=positions, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lens_model = getattr(source_pix_result_1.model.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.PowerLaw),\n", + " mass_result=lens_model.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lens_model.redshift,\n", + " mass=mass,\n", + " shear=lens_model.shear,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " n_live = 200 + 100 * (n_lenses - 1)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "pixel_scale = 0.1\n", + "mask_radius = 7.5\n", + "redshift_lens = 0.5\n", + "redshift_source = 1.0\n", + "n_batch = 20\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=pixel_scale,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "main_lens_centres.json \u2014 required; determines the number of main lenses.\n", + "extra_galaxies_centres.json \u2014 optional; empty list if absent.\n", + "\n", + "All files contain a list of [y, x] arcsecond coordinates." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = _load_centres(dataset_path / \"main_lens_centres.json\")\n", + "extra_lens_centres = _load_centres(dataset_path / \"extra_galaxies_centres.json\")\n", + "\n", + "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define the 2D mask applied to the dataset for the model-fit. We use a 7.5 arcsecond circular mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"group\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE for the no-lens-light group case.\n", + "\n", + "The pipeline is:\n", + "\n", + " 1. `source_lp`: Fit mass + MGE source directly \u2014 no lens light for any galaxy.\n", + " 2. `source_pix_1`: Pixelized source, mass carried forward.\n", + " 3. `source_pix_2`: Refined pixelized source.\n", + " 4. `mass_total`: Final PowerLaw mass fit.\n", + "\n", + "There is **no** `source_lp_0` (light-only) stage and **no** `light_lp` stage because there is no lens light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " main_lens_centres=main_lens_centres,\n", + " extra_lens_centres=extra_lens_centres,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1, dataset, adapt_images = source_pix_1(\n", + " dataset=dataset,\n", + " mask=mask,\n", + " settings_search=settings_search,\n", + " source_lp_result=source_lp_result,\n", + " pixel_scale=pixel_scale,\n", + " mask_radius=mask_radius,\n", + " positions=positions,\n", + " n_batch=n_batch,\n", + ")\n", + "\n", + "source_pix_result_2, dataset, adapt_images = source_pix_2(\n", + " dataset=dataset,\n", + " mask=mask,\n", + " settings_search=settings_search,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " pixel_scale=pixel_scale,\n", + " mask_radius=mask_radius,\n", + " n_batch=n_batch,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " adapt_images=adapt_images,\n", + " positions=positions,\n", + " redshift_lens=redshift_lens,\n", + " n_batch=n_batch,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/pixelization/adaptive.ipynb b/notebooks/group/features/pixelization/adaptive.ipynb index 909eee32e..d12ebc293 100644 --- a/notebooks/group/features/pixelization/adaptive.ipynb +++ b/notebooks/group/features/pixelization/adaptive.ipynb @@ -1,567 +1,604 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Adaptive Pixelization (Group)\n", - "=============================\n", - "\n", - "This script demonstrates how adaptive pixelization features work for group-scale strong lenses.\n", - "\n", - "Adaptive pixelizations adapt the mesh density and regularization strength to the source galaxy's unlensed\n", - "morphology. More source pixels are placed where the source is brightest, and regularization is relaxed\n", - "in bright regions to preserve detail while remaining strong in faint regions to suppress noise.\n", - "\n", - "For group-scale lenses, correct lens-light subtraction from ALL galaxies (main + extra) is critical\n", - "for producing good adapt_data, which drives the adaptive mesh. If the lens light subtraction is poor,\n", - "the adaptive mesh may concentrate pixels on residuals rather than genuine source emission.\n", - "\n", - "This script uses search chaining: an initial parametric fit establishes the lens model and source\n", - "morphology, then subsequent searches use adaptive pixelization features.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", - "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", - "- **Search 1:** Fit a parametric model to establish the lens model and source morphology.\n", - "- **Search 2:** Introduce a pixelization with constant regularization.\n", - "- **Search 3:** Use adaptive mesh and regularization driven by adapt_data from search 2.\n", - "- **Adapt Images:** How adapt_data is constructed from the lens-subtracted source image.\n", - "\n", - "__Adaptive Features__\n", - "\n", - "Two key adaptive classes are used:\n", - "\n", - " - `RectangularAdaptImage` mesh: adapts the rectangular source-pixel upsampling to the source's unlensed\n", - " morphology. More rectangular pixels are placed where the source is located, even in low magnification\n", - " regions.\n", - "\n", - " - `Adapt` regularization: adapts the regularization coefficient to the source's unlensed morphology.\n", - " Bright regions are regularized less (preserving detail), faint regions are regularized more\n", - " (suppressing noise).\n", - "\n", - "For group lenses, the adapt_data is the lens-subtracted image, where the light of ALL group galaxies\n", - "has been removed. This means accurate modeling of every galaxy's light profile is essential." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"group\") / \"features\" / \"pixelization\" / \"adaptive\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Search 1 fits a parametric group model to establish the lens mass model and source morphology.\n", - "\n", - " - Main lens galaxy: MGE light + Isothermal mass + ExternalShear.\n", - " - Extra galaxies: MGE light + IsothermalSph mass (fixed centres, bounded Einstein radii).\n", - " - Source: MGE light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source (parametric):\n", - "\n", - "source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model_1 = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")\n", - "\n", - "print(model_1.info)\n", - "\n", - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__parametric\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_like_max=500,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "Search 2 introduces a pixelization with constant regularization. The lens mass model is taken from\n", - "search 1 as a model (with priors from the previous result)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict_2 = {}\n", - "for i, _ in enumerate(main_lens_centres):\n", - " lens_dict_2[f\"lens_{i}\"] = getattr(result_1.model.galaxies, f\"lens_{i}\")\n", - "\n", - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", - " regularization=al.reg.Constant,\n", - ")\n", - "\n", - "source_2 = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "model_2 = af.Collection(\n", - " galaxies=af.Collection(**lens_dict_2, source=source_2),\n", - " extra_galaxies=result_1.model.extra_galaxies,\n", - ")\n", - "\n", - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__pixelization_setup\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_like_max=500,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[\n", - " result_1.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " use_jax=True,\n", - ")\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Adaptive Pixelization (Search 3)__\n", - "\n", - "Search 3 uses the adaptive pixelization classes:\n", - "\n", - " - `RectangularAdaptImage` mesh: adapts pixel density to the source morphology.\n", - " - `Adapt` regularization: adapts smoothing strength to the source brightness.\n", - "\n", - "The lens mass is fixed from search 2 to ensure the adaptation is performed quickly.\n", - "\n", - "__Adapt Images__\n", - "\n", - "The `adapt_images` are constructed from the lens-subtracted source image of search 2. For group lenses,\n", - "this means the light of ALL galaxies (main + extra) has been subtracted, leaving only the lensed source.\n", - "\n", - "If the lens light subtraction is poor (e.g. because an extra galaxy's light was not modeled), the\n", - "adapt_data will contain residuals that confuse the adaptive mesh." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result_2)\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "# Fix the lens mass model from search 2.\n", - "lens_dict_3 = {}\n", - "for i, _ in enumerate(main_lens_centres):\n", - " lens_dict_3[f\"lens_{i}\"] = getattr(result_2.instance.galaxies, f\"lens_{i}\")\n", - "\n", - "pixelization_3 = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.RectangularAdaptImage(shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "source_3 = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization_3)\n", - "\n", - "model_3 = af.Collection(\n", - " galaxies=af.Collection(**lens_dict_3, source=source_3),\n", - " extra_galaxies=result_2.instance.extra_galaxies,\n", - ")\n", - "\n", - "search_3 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[3]__adaptive_pixelization\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - " n_like_max=500,\n", - ")\n", - "\n", - "analysis_3 = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " result_2.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " use_jax=True,\n", - ")\n", - "\n", - "result_3 = search_3.fit(model=model_3, analysis=analysis_3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 3)__\n", - "\n", - "The adaptive pixelization result should show significantly higher likelihood than search 2, with\n", - "the source reconstruction concentrating pixels in bright source regions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_3.info)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result_3.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search 4: Free Mass Model__\n", - "\n", - "Finally, we refit the lens mass model with the adaptive pixelization fixed from search 3." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict_4 = {}\n", - "for i, _ in enumerate(main_lens_centres):\n", - " bulge_i = getattr(result_2.instance.galaxies, f\"lens_{i}\").bulge\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens_dict_4[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge_i,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - "extra_galaxies_4_list = []\n", - "for centre in extra_galaxies_centres:\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - " extra_galaxies_4_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies_4 = af.Collection(extra_galaxies_4_list)\n", - "\n", - "source_4 = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " pixelization=result_3.instance.galaxies.source.pixelization,\n", - ")\n", - "\n", - "model_4 = af.Collection(\n", - " galaxies=af.Collection(**lens_dict_4, source=source_4),\n", - " extra_galaxies=extra_galaxies_4,\n", - ")\n", - "\n", - "search_4 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[4]__adapt_free_mass\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_4 = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - ")\n", - "\n", - "result_4 = search_4.fit(model=model_4, analysis=analysis_4)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated adaptive pixelization for group-scale lenses.\n", - "\n", - "Key points:\n", - " - Adaptive pixelizations are set up via search chaining: parametric fit -> constant pixelization -> adaptive.\n", - " - The adapt_data is the lens-subtracted image, so accurate light modeling of ALL group galaxies is essential.\n", - " - `RectangularAdaptImage` concentrates source pixels where the source is brightest.\n", - " - `Adapt` regularization varies smoothing based on source brightness.\n", - " - The SLaM pipeline (see `group/features/pixelization/slam.py`) automates this entire process." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Adaptive Pixelization (Group)\n", + "=============================\n", + "\n", + "This script demonstrates how adaptive pixelization features work for group-scale strong lenses.\n", + "\n", + "Adaptive pixelizations adapt the mesh density and regularization strength to the source galaxy's unlensed\n", + "morphology. More source pixels are placed where the source is brightest, and regularization is relaxed\n", + "in bright regions to preserve detail while remaining strong in faint regions to suppress noise.\n", + "\n", + "For group-scale lenses, correct lens-light subtraction from ALL galaxies (main + extra) is critical\n", + "for producing good adapt_data, which drives the adaptive mesh. If the lens light subtraction is poor,\n", + "the adaptive mesh may concentrate pixels on residuals rather than genuine source emission.\n", + "\n", + "This script uses search chaining: an initial parametric fit establishes the lens model and source\n", + "morphology, then subsequent searches use adaptive pixelization features.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", + "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", + "- **Search 1:** Fit a parametric model to establish the lens model and source morphology.\n", + "- **Search 2:** Introduce a pixelization with constant regularization.\n", + "- **Search 3:** Use adaptive mesh and regularization driven by adapt_data from search 2.\n", + "- **Adapt Images:** How adapt_data is constructed from the lens-subtracted source image.\n", + "\n", + "__Adaptive Features__\n", + "\n", + "Two key adaptive classes are used:\n", + "\n", + " - `RectangularAdaptImage` mesh: adapts the rectangular source-pixel upsampling to the source's unlensed\n", + " morphology. More rectangular pixels are placed where the source is located, even in low magnification\n", + " regions.\n", + "\n", + " - `Adapt` regularization: adapts the regularization coefficient to the source's unlensed morphology.\n", + " Bright regions are regularized less (preserving detail), faint regions are regularized more\n", + " (suppressing noise).\n", + "\n", + "For group lenses, the adapt_data is the lens-subtracted image, where the light of ALL group galaxies\n", + "has been removed. This means accurate modeling of every galaxy's light profile is essential." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"group\") / \"features\" / \"pixelization\" / \"adaptive\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Search 1 fits a parametric group model to establish the lens mass model and source morphology.\n", + "\n", + " - Main lens galaxy: MGE light + Isothermal mass + ExternalShear.\n", + " - Extra galaxies: MGE light + IsothermalSph mass (fixed centres, bounded Einstein radii).\n", + " - Source: MGE light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source (parametric):\n", + "\n", + "source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model_1 = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")\n", + "\n", + "print(model_1.info)\n", + "\n", + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__parametric\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_like_max=500,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "Search 2 introduces a pixelization with constant regularization. The lens mass model is taken from\n", + "search 1 as a model (with priors from the previous result)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict_2 = {}\n", + "for i, _ in enumerate(main_lens_centres):\n", + " lens_dict_2[f\"lens_{i}\"] = getattr(result_1.model.galaxies, f\"lens_{i}\")\n", + "\n", + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", + " regularization=al.reg.Constant,\n", + ")\n", + "\n", + "source_2 = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "model_2 = af.Collection(\n", + " galaxies=af.Collection(**lens_dict_2, source=source_2),\n", + " extra_galaxies=result_1.model.extra_galaxies,\n", + ")\n", + "\n", + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__pixelization_setup\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_like_max=500,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[\n", + " result_1.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " use_jax=True,\n", + ")\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Adaptive Pixelization (Search 3)__\n", + "\n", + "Search 3 uses the adaptive pixelization classes:\n", + "\n", + " - `RectangularAdaptImage` mesh: adapts pixel density to the source morphology.\n", + " - `Adapt` regularization: adapts smoothing strength to the source brightness.\n", + "\n", + "The lens mass is fixed from search 2 to ensure the adaptation is performed quickly.\n", + "\n", + "__Adapt Images__\n", + "\n", + "The `adapt_images` are constructed from the lens-subtracted source image of search 2. For group lenses,\n", + "this means the light of ALL galaxies (main + extra) has been subtracted, leaving only the lensed source.\n", + "\n", + "If the lens light subtraction is poor (e.g. because an extra galaxy's light was not modeled), the\n", + "adapt_data will contain residuals that confuse the adaptive mesh." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result_2)\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "# Fix the lens mass model from search 2.\n", + "lens_dict_3 = {}\n", + "for i, _ in enumerate(main_lens_centres):\n", + " lens_dict_3[f\"lens_{i}\"] = getattr(result_2.instance.galaxies, f\"lens_{i}\")\n", + "\n", + "pixelization_3 = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.RectangularAdaptImage(shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "source_3 = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization_3)\n", + "\n", + "model_3 = af.Collection(\n", + " galaxies=af.Collection(**lens_dict_3, source=source_3),\n", + " extra_galaxies=result_2.instance.extra_galaxies,\n", + ")\n", + "\n", + "search_3 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[3]__adaptive_pixelization\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + " n_like_max=500,\n", + ")\n", + "\n", + "analysis_3 = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " result_2.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " use_jax=True,\n", + ")\n", + "\n", + "result_3 = search_3.fit(model=model_3, analysis=analysis_3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 3)__\n", + "\n", + "The adaptive pixelization result should show significantly higher likelihood than search 2, with\n", + "the source reconstruction concentrating pixels in bright source regions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_3.info)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result_3.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search 4: Free Mass Model__\n", + "\n", + "Finally, we refit the lens mass model with the adaptive pixelization fixed from search 3." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict_4 = {}\n", + "for i, _ in enumerate(main_lens_centres):\n", + " bulge_i = getattr(result_2.instance.galaxies, f\"lens_{i}\").bulge\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens_dict_4[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge_i,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + "extra_galaxies_4_list = []\n", + "for centre in extra_galaxies_centres:\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + " extra_galaxies_4_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies_4 = af.Collection(extra_galaxies_4_list)\n", + "\n", + "source_4 = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " pixelization=result_3.instance.galaxies.source.pixelization,\n", + ")\n", + "\n", + "model_4 = af.Collection(\n", + " galaxies=af.Collection(**lens_dict_4, source=source_4),\n", + " extra_galaxies=extra_galaxies_4,\n", + ")\n", + "\n", + "search_4 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[4]__adapt_free_mass\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_4 = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + ")\n", + "\n", + "result_4 = search_4.fit(model=model_4, analysis=analysis_4)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated adaptive pixelization for group-scale lenses.\n", + "\n", + "Key points:\n", + " - Adaptive pixelizations are set up via search chaining: parametric fit -> constant pixelization -> adaptive.\n", + " - The adapt_data is the lens-subtracted image, so accurate light modeling of ALL group galaxies is essential.\n", + " - `RectangularAdaptImage` concentrates source pixels where the source is brightest.\n", + " - `Adapt` regularization varies smoothing based on source brightness.\n", + " - The SLaM pipeline (see `group/features/pixelization/slam.py`) automates this entire process." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/pixelization/cpu_fast_modeling.ipynb b/notebooks/group/features/pixelization/cpu_fast_modeling.ipynb index dc05c3335..382bb1dc3 100644 --- a/notebooks/group/features/pixelization/cpu_fast_modeling.ipynb +++ b/notebooks/group/features/pixelization/cpu_fast_modeling.ipynb @@ -1,465 +1,502 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "CPU Fast Modeling: Pixelization (Group)\n", - "=======================================\n", - "\n", - "This script demonstrates how to achieve fast pixelization performance on a CPU without JAX for group-scale\n", - "strong lenses, by combining:\n", - "\n", - " - `numba` for optimized numerical routines.\n", - " - Python `multiprocessing` to exploit multiple CPU cores.\n", - " - Sparse operator formalism for efficient linear algebra.\n", - "\n", - "For group-scale lenses, the larger 7.5\" mask means significantly more image pixels than galaxy-scale\n", - "lenses, which increases the size of the matrices used in the pixelized source reconstruction. CPU\n", - "optimization is therefore even more important for group lenses, as the matrix operations scale with\n", - "the number of image pixels.\n", - "\n", - "On machines with many CPU cores (e.g. HPC clusters with >10 CPUs), this method can outperform JAX GPU\n", - "acceleration for pixelized source modeling, because pixelizations rely on sparse linear algebra which\n", - "is not currently optimized in JAX.\n", - "\n", - "> Note: This performance advantage applies only to pixelized sources. For parametric sources or\n", - "> multi-Gaussian models, JAX (especially with a GPU) is significantly faster.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", - "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", - "- **Sparse Operators:** Pre-compute sparse matrices for CPU-accelerated pixelization.\n", - "- **Fit:** Fit the group lens with a pixelized source on CPU.\n", - "- **Model:** Full model-fit with CPU parallelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "try:\n", - " import numba\n", - "except ModuleNotFoundError:\n", - " input(\n", - " \"##################\\n\"\n", - " \"##### NUMBA ######\\n\"\n", - " \"##################\\n\\n\"\n", - " \"\"\"\n", - " Numba is not currently installed.\n", - "\n", - " Numba is a library which makes PyAutoLens run a lot faster. Certain functionality is disabled without numba\n", - " and will raise an exception if it is used.\n", - "\n", - " If you have not tried installing numba, I recommend you try and do so now by running the following\n", - " commands in your command line / bash terminal now:\n", - "\n", - " pip install --upgrade pip\n", - " pip install numba\n", - "\n", - " [Press Enter to continue]\n", - " \"\"\"\n", - " )\n", - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=4,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "Pixelized source modeling requires dense linear algebra operations. These calculations can be greatly\n", - "accelerated using the sparse operator formalism.\n", - "\n", - "For group-scale lenses, this is especially important because the 7.5\" mask contains many more pixels\n", - "than a typical galaxy-scale 3.0\" mask, making the matrices much larger.\n", - "\n", - "Computing the operator matrices takes anywhere from a few seconds to a few minutes, depending on the\n", - "dataset size. After it is computed once, every model-fit using pixelization becomes substantially faster." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_sparse_operator_cpu()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We first demonstrate a single fit using the CPU-accelerated pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", - "regularization = al.reg.Constant(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(f\"Log Likelihood: {fit.log_likelihood}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We now perform a full model-fit using the sparse operator formalism on the CPU.\n", - "\n", - "There are two key differences from JAX-based pixelization examples:\n", - "\n", - " - **JAX is disabled**: The `AnalysisImaging` class is created with `use_jax=False`.\n", - " - **CPU parallelization**: The non-linear search uses `number_of_cores` to parallelize likelihood\n", - " evaluations using Python's `multiprocessing`.\n", - "\n", - "For group-scale lenses, CPU parallelization is especially beneficial because each likelihood evaluation\n", - "is more expensive (due to the larger mask and more galaxies)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source: Rectangular pixelization with constant regularization.\n", - "\n", - "pix = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", - " regularization=al.reg.Constant,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pix)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\" / \"pixelization\",\n", - " name=\"cpu_fast_modeling\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " number_of_cores=2, # CPU specific: parallelize likelihood evaluations\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=False, # CPU specific: disable JAX compilation\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated CPU-optimized pixelization modeling for group-scale lenses.\n", - "\n", - "Key points for group lenses:\n", - " - The larger 7.5\" mask means more image pixels, making sparse operators especially valuable.\n", - " - CPU parallelization via `number_of_cores` provides significant speedup for the expensive likelihood\n", - " evaluations of group-scale fits.\n", - " - On HPC clusters with many cores, this approach can outperform JAX GPU acceleration for pixelized sources.\n", - " - For parametric sources (MGE, Sersic), JAX with GPU remains faster." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "CPU Fast Modeling: Pixelization (Group)\n", + "=======================================\n", + "\n", + "This script demonstrates how to achieve fast pixelization performance on a CPU without JAX for group-scale\n", + "strong lenses, by combining:\n", + "\n", + " - `numba` for optimized numerical routines.\n", + " - Python `multiprocessing` to exploit multiple CPU cores.\n", + " - Sparse operator formalism for efficient linear algebra.\n", + "\n", + "For group-scale lenses, the larger 7.5\" mask means significantly more image pixels than galaxy-scale\n", + "lenses, which increases the size of the matrices used in the pixelized source reconstruction. CPU\n", + "optimization is therefore even more important for group lenses, as the matrix operations scale with\n", + "the number of image pixels.\n", + "\n", + "On machines with many CPU cores (e.g. HPC clusters with >10 CPUs), this method can outperform JAX GPU\n", + "acceleration for pixelized source modeling, because pixelizations rely on sparse linear algebra which\n", + "is not currently optimized in JAX.\n", + "\n", + "> Note: This performance advantage applies only to pixelized sources. For parametric sources or\n", + "> multi-Gaussian models, JAX (especially with a GPU) is significantly faster.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", + "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", + "- **Sparse Operators:** Pre-compute sparse matrices for CPU-accelerated pixelization.\n", + "- **Fit:** Fit the group lens with a pixelized source on CPU.\n", + "- **Model:** Full model-fit with CPU parallelization." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "try:\n", + " import numba\n", + "except ModuleNotFoundError:\n", + " input(\n", + " \"##################\\n\"\n", + " \"##### NUMBA ######\\n\"\n", + " \"##################\\n\\n\"\n", + " \"\"\"\n", + " Numba is not currently installed.\n", + "\n", + " Numba is a library which makes PyAutoLens run a lot faster. Certain functionality is disabled without numba\n", + " and will raise an exception if it is used.\n", + "\n", + " If you have not tried installing numba, I recommend you try and do so now by running the following\n", + " commands in your command line / bash terminal now:\n", + "\n", + " pip install --upgrade pip\n", + " pip install numba\n", + "\n", + " [Press Enter to continue]\n", + " \"\"\"\n", + " )\n", + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=4,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "Pixelized source modeling requires dense linear algebra operations. These calculations can be greatly\n", + "accelerated using the sparse operator formalism.\n", + "\n", + "For group-scale lenses, this is especially important because the 7.5\" mask contains many more pixels\n", + "than a typical galaxy-scale 3.0\" mask, making the matrices much larger.\n", + "\n", + "Computing the operator matrices takes anywhere from a few seconds to a few minutes, depending on the\n", + "dataset size. After it is computed once, every model-fit using pixelization becomes substantially faster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_sparse_operator_cpu()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We first demonstrate a single fit using the CPU-accelerated pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", + "regularization = al.reg.Constant(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(f\"Log Likelihood: {fit.log_likelihood}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We now perform a full model-fit using the sparse operator formalism on the CPU.\n", + "\n", + "There are two key differences from JAX-based pixelization examples:\n", + "\n", + " - **JAX is disabled**: The `AnalysisImaging` class is created with `use_jax=False`.\n", + " - **CPU parallelization**: The non-linear search uses `number_of_cores` to parallelize likelihood\n", + " evaluations using Python's `multiprocessing`.\n", + "\n", + "For group-scale lenses, CPU parallelization is especially beneficial because each likelihood evaluation\n", + "is more expensive (due to the larger mask and more galaxies)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source: Rectangular pixelization with constant regularization.\n", + "\n", + "pix = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", + " regularization=al.reg.Constant,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pix)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\" / \"pixelization\",\n", + " name=\"cpu_fast_modeling\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " number_of_cores=2, # CPU specific: parallelize likelihood evaluations\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=False, # CPU specific: disable JAX compilation\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated CPU-optimized pixelization modeling for group-scale lenses.\n", + "\n", + "Key points for group lenses:\n", + " - The larger 7.5\" mask means more image pixels, making sparse operators especially valuable.\n", + " - CPU parallelization via `number_of_cores` provides significant speedup for the expensive likelihood\n", + " evaluations of group-scale fits.\n", + " - On HPC clusters with many cores, this approach can outperform JAX GPU acceleration for pixelized sources.\n", + " - For parametric sources (MGE, Sersic), JAX with GPU remains faster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/pixelization/delaunay.ipynb b/notebooks/group/features/pixelization/delaunay.ipynb index e9d62e986..25701c6bb 100644 --- a/notebooks/group/features/pixelization/delaunay.ipynb +++ b/notebooks/group/features/pixelization/delaunay.ipynb @@ -1,587 +1,624 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Delaunay Pixelization (Group)\n", - "=============================\n", - "\n", - "This script demonstrates using a Delaunay triangulation mesh for source reconstruction in a group-scale\n", - "strong lens, as an alternative to the rectangular mesh used in other examples.\n", - "\n", - "The Delaunay mesh has several unique properties:\n", - "\n", - " - **Adaptive Mesh**: In the source plane, the Delaunay mesh uses irregularly shaped triangles rather than\n", - " uniform rectangular pixels. This allows the mesh to better adapt to irregular and asymmetric source\n", - " morphologies.\n", - "\n", - " - **Image Mesh**: The vertices of the Delaunay triangles are computed by overlaying a coarse uniform grid\n", - " in the image plane and ray-tracing these coordinates to the source plane. This helps the mesh adapt\n", - " to the magnification pattern of the lens.\n", - "\n", - " - **Interpolation**: The Delaunay mesh uses barycentric interpolation within each triangle, which is\n", - " different from the bilinear interpolation used by rectangular meshes.\n", - "\n", - "For group-scale lenses, the Delaunay mesh is particularly advantageous because the complex mass\n", - "distribution (from multiple galaxies) creates an irregular magnification pattern, and the Delaunay\n", - "triangles naturally adapt to follow the source morphology in this environment.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", - "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", - "- **Fit:** Create a FitImaging with a Delaunay pixelized source for a group lens.\n", - "- **Model:** Compose a group lens model with a Delaunay pixelized source for modeling.\n", - "- **Advantages for Group Lenses:** Why Delaunay meshes are well-suited for group-scale lensing." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autoarray.inversion.plot.inversion_plots import subplot_of_mapper, subplot_mappings" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=4,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image-Plane Mesh Grid__\n", - "\n", - "The Delaunay mesh requires an image-plane mesh grid whose ray-traced positions become the Delaunay\n", - "vertices in the source plane. We build it via an `Overlay` image-mesh covering the masked field, then\n", - "append a ring of edge pixels at the mask boundary so the linear inversion can zero them out. The same\n", - "image-plane mesh grid is reused below for both the concrete fit and the modeling section." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "edge_pixels_total = 30\n", - "\n", - "image_mesh = al.image_mesh.Overlay(shape=(22, 22))\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=dataset.mask)\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We create a fit using a Delaunay mesh whose source-pixel count matches the image-plane mesh grid\n", - "computed above, with constant regularization.\n", - "\n", - "For the group lens, we include the main lens galaxy and extra galaxies with their light and mass\n", - "profiles, and a source galaxy with the Delaunay pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(\n", - " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid}\n", - ")\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer, adapt_images=adapt_images)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit subplot shows the pixelized source does a good job at capturing the appearance of the lensed\n", - "source galaxy across the wide field of the group lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(f\"Log Likelihood: {fit.log_likelihood}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Inversion Visualization__\n", - "\n", - "The Delaunay source reconstruction can be visualized using bespoke inversion plots." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "\n", - "subplot_of_mapper(inversion=inversion, mapper_index=0)\n", - "subplot_mappings(inversion=inversion, pixelization_index=0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstructed Source__\n", - "\n", - "The reconstructed source pixel fluxes and their (y,x) positions in the source plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = inversion.linear_obj_list[0]\n", - "\n", - "reconstruction = inversion.reconstruction\n", - "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", - "\n", - "print(f\"Number of source pixels: {len(reconstruction)}\")\n", - "print(f\"Total source flux: {np.sum(reconstruction)} e- s^-1\")\n", - "print(f\"Source plane mesh grid: {source_plane_mesh_grid}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Advantages for Group Lenses__\n", - "\n", - "The Delaunay mesh is particularly well-suited for group-scale lensing because:\n", - "\n", - " 1. **Irregular magnification**: The combined mass of multiple galaxies creates a complex magnification\n", - " pattern with multiple critical curves. The Delaunay mesh naturally adapts to this by placing more\n", - " triangles in regions of high magnification (where many image pixels map to a small source-plane area).\n", - "\n", - " 2. **Extended arcs**: Group lenses often produce extended arcs that span a large fraction of the image.\n", - " The Delaunay mesh can efficiently cover these extended structures with triangles of varying size.\n", - "\n", - " 3. **Multiple images**: The mass distribution of a group lens can produce many distinct images of the\n", - " source. The Delaunay mesh places vertices based on the image-plane grid, so it naturally creates\n", - " source pixels wherever images of the source appear.\n", - "\n", - " 4. **Source complexity**: The lensed sources in group systems are often complex (e.g. merging galaxies,\n", - " star-forming clumps) and benefit from the adaptive triangle sizes of the Delaunay mesh." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Model__\n", - "\n", - "We compose a group lens model with a Delaunay source for full modeling.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source: Delaunay pixelization with ConstantSplit regularization.\n", - "#\n", - "# The Delaunay mesh is constructed as a concrete instance (not `af.Model`) because its `pixels` count\n", - "# is fixed by the image-plane mesh grid built above. JAX requires this to be static across samples.\n", - "#\n", - "# `ConstantSplit` is used for this first-pass model because adapt data (per-galaxy images from a\n", - "# previous search) is not yet available. A SLaM pipeline can later upgrade to `AdaptSplit` once the\n", - "# source has been imaged \u2014 see `scripts/group/slam.py` for the canonical chained pattern.\n", - "\n", - "pix = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.ConstantSplit,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pix)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\" / \"pixelization\",\n", - " name=\"delaunay\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20,\n", - " iterations_per_quick_update=10000,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "The image-plane mesh grid is paired with the source galaxy via `AdaptImages`, keyed by the model\n", - "path so it resolves at instance time during the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "adapt_images = al.AdaptImages(\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " settings=al.Settings(use_mixed_precision=True),\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated Delaunay pixelization for group-scale lenses. The Delaunay mesh provides an\n", - "irregular, adaptive triangulation that follows the source morphology, making it well-suited for the\n", - "complex magnification patterns produced by group-scale mass distributions.\n", - "\n", - "For automated modeling, the SLaM pipeline uses a Hilbert mesh by default, but the Delaunay mesh can be\n", - "substituted by changing the mesh class in the source pixelization pipeline stages." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Delaunay Pixelization (Group)\n", + "=============================\n", + "\n", + "This script demonstrates using a Delaunay triangulation mesh for source reconstruction in a group-scale\n", + "strong lens, as an alternative to the rectangular mesh used in other examples.\n", + "\n", + "The Delaunay mesh has several unique properties:\n", + "\n", + " - **Adaptive Mesh**: In the source plane, the Delaunay mesh uses irregularly shaped triangles rather than\n", + " uniform rectangular pixels. This allows the mesh to better adapt to irregular and asymmetric source\n", + " morphologies.\n", + "\n", + " - **Image Mesh**: The vertices of the Delaunay triangles are computed by overlaying a coarse uniform grid\n", + " in the image plane and ray-tracing these coordinates to the source plane. This helps the mesh adapt\n", + " to the magnification pattern of the lens.\n", + "\n", + " - **Interpolation**: The Delaunay mesh uses barycentric interpolation within each triangle, which is\n", + " different from the bilinear interpolation used by rectangular meshes.\n", + "\n", + "For group-scale lenses, the Delaunay mesh is particularly advantageous because the complex mass\n", + "distribution (from multiple galaxies) creates an irregular magnification pattern, and the Delaunay\n", + "triangles naturally adapt to follow the source morphology in this environment.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", + "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", + "- **Fit:** Create a FitImaging with a Delaunay pixelized source for a group lens.\n", + "- **Model:** Compose a group lens model with a Delaunay pixelized source for modeling.\n", + "- **Advantages for Group Lenses:** Why Delaunay meshes are well-suited for group-scale lensing." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autoarray.inversion.plot.inversion_plots import subplot_of_mapper, subplot_mappings" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=4,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image-Plane Mesh Grid__\n", + "\n", + "The Delaunay mesh requires an image-plane mesh grid whose ray-traced positions become the Delaunay\n", + "vertices in the source plane. We build it via an `Overlay` image-mesh covering the masked field, then\n", + "append a ring of edge pixels at the mask boundary so the linear inversion can zero them out. The same\n", + "image-plane mesh grid is reused below for both the concrete fit and the modeling section." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "edge_pixels_total = 30\n", + "\n", + "image_mesh = al.image_mesh.Overlay(shape=(22, 22))\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=dataset.mask)\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We create a fit using a Delaunay mesh whose source-pixel count matches the image-plane mesh grid\n", + "computed above, with constant regularization.\n", + "\n", + "For the group lens, we include the main lens galaxy and extra galaxies with their light and mass\n", + "profiles, and a source galaxy with the Delaunay pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(\n", + " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid}\n", + ")\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer, adapt_images=adapt_images)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit subplot shows the pixelized source does a good job at capturing the appearance of the lensed\n", + "source galaxy across the wide field of the group lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(f\"Log Likelihood: {fit.log_likelihood}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Inversion Visualization__\n", + "\n", + "The Delaunay source reconstruction can be visualized using bespoke inversion plots." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "\n", + "subplot_of_mapper(inversion=inversion, mapper_index=0)\n", + "subplot_mappings(inversion=inversion, pixelization_index=0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstructed Source__\n", + "\n", + "The reconstructed source pixel fluxes and their (y,x) positions in the source plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = inversion.linear_obj_list[0]\n", + "\n", + "reconstruction = inversion.reconstruction\n", + "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", + "\n", + "print(f\"Number of source pixels: {len(reconstruction)}\")\n", + "print(f\"Total source flux: {np.sum(reconstruction)} e- s^-1\")\n", + "print(f\"Source plane mesh grid: {source_plane_mesh_grid}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Advantages for Group Lenses__\n", + "\n", + "The Delaunay mesh is particularly well-suited for group-scale lensing because:\n", + "\n", + " 1. **Irregular magnification**: The combined mass of multiple galaxies creates a complex magnification\n", + " pattern with multiple critical curves. The Delaunay mesh naturally adapts to this by placing more\n", + " triangles in regions of high magnification (where many image pixels map to a small source-plane area).\n", + "\n", + " 2. **Extended arcs**: Group lenses often produce extended arcs that span a large fraction of the image.\n", + " The Delaunay mesh can efficiently cover these extended structures with triangles of varying size.\n", + "\n", + " 3. **Multiple images**: The mass distribution of a group lens can produce many distinct images of the\n", + " source. The Delaunay mesh places vertices based on the image-plane grid, so it naturally creates\n", + " source pixels wherever images of the source appear.\n", + "\n", + " 4. **Source complexity**: The lensed sources in group systems are often complex (e.g. merging galaxies,\n", + " star-forming clumps) and benefit from the adaptive triangle sizes of the Delaunay mesh." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Model__\n", + "\n", + "We compose a group lens model with a Delaunay source for full modeling.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source: Delaunay pixelization with ConstantSplit regularization.\n", + "#\n", + "# The Delaunay mesh is constructed as a concrete instance (not `af.Model`) because its `pixels` count\n", + "# is fixed by the image-plane mesh grid built above. JAX requires this to be static across samples.\n", + "#\n", + "# `ConstantSplit` is used for this first-pass model because adapt data (per-galaxy images from a\n", + "# previous search) is not yet available. A SLaM pipeline can later upgrade to `AdaptSplit` once the\n", + "# source has been imaged \u2014 see `scripts/group/slam.py` for the canonical chained pattern.\n", + "\n", + "pix = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.ConstantSplit,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pix)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\" / \"pixelization\",\n", + " name=\"delaunay\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20,\n", + " iterations_per_quick_update=10000,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "The image-plane mesh grid is paired with the source galaxy via `AdaptImages`, keyed by the model\n", + "path so it resolves at instance time during the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "adapt_images = al.AdaptImages(\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " settings=al.Settings(use_mixed_precision=True),\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated Delaunay pixelization for group-scale lenses. The Delaunay mesh provides an\n", + "irregular, adaptive triangulation that follows the source morphology, making it well-suited for the\n", + "complex magnification patterns produced by group-scale mass distributions.\n", + "\n", + "For automated modeling, the SLaM pipeline uses a Hilbert mesh by default, but the Delaunay mesh can be\n", + "substituted by changing the mesh class in the source pixelization pipeline stages." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/pixelization/fit.ipynb b/notebooks/group/features/pixelization/fit.ipynb index 2784a7926..93f0d4d39 100644 --- a/notebooks/group/features/pixelization/fit.ipynb +++ b/notebooks/group/features/pixelization/fit.ipynb @@ -1,495 +1,532 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fit Features: Pixelization (Group)\n", - "==================================\n", - "\n", - "This script demonstrates how to create a `FitImaging` object for a group-scale strong lens where the source\n", - "galaxy is reconstructed using a pixelized mesh, rather than parametric light profiles.\n", - "\n", - "For group-scale lenses, the fit automatically handles contributions from all galaxies: the main lens galaxy,\n", - "extra galaxies, and the pixelized source. The combined deflection field from all mass profiles is used to\n", - "ray-trace image pixels to the source plane, where the Delaunay mesh reconstructs the source emission.\n", - "\n", - "This example uses concrete (non-model) galaxy objects to illustrate the API, rather than performing a\n", - "full model-fit.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", - "- **Galaxy Centres:** Load centres for main lens and extra galaxies from JSON files.\n", - "- **Fitting:** Create a FitImaging with a pixelized source for a group lens.\n", - "- **Inversion:** Inspect the pixelized source reconstruction via the inversion object.\n", - "- **Over Sampling:** Adaptive over-sampling at all galaxy centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autoarray.inversion.plot.inversion_plots import subplot_of_mapper, subplot_mappings" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask for the group-scale lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Load centres for the main lens galaxies and extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "We apply adaptive over-sampling at all galaxy centres for the lens light profiles, and a uniform\n", - "over-sampling for the pixelization grid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=4,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image-Plane Mesh Grid__\n", - "\n", - "The Delaunay mesh requires an image-plane mesh grid whose ray-traced positions become the Delaunay\n", - "vertices in the source plane. We build it via an `Overlay` image-mesh covering the masked field, then\n", - "append a ring of edge pixels at the mask boundary so the linear inversion can zero them out." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "edge_pixels_total = 30\n", - "\n", - "image_mesh = al.image_mesh.Overlay(shape=(22, 22))\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=dataset.mask)\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fitting__\n", - "\n", - "We create concrete galaxy objects for the group lens system. The main lens and extra galaxies have\n", - "light and mass profiles, while the source galaxy uses a pixelized reconstruction.\n", - "\n", - "The `Pixelization` uses a `Delaunay` mesh whose pixel count matches the image-plane mesh grid\n", - "computed above (with `edge_pixels_total` boundary vertices reserved for zeroing) and `Constant`\n", - "regularization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We create a tracer including all group galaxies and the pixelized source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `FitImaging` object handles the full group-scale fit automatically: it sums the lens light from all\n", - "galaxies, computes the combined deflection field from all mass profiles, ray-traces to the source plane,\n", - "performs the pixelized source reconstruction, convolves with the PSF, and computes the log likelihood.\n", - "\n", - "The image-plane mesh grid is supplied via an `AdaptImages` object, which pairs it with the source\n", - "galaxy so the Delaunay vertices can be ray-traced to the source plane during the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "adapt_images = al.AdaptImages(\n", - " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid}\n", - ")\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer, adapt_images=adapt_images)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(f\"Log Likelihood: {fit.log_likelihood}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Inversion__\n", - "\n", - "The pixelized source reconstruction is stored in the `inversion` attribute of the fit. This contains the\n", - "reconstructed source pixel fluxes, the mapping between image and source pixels, and diagnostic quantities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "\n", - "print(f\"Inversion Object: {inversion}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The reconstructed source pixel fluxes are available as a 1D array." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = inversion.reconstruction\n", - "\n", - "print(f\"Reconstructed Source Pixel Fluxes: {reconstruction}\")\n", - "print(f\"Total Source Flux: {np.sum(reconstruction)} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Bespoke pixelization visualizations show the source reconstruction and image-source mapping." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "subplot_of_mapper(inversion=inversion, mapper_index=0)\n", - "subplot_mappings(inversion=inversion, pixelization_index=0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Algebra Matrices__\n", - "\n", - "The inversion exposes the linear algebra matrices used for the source reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.curvature_matrix)\n", - "print(inversion.regularization_matrix)\n", - "print(inversion.curvature_reg_matrix)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Plane Quantities__\n", - "\n", - "For a group-scale fit with a pixelized source, the model images of each plane are accessible. The first\n", - "plane contains all lens galaxies (main + extra), the second plane contains the pixelized source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_images_of_planes_list[0].slim)\n", - "print(fit.model_images_of_planes_list[1].slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Subtracted images isolate the emission of each plane. For example, the second subtracted image has all\n", - "lens galaxy light removed, leaving only the lensed source emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.subtracted_images_of_planes_list[0].slim)\n", - "print(fit.subtracted_images_of_planes_list[1].slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Figures of Merit__\n", - "\n", - "The fit provides single-valued figures of merit. For a pixelized source, the log likelihood includes\n", - "the Bayesian evidence from the regularization, which penalizes overly complex source reconstructions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Chi Squared: {fit.chi_squared}\")\n", - "print(f\"Noise Normalization: {fit.noise_normalization}\")\n", - "print(f\"Log Likelihood: {fit.log_likelihood}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated how to create a FitImaging with a pixelized source for a group-scale lens.\n", - "\n", - "The key points are:\n", - "\n", - " - The FitImaging automatically handles all group galaxies when computing the fit.\n", - " - The combined deflection field from all mass profiles is used for ray-tracing to the source plane.\n", - " - The inversion object provides access to the source reconstruction, mapping matrices, and evidence terms.\n", - " - Plane quantities allow isolating the contributions of lens galaxies vs. the pixelized source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fit Features: Pixelization (Group)\n", + "==================================\n", + "\n", + "This script demonstrates how to create a `FitImaging` object for a group-scale strong lens where the source\n", + "galaxy is reconstructed using a pixelized mesh, rather than parametric light profiles.\n", + "\n", + "For group-scale lenses, the fit automatically handles contributions from all galaxies: the main lens galaxy,\n", + "extra galaxies, and the pixelized source. The combined deflection field from all mass profiles is used to\n", + "ray-trace image pixels to the source plane, where the Delaunay mesh reconstructs the source emission.\n", + "\n", + "This example uses concrete (non-model) galaxy objects to illustrate the API, rather than performing a\n", + "full model-fit.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", + "- **Galaxy Centres:** Load centres for main lens and extra galaxies from JSON files.\n", + "- **Fitting:** Create a FitImaging with a pixelized source for a group lens.\n", + "- **Inversion:** Inspect the pixelized source reconstruction via the inversion object.\n", + "- **Over Sampling:** Adaptive over-sampling at all galaxy centres." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autoarray.inversion.plot.inversion_plots import subplot_of_mapper, subplot_mappings" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask for the group-scale lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Load centres for the main lens galaxies and extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "We apply adaptive over-sampling at all galaxy centres for the lens light profiles, and a uniform\n", + "over-sampling for the pixelization grid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=4,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image-Plane Mesh Grid__\n", + "\n", + "The Delaunay mesh requires an image-plane mesh grid whose ray-traced positions become the Delaunay\n", + "vertices in the source plane. We build it via an `Overlay` image-mesh covering the masked field, then\n", + "append a ring of edge pixels at the mask boundary so the linear inversion can zero them out." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "edge_pixels_total = 30\n", + "\n", + "image_mesh = al.image_mesh.Overlay(shape=(22, 22))\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=dataset.mask)\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fitting__\n", + "\n", + "We create concrete galaxy objects for the group lens system. The main lens and extra galaxies have\n", + "light and mass profiles, while the source galaxy uses a pixelized reconstruction.\n", + "\n", + "The `Pixelization` uses a `Delaunay` mesh whose pixel count matches the image-plane mesh grid\n", + "computed above (with `edge_pixels_total` boundary vertices reserved for zeroing) and `Constant`\n", + "regularization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We create a tracer including all group galaxies and the pixelized source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `FitImaging` object handles the full group-scale fit automatically: it sums the lens light from all\n", + "galaxies, computes the combined deflection field from all mass profiles, ray-traces to the source plane,\n", + "performs the pixelized source reconstruction, convolves with the PSF, and computes the log likelihood.\n", + "\n", + "The image-plane mesh grid is supplied via an `AdaptImages` object, which pairs it with the source\n", + "galaxy so the Delaunay vertices can be ray-traced to the source plane during the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "adapt_images = al.AdaptImages(\n", + " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid}\n", + ")\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer, adapt_images=adapt_images)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(f\"Log Likelihood: {fit.log_likelihood}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Inversion__\n", + "\n", + "The pixelized source reconstruction is stored in the `inversion` attribute of the fit. This contains the\n", + "reconstructed source pixel fluxes, the mapping between image and source pixels, and diagnostic quantities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "\n", + "print(f\"Inversion Object: {inversion}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The reconstructed source pixel fluxes are available as a 1D array." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = inversion.reconstruction\n", + "\n", + "print(f\"Reconstructed Source Pixel Fluxes: {reconstruction}\")\n", + "print(f\"Total Source Flux: {np.sum(reconstruction)} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Bespoke pixelization visualizations show the source reconstruction and image-source mapping." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "subplot_of_mapper(inversion=inversion, mapper_index=0)\n", + "subplot_mappings(inversion=inversion, pixelization_index=0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Algebra Matrices__\n", + "\n", + "The inversion exposes the linear algebra matrices used for the source reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.curvature_matrix)\n", + "print(inversion.regularization_matrix)\n", + "print(inversion.curvature_reg_matrix)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Plane Quantities__\n", + "\n", + "For a group-scale fit with a pixelized source, the model images of each plane are accessible. The first\n", + "plane contains all lens galaxies (main + extra), the second plane contains the pixelized source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_images_of_planes_list[0].slim)\n", + "print(fit.model_images_of_planes_list[1].slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Subtracted images isolate the emission of each plane. For example, the second subtracted image has all\n", + "lens galaxy light removed, leaving only the lensed source emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.subtracted_images_of_planes_list[0].slim)\n", + "print(fit.subtracted_images_of_planes_list[1].slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Figures of Merit__\n", + "\n", + "The fit provides single-valued figures of merit. For a pixelized source, the log likelihood includes\n", + "the Bayesian evidence from the regularization, which penalizes overly complex source reconstructions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Chi Squared: {fit.chi_squared}\")\n", + "print(f\"Noise Normalization: {fit.noise_normalization}\")\n", + "print(f\"Log Likelihood: {fit.log_likelihood}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated how to create a FitImaging with a pixelized source for a group-scale lens.\n", + "\n", + "The key points are:\n", + "\n", + " - The FitImaging automatically handles all group galaxies when computing the fit.\n", + " - The combined deflection field from all mass profiles is used for ray-tracing to the source plane.\n", + " - The inversion object provides access to the source reconstruction, mapping matrices, and evidence terms.\n", + " - Plane quantities allow isolating the contributions of lens galaxies vs. the pixelized source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/pixelization/likelihood_function.ipynb b/notebooks/group/features/pixelization/likelihood_function.ipynb index 37b4cabec..1b36cd8d5 100644 --- a/notebooks/group/features/pixelization/likelihood_function.ipynb +++ b/notebooks/group/features/pixelization/likelihood_function.ipynb @@ -1,592 +1,629 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Pixelization (Group)__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data of\n", - "a group-scale strong lens where the source galaxy is reconstructed using a pixelized mesh.\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers, so that readers can understand the pixelized\n", - " likelihood function for group-scale lenses without having to read the source code.\n", - "\n", - " - To illustrate how group-scale lensing with a pixelized source differs from galaxy-scale lensing: multiple\n", - " mass profiles from multiple galaxies contribute to the deflection angles used for ray-tracing, and the\n", - " pixelized source reconstruction involves linear algebra (the mapping matrix, regularization, etc.).\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", - "- **Lens Galaxies:** Define the main lens and extra galaxies with light and mass profiles.\n", - "- **Source Galaxy Pixelization:** The source uses a Delaunay mesh with constant regularization.\n", - "- **Lens Light:** Compute the total lens light from all galaxies.\n", - "- **Deflection Angles:** Compute deflection angles from all mass profiles.\n", - "- **Ray Tracing:** Ray-trace image pixels to the source plane using combined deflections.\n", - "- **Pixelized Source Reconstruction:** The linear algebra inversion step.\n", - "- **Likelihood Function:** Compute the log likelihood including regularization evidence terms.\n", - "- **Fit:** Confirm the step-by-step calculation matches the FitImaging object.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function of a pixelization builds on that used for standard light profiles and\n", - "linear light profiles. You should first read:\n", - "\n", - "- `group/likelihood_function.py` (group-scale parametric likelihood).\n", - "- `imaging/features/pixelization/likelihood_function.py` (galaxy-scale pixelized likelihood)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autoarray as aa\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the group-scale strong lens dataset with 0.1 arcsecond-per-pixel resolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"group\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5\" circular mask for the group-scale lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=masked_dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "For simplicity in this step-by-step guide, we disable over sampling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_dataset = masked_dataset.apply_over_sampling(\n", - " over_sample_size_lp=1,\n", - " over_sample_size_pixelization=1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxy__\n", - "\n", - "The main lens galaxy has a spherical Sersic light profile and an isothermal mass profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "The extra galaxies each have their own light and mass profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image-Plane Mesh Grid__\n", - "\n", - "The Delaunay mesh requires an image-plane mesh grid whose ray-traced positions become the Delaunay\n", - "vertices in the source plane. We build it via an `Overlay` image-mesh covering the masked field, then\n", - "append a ring of edge pixels at the mask boundary so the linear inversion can zero them out." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "edge_pixels_total = 30\n", - "\n", - "image_mesh = al.image_mesh.Overlay(shape=(22, 22))\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=masked_dataset.mask)\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Pixelization__\n", - "\n", - "The source galaxy is reconstructed using a Delaunay mesh with constant regularization, rather than\n", - "an analytic light profile.\n", - "\n", - "The `Pixelization` consists of:\n", - "\n", - " - `mesh`: A `Delaunay` triangulation whose source-pixel count matches the image-plane mesh grid\n", - " computed above (with `edge_pixels_total` boundary vertices reserved for zeroing). The triangle\n", - " vertices are determined by ray-tracing this image-plane grid to the source plane.\n", - "\n", - " - `regularization`: A `Constant` scheme that applies uniform smoothing across all source pixels, with\n", - " a single free regularization coefficient." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Light__\n", - "\n", - "Compute the total lens light from ALL galaxies (main + extra) in the group.\n", - "\n", - "For group-scale lenses, the total lens light is the sum of images from every lens galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", - "extra_0_image_2d = extra_galaxy_0.image_2d_from(grid=masked_dataset.grid)\n", - "extra_1_image_2d = extra_galaxy_1.image_2d_from(grid=masked_dataset.grid)\n", - "\n", - "total_lens_image_2d = lens_image_2d + extra_0_image_2d + extra_1_image_2d\n", - "\n", - "aplt.plot_array(array=total_lens_image_2d, title=\"Total Lens Light (All Galaxies)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We also compute blurring images for PSF convolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n", - "extra_0_blurring_image_2d = extra_galaxy_0.image_2d_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")\n", - "extra_1_blurring_image_2d = extra_galaxy_1.image_2d_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")\n", - "\n", - "total_lens_blurring_image_2d = (\n", - " lens_blurring_image_2d + extra_0_blurring_image_2d + extra_1_blurring_image_2d\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Deflection Angles__\n", - "\n", - "We compute deflection angles from ALL mass profiles in the group.\n", - "\n", - "For group-scale lensing, the total deflection is:\n", - "\n", - " alpha_total = alpha_lens + alpha_extra_0 + alpha_extra_1\n", - "\n", - "Each galaxy's mass profile contributes to the total deflection, and errors in any of them lead\n", - "to incorrect source-plane coordinates." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_lens = lens_galaxy.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_extra_0 = extra_galaxy_0.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_extra_1 = extra_galaxy_1.deflections_yx_2d_from(grid=masked_dataset.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Ray-trace every 2D (y,x) coordinate from the image-plane to the source-plane using the summed\n", - "deflection angles from ALL galaxies:\n", - "\n", - " beta = theta - alpha_total(theta)\n", - "\n", - "The `Tracer` object handles this automatically." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", - "\n", - "aplt.plot_grid(grid=traced_grid, title=\"Source Plane Grid (Traced)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixelized Source Reconstruction__\n", - "\n", - "Unlike a parametric source, a pixelized source does not have a closed-form image. Instead, the source\n", - "is reconstructed via a linear algebra inversion.\n", - "\n", - "The key steps are:\n", - "\n", - " 1. **Mesh construction**: The Delaunay triangulation is built in the source plane from the ray-traced\n", - " positions of a coarse image-plane grid.\n", - "\n", - " 2. **Mapping matrix**: A matrix F is constructed where F_ij describes the fractional contribution of\n", - " source pixel j to image pixel i. For a Delaunay mesh, this uses barycentric interpolation within\n", - " each triangle.\n", - "\n", - " 3. **Data vector**: d = F^T (D / sigma^2), where D is the (lens-subtracted) data and sigma is the\n", - " noise map.\n", - "\n", - " 4. **Curvature matrix**: C = F^T diag(1/sigma^2) F, which encodes how image pixels constrain source\n", - " pixels.\n", - "\n", - " 5. **Regularization matrix**: H encodes the smoothness prior on the source. For constant regularization,\n", - " H = lambda * G, where G penalizes flux differences between neighboring source pixels.\n", - "\n", - " 6. **Inversion**: Solve (C + H) s = d for the source pixel fluxes s. A positive-only solver is used\n", - " to ensure all reconstructed fluxes are physical (non-negative).\n", - "\n", - " 7. **Model image**: The reconstructed source fluxes are mapped back to the image plane via the mapping\n", - " matrix to produce the model source image, which is then added to the lens light and convolved\n", - " with the PSF.\n", - "\n", - "For group-scale lenses, the larger mask means more image pixels, making the mapping matrix larger\n", - "and the inversion more computationally expensive." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Likelihood Function__\n", - "\n", - "The log likelihood for a pixelized source has additional terms compared to a parametric source:\n", - "\n", - " -2 ln L = chi^2 + s^T H s - log|H| + log|C + H| + N ln(2 pi sigma^2)\n", - "\n", - "Where:\n", - "\n", - " - chi^2 = sum((data - model)^2 / sigma^2) is the standard goodness-of-fit term\n", - " - s^T H s is the regularization penalty (source smoothness)\n", - " - log|H| and log|C + H| are evidence terms that balance fit quality vs. source complexity\n", - " - N ln(2 pi sigma^2) is the noise normalization\n", - "\n", - "These extra terms implement Bayesian regularization: simpler (smoother) source reconstructions are\n", - "preferred unless the data demands more complexity.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Fit__\n", - "\n", - "The `FitImaging` object performs all of the above steps automatically. We verify that it produces the\n", - "correct result for the group-scale pixelized fit.\n", - "\n", - "The image-plane mesh grid is supplied via an `AdaptImages` object, which pairs it with the source\n", - "galaxy so the Delaunay vertices can be ray-traced to the source plane during the fit.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "adapt_images = al.AdaptImages(\n", - " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid}\n", - ")\n", - "\n", - "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer, adapt_images=adapt_images)\n", - "fit_figure_of_merit = fit.figure_of_merit\n", - "print(f\"Log Likelihood: {fit_figure_of_merit}\")\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Inversion Details__\n", - "\n", - "The inversion object provides access to the individual terms of the evidence-based likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "\n", - "print(f\"Regularization Term (s^T H s): {inversion.regularization_term}\")\n", - "print(f\"log|H|: {inversion.log_det_regularization_matrix_term}\")\n", - "print(f\"log|C + H|: {inversion.log_det_curvature_reg_matrix_term}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "We have presented a step-by-step guide to the group-scale pixelized source likelihood function.\n", - "\n", - "The key differences from the galaxy-scale pixelized likelihood are:\n", - "\n", - " - Multiple lens galaxies (main + extra) each contribute light and mass profiles.\n", - " - The total deflection field is the sum of deflections from ALL galaxies in the group.\n", - " - The larger mask (7.5\") means more image pixels, making the mapping matrix and inversion larger.\n", - " - Accurate lens light subtraction from all group members is critical for a clean source reconstruction.\n", - "\n", - "The key differences from the group-scale parametric likelihood are:\n", - "\n", - " - The source is reconstructed via linear algebra (mapping matrix, regularization) rather than evaluated\n", - " from a parametric profile.\n", - " - The likelihood includes Bayesian evidence terms that penalize overly complex source reconstructions.\n", - " - A positive-only solver ensures physical (non-negative) source pixel fluxes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Pixelization (Group)__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data of\n", + "a group-scale strong lens where the source galaxy is reconstructed using a pixelized mesh.\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers, so that readers can understand the pixelized\n", + " likelihood function for group-scale lenses without having to read the source code.\n", + "\n", + " - To illustrate how group-scale lensing with a pixelized source differs from galaxy-scale lensing: multiple\n", + " mass profiles from multiple galaxies contribute to the deflection angles used for ray-tracing, and the\n", + " pixelized source reconstruction involves linear algebra (the mapping matrix, regularization, etc.).\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", + "- **Lens Galaxies:** Define the main lens and extra galaxies with light and mass profiles.\n", + "- **Source Galaxy Pixelization:** The source uses a Delaunay mesh with constant regularization.\n", + "- **Lens Light:** Compute the total lens light from all galaxies.\n", + "- **Deflection Angles:** Compute deflection angles from all mass profiles.\n", + "- **Ray Tracing:** Ray-trace image pixels to the source plane using combined deflections.\n", + "- **Pixelized Source Reconstruction:** The linear algebra inversion step.\n", + "- **Likelihood Function:** Compute the log likelihood including regularization evidence terms.\n", + "- **Fit:** Confirm the step-by-step calculation matches the FitImaging object.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function of a pixelization builds on that used for standard light profiles and\n", + "linear light profiles. You should first read:\n", + "\n", + "- `group/likelihood_function.py` (group-scale parametric likelihood).\n", + "- `imaging/features/pixelization/likelihood_function.py` (galaxy-scale pixelized likelihood)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autoarray as aa\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the group-scale strong lens dataset with 0.1 arcsecond-per-pixel resolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"group\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5\" circular mask for the group-scale lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=masked_dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "For simplicity in this step-by-step guide, we disable over sampling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_dataset = masked_dataset.apply_over_sampling(\n", + " over_sample_size_lp=1,\n", + " over_sample_size_pixelization=1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxy__\n", + "\n", + "The main lens galaxy has a spherical Sersic light profile and an isothermal mass profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "The extra galaxies each have their own light and mass profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image-Plane Mesh Grid__\n", + "\n", + "The Delaunay mesh requires an image-plane mesh grid whose ray-traced positions become the Delaunay\n", + "vertices in the source plane. We build it via an `Overlay` image-mesh covering the masked field, then\n", + "append a ring of edge pixels at the mask boundary so the linear inversion can zero them out." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "edge_pixels_total = 30\n", + "\n", + "image_mesh = al.image_mesh.Overlay(shape=(22, 22))\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=masked_dataset.mask)\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Pixelization__\n", + "\n", + "The source galaxy is reconstructed using a Delaunay mesh with constant regularization, rather than\n", + "an analytic light profile.\n", + "\n", + "The `Pixelization` consists of:\n", + "\n", + " - `mesh`: A `Delaunay` triangulation whose source-pixel count matches the image-plane mesh grid\n", + " computed above (with `edge_pixels_total` boundary vertices reserved for zeroing). The triangle\n", + " vertices are determined by ray-tracing this image-plane grid to the source plane.\n", + "\n", + " - `regularization`: A `Constant` scheme that applies uniform smoothing across all source pixels, with\n", + " a single free regularization coefficient." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Light__\n", + "\n", + "Compute the total lens light from ALL galaxies (main + extra) in the group.\n", + "\n", + "For group-scale lenses, the total lens light is the sum of images from every lens galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", + "extra_0_image_2d = extra_galaxy_0.image_2d_from(grid=masked_dataset.grid)\n", + "extra_1_image_2d = extra_galaxy_1.image_2d_from(grid=masked_dataset.grid)\n", + "\n", + "total_lens_image_2d = lens_image_2d + extra_0_image_2d + extra_1_image_2d\n", + "\n", + "aplt.plot_array(array=total_lens_image_2d, title=\"Total Lens Light (All Galaxies)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We also compute blurring images for PSF convolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n", + "extra_0_blurring_image_2d = extra_galaxy_0.image_2d_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")\n", + "extra_1_blurring_image_2d = extra_galaxy_1.image_2d_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")\n", + "\n", + "total_lens_blurring_image_2d = (\n", + " lens_blurring_image_2d + extra_0_blurring_image_2d + extra_1_blurring_image_2d\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Deflection Angles__\n", + "\n", + "We compute deflection angles from ALL mass profiles in the group.\n", + "\n", + "For group-scale lensing, the total deflection is:\n", + "\n", + " alpha_total = alpha_lens + alpha_extra_0 + alpha_extra_1\n", + "\n", + "Each galaxy's mass profile contributes to the total deflection, and errors in any of them lead\n", + "to incorrect source-plane coordinates." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_lens = lens_galaxy.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_extra_0 = extra_galaxy_0.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_extra_1 = extra_galaxy_1.deflections_yx_2d_from(grid=masked_dataset.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Ray-trace every 2D (y,x) coordinate from the image-plane to the source-plane using the summed\n", + "deflection angles from ALL galaxies:\n", + "\n", + " beta = theta - alpha_total(theta)\n", + "\n", + "The `Tracer` object handles this automatically." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", + "\n", + "aplt.plot_grid(grid=traced_grid, title=\"Source Plane Grid (Traced)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixelized Source Reconstruction__\n", + "\n", + "Unlike a parametric source, a pixelized source does not have a closed-form image. Instead, the source\n", + "is reconstructed via a linear algebra inversion.\n", + "\n", + "The key steps are:\n", + "\n", + " 1. **Mesh construction**: The Delaunay triangulation is built in the source plane from the ray-traced\n", + " positions of a coarse image-plane grid.\n", + "\n", + " 2. **Mapping matrix**: A matrix F is constructed where F_ij describes the fractional contribution of\n", + " source pixel j to image pixel i. For a Delaunay mesh, this uses barycentric interpolation within\n", + " each triangle.\n", + "\n", + " 3. **Data vector**: d = F^T (D / sigma^2), where D is the (lens-subtracted) data and sigma is the\n", + " noise map.\n", + "\n", + " 4. **Curvature matrix**: C = F^T diag(1/sigma^2) F, which encodes how image pixels constrain source\n", + " pixels.\n", + "\n", + " 5. **Regularization matrix**: H encodes the smoothness prior on the source. For constant regularization,\n", + " H = lambda * G, where G penalizes flux differences between neighboring source pixels.\n", + "\n", + " 6. **Inversion**: Solve (C + H) s = d for the source pixel fluxes s. A positive-only solver is used\n", + " to ensure all reconstructed fluxes are physical (non-negative).\n", + "\n", + " 7. **Model image**: The reconstructed source fluxes are mapped back to the image plane via the mapping\n", + " matrix to produce the model source image, which is then added to the lens light and convolved\n", + " with the PSF.\n", + "\n", + "For group-scale lenses, the larger mask means more image pixels, making the mapping matrix larger\n", + "and the inversion more computationally expensive." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Likelihood Function__\n", + "\n", + "The log likelihood for a pixelized source has additional terms compared to a parametric source:\n", + "\n", + " -2 ln L = chi^2 + s^T H s - log|H| + log|C + H| + N ln(2 pi sigma^2)\n", + "\n", + "Where:\n", + "\n", + " - chi^2 = sum((data - model)^2 / sigma^2) is the standard goodness-of-fit term\n", + " - s^T H s is the regularization penalty (source smoothness)\n", + " - log|H| and log|C + H| are evidence terms that balance fit quality vs. source complexity\n", + " - N ln(2 pi sigma^2) is the noise normalization\n", + "\n", + "These extra terms implement Bayesian regularization: simpler (smoother) source reconstructions are\n", + "preferred unless the data demands more complexity.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Fit__\n", + "\n", + "The `FitImaging` object performs all of the above steps automatically. We verify that it produces the\n", + "correct result for the group-scale pixelized fit.\n", + "\n", + "The image-plane mesh grid is supplied via an `AdaptImages` object, which pairs it with the source\n", + "galaxy so the Delaunay vertices can be ray-traced to the source plane during the fit.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "adapt_images = al.AdaptImages(\n", + " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid}\n", + ")\n", + "\n", + "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer, adapt_images=adapt_images)\n", + "fit_figure_of_merit = fit.figure_of_merit\n", + "print(f\"Log Likelihood: {fit_figure_of_merit}\")\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Inversion Details__\n", + "\n", + "The inversion object provides access to the individual terms of the evidence-based likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "\n", + "print(f\"Regularization Term (s^T H s): {inversion.regularization_term}\")\n", + "print(f\"log|H|: {inversion.log_det_regularization_matrix_term}\")\n", + "print(f\"log|C + H|: {inversion.log_det_curvature_reg_matrix_term}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "We have presented a step-by-step guide to the group-scale pixelized source likelihood function.\n", + "\n", + "The key differences from the galaxy-scale pixelized likelihood are:\n", + "\n", + " - Multiple lens galaxies (main + extra) each contribute light and mass profiles.\n", + " - The total deflection field is the sum of deflections from ALL galaxies in the group.\n", + " - The larger mask (7.5\") means more image pixels, making the mapping matrix and inversion larger.\n", + " - Accurate lens light subtraction from all group members is critical for a clean source reconstruction.\n", + "\n", + "The key differences from the group-scale parametric likelihood are:\n", + "\n", + " - The source is reconstructed via linear algebra (mapping matrix, regularization) rather than evaluated\n", + " from a parametric profile.\n", + " - The likelihood includes Bayesian evidence terms that penalize overly complex source reconstructions.\n", + " - A positive-only solver ensures physical (non-negative) source pixel fluxes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/pixelization/modeling.ipynb b/notebooks/group/features/pixelization/modeling.ipynb index a930c930d..1c202c285 100644 --- a/notebooks/group/features/pixelization/modeling.ipynb +++ b/notebooks/group/features/pixelization/modeling.ipynb @@ -1,511 +1,548 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Pixelization (Group)\n", - "=======================================\n", - "\n", - "This script fits a group-scale strong lens where the source galaxy is reconstructed using a pixelized mesh\n", - "(Delaunay triangulation) with adaptive regularization, rather than parametric light profiles.\n", - "\n", - "For group-scale lenses, pixelized source reconstructions are especially valuable because the lensed source\n", - "morphology is often complex, with multiple extended arcs and counter-images produced by the combined mass of\n", - "several galaxies. A parametric source model (e.g. Sersic or MGE) may struggle to capture these features,\n", - "whereas a pixelized mesh can reconstruct arbitrarily complex source-plane emission.\n", - "\n", - "The main lens galaxies and extra galaxies are modeled with MGE light profiles (via ``al.model_util.mge_model_from``)\n", - "and Isothermal mass profiles, following the standard group modeling pattern. The source galaxy uses a\n", - "Delaunay pixelization with ConstantSplit regularization (`AdaptSplit` requires adapt data from a\n", - "prior search and is wired up later by the SLaM pipeline \u2014 see `scripts/group/slam.py`).\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", - "- **Galaxy Centres:** Load centres for main lens and extra galaxies from JSON files.\n", - "- **Model:** Compose the group lens model with a pixelized source.\n", - "- **Over Sampling:** Adaptive over-sampling at all galaxy centres.\n", - "- **Search:** Configure the non-linear search.\n", - "- **Analysis:** Create the Analysis object with pixelization settings.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Example__\n", - "\n", - "This script fits an `Imaging` dataset of a 'group-scale' strong lens where:\n", - "\n", - " - There is a main lens galaxy whose light is an MGE and total mass is `Isothermal` with `ExternalShear`.\n", - " - There are two extra lens galaxies with MGE light and `IsothermalSph` mass, centres fixed.\n", - " - The source galaxy's light is reconstructed using a `Delaunay` mesh with `ConstantSplit` regularization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`, which is the dataset we will use to perform lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5 arcsecond circular mask, which is larger than galaxy-scale lenses because group-scale systems\n", - "have lensed emission spread over a wider area." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Load the centres of the main lens galaxies and extra galaxies from JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image-Plane Mesh Grid__\n", - "\n", - "The Delaunay mesh requires an image-plane mesh grid whose ray-traced positions become the Delaunay\n", - "vertices in the source plane. We build it via an `Overlay` image-mesh covering the masked field, then\n", - "append a ring of edge pixels at the mask boundary so the linear inversion can zero them out.\n", - "\n", - "This grid is reused below: its shape sets the number of source pixels in the mesh, and the grid\n", - "itself is paired with the source galaxy via `AdaptImages` when the analysis is constructed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "edge_pixels_total = 30\n", - "\n", - "image_mesh = al.image_mesh.Overlay(shape=(22, 22))\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=dataset.mask)\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a group lens model where:\n", - "\n", - " - Each main lens galaxy has MGE light and Isothermal mass. Only lens_0 carries ExternalShear.\n", - " - Each extra galaxy has MGE light and IsothermalSph mass with fixed centres and bounded Einstein radii.\n", - " - The source galaxy uses a Delaunay pixelization with ConstantSplit regularization. `ConstantSplit`\n", - " is used for this first-pass model because adapt data (per-galaxy images from a previous search)\n", - " is not yet available; a SLaM pipeline can later upgrade to `AdaptSplit` once the source has been\n", - " imaged \u2014 see `scripts/group/slam.py` for the chained pattern.\n", - "\n", - "The pixelized source captures complex lensed morphologies far better than parametric profiles, which is\n", - "especially important for group lenses where the extended mass distribution creates intricate arc structures." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = lens\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source:\n", - "#\n", - "# The Delaunay mesh is built as a concrete instance (not `af.Model`) because its `pixels` count is\n", - "# fixed by the image-plane mesh grid built above. JAX requires this to be static across samples.\n", - "\n", - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.ConstantSplit,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format, confirming the pixelized source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) ensures the lens light\n", - "profiles are accurately evaluated across the full field of the group.\n", - "\n", - "For the pixelization, a separate uniform over-sampling is applied via `over_sample_size_pixelization`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=4,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "We use Nautilus with `n_live=100` and `n_batch=20`, suitable for a group-scale pixelization model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"pixelization\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "The `AnalysisImaging` object defines how the model is fitted to the data.\n", - "\n", - "For pixelized source fits, we enable mixed precision to speed up GPU run times on consumer hardware.\n", - "\n", - "The image-plane mesh grid is paired with the source galaxy via `AdaptImages`, keyed by the model\n", - "path so it resolves at instance time during the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "adapt_images = al.AdaptImages(\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " settings=al.Settings(use_mixed_precision=True),\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "Group-scale pixelization fits are more computationally expensive than galaxy-scale fits because the larger\n", - "7.5\" mask contains many more image pixels, all of which must be mapped to source pixels. GPU acceleration\n", - "is strongly recommended." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - " \"\"\"\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result contains the best-fit pixelized source reconstruction alongside the group lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated how to model a group-scale strong lens with a pixelized source reconstruction.\n", - "\n", - "The pixelized source is particularly powerful for group lenses because:\n", - "\n", - " - The combined mass of multiple galaxies creates complex arc structures that parametric profiles cannot capture.\n", - " - The Delaunay mesh adapts to the source morphology, placing more triangles where the source is brightest.\n", - " - `ConstantSplit` regularization is used as a first-pass scheme; chained pipelines can upgrade to\n", - " `AdaptSplit` once an adapt image of the source is available, allowing different smoothing in\n", - " bright vs. faint source regions.\n", - "\n", - "For automated modeling of large samples, the SLaM pipeline (see `group/features/pixelization/slam.py`) provides\n", - "a robust framework that chains parametric and pixelized source fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Pixelization (Group)\n", + "=======================================\n", + "\n", + "This script fits a group-scale strong lens where the source galaxy is reconstructed using a pixelized mesh\n", + "(Delaunay triangulation) with adaptive regularization, rather than parametric light profiles.\n", + "\n", + "For group-scale lenses, pixelized source reconstructions are especially valuable because the lensed source\n", + "morphology is often complex, with multiple extended arcs and counter-images produced by the combined mass of\n", + "several galaxies. A parametric source model (e.g. Sersic or MGE) may struggle to capture these features,\n", + "whereas a pixelized mesh can reconstruct arbitrarily complex source-plane emission.\n", + "\n", + "The main lens galaxies and extra galaxies are modeled with MGE light profiles (via ``al.model_util.mge_model_from``)\n", + "and Isothermal mass profiles, following the standard group modeling pattern. The source galaxy uses a\n", + "Delaunay pixelization with ConstantSplit regularization (`AdaptSplit` requires adapt data from a\n", + "prior search and is wired up later by the SLaM pipeline \u2014 see `scripts/group/slam.py`).\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", + "- **Galaxy Centres:** Load centres for main lens and extra galaxies from JSON files.\n", + "- **Model:** Compose the group lens model with a pixelized source.\n", + "- **Over Sampling:** Adaptive over-sampling at all galaxy centres.\n", + "- **Search:** Configure the non-linear search.\n", + "- **Analysis:** Create the Analysis object with pixelization settings.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Example__\n", + "\n", + "This script fits an `Imaging` dataset of a 'group-scale' strong lens where:\n", + "\n", + " - There is a main lens galaxy whose light is an MGE and total mass is `Isothermal` with `ExternalShear`.\n", + " - There are two extra lens galaxies with MGE light and `IsothermalSph` mass, centres fixed.\n", + " - The source galaxy's light is reconstructed using a `Delaunay` mesh with `ConstantSplit` regularization." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`, which is the dataset we will use to perform lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5 arcsecond circular mask, which is larger than galaxy-scale lenses because group-scale systems\n", + "have lensed emission spread over a wider area." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Load the centres of the main lens galaxies and extra galaxies from JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image-Plane Mesh Grid__\n", + "\n", + "The Delaunay mesh requires an image-plane mesh grid whose ray-traced positions become the Delaunay\n", + "vertices in the source plane. We build it via an `Overlay` image-mesh covering the masked field, then\n", + "append a ring of edge pixels at the mask boundary so the linear inversion can zero them out.\n", + "\n", + "This grid is reused below: its shape sets the number of source pixels in the mesh, and the grid\n", + "itself is paired with the source galaxy via `AdaptImages` when the analysis is constructed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "edge_pixels_total = 30\n", + "\n", + "image_mesh = al.image_mesh.Overlay(shape=(22, 22))\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=dataset.mask)\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a group lens model where:\n", + "\n", + " - Each main lens galaxy has MGE light and Isothermal mass. Only lens_0 carries ExternalShear.\n", + " - Each extra galaxy has MGE light and IsothermalSph mass with fixed centres and bounded Einstein radii.\n", + " - The source galaxy uses a Delaunay pixelization with ConstantSplit regularization. `ConstantSplit`\n", + " is used for this first-pass model because adapt data (per-galaxy images from a previous search)\n", + " is not yet available; a SLaM pipeline can later upgrade to `AdaptSplit` once the source has been\n", + " imaged \u2014 see `scripts/group/slam.py` for the chained pattern.\n", + "\n", + "The pixelized source captures complex lensed morphologies far better than parametric profiles, which is\n", + "especially important for group lenses where the extended mass distribution creates intricate arc structures." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = lens\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source:\n", + "#\n", + "# The Delaunay mesh is built as a concrete instance (not `af.Model`) because its `pixels` count is\n", + "# fixed by the image-plane mesh grid built above. JAX requires this to be static across samples.\n", + "\n", + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.ConstantSplit,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format, confirming the pixelized source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) ensures the lens light\n", + "profiles are accurately evaluated across the full field of the group.\n", + "\n", + "For the pixelization, a separate uniform over-sampling is applied via `over_sample_size_pixelization`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=4,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "We use Nautilus with `n_live=100` and `n_batch=20`, suitable for a group-scale pixelization model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"pixelization\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "The `AnalysisImaging` object defines how the model is fitted to the data.\n", + "\n", + "For pixelized source fits, we enable mixed precision to speed up GPU run times on consumer hardware.\n", + "\n", + "The image-plane mesh grid is paired with the source galaxy via `AdaptImages`, keyed by the model\n", + "path so it resolves at instance time during the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "adapt_images = al.AdaptImages(\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " settings=al.Settings(use_mixed_precision=True),\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "Group-scale pixelization fits are more computationally expensive than galaxy-scale fits because the larger\n", + "7.5\" mask contains many more image pixels, all of which must be mapped to source pixels. GPU acceleration\n", + "is strongly recommended." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result contains the best-fit pixelized source reconstruction alongside the group lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated how to model a group-scale strong lens with a pixelized source reconstruction.\n", + "\n", + "The pixelized source is particularly powerful for group lenses because:\n", + "\n", + " - The combined mass of multiple galaxies creates complex arc structures that parametric profiles cannot capture.\n", + " - The Delaunay mesh adapts to the source morphology, placing more triangles where the source is brightest.\n", + " - `ConstantSplit` regularization is used as a first-pass scheme; chained pipelines can upgrade to\n", + " `AdaptSplit` once an adapt image of the source is available, allowing different smoothing in\n", + " bright vs. faint source regions.\n", + "\n", + "For automated modeling of large samples, the SLaM pipeline (see `group/features/pixelization/slam.py`) provides\n", + "a robust framework that chains parametric and pixelized source fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/pixelization/slam.ipynb b/notebooks/group/features/pixelization/slam.ipynb index 94c02e67c..553dc6844 100644 --- a/notebooks/group/features/pixelization/slam.ipynb +++ b/notebooks/group/features/pixelization/slam.ipynb @@ -1,1021 +1,1058 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: Group SLaM\n", - "=========================\n", - "\n", - "This script uses the SLaM (Source, Light and Mass) pipelines to fit a group-scale strong lens where the\n", - "source galaxy is reconstructed using a pixelized mesh with adaptive regularization.\n", - "\n", - "This is essentially the same as the main `group/slam.py` script, since the group SLaM pipeline already\n", - "uses a Delaunay pixelization as its default source model. This feature script documents the\n", - "pixelization-specific choices and parameters used in the pipeline.\n", - "\n", - "__Pixelization Choices__\n", - "\n", - "The SLaM pipeline makes the following pixelization-specific decisions:\n", - "\n", - " - **Hilbert mesh**: The SOURCE PIX pipelines use a `Hilbert` image mesh, which distributes source pixels\n", - " along a space-filling Hilbert curve. The pixel count is set automatically based on the data's pixel scale\n", - " via `al.model_util.hilbert_pixels_from_pixel_scale`. This ensures the mesh resolution scales with data quality.\n", - "\n", - " - **Adapt regularization**: Rather than constant regularization, `Adapt` adapts the smoothing\n", - " to the source morphology. Bright source regions receive less smoothing (preserving detail) while faint\n", - " regions receive more smoothing (suppressing noise). Rectangular pixelizations use `Adapt` (not\n", - " `AdaptSplit`, which is only appropriate for irregular meshes such as Delaunay or Voronoi).\n", - "\n", - " - **Edge pixels**: The Hilbert mesh includes `edge_pixels_total=30` padding pixels at the border of the\n", - " source-plane reconstruction, preventing edge artifacts.\n", - "\n", - " - **Two-stage pixelization**: SOURCE PIX 1 establishes the pixelization with initial adapt data, then\n", - " SOURCE PIX 2 refines it using capped adapt data from the first stage. This two-stage approach\n", - " ensures the adaptive features converge to robust solutions.\n", - "\n", - " - **Signal-adaptive over-sampling**: Pixels above a signal-to-noise threshold use higher over-sampling\n", - " (sub_size=4) than faint pixels (sub_size=2), balancing accuracy and speed.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with group/slam.py and the\n", - " pixelization feature scripts.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", - "- **Mask:** Define the 2D mask applied to the dataset.\n", - "- **SLaM Pipeline:** The full SLaM pipeline with pixelization-specific documentation.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`): Introduction to SLaM pipeline structure.\n", - "- **Group SLaM** (`group/slam`): How group-scale SLaM handles extra and scaling galaxies.\n", - "- **Pixelization Modeling** (`group/features/pixelization/modeling`): Group pixelization basics.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE (two searches), SOURCE PIX PIPELINE (two searches), LIGHT LP PIPELINE\n", - "and TOTAL MASS PIPELINE this SLaM modeling script fits `Imaging` data of a group-scale strong lens\n", - "where in the final model:\n", - "\n", - " - Each main lens galaxy has a free MGE bulge and a `PowerLaw` total mass.\n", - " - Each extra galaxy has a free MGE bulge and a luminosity-bounded `Isothermal` mass.\n", - " - The source galaxy's light is a rectangular adaptive `Pixelization` with `Adapt` regularization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "\n", - "def _load_centres(path):\n", - " \"\"\"Load a centres JSON file, returning an empty list if the file is absent.\"\"\"\n", - " try:\n", - " return al.Grid2DIrregular(al.from_json(file_path=path))\n", - " except FileNotFoundError:\n", - " return al.Grid2DIrregular([])\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 0__\n", - "\n", - "Fits light only -- no mass, no source -- for every galaxy simultaneously, giving the next search\n", - "clean fixed light models to build on.\n", - "\n", - "This is identical to the group/slam.py SOURCE LP 0 pipeline." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_0(\n", - " dataset,\n", - " settings_search,\n", - " main_lens_centres,\n", - " extra_lens_centres,\n", - " mask_radius,\n", - " redshift_lens,\n", - " n_batch=50,\n", - "):\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " # --- main lens light models (one per centre, light only) ---\n", - " lens_dict = {}\n", - " for i, centre in enumerate(main_lens_centres):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=False,\n", - " centre=(centre[0], centre[1]),\n", - " centre_sigma=0.1,\n", - " )\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=bulge, disk=None, point=None\n", - " )\n", - "\n", - " # --- extra lens galaxy light models ---\n", - " extra_light_models = []\n", - " for centre in extra_lens_centres:\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_prior_is_uniform=True,\n", - " centre=(centre[0], centre[1]),\n", - " ell_comps_prior_is_uniform=True,\n", - " )\n", - " extra_light_models.append(\n", - " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", - "\n", - " n_extra = len(extra_galaxies) if extra_galaxies is not None else 0\n", - " n_live = 100 + 30 * len(lens_dict) + 30 * n_extra\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[0]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " n_like_max=1000000,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 1__\n", - "\n", - "Equivalent to `source_lp` in `slam_start_here.py`, except lens light is fixed from `source_lp[0]`\n", - "rather than free, and mass and source are introduced here for the first time.\n", - "\n", - "Multiple main-lens galaxies each get an `Isothermal` mass; only `lens_0` carries an `ExternalShear`.\n", - "Extra-galaxy Einstein radii are bounded by a luminosity-derived prior." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_1(\n", - " dataset,\n", - " settings_search,\n", - " source_lp_result_0,\n", - " positions,\n", - " pixel_scale,\n", - " redshift_lens,\n", - " redshift_source,\n", - " source_mge_radius,\n", - " n_batch=50,\n", - "):\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(source_lp_result_0.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - " n_extra = (\n", - " len(list(source_lp_result_0.instance.extra_galaxies))\n", - " if source_lp_result_0.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " tracer = (\n", - " source_lp_result_0.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", - " )\n", - "\n", - " # Source MGE centred on primary lens bulge from stage 0.\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=source_mge_radius,\n", - " total_gaussians=30,\n", - " centre=source_lp_result_0.instance.galaxies.lens_0.bulge.centre,\n", - " centre_prior_is_uniform=False,\n", - " centre_sigma=0.6,\n", - " )\n", - "\n", - " # --- main lens full models (light fixed from stage 0, mass + shear free) ---\n", - " lens_dict = {}\n", - " for i in range(n_main):\n", - " lp0_lens = getattr(source_lp_result_0.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = lp0_lens.bulge.centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=5.0)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lp0_lens.bulge,\n", - " disk=lp0_lens.disk,\n", - " point=lp0_lens.point,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " # --- extra lens galaxy models (light fixed, mass bounded by luminosity) ---\n", - " extra_mass_models = []\n", - " for i in range(n_extra):\n", - " lp0_extra = source_lp_result_0.instance.extra_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = lp0_extra.bulge.centre\n", - " mass.ell_comps = lp0_extra.bulge.ell_comps\n", - "\n", - " luminosity_per_gaussian_list = [\n", - " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", - " for g in tracer.galaxies[n_main + i].bulge.profile_list\n", - " ]\n", - " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", - " luminosity_cap = 5 * 0.5 * total_luminosity**0.6\n", - " upper_limit = min(luminosity_cap, 5.0) if luminosity_cap > 0 else 5.0\n", - " mass.einstein_radius = af.UniformPrior(\n", - " lower_limit=0.0,\n", - " upper_limit=upper_limit,\n", - " )\n", - "\n", - " extra_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=lp0_extra.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", - " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " n_extra_model = len(extra_galaxies) if extra_galaxies is not None else 0\n", - " n_live = 150 + 30 * n_main + 30 * n_extra_model\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " n_like_max=200000,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "This is where the pixelization is first introduced. The source switches from a parametric MGE to a\n", - "rectangular adaptive pixelization with `Adapt` regularization.\n", - "\n", - "Key pixelization choices:\n", - " - `al.mesh.RectangularAdaptDensity`: a rectangular mesh whose cell density adapts to the source brightness.\n", - " - Fixed `mesh_shape = (28, 28)` pixels, chosen to match data resolution.\n", - " - `al.reg.Adapt`: adaptive regularization that varies smoothing based on source brightness.\n", - " Rectangular meshes use `Adapt` (not `AdaptSplit`, which is reserved for irregular meshes).\n", - "\n", - "Signal-adaptive over-sampling is applied: pixels above the S/N threshold use sub_size=4, others sub_size=2." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " dataset,\n", - " mask,\n", - " settings_search,\n", - " source_lp_result_1,\n", - " over_sample_size,\n", - " pixel_scale,\n", - " mask_radius,\n", - " positions,\n", - " n_batch=20,\n", - "):\n", - " mesh_shape = (28, 28)\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=4,\n", - " )\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result_1\n", - " )\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result_1.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(source_lp_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " # Main lens galaxies: mass as model from previous result.\n", - " lens_dict = {}\n", - " for i in range(n_main):\n", - " prev_lens = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=getattr(source_lp_result_1.model.galaxies, f\"lens_{i}\").mass,\n", - " mass_result=getattr(source_lp_result_1.model.galaxies, f\"lens_{i}\").mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=prev_lens.redshift,\n", - " bulge=prev_lens.bulge,\n", - " disk=prev_lens.disk,\n", - " mass=mass,\n", - " shear=source_lp_result_1.model.galaxies.lens_0.shear if i == 0 else None,\n", - " )\n", - "\n", - " # Extra galaxies: carried forward as model parameters.\n", - " extra_galaxies = (\n", - " source_lp_result_1.model.extra_galaxies\n", - " if source_lp_result_1.model.extra_galaxies is not None\n", - " else None\n", - " )\n", - "\n", - " # Source: RectangularAdaptDensity mesh with Adapt regularization.\n", - " pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization=af.Model(al.reg.Adapt),\n", - " )\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", - " pixelization=pixelization,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - " return result, dataset, adapt_images\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to SOURCE PIX 1 but uses capped adapt data from the first pixelization stage, ensuring\n", - "the adaptive mesh converges to a robust solution. The lens mass is fixed from SOURCE PIX 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " dataset,\n", - " settings_search,\n", - " source_lp_result_1,\n", - " source_pix_result_1,\n", - " adapt_images,\n", - " pixel_scale,\n", - " n_batch=20,\n", - "):\n", - " mesh_shape = (28, 28)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " # Cap the adapt data to prevent extreme values from dominating the mesh.\n", - " adapt_images_capped = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images_capped,\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " # Main lens galaxies: mass fixed from SOURCE PIX 1.\n", - " lens_dict = {}\n", - " for i in range(n_main):\n", - " prev_lens = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", - " lens_dict[f\"lens_{i}\"] = prev_lens\n", - "\n", - " extra_galaxies = source_pix_result_1.instance.extra_galaxies\n", - "\n", - " pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=af.Model(al.reg.Adapt),\n", - " )\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", - " pixelization=pixelization,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Fits the lens galaxy light with mass and pixelized source fixed from the SOURCE PIX pipelines.\n", - "Extra galaxies receive a fresh free MGE light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " dataset,\n", - " settings_search,\n", - " source_pix_result_1,\n", - " source_pix_result_2,\n", - " adapt_images,\n", - " mask_radius,\n", - " redshift_lens,\n", - " n_batch=20,\n", - "):\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_main):\n", - " prev_lens = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=False,\n", - " centre=prev_lens.bulge.centre,\n", - " centre_sigma=0.1,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=bulge,\n", - " mass=prev_lens.mass,\n", - " shear=prev_lens.shear if i == 0 else None,\n", - " )\n", - "\n", - " # Extra galaxies: fresh MGE light, mass fixed.\n", - " n_extra = (\n", - " len(list(source_pix_result_1.instance.extra_galaxies))\n", - " if source_pix_result_1.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " extra_models = []\n", - " for i in range(n_extra):\n", - " prev_extra = source_pix_result_1.instance.extra_galaxies[i]\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_prior_is_uniform=True,\n", - " centre=prev_extra.bulge.centre,\n", - " ell_comps_prior_is_uniform=True,\n", - " )\n", - "\n", - " extra_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=prev_extra.mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_models) if extra_models else None\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_pix_result_2, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Fits a PowerLaw total mass model with priors from the SOURCE PIX pipeline and lens light fixed from\n", - "the LIGHT LP pipeline. The pixelized source is carried forward." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " dataset,\n", - " settings_search,\n", - " source_pix_result_1,\n", - " source_pix_result_2,\n", - " light_result,\n", - " adapt_images,\n", - " redshift_lens,\n", - " n_batch=20,\n", - "):\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_2.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(light_result.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_main):\n", - " light_lens = getattr(light_result.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = af.Model(al.mp.PowerLaw)\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=getattr(source_pix_result_1.model.galaxies, f\"lens_{i}\").mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=light_lens.bulge,\n", - " mass=mass,\n", - " shear=source_pix_result_1.model.galaxies.lens_0.shear if i == 0 else None,\n", - " )\n", - "\n", - " # Extra galaxies: fresh mass, light fixed from LIGHT LP.\n", - " n_extra = (\n", - " len(list(light_result.instance.extra_galaxies))\n", - " if light_result.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " extra_models = []\n", - " for i in range(n_extra):\n", - " light_extra = light_result.instance.extra_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = light_extra.mass.centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " extra_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=light_extra.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_models) if extra_models else None\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens group dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Load centres for main lens and extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = _load_centres(dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = _load_centres(dataset_path / \"extra_galaxies_centres.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a 7.5\" circular mask for the group-scale lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0\n", - "pixel_scale = 0.1\n", - "\n", - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"group\") / \"features\" / \"pixelization\" / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Positions are computed from the SOURCE LP result and used in subsequent pixelization searches\n", - "to prevent demagnified source solutions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__SLaM Pipeline__\n", - "\n", - "The full SLaM pipeline is executed below. Each stage is documented with its pixelization-specific choices.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# --- SOURCE LP PIPELINE 0: Light only ---\n", - "\n", - "source_lp_result_0 = source_lp_0(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " main_lens_centres=main_lens_centres,\n", - " extra_lens_centres=extra_galaxies_centres,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - ")\n", - "\n", - "# --- SOURCE LP PIPELINE 1: Add mass + parametric source ---\n", - "\n", - "positions = source_lp_result_0.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - ").positions\n", - "\n", - "source_mge_radius = mask_radius\n", - "\n", - "source_lp_result_1 = source_lp_1(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " source_lp_result_0=source_lp_result_0,\n", - " positions=positions,\n", - " pixel_scale=pixel_scale,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - " source_mge_radius=source_mge_radius,\n", - ")\n", - "\n", - "# --- SOURCE PIX PIPELINE 1: Introduce Hilbert pixelization ---\n", - "\n", - "source_pix_result_1, dataset_pix, adapt_images = source_pix_1(\n", - " dataset=dataset,\n", - " mask=mask,\n", - " settings_search=settings_search,\n", - " source_lp_result_1=source_lp_result_1,\n", - " over_sample_size=over_sample_size,\n", - " pixel_scale=pixel_scale,\n", - " mask_radius=mask_radius,\n", - " positions=positions,\n", - ")\n", - "\n", - "# --- SOURCE PIX PIPELINE 2: Refine with capped adapt data ---\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " dataset=dataset_pix,\n", - " settings_search=settings_search,\n", - " source_lp_result_1=source_lp_result_1,\n", - " source_pix_result_1=source_pix_result_1,\n", - " adapt_images=adapt_images,\n", - " pixel_scale=pixel_scale,\n", - ")\n", - "\n", - "# --- LIGHT LP PIPELINE: Refit lens light ---\n", - "\n", - "light_result = light_lp(\n", - " dataset=dataset_pix,\n", - " settings_search=settings_search,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " adapt_images=adapt_images,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - ")\n", - "\n", - "# --- MASS TOTAL PIPELINE: Fit PowerLaw mass ---\n", - "\n", - "mass_result = mass_total(\n", - " dataset=dataset_pix,\n", - " settings_search=settings_search,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " light_result=light_result,\n", - " adapt_images=adapt_images,\n", - " redshift_lens=redshift_lens,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The final result contains the full group lens model with a pixelized source reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(mass_result.info)\n", - "\n", - "aplt.subplot_fit_imaging(fit=mass_result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated the full SLaM pipeline for group-scale lenses with a pixelized source.\n", - "\n", - "The key pixelization choices are:\n", - " - Rectangular adaptive mesh with a fixed shape (28, 28).\n", - " - `Adapt` regularization for brightness-dependent smoothing.\n", - " - Two-stage pixelization refinement with capped adapt data.\n", - " - Fixed pixelization over-sampling of 4.\n", - "\n", - "These choices are the recommended defaults for science-grade group-scale lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: Group SLaM\n", + "=========================\n", + "\n", + "This script uses the SLaM (Source, Light and Mass) pipelines to fit a group-scale strong lens where the\n", + "source galaxy is reconstructed using a pixelized mesh with adaptive regularization.\n", + "\n", + "This is essentially the same as the main `group/slam.py` script, since the group SLaM pipeline already\n", + "uses a Delaunay pixelization as its default source model. This feature script documents the\n", + "pixelization-specific choices and parameters used in the pipeline.\n", + "\n", + "__Pixelization Choices__\n", + "\n", + "The SLaM pipeline makes the following pixelization-specific decisions:\n", + "\n", + " - **Hilbert mesh**: The SOURCE PIX pipelines use a `Hilbert` image mesh, which distributes source pixels\n", + " along a space-filling Hilbert curve. The pixel count is set automatically based on the data's pixel scale\n", + " via `al.model_util.hilbert_pixels_from_pixel_scale`. This ensures the mesh resolution scales with data quality.\n", + "\n", + " - **Adapt regularization**: Rather than constant regularization, `Adapt` adapts the smoothing\n", + " to the source morphology. Bright source regions receive less smoothing (preserving detail) while faint\n", + " regions receive more smoothing (suppressing noise). Rectangular pixelizations use `Adapt` (not\n", + " `AdaptSplit`, which is only appropriate for irregular meshes such as Delaunay or Voronoi).\n", + "\n", + " - **Edge pixels**: The Hilbert mesh includes `edge_pixels_total=30` padding pixels at the border of the\n", + " source-plane reconstruction, preventing edge artifacts.\n", + "\n", + " - **Two-stage pixelization**: SOURCE PIX 1 establishes the pixelization with initial adapt data, then\n", + " SOURCE PIX 2 refines it using capped adapt data from the first stage. This two-stage approach\n", + " ensures the adaptive features converge to robust solutions.\n", + "\n", + " - **Signal-adaptive over-sampling**: Pixels above a signal-to-noise threshold use higher over-sampling\n", + " (sub_size=4) than faint pixels (sub_size=2), balancing accuracy and speed.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with group/slam.py and the\n", + " pixelization feature scripts.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", + "- **Mask:** Define the 2D mask applied to the dataset.\n", + "- **SLaM Pipeline:** The full SLaM pipeline with pixelization-specific documentation.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`): Introduction to SLaM pipeline structure.\n", + "- **Group SLaM** (`group/slam`): How group-scale SLaM handles extra and scaling galaxies.\n", + "- **Pixelization Modeling** (`group/features/pixelization/modeling`): Group pixelization basics.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE (two searches), SOURCE PIX PIPELINE (two searches), LIGHT LP PIPELINE\n", + "and TOTAL MASS PIPELINE this SLaM modeling script fits `Imaging` data of a group-scale strong lens\n", + "where in the final model:\n", + "\n", + " - Each main lens galaxy has a free MGE bulge and a `PowerLaw` total mass.\n", + " - Each extra galaxy has a free MGE bulge and a luminosity-bounded `Isothermal` mass.\n", + " - The source galaxy's light is a rectangular adaptive `Pixelization` with `Adapt` regularization." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "\n", + "def _load_centres(path):\n", + " \"\"\"Load a centres JSON file, returning an empty list if the file is absent.\"\"\"\n", + " try:\n", + " return al.Grid2DIrregular(al.from_json(file_path=path))\n", + " except FileNotFoundError:\n", + " return al.Grid2DIrregular([])\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 0__\n", + "\n", + "Fits light only -- no mass, no source -- for every galaxy simultaneously, giving the next search\n", + "clean fixed light models to build on.\n", + "\n", + "This is identical to the group/slam.py SOURCE LP 0 pipeline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_0(\n", + " dataset,\n", + " settings_search,\n", + " main_lens_centres,\n", + " extra_lens_centres,\n", + " mask_radius,\n", + " redshift_lens,\n", + " n_batch=50,\n", + "):\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " # --- main lens light models (one per centre, light only) ---\n", + " lens_dict = {}\n", + " for i, centre in enumerate(main_lens_centres):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=False,\n", + " centre=(centre[0], centre[1]),\n", + " centre_sigma=0.1,\n", + " )\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=bulge, disk=None, point=None\n", + " )\n", + "\n", + " # --- extra lens galaxy light models ---\n", + " extra_light_models = []\n", + " for centre in extra_lens_centres:\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_prior_is_uniform=True,\n", + " centre=(centre[0], centre[1]),\n", + " ell_comps_prior_is_uniform=True,\n", + " )\n", + " extra_light_models.append(\n", + " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", + "\n", + " n_extra = len(extra_galaxies) if extra_galaxies is not None else 0\n", + " n_live = 100 + 30 * len(lens_dict) + 30 * n_extra\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[0]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " n_like_max=1000000,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 1__\n", + "\n", + "Equivalent to `source_lp` in `slam_start_here.py`, except lens light is fixed from `source_lp[0]`\n", + "rather than free, and mass and source are introduced here for the first time.\n", + "\n", + "Multiple main-lens galaxies each get an `Isothermal` mass; only `lens_0` carries an `ExternalShear`.\n", + "Extra-galaxy Einstein radii are bounded by a luminosity-derived prior." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_1(\n", + " dataset,\n", + " settings_search,\n", + " source_lp_result_0,\n", + " positions,\n", + " pixel_scale,\n", + " redshift_lens,\n", + " redshift_source,\n", + " source_mge_radius,\n", + " n_batch=50,\n", + "):\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(source_lp_result_0.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + " n_extra = (\n", + " len(list(source_lp_result_0.instance.extra_galaxies))\n", + " if source_lp_result_0.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " tracer = (\n", + " source_lp_result_0.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", + " )\n", + "\n", + " # Source MGE centred on primary lens bulge from stage 0.\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=source_mge_radius,\n", + " total_gaussians=30,\n", + " centre=source_lp_result_0.instance.galaxies.lens_0.bulge.centre,\n", + " centre_prior_is_uniform=False,\n", + " centre_sigma=0.6,\n", + " )\n", + "\n", + " # --- main lens full models (light fixed from stage 0, mass + shear free) ---\n", + " lens_dict = {}\n", + " for i in range(n_main):\n", + " lp0_lens = getattr(source_lp_result_0.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = lp0_lens.bulge.centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=5.0)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lp0_lens.bulge,\n", + " disk=lp0_lens.disk,\n", + " point=lp0_lens.point,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " # --- extra lens galaxy models (light fixed, mass bounded by luminosity) ---\n", + " extra_mass_models = []\n", + " for i in range(n_extra):\n", + " lp0_extra = source_lp_result_0.instance.extra_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = lp0_extra.bulge.centre\n", + " mass.ell_comps = lp0_extra.bulge.ell_comps\n", + "\n", + " luminosity_per_gaussian_list = [\n", + " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", + " for g in tracer.galaxies[n_main + i].bulge.profile_list\n", + " ]\n", + " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", + " luminosity_cap = 5 * 0.5 * total_luminosity**0.6\n", + " upper_limit = min(luminosity_cap, 5.0) if luminosity_cap > 0 else 5.0\n", + " mass.einstein_radius = af.UniformPrior(\n", + " lower_limit=0.0,\n", + " upper_limit=upper_limit,\n", + " )\n", + "\n", + " extra_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=lp0_extra.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", + " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " n_extra_model = len(extra_galaxies) if extra_galaxies is not None else 0\n", + " n_live = 150 + 30 * n_main + 30 * n_extra_model\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " n_like_max=200000,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "This is where the pixelization is first introduced. The source switches from a parametric MGE to a\n", + "rectangular adaptive pixelization with `Adapt` regularization.\n", + "\n", + "Key pixelization choices:\n", + " - `al.mesh.RectangularAdaptDensity`: a rectangular mesh whose cell density adapts to the source brightness.\n", + " - Fixed `mesh_shape = (28, 28)` pixels, chosen to match data resolution.\n", + " - `al.reg.Adapt`: adaptive regularization that varies smoothing based on source brightness.\n", + " Rectangular meshes use `Adapt` (not `AdaptSplit`, which is reserved for irregular meshes).\n", + "\n", + "Signal-adaptive over-sampling is applied: pixels above the S/N threshold use sub_size=4, others sub_size=2." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " dataset,\n", + " mask,\n", + " settings_search,\n", + " source_lp_result_1,\n", + " over_sample_size,\n", + " pixel_scale,\n", + " mask_radius,\n", + " positions,\n", + " n_batch=20,\n", + "):\n", + " mesh_shape = (28, 28)\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=4,\n", + " )\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result_1\n", + " )\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result_1.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(source_lp_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " # Main lens galaxies: mass as model from previous result.\n", + " lens_dict = {}\n", + " for i in range(n_main):\n", + " prev_lens = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=getattr(source_lp_result_1.model.galaxies, f\"lens_{i}\").mass,\n", + " mass_result=getattr(source_lp_result_1.model.galaxies, f\"lens_{i}\").mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=prev_lens.redshift,\n", + " bulge=prev_lens.bulge,\n", + " disk=prev_lens.disk,\n", + " mass=mass,\n", + " shear=source_lp_result_1.model.galaxies.lens_0.shear if i == 0 else None,\n", + " )\n", + "\n", + " # Extra galaxies: carried forward as model parameters.\n", + " extra_galaxies = (\n", + " source_lp_result_1.model.extra_galaxies\n", + " if source_lp_result_1.model.extra_galaxies is not None\n", + " else None\n", + " )\n", + "\n", + " # Source: RectangularAdaptDensity mesh with Adapt regularization.\n", + " pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization=af.Model(al.reg.Adapt),\n", + " )\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", + " pixelization=pixelization,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + " return result, dataset, adapt_images\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to SOURCE PIX 1 but uses capped adapt data from the first pixelization stage, ensuring\n", + "the adaptive mesh converges to a robust solution. The lens mass is fixed from SOURCE PIX 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " dataset,\n", + " settings_search,\n", + " source_lp_result_1,\n", + " source_pix_result_1,\n", + " adapt_images,\n", + " pixel_scale,\n", + " n_batch=20,\n", + "):\n", + " mesh_shape = (28, 28)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " # Cap the adapt data to prevent extreme values from dominating the mesh.\n", + " adapt_images_capped = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images_capped,\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " # Main lens galaxies: mass fixed from SOURCE PIX 1.\n", + " lens_dict = {}\n", + " for i in range(n_main):\n", + " prev_lens = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", + " lens_dict[f\"lens_{i}\"] = prev_lens\n", + "\n", + " extra_galaxies = source_pix_result_1.instance.extra_galaxies\n", + "\n", + " pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=af.Model(al.reg.Adapt),\n", + " )\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", + " pixelization=pixelization,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Fits the lens galaxy light with mass and pixelized source fixed from the SOURCE PIX pipelines.\n", + "Extra galaxies receive a fresh free MGE light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " dataset,\n", + " settings_search,\n", + " source_pix_result_1,\n", + " source_pix_result_2,\n", + " adapt_images,\n", + " mask_radius,\n", + " redshift_lens,\n", + " n_batch=20,\n", + "):\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_main):\n", + " prev_lens = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=False,\n", + " centre=prev_lens.bulge.centre,\n", + " centre_sigma=0.1,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=bulge,\n", + " mass=prev_lens.mass,\n", + " shear=prev_lens.shear if i == 0 else None,\n", + " )\n", + "\n", + " # Extra galaxies: fresh MGE light, mass fixed.\n", + " n_extra = (\n", + " len(list(source_pix_result_1.instance.extra_galaxies))\n", + " if source_pix_result_1.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " extra_models = []\n", + " for i in range(n_extra):\n", + " prev_extra = source_pix_result_1.instance.extra_galaxies[i]\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_prior_is_uniform=True,\n", + " centre=prev_extra.bulge.centre,\n", + " ell_comps_prior_is_uniform=True,\n", + " )\n", + "\n", + " extra_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=prev_extra.mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_models) if extra_models else None\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_pix_result_2, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Fits a PowerLaw total mass model with priors from the SOURCE PIX pipeline and lens light fixed from\n", + "the LIGHT LP pipeline. The pixelized source is carried forward." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " dataset,\n", + " settings_search,\n", + " source_pix_result_1,\n", + " source_pix_result_2,\n", + " light_result,\n", + " adapt_images,\n", + " redshift_lens,\n", + " n_batch=20,\n", + "):\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_2.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(light_result.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_main):\n", + " light_lens = getattr(light_result.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = af.Model(al.mp.PowerLaw)\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=getattr(source_pix_result_1.model.galaxies, f\"lens_{i}\").mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=light_lens.bulge,\n", + " mass=mass,\n", + " shear=source_pix_result_1.model.galaxies.lens_0.shear if i == 0 else None,\n", + " )\n", + "\n", + " # Extra galaxies: fresh mass, light fixed from LIGHT LP.\n", + " n_extra = (\n", + " len(list(light_result.instance.extra_galaxies))\n", + " if light_result.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " extra_models = []\n", + " for i in range(n_extra):\n", + " light_extra = light_result.instance.extra_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = light_extra.mass.centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " extra_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=light_extra.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_models) if extra_models else None\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens group dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Load centres for main lens and extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = _load_centres(dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = _load_centres(dataset_path / \"extra_galaxies_centres.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a 7.5\" circular mask for the group-scale lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0\n", + "pixel_scale = 0.1\n", + "\n", + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"group\") / \"features\" / \"pixelization\" / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Positions are computed from the SOURCE LP result and used in subsequent pixelization searches\n", + "to prevent demagnified source solutions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__SLaM Pipeline__\n", + "\n", + "The full SLaM pipeline is executed below. Each stage is documented with its pixelization-specific choices.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# --- SOURCE LP PIPELINE 0: Light only ---\n", + "\n", + "source_lp_result_0 = source_lp_0(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " main_lens_centres=main_lens_centres,\n", + " extra_lens_centres=extra_galaxies_centres,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + ")\n", + "\n", + "# --- SOURCE LP PIPELINE 1: Add mass + parametric source ---\n", + "\n", + "positions = source_lp_result_0.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + ").positions\n", + "\n", + "source_mge_radius = mask_radius\n", + "\n", + "source_lp_result_1 = source_lp_1(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " source_lp_result_0=source_lp_result_0,\n", + " positions=positions,\n", + " pixel_scale=pixel_scale,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + " source_mge_radius=source_mge_radius,\n", + ")\n", + "\n", + "# --- SOURCE PIX PIPELINE 1: Introduce Hilbert pixelization ---\n", + "\n", + "source_pix_result_1, dataset_pix, adapt_images = source_pix_1(\n", + " dataset=dataset,\n", + " mask=mask,\n", + " settings_search=settings_search,\n", + " source_lp_result_1=source_lp_result_1,\n", + " over_sample_size=over_sample_size,\n", + " pixel_scale=pixel_scale,\n", + " mask_radius=mask_radius,\n", + " positions=positions,\n", + ")\n", + "\n", + "# --- SOURCE PIX PIPELINE 2: Refine with capped adapt data ---\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " dataset=dataset_pix,\n", + " settings_search=settings_search,\n", + " source_lp_result_1=source_lp_result_1,\n", + " source_pix_result_1=source_pix_result_1,\n", + " adapt_images=adapt_images,\n", + " pixel_scale=pixel_scale,\n", + ")\n", + "\n", + "# --- LIGHT LP PIPELINE: Refit lens light ---\n", + "\n", + "light_result = light_lp(\n", + " dataset=dataset_pix,\n", + " settings_search=settings_search,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " adapt_images=adapt_images,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + ")\n", + "\n", + "# --- MASS TOTAL PIPELINE: Fit PowerLaw mass ---\n", + "\n", + "mass_result = mass_total(\n", + " dataset=dataset_pix,\n", + " settings_search=settings_search,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " light_result=light_result,\n", + " adapt_images=adapt_images,\n", + " redshift_lens=redshift_lens,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The final result contains the full group lens model with a pixelized source reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(mass_result.info)\n", + "\n", + "aplt.subplot_fit_imaging(fit=mass_result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated the full SLaM pipeline for group-scale lenses with a pixelized source.\n", + "\n", + "The key pixelization choices are:\n", + " - Rectangular adaptive mesh with a fixed shape (28, 28).\n", + " - `Adapt` regularization for brightness-dependent smoothing.\n", + " - Two-stage pixelization refinement with capped adapt data.\n", + " - Fixed pixelization over-sampling of 4.\n", + "\n", + "These choices are the recommended defaults for science-grade group-scale lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/pixelization/source_science.ipynb b/notebooks/group/features/pixelization/source_science.ipynb index 9f3166c79..04df6d711 100644 --- a/notebooks/group/features/pixelization/source_science.ipynb +++ b/notebooks/group/features/pixelization/source_science.ipynb @@ -1,623 +1,660 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Source Science: Pixelization (Group)\n", - "====================================\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy.\n", - "\n", - "Using a pixelized source reconstruction from a group-scale lens, we can compute key quantities such as:\n", - "\n", - " - The total flux of the reconstructed source.\n", - " - The magnification of the source due to the combined mass of all group galaxies.\n", - " - The intrinsic size and morphology of the source.\n", - "\n", - "For group-scale lenses, ALL mass profiles in the group (main lens + extra galaxies) contribute to the\n", - "magnification. Omitting any galaxy's mass would give an incorrect magnification estimate.\n", - "\n", - "This script also compares the pixelized source flux to a parametric estimate, demonstrating that the\n", - "pixelized reconstruction can recover source properties that parametric models may miss.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", - "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", - "- **Model Fit:** Create a FitImaging with a pixelized source.\n", - "- **Source Flux:** Compute the total flux from the pixelized source reconstruction.\n", - "- **Source Magnification:** Compute the magnification using all group mass profiles.\n", - "- **Impact of Extra Galaxies:** Demonstrate how omitting extra galaxies affects magnification.\n", - "- **Interpolated Source:** Interpolate the pixelized source to a uniform grid.\n", - "- **Parametric Comparison:** Compare pixelized source flux to a parametric estimate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=4,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "We create a fit with a pixelized source using concrete galaxy objects.\n", - "\n", - "For the group lens, we include the main lens galaxy, extra galaxies, and a pixelized source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", - "\n", - "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", - "regularization = al.reg.Constant(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Inversion__\n", - "\n", - "The inversion contains all information about the pixelized source reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "\n", - "mapper = inversion.cls_list_from(cls=al.Mapper)[0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Flux__\n", - "\n", - "The total flux of the pixelized source reconstruction is the sum of all source pixel fluxes.\n", - "\n", - "The units are the same as the data (typically electrons per second, e- s^-1)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = inversion.reconstruction\n", - "\n", - "total_source_flux = np.sum(reconstruction)\n", - "\n", - "print(f\"Total Source Flux via Pixelization: {total_source_flux} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The source pixel positions in the source plane are also available." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", - "\n", - "print(f\"Source Plane Mesh Grid: {source_plane_mesh_grid}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Magnification__\n", - "\n", - "The overall magnification of the source is the ratio of total flux in the image plane to total flux\n", - "in the source plane.\n", - "\n", - "The image-plane reconstruction (mapped back from source pixels) is available from the inversion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_operated_data = inversion.mapped_reconstructed_operated_data" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To compute the magnification, we need the areas of source and image pixels.\n", - "\n", - "The image-plane pixels have a uniform area defined by the dataset pixel scale." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.interpolate import griddata\n", - "\n", - "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", - "\n", - "interpolated_reconstruction = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", - ")\n", - "\n", - "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_ndarray,\n", - " pixel_scales=interpolation_grid.pixel_scales,\n", - ")\n", - "\n", - "magnification = np.sum(\n", - " mapped_reconstructed_operated_data * mapped_reconstructed_operated_data.pixel_area\n", - ") / np.sum(interpolated_reconstruction * interpolated_reconstruction.pixel_area)\n", - "\n", - "print(f\"Source Magnification (all group galaxies): {magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Impact of Extra Galaxies__\n", - "\n", - "For group-scale lenses, ALL mass profiles contribute to the magnification. Omitting extra galaxies\n", - "gives an incorrect estimate.\n", - "\n", - "Below, we compute the magnification using only the main lens galaxy for comparison." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_main_only = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit_main_only = al.FitImaging(dataset=dataset, tracer=tracer_main_only)\n", - "\n", - "inversion_main_only = fit_main_only.inversion\n", - "reconstruction_main_only = inversion_main_only.reconstruction\n", - "mapped_main_only = inversion_main_only.mapped_reconstructed_operated_data\n", - "\n", - "mapper_main_only = inversion_main_only.cls_list_from(cls=al.Mapper)[0]\n", - "source_plane_mesh_grid_main_only = mapper_main_only.source_plane_mesh_grid\n", - "\n", - "interpolated_main_only = griddata(\n", - " points=source_plane_mesh_grid_main_only,\n", - " values=reconstruction_main_only,\n", - " xi=interpolation_grid,\n", - ")\n", - "\n", - "interpolated_main_only_ndarray = interpolated_main_only.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_main_only = al.Array2D.no_mask(\n", - " values=interpolated_main_only_ndarray,\n", - " pixel_scales=interpolation_grid.pixel_scales,\n", - ")\n", - "\n", - "magnification_main_only = np.sum(\n", - " mapped_main_only * mapped_main_only.pixel_area\n", - ") / np.sum(interpolated_main_only * interpolated_main_only.pixel_area)\n", - "\n", - "print(f\"Source Magnification (main lens only): {magnification_main_only}\")\n", - "print(\n", - " f\"Magnification difference when omitting extra galaxies: \"\n", - " f\"{magnification - magnification_main_only:.4f}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolated Source__\n", - "\n", - "For detailed source science, the pixelized source can be interpolated to a uniform 2D grid.\n", - "This allows standard image analysis tools to be used." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=interpolated_reconstruction, title=\"Interpolated Source\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Zoom__\n", - "\n", - "We can zoom in on the source region for higher resolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extent = (-1.0, 1.0, -1.0, 1.0)\n", - "shape_native = (401, 401)\n", - "\n", - "interpolation_grid_zoom = al.Grid2D.from_extent(\n", - " extent=extent,\n", - " shape_native=shape_native,\n", - ")\n", - "\n", - "interpolated_reconstruction_zoom = griddata(\n", - " points=source_plane_mesh_grid,\n", - " values=reconstruction,\n", - " xi=interpolation_grid_zoom,\n", - ")\n", - "\n", - "interpolated_reconstruction_zoom_ndarray = interpolated_reconstruction_zoom.reshape(\n", - " interpolation_grid_zoom.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction_zoom = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_zoom_ndarray,\n", - " pixel_scales=interpolation_grid_zoom.pixel_scales,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=interpolated_reconstruction_zoom, title=\"Zoomed Interpolated Source\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors__\n", - "\n", - "The reconstruction noise map provides errors on each source pixel, enabling uncertainty propagation\n", - "for source science calculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction_noise_map = inversion.reconstruction_noise_map\n", - "\n", - "interpolated_noise_map = griddata(\n", - " points=source_plane_mesh_grid,\n", - " values=reconstruction_noise_map,\n", - " xi=interpolation_grid,\n", - ")\n", - "\n", - "interpolated_noise_map_ndarray = interpolated_noise_map.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_noise_map = al.Array2D.no_mask(\n", - " values=interpolated_noise_map_ndarray, pixel_scales=interpolation_grid.pixel_scales\n", - ")\n", - "\n", - "aplt.plot_array(array=interpolated_noise_map, title=\"Source Reconstruction Noise Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parametric Comparison__\n", - "\n", - "For comparison, we compute the source flux using a parametric source model (Sersic) that\n", - "approximates the true source used to simulate the dataset.\n", - "\n", - "The pixelized reconstruction may recover more flux than the parametric model if the source has\n", - "complex morphology that a Sersic profile cannot capture." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_parametric = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "parametric_image = source_parametric.bulge.image_2d_from(grid=grid)\n", - "total_parametric_flux = np.sum(parametric_image)\n", - "\n", - "print(f\"Total Source Flux (Parametric Sersic): {total_parametric_flux} e- s^-1\")\n", - "print(f\"Total Source Flux (Pixelized): {total_source_flux} e- s^-1\")\n", - "print(\n", - " f\"Flux Difference (Pixelized - Parametric): \"\n", - " f\"{total_source_flux - total_parametric_flux:.4f} e- s^-1\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Magnification via Mesh__\n", - "\n", - "For more accurate magnification, we can use the mesh pixel areas directly rather than interpolating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_areas = mapper.mesh_geometry.areas_for_magnification\n", - "\n", - "magnification_mesh = np.sum(\n", - " mapped_reconstructed_operated_data * mapped_reconstructed_operated_data.pixel_area\n", - ") / np.sum(reconstruction * mesh_areas)\n", - "\n", - "print(f\"Magnification via Mesh Areas: {magnification_mesh}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated source science calculations using a pixelized source reconstruction for a\n", - "group-scale lens.\n", - "\n", - "Key points:\n", - " - The total source flux is the sum of all reconstructed source pixel values.\n", - " - Magnification is computed as the ratio of image-plane to source-plane flux.\n", - " - ALL group mass profiles must be included for accurate magnification.\n", - " - The interpolated source can be analyzed using standard image tools.\n", - " - Pixelized reconstructions may recover more flux than parametric models for complex sources." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Source Science: Pixelization (Group)\n", + "====================================\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy.\n", + "\n", + "Using a pixelized source reconstruction from a group-scale lens, we can compute key quantities such as:\n", + "\n", + " - The total flux of the reconstructed source.\n", + " - The magnification of the source due to the combined mass of all group galaxies.\n", + " - The intrinsic size and morphology of the source.\n", + "\n", + "For group-scale lenses, ALL mass profiles in the group (main lens + extra galaxies) contribute to the\n", + "magnification. Omitting any galaxy's mass would give an incorrect magnification estimate.\n", + "\n", + "This script also compares the pixelized source flux to a parametric estimate, demonstrating that the\n", + "pixelized reconstruction can recover source properties that parametric models may miss.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the group dataset and 7.5\" mask.\n", + "- **Galaxy Centres:** Load centres for main lens and extra galaxies.\n", + "- **Model Fit:** Create a FitImaging with a pixelized source.\n", + "- **Source Flux:** Compute the total flux from the pixelized source reconstruction.\n", + "- **Source Magnification:** Compute the magnification using all group mass profiles.\n", + "- **Impact of Extra Galaxies:** Demonstrate how omitting extra galaxies affects magnification.\n", + "- **Interpolated Source:** Interpolate the pixelized source to a uniform grid.\n", + "- **Parametric Comparison:** Compare pixelized source flux to a parametric estimate." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=4,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "We create a fit with a pixelized source using concrete galaxy objects.\n", + "\n", + "For the group lens, we include the main lens galaxy, extra galaxies, and a pixelized source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", + "\n", + "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", + "regularization = al.reg.Constant(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Inversion__\n", + "\n", + "The inversion contains all information about the pixelized source reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "\n", + "mapper = inversion.cls_list_from(cls=al.Mapper)[0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Flux__\n", + "\n", + "The total flux of the pixelized source reconstruction is the sum of all source pixel fluxes.\n", + "\n", + "The units are the same as the data (typically electrons per second, e- s^-1)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = inversion.reconstruction\n", + "\n", + "total_source_flux = np.sum(reconstruction)\n", + "\n", + "print(f\"Total Source Flux via Pixelization: {total_source_flux} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The source pixel positions in the source plane are also available." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", + "\n", + "print(f\"Source Plane Mesh Grid: {source_plane_mesh_grid}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Magnification__\n", + "\n", + "The overall magnification of the source is the ratio of total flux in the image plane to total flux\n", + "in the source plane.\n", + "\n", + "The image-plane reconstruction (mapped back from source pixels) is available from the inversion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_operated_data = inversion.mapped_reconstructed_operated_data" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To compute the magnification, we need the areas of source and image pixels.\n", + "\n", + "The image-plane pixels have a uniform area defined by the dataset pixel scale." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.interpolate import griddata\n", + "\n", + "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", + "\n", + "interpolated_reconstruction = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", + ")\n", + "\n", + "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_ndarray,\n", + " pixel_scales=interpolation_grid.pixel_scales,\n", + ")\n", + "\n", + "magnification = np.sum(\n", + " mapped_reconstructed_operated_data * mapped_reconstructed_operated_data.pixel_area\n", + ") / np.sum(interpolated_reconstruction * interpolated_reconstruction.pixel_area)\n", + "\n", + "print(f\"Source Magnification (all group galaxies): {magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Impact of Extra Galaxies__\n", + "\n", + "For group-scale lenses, ALL mass profiles contribute to the magnification. Omitting extra galaxies\n", + "gives an incorrect estimate.\n", + "\n", + "Below, we compute the magnification using only the main lens galaxy for comparison." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_main_only = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit_main_only = al.FitImaging(dataset=dataset, tracer=tracer_main_only)\n", + "\n", + "inversion_main_only = fit_main_only.inversion\n", + "reconstruction_main_only = inversion_main_only.reconstruction\n", + "mapped_main_only = inversion_main_only.mapped_reconstructed_operated_data\n", + "\n", + "mapper_main_only = inversion_main_only.cls_list_from(cls=al.Mapper)[0]\n", + "source_plane_mesh_grid_main_only = mapper_main_only.source_plane_mesh_grid\n", + "\n", + "interpolated_main_only = griddata(\n", + " points=source_plane_mesh_grid_main_only,\n", + " values=reconstruction_main_only,\n", + " xi=interpolation_grid,\n", + ")\n", + "\n", + "interpolated_main_only_ndarray = interpolated_main_only.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_main_only = al.Array2D.no_mask(\n", + " values=interpolated_main_only_ndarray,\n", + " pixel_scales=interpolation_grid.pixel_scales,\n", + ")\n", + "\n", + "magnification_main_only = np.sum(\n", + " mapped_main_only * mapped_main_only.pixel_area\n", + ") / np.sum(interpolated_main_only * interpolated_main_only.pixel_area)\n", + "\n", + "print(f\"Source Magnification (main lens only): {magnification_main_only}\")\n", + "print(\n", + " f\"Magnification difference when omitting extra galaxies: \"\n", + " f\"{magnification - magnification_main_only:.4f}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolated Source__\n", + "\n", + "For detailed source science, the pixelized source can be interpolated to a uniform 2D grid.\n", + "This allows standard image analysis tools to be used." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=interpolated_reconstruction, title=\"Interpolated Source\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Zoom__\n", + "\n", + "We can zoom in on the source region for higher resolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extent = (-1.0, 1.0, -1.0, 1.0)\n", + "shape_native = (401, 401)\n", + "\n", + "interpolation_grid_zoom = al.Grid2D.from_extent(\n", + " extent=extent,\n", + " shape_native=shape_native,\n", + ")\n", + "\n", + "interpolated_reconstruction_zoom = griddata(\n", + " points=source_plane_mesh_grid,\n", + " values=reconstruction,\n", + " xi=interpolation_grid_zoom,\n", + ")\n", + "\n", + "interpolated_reconstruction_zoom_ndarray = interpolated_reconstruction_zoom.reshape(\n", + " interpolation_grid_zoom.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction_zoom = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_zoom_ndarray,\n", + " pixel_scales=interpolation_grid_zoom.pixel_scales,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=interpolated_reconstruction_zoom, title=\"Zoomed Interpolated Source\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors__\n", + "\n", + "The reconstruction noise map provides errors on each source pixel, enabling uncertainty propagation\n", + "for source science calculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction_noise_map = inversion.reconstruction_noise_map\n", + "\n", + "interpolated_noise_map = griddata(\n", + " points=source_plane_mesh_grid,\n", + " values=reconstruction_noise_map,\n", + " xi=interpolation_grid,\n", + ")\n", + "\n", + "interpolated_noise_map_ndarray = interpolated_noise_map.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_noise_map = al.Array2D.no_mask(\n", + " values=interpolated_noise_map_ndarray, pixel_scales=interpolation_grid.pixel_scales\n", + ")\n", + "\n", + "aplt.plot_array(array=interpolated_noise_map, title=\"Source Reconstruction Noise Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parametric Comparison__\n", + "\n", + "For comparison, we compute the source flux using a parametric source model (Sersic) that\n", + "approximates the true source used to simulate the dataset.\n", + "\n", + "The pixelized reconstruction may recover more flux than the parametric model if the source has\n", + "complex morphology that a Sersic profile cannot capture." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_parametric = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "parametric_image = source_parametric.bulge.image_2d_from(grid=grid)\n", + "total_parametric_flux = np.sum(parametric_image)\n", + "\n", + "print(f\"Total Source Flux (Parametric Sersic): {total_parametric_flux} e- s^-1\")\n", + "print(f\"Total Source Flux (Pixelized): {total_source_flux} e- s^-1\")\n", + "print(\n", + " f\"Flux Difference (Pixelized - Parametric): \"\n", + " f\"{total_source_flux - total_parametric_flux:.4f} e- s^-1\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Magnification via Mesh__\n", + "\n", + "For more accurate magnification, we can use the mesh pixel areas directly rather than interpolating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_areas = mapper.mesh_geometry.areas_for_magnification\n", + "\n", + "magnification_mesh = np.sum(\n", + " mapped_reconstructed_operated_data * mapped_reconstructed_operated_data.pixel_area\n", + ") / np.sum(reconstruction * mesh_areas)\n", + "\n", + "print(f\"Magnification via Mesh Areas: {magnification_mesh}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated source science calculations using a pixelized source reconstruction for a\n", + "group-scale lens.\n", + "\n", + "Key points:\n", + " - The total source flux is the sum of all reconstructed source pixel values.\n", + " - Magnification is computed as the ratio of image-plane to source-plane flux.\n", + " - ALL group mass profiles must be included for accurate magnification.\n", + " - The interpolated source can be analyzed using standard image tools.\n", + " - Pixelized reconstructions may recover more flux than parametric models for complex sources." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/scaling_relation/fit.ipynb b/notebooks/group/features/scaling_relation/fit.ipynb index 6ff6d8f08..9fe61419f 100644 --- a/notebooks/group/features/scaling_relation/fit.ipynb +++ b/notebooks/group/features/scaling_relation/fit.ipynb @@ -1,530 +1,570 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Group Scaling Relation Fit\n", - "====================================\n", - "\n", - "A group-scale strong lens often has many foreground galaxies near the line of sight to the source, on top of one\n", - "or more primary lens galaxies. The **three-tier API** splits these galaxies into populations that the lens model\n", - "treats differently:\n", - "\n", - " - **Main lens galaxies** (`main_lens_centres.json`): the primary lens(es) \u2014 each modelled with its own free\n", - " mass parameters via the group `lens_dict` API.\n", - " - **Extra galaxies** (`extra_galaxies_centres.json`): individually-modelled companions, each with its own free\n", - " `einstein_radius`. Use this tier for the brighter / closer companions that contribute non-trivially to the\n", - " lensing on their own.\n", - " - **Scaling galaxies** (`scaling_galaxies_centres.json` + `scaling_galaxies.csv`): the long tail of fainter\n", - " companions whose Einstein radii are tied together via a shared two-parameter relation\n", - " `einstein_radius = scaling_factor * luminosity ** scaling_exponent`. Adding more galaxies to this tier does\n", - " not grow the model.\n", - "\n", - "This script illustrates the API for performing a fit to a group-scale strong lens with all three tiers active,\n", - "via the standard `Tracer` and `FitImaging` objects, without invoking a non-linear search.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres + Luminosities:** Load main, extras and scaling-tier centres + luminosities.\n", - "- **Over Sampling:** Adaptive over-sampling at every galaxy centre.\n", - "- **MGE Basis:** Build a `Basis` of linear Gaussians for the source.\n", - "- **Galaxies:** Concrete composition \u2014 `lens_dict` + extras + scaling-tier + source.\n", - "- **Tracer:** Build the `Tracer` and fit the dataset.\n", - "- **Three-Tier Deflection Tour:** Per-tier deflection sums into the tracer's total deflection.\n", - "- **Intensities:** The solved-for linear light profile `intensity` values for each MGE Gaussian.\n", - "- **Wrap Up:** Summary and next steps.\n", - "\n", - "__Prerequisites__\n", - "\n", - "This script focuses on the API specific to a group-scale three-tier extras population. For background:\n", - "\n", - " - `autolens_workspace/scripts/imaging/features/scaling_relation/fit.py` \u2014 the single-main-lens version of this\n", - " script. The three-tier walkthrough below generalises that example across multiple main lens galaxies.\n", - " - `autolens_workspace/scripts/group/start_here.py` \u2014 the group-scale `lens_dict` API, including how\n", - " `main_lens_centres.json` is loaded.\n", - " - `autolens_workspace/scripts/group/features/scaling_relation/modeling.py` \u2014 the search-based version of this\n", - " script, which composes the same model via `af.Model` with free `scaling_factor` and `scaling_exponent` priors.\n", - "\n", - "The group simulator here has only ONE main lens galaxy, so the `lens_dict` has a single entry `lens_0`. The\n", - "pattern generalises naturally to groups with multiple main lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the group-scale strong lens dataset `scaling_relation` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"scaling_relation\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/scaling_relation/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres + Luminosities__\n", - "\n", - "Load each tier's data:\n", - " - main lens centres from JSON,\n", - " - individually-modelled extras centres from JSON,\n", - " - scaling-tier centres AND luminosities from CSV." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "scaling_table = al.galaxy_table_from_csv(\n", - " file_path=dataset_path / \"scaling_galaxies.csv\"\n", - ")\n", - "scaling_galaxies_centres = scaling_table.centres\n", - "scaling_galaxies_luminosities = scaling_table.luminosities\n", - "\n", - "print(f\"Main lens centres: {list(main_lens_centres)}\")\n", - "print(f\"Individually-modelled extras centres: {list(extra_galaxies_centres)}\")\n", - "print(f\"Scaling-tier extras centres: {list(scaling_galaxies_centres)}\")\n", - "print(f\"Scaling-tier extras luminosities: {scaling_galaxies_luminosities}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define an 8.0\" circular mask, large enough to include the main lens and all extra + scaling galaxies (the most\n", - "distant sits at radius ~7.5\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 8.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling at every galaxy centre across all three tiers." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_galaxy_centres = (\n", - " [tuple(c) for c in main_lens_centres]\n", - " + [tuple(c) for c in extra_galaxies_centres]\n", - " + [tuple(c) for c in scaling_galaxies_centres]\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_galaxy_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE Basis__\n", - "\n", - "A `Basis` of 30 linear Gaussians for the source galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "source_bulge = build_source_basis(centre=(0.0, 0.1))\n", - "\n", - "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "Three-tier concrete composition:\n", - "\n", - " - `lens_dict` (z=0.5): one `Galaxy` per main lens centre, each with `SersicSph` light + `IsothermalSph` mass.\n", - " The simulator here has a single main lens with `einstein_radius=4.0`.\n", - " - `individual_extras` (z=0.5): two individually-modelled companions with simulator-true Einstein radii 0.8\n", - " and 1.0.\n", - " - `scaling_extras` (z=0.5): two scaling-tier companions whose Einstein radii are derived from\n", - " `einstein_radius = 0.3 * luminosity ** 1.0` (simulator truth: 0.135 each from luminosity 0.45).\n", - " - `source` (z=1.0): the MGE basis above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_truth = [\n", - " dict(intensity=0.7, effective_radius=2.0, sersic_index=4.0, einstein_radius=4.0),\n", - "]\n", - "\n", - "lens_dict = {}\n", - "for i, (centre, truth) in enumerate(zip(main_lens_centres, main_lens_truth)):\n", - " lens_dict[f\"lens_{i}\"] = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=tuple(centre),\n", - " intensity=truth[\"intensity\"],\n", - " effective_radius=truth[\"effective_radius\"],\n", - " sersic_index=truth[\"sersic_index\"],\n", - " ),\n", - " mass=al.mp.IsothermalSph(\n", - " centre=tuple(centre), einstein_radius=truth[\"einstein_radius\"]\n", - " ),\n", - " )\n", - "\n", - "extra_truth = [\n", - " dict(intensity=0.9, effective_radius=0.8, sersic_index=3.0, einstein_radius=0.8),\n", - " dict(intensity=0.9, effective_radius=0.8, sersic_index=3.0, einstein_radius=1.0),\n", - "]\n", - "\n", - "individual_extras = []\n", - "for centre, truth in zip(extra_galaxies_centres, extra_truth):\n", - " individual_extras.append(\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=tuple(centre),\n", - " intensity=truth[\"intensity\"],\n", - " effective_radius=truth[\"effective_radius\"],\n", - " sersic_index=truth[\"sersic_index\"],\n", - " ),\n", - " mass=al.mp.IsothermalSph(\n", - " centre=tuple(centre), einstein_radius=truth[\"einstein_radius\"]\n", - " ),\n", - " )\n", - " )\n", - "\n", - "scaling_factor = 0.3\n", - "scaling_exponent = 1.0\n", - "\n", - "scaling_extras = []\n", - "scaling_extras_einstein_radii = []\n", - "for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities):\n", - " einstein_radius = scaling_factor * luminosity**scaling_exponent\n", - " scaling_extras_einstein_radii.append(einstein_radius)\n", - " scaling_extras.append(\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=tuple(centre),\n", - " intensity=luminosity,\n", - " effective_radius=0.6,\n", - " sersic_index=2.5,\n", - " ),\n", - " mass=al.mp.IsothermalSph(\n", - " centre=tuple(centre), einstein_radius=einstein_radius\n", - " ),\n", - " )\n", - " )\n", - "\n", - "source = al.Galaxy(redshift=1.0, bulge=source_bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "The `Tracer` queries every mass profile across all three tiers and sums their deflections." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=list(lens_dict.values()) + individual_extras + scaling_extras + [source]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Three-Tier Deflection Tour__\n", - "\n", - "The lens-plane total deflection is the sum of three tier-wise contributions:\n", - "\n", - " alpha_lens_total(theta) = sum_i alpha_lens_i(theta)\n", - " + sum_j alpha_extra_individual_j(theta)\n", - " + sum_k alpha_extra_scaling_k(theta)\n", - "\n", - "where the scaling tier contributions come from mass profiles whose Einstein radii are derived from the shared\n", - "relation. We verify this by computing each tier explicitly and confirming the grand sum equals what the `Tracer`\n", - "returns." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = dataset.grid\n", - "\n", - "alpha_main_per_lens = [\n", - " g.mass.deflections_yx_2d_from(grid=grid) for g in lens_dict.values()\n", - "]\n", - "alpha_individual = [g.mass.deflections_yx_2d_from(grid=grid) for g in individual_extras]\n", - "alpha_scaling = [g.mass.deflections_yx_2d_from(grid=grid) for g in scaling_extras]\n", - "\n", - "alpha_main_total = sum(alpha_main_per_lens)\n", - "alpha_individual_total = sum(alpha_individual)\n", - "alpha_scaling_total = sum(alpha_scaling)\n", - "\n", - "print(f\"alpha_main_lens (tier sum, first coord) : {alpha_main_total[0]}\")\n", - "print(f\"alpha_individual (tier sum, first coord) : {alpha_individual_total[0]}\")\n", - "print(f\"alpha_scaling (tier sum, first coord) : {alpha_scaling_total[0]}\")\n", - "\n", - "for centre, luminosity, er in zip(\n", - " scaling_galaxies_centres,\n", - " scaling_galaxies_luminosities,\n", - " scaling_extras_einstein_radii,\n", - "):\n", - " print(\n", - " f\" scaling galaxy @ {tuple(centre)}: \"\n", - " f\"einstein_radius = {scaling_factor:.2f} * {luminosity:.3f} ** {scaling_exponent:.1f} = {er:.4f}\"\n", - " )\n", - "\n", - "alpha_total_summed = alpha_main_total + alpha_individual_total + alpha_scaling_total\n", - "\n", - "traced_grids = tracer.traced_grid_2d_list_from(grid=grid)\n", - "alpha_total_tracer = grid - traced_grids[1]\n", - "\n", - "print(f\"\\nalpha_total (summed by hand, first 3): {alpha_total_summed[:3]}\")\n", - "print(f\"alpha_total (from tracer, first 3): {alpha_total_tracer[:3]}\")\n", - "\n", - "assert np.allclose(np.asarray(alpha_total_summed), np.asarray(alpha_total_tracer))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "After the fit, every linear Gaussian in the source MGE basis has been assigned an `intensity` via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"\\nFirst Gaussian intensity, source = \"\n", - " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", - ")\n", - "\n", - "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer_fitted.galaxies[-1].bulge, grid=plot_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated the group-scale three-tier API and the per-tier deflection composition, without\n", - "invoking a non-linear search. The scaling relation collapses what would otherwise be N free `einstein_radius`\n", - "parameters into 2 shared parameters (`scaling_factor` and `scaling_exponent`), letting the model dimensionality\n", - "stay constant as galaxy count grows.\n", - "\n", - "In a real modeling workflow:\n", - "\n", - " - `modeling.py` runs the search-based version, where `scaling_factor` and `scaling_exponent` are free `af.Model`\n", - " parameters with `UniformPrior`s.\n", - " - `modeling_for_luminosities.py` is the standalone light-only fit that produces the luminosities consumed by\n", - " the scaling relation. In production this stage is the `source_lp[0]` step of a SLaM pipeline.\n", - " - `autolens_workspace/scripts/group/slam.py` is the full SLaM pipeline, which already implements scaling\n", - " galaxies via this same composition under the hood.\n", - "\n", - "The key takeaway is that the group `lens_dict + extra_galaxies + scaling_galaxies` collection structure scales\n", - "naturally to any number of main lens galaxies and any number of foreground galaxies \u2014 the lens-plane deflection\n", - "is still a simple per-galaxy sum, but the scaling-tier contributions are parameterized through luminosity rather\n", - "than per-galaxy free parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Group Scaling Relation Fit\n", + "====================================\n", + "\n", + "A group-scale strong lens often has many foreground galaxies near the line of sight to the source, on top of one\n", + "or more primary lens galaxies. The **three-tier API** splits these galaxies into populations that the lens model\n", + "treats differently:\n", + "\n", + " - **Main lens galaxies** (`main_lens_centres.json`): the primary lens(es) \u2014 each modelled with its own free\n", + " mass parameters via the group `lens_dict` API.\n", + " - **Extra galaxies** (`extra_galaxies_centres.json`): individually-modelled companions, each with its own free\n", + " `einstein_radius`. Use this tier for the brighter / closer companions that contribute non-trivially to the\n", + " lensing on their own.\n", + " - **Scaling galaxies** (`scaling_galaxies_centres.json` + `scaling_galaxies.csv`): the long tail of fainter\n", + " companions whose Einstein radii are tied together via a shared reference-anchored relation\n", + " `einstein_radius = einstein_radius_ref * (luminosity / luminosity_ref) ** 0.5` (exponent fixed at the\n", + " Faber-Jackson value; the Lenstool convention). Adding more galaxies to this tier does not grow the model.\n", + "\n", + "This script illustrates the API for performing a fit to a group-scale strong lens with all three tiers active,\n", + "via the standard `Tracer` and `FitImaging` objects, without invoking a non-linear search.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres + Luminosities:** Load main, extras and scaling-tier centres + luminosities.\n", + "- **Over Sampling:** Adaptive over-sampling at every galaxy centre.\n", + "- **MGE Basis:** Build a `Basis` of linear Gaussians for the source.\n", + "- **Galaxies:** Concrete composition \u2014 `lens_dict` + extras + scaling-tier + source.\n", + "- **Tracer:** Build the `Tracer` and fit the dataset.\n", + "- **Three-Tier Deflection Tour:** Per-tier deflection sums into the tracer's total deflection.\n", + "- **Intensities:** The solved-for linear light profile `intensity` values for each MGE Gaussian.\n", + "- **Wrap Up:** Summary and next steps.\n", + "\n", + "__Prerequisites__\n", + "\n", + "This script focuses on the API specific to a group-scale three-tier extras population. For background:\n", + "\n", + " - `autolens_workspace/scripts/imaging/features/scaling_relation/fit.py` \u2014 the single-main-lens version of this\n", + " script. The three-tier walkthrough below generalises that example across multiple main lens galaxies.\n", + " - `autolens_workspace/scripts/group/start_here.py` \u2014 the group-scale `lens_dict` API, including how\n", + " `main_lens_centres.json` is loaded.\n", + " - `autolens_workspace/scripts/group/features/scaling_relation/modeling.py` \u2014 the search-based version of this\n", + " script, which composes the same model via `af.Model` with a free `einstein_radius_ref` prior.\n", + "\n", + "The group simulator here has only ONE main lens galaxy, so the `lens_dict` has a single entry `lens_0`. The\n", + "pattern generalises naturally to groups with multiple main lens galaxies." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the group-scale strong lens dataset `scaling_relation` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"scaling_relation\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/scaling_relation/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres + Luminosities__\n", + "\n", + "Load each tier's data:\n", + " - main lens centres from JSON,\n", + " - individually-modelled extras centres from JSON,\n", + " - scaling-tier centres AND luminosities from CSV." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "scaling_table = al.galaxy_table_from_csv(\n", + " file_path=dataset_path / \"scaling_galaxies.csv\"\n", + ")\n", + "scaling_galaxies_centres = scaling_table.centres\n", + "scaling_galaxies_luminosities = scaling_table.luminosities\n", + "\n", + "print(f\"Main lens centres: {list(main_lens_centres)}\")\n", + "print(f\"Individually-modelled extras centres: {list(extra_galaxies_centres)}\")\n", + "print(f\"Scaling-tier extras centres: {list(scaling_galaxies_centres)}\")\n", + "print(f\"Scaling-tier extras luminosities: {scaling_galaxies_luminosities}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define an 8.0\" circular mask, large enough to include the main lens and all extra + scaling galaxies (the most\n", + "distant sits at radius ~7.5\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 8.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling at every galaxy centre across all three tiers." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_galaxy_centres = (\n", + " [tuple(c) for c in main_lens_centres]\n", + " + [tuple(c) for c in extra_galaxies_centres]\n", + " + [tuple(c) for c in scaling_galaxies_centres]\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_galaxy_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE Basis__\n", + "\n", + "A `Basis` of 30 linear Gaussians for the source galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "source_bulge = build_source_basis(centre=(0.0, 0.1))\n", + "\n", + "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "Three-tier concrete composition:\n", + "\n", + " - `lens_dict` (z=0.5): one `Galaxy` per main lens centre, each with `SersicSph` light + `IsothermalSph` mass.\n", + " The simulator here has a single main lens with `einstein_radius=4.0`.\n", + " - `individual_extras` (z=0.5): two individually-modelled companions with simulator-true Einstein radii 0.8\n", + " and 1.0.\n", + " - `scaling_extras` (z=0.5): two scaling-tier companions whose Einstein radii are derived from\n", + " `einstein_radius = 0.3 * luminosity ** 1.0` (simulator truth: 0.135 each from luminosity 0.45).\n", + " - `source` (z=1.0): the MGE basis above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_truth = [\n", + " dict(intensity=0.7, effective_radius=2.0, sersic_index=4.0, einstein_radius=4.0),\n", + "]\n", + "\n", + "lens_dict = {}\n", + "for i, (centre, truth) in enumerate(zip(main_lens_centres, main_lens_truth)):\n", + " lens_dict[f\"lens_{i}\"] = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=tuple(centre),\n", + " intensity=truth[\"intensity\"],\n", + " effective_radius=truth[\"effective_radius\"],\n", + " sersic_index=truth[\"sersic_index\"],\n", + " ),\n", + " mass=al.mp.IsothermalSph(\n", + " centre=tuple(centre), einstein_radius=truth[\"einstein_radius\"]\n", + " ),\n", + " )\n", + "\n", + "extra_truth = [\n", + " dict(intensity=0.9, effective_radius=0.8, sersic_index=3.0, einstein_radius=0.8),\n", + " dict(intensity=0.9, effective_radius=0.8, sersic_index=3.0, einstein_radius=1.0),\n", + "]\n", + "\n", + "individual_extras = []\n", + "for centre, truth in zip(extra_galaxies_centres, extra_truth):\n", + " individual_extras.append(\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=tuple(centre),\n", + " intensity=truth[\"intensity\"],\n", + " effective_radius=truth[\"effective_radius\"],\n", + " sersic_index=truth[\"sersic_index\"],\n", + " ),\n", + " mass=al.mp.IsothermalSph(\n", + " centre=tuple(centre), einstein_radius=truth[\"einstein_radius\"]\n", + " ),\n", + " )\n", + " )\n", + "\n", + "einstein_radius_ref = 0.135\n", + "scaling_exponent = 0.5\n", + "luminosity_ref = max(scaling_galaxies_luminosities)\n", + "\n", + "scaling_extras = []\n", + "scaling_extras_einstein_radii = []\n", + "for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities):\n", + " einstein_radius = (\n", + " einstein_radius_ref * (luminosity / luminosity_ref) ** scaling_exponent\n", + " )\n", + " scaling_extras_einstein_radii.append(einstein_radius)\n", + " scaling_extras.append(\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=tuple(centre),\n", + " intensity=luminosity,\n", + " effective_radius=0.6,\n", + " sersic_index=2.5,\n", + " ),\n", + " mass=al.mp.IsothermalSph(\n", + " centre=tuple(centre), einstein_radius=einstein_radius\n", + " ),\n", + " )\n", + " )\n", + "\n", + "source = al.Galaxy(redshift=1.0, bulge=source_bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "The `Tracer` queries every mass profile across all three tiers and sums their deflections." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=list(lens_dict.values()) + individual_extras + scaling_extras + [source]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Three-Tier Deflection Tour__\n", + "\n", + "The lens-plane total deflection is the sum of three tier-wise contributions:\n", + "\n", + " alpha_lens_total(theta) = sum_i alpha_lens_i(theta)\n", + " + sum_j alpha_extra_individual_j(theta)\n", + " + sum_k alpha_extra_scaling_k(theta)\n", + "\n", + "where the scaling tier contributions come from mass profiles whose Einstein radii are derived from the shared\n", + "relation. We verify this by computing each tier explicitly and confirming the grand sum equals what the `Tracer`\n", + "returns." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = dataset.grid\n", + "\n", + "alpha_main_per_lens = [\n", + " g.mass.deflections_yx_2d_from(grid=grid) for g in lens_dict.values()\n", + "]\n", + "alpha_individual = [g.mass.deflections_yx_2d_from(grid=grid) for g in individual_extras]\n", + "alpha_scaling = [g.mass.deflections_yx_2d_from(grid=grid) for g in scaling_extras]\n", + "\n", + "alpha_main_total = sum(alpha_main_per_lens)\n", + "alpha_individual_total = sum(alpha_individual)\n", + "alpha_scaling_total = sum(alpha_scaling)\n", + "\n", + "print(f\"alpha_main_lens (tier sum, first coord) : {alpha_main_total[0]}\")\n", + "print(f\"alpha_individual (tier sum, first coord) : {alpha_individual_total[0]}\")\n", + "print(f\"alpha_scaling (tier sum, first coord) : {alpha_scaling_total[0]}\")\n", + "\n", + "for centre, luminosity, er in zip(\n", + " scaling_galaxies_centres,\n", + " scaling_galaxies_luminosities,\n", + " scaling_extras_einstein_radii,\n", + "):\n", + " print(\n", + " f\" scaling galaxy @ {tuple(centre)}: \"\n", + " f\"einstein_radius = {einstein_radius_ref:.3f} * ({luminosity:.3f} / {luminosity_ref:.3f}) ** {scaling_exponent:.1f} = {er:.4f}\"\n", + " )\n", + "\n", + "alpha_total_summed = alpha_main_total + alpha_individual_total + alpha_scaling_total\n", + "\n", + "traced_grids = tracer.traced_grid_2d_list_from(grid=grid)\n", + "alpha_total_tracer = grid - traced_grids[1]\n", + "\n", + "print(f\"\\nalpha_total (summed by hand, first 3): {alpha_total_summed[:3]}\")\n", + "print(f\"alpha_total (from tracer, first 3): {alpha_total_tracer[:3]}\")\n", + "\n", + "assert np.allclose(np.asarray(alpha_total_summed), np.asarray(alpha_total_tracer))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "After the fit, every linear Gaussian in the source MGE basis has been assigned an `intensity` via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"\\nFirst Gaussian intensity, source = \"\n", + " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", + ")\n", + "\n", + "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer_fitted.galaxies[-1].bulge, grid=plot_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated the group-scale three-tier API and the per-tier deflection composition, without\n", + "invoking a non-linear search. The scaling relation collapses what would otherwise be N free `einstein_radius`\n", + "parameters into a single shared normalization (`einstein_radius_ref`, the brightest member's Einstein radius,\n", + "with the exponent fixed at 0.5), letting the model dimensionality stay constant as galaxy count grows.\n", + "\n", + "In a real modeling workflow:\n", + "\n", + " - `modeling.py` runs the search-based version, where `einstein_radius_ref` is a free `af.Model` parameter with\n", + " a `UniformPrior`.\n", + " - `modeling_for_luminosities.py` is the standalone light-only fit that produces the luminosities consumed by\n", + " the scaling relation. In production this stage is the `source_lp[0]` step of a SLaM pipeline.\n", + " - `autolens_workspace/scripts/group/slam.py` is the full SLaM pipeline, which already implements scaling\n", + " galaxies via this same composition under the hood.\n", + "\n", + "The key takeaway is that the group `lens_dict + extra_galaxies + scaling_galaxies` collection structure scales\n", + "naturally to any number of main lens galaxies and any number of foreground galaxies \u2014 the lens-plane deflection\n", + "is still a simple per-galaxy sum, but the scaling-tier contributions are parameterized through luminosity rather\n", + "than per-galaxy free parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/scaling_relation/likelihood_function.ipynb b/notebooks/group/features/scaling_relation/likelihood_function.ipynb index 926264c58..10981ab08 100644 --- a/notebooks/group/features/scaling_relation/likelihood_function.ipynb +++ b/notebooks/group/features/scaling_relation/likelihood_function.ipynb @@ -1,420 +1,460 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Group Scaling Relation__\n", - "\n", - "This script describes the additional steps required to compute the `log_likelihood` for a group-scale strong\n", - "lens whose foreground galaxy population is split across three tiers \u2014 main lens galaxies (modelled via the\n", - "group `lens_dict` API), individually-modelled extras (each with its own free `einstein_radius`), and\n", - "scaling-tier extras (whose Einstein radii are derived from a shared two-parameter relation\n", - "`einstein_radius = scaling_factor * luminosity ** scaling_exponent`).\n", - "\n", - "This script does NOT repeat the steps shared with single-plane lensing (mask, image-plane grid, PSF convolution,\n", - "chi-squared, noise normalization, linear-algebra solver for MGE source intensities). It documents only the part\n", - "of the likelihood function which is specific to the group three-tier API: the multi-tier lens-plane deflection\n", - "composition.\n", - "\n", - "__Prerequisites__\n", - "\n", - " - `autolens_workspace/scripts/imaging/features/scaling_relation/likelihood_function.py` \u2014 the single-main-lens\n", - " walkthrough. The three-tier version below generalises that example across the `lens_dict` API.\n", - " - `autolens_workspace/scripts/group/start_here.py` \u2014 the group `lens_dict` API, including how\n", - " `main_lens_centres.json` is loaded.\n", - " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 canonical single-plane walkthrough.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a `Basis`\n", - " of linear Gaussians is solved for via linear algebra.\n", - " - `autolens_workspace/scripts/group/features/scaling_relation/modeling.py` \u2014 search-based three-tier version.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order (see above).\n", - "- **Dataset & Mask.**\n", - "- **Centres + Luminosities:** Load each tier's data.\n", - "- **Galaxies:** lens_dict + extras + scaling-tier + MGE source.\n", - "- **Three-Tier Deflection:** Per-tier deflection sums plus the scaling-relation evaluation.\n", - "- **Manual Ray-Tracing:** Hand-compute the source-plane grid and confirm it matches `Tracer`.\n", - "- **Source-Plane Image:** Source MGE evaluated at the ray-traced grid.\n", - "- **Fit Check.**\n", - "- **Wrap Up.**\n", - "\n", - "__What Changes For A Group Three-Tier Scaling Relation__\n", - "\n", - "For a single-galaxy lens with one mass profile, the lens-plane deflection is just one profile evaluated at\n", - "each image-plane coordinate:\n", - "\n", - " alpha_lens(theta) = alpha_total(theta ; Isothermal parameters)\n", - "\n", - "For a group-scale lens with the three-tier API, the lens-plane deflection is the SUM across every main lens\n", - "galaxy AND every extras / scaling-tier galaxy:\n", - "\n", - " alpha_lens(theta) = sum_i alpha_main_lens_i(theta)\n", - " + sum_j alpha_extra_individual_j(theta)\n", - " + sum_k alpha_extra_scaling_k(theta)\n", - "\n", - " where alpha_extra_scaling_k is the deflection of a mass profile whose\n", - " einstein_radius_k = scaling_factor * luminosity_k ** scaling_exponent.\n", - "\n", - "The model gains exactly 2 free parameters (`scaling_factor`, `scaling_exponent`) regardless of how many galaxies\n", - "sit on the scaling tier. Every other step of the likelihood (PSF convolution, chi-squared, noise normalization,\n", - "MGE linear-algebra solver) is unchanged." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the group scaling_relation dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"scaling_relation\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/group/features/scaling_relation/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 8.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres + Luminosities__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "scaling_table = al.galaxy_table_from_csv(\n", - " file_path=dataset_path / \"scaling_galaxies.csv\"\n", - ")\n", - "scaling_galaxies_centres = scaling_table.centres\n", - "scaling_galaxies_luminosities = scaling_table.luminosities" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "Three populations participate in the ray-tracing:\n", - "\n", - " - `lens_dict` (z=0.5): one `IsothermalSph` mass per main lens centre \u2014 here, just one with `einstein_radius=4.0`.\n", - " - `individual_extras` (z=0.5): two `IsothermalSph` masses with simulator-true Einstein radii 0.8 and 1.0.\n", - " - `scaling_extras` (z=0.5): two `IsothermalSph` masses with Einstein radii from the scaling relation.\n", - " - `source` (z=1.0): a small MGE light component (10 linear Gaussians)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 10\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "main_lens_einstein_radii = [4.0]\n", - "\n", - "lens_dict = {}\n", - "for i, (centre, er) in enumerate(zip(main_lens_centres, main_lens_einstein_radii)):\n", - " lens_dict[f\"lens_{i}\"] = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=tuple(centre), einstein_radius=er),\n", - " )\n", - "\n", - "individual_extras_einstein_radii = [0.8, 1.0]\n", - "individual_extras = [\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=tuple(centre), einstein_radius=er),\n", - " )\n", - " for centre, er in zip(extra_galaxies_centres, individual_extras_einstein_radii)\n", - "]\n", - "\n", - "scaling_factor = 0.3\n", - "scaling_exponent = 1.0\n", - "\n", - "scaling_extras = []\n", - "for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities):\n", - " einstein_radius = scaling_factor * luminosity**scaling_exponent\n", - " scaling_extras.append(\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(\n", - " centre=tuple(centre), einstein_radius=einstein_radius\n", - " ),\n", - " )\n", - " )\n", - "\n", - "source = al.Galaxy(redshift=1.0, bulge=build_source_basis(centre=(0.0, 0.1)))\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=list(lens_dict.values()) + individual_extras + scaling_extras + [source]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Three-Tier Deflection__\n", - "\n", - "The `Tracer.traced_grid_2d_list_from(...)` call performs the standard image-plane \u2192 source-plane ray-tracing.\n", - "Internally it queries every mass profile across all three tiers and sums their deflections.\n", - "\n", - "To make the decomposition concrete we re-compute the same source-plane grid by hand. Each profile exposes its\n", - "own `deflections_yx_2d_from`; the SUM of per-galaxy contributions across the three tiers is what the tracer\n", - "applies internally:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_grid = dataset.grid\n", - "\n", - "alpha_main_per_lens = [\n", - " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in lens_dict.values()\n", - "]\n", - "alpha_individual = [\n", - " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in individual_extras\n", - "]\n", - "alpha_scaling = [\n", - " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in scaling_extras\n", - "]\n", - "\n", - "alpha_main_total = sum(alpha_main_per_lens)\n", - "alpha_individual_total = sum(alpha_individual)\n", - "alpha_scaling_total = sum(alpha_scaling)\n", - "\n", - "alpha_total = alpha_main_total + alpha_individual_total + alpha_scaling_total\n", - "\n", - "print(f\"alpha_main_lens (tier sum, first coord) : {alpha_main_total[0]}\")\n", - "print(f\"alpha_individual (tier sum, first coord) : {alpha_individual_total[0]}\")\n", - "print(f\"alpha_scaling (tier sum, first coord) : {alpha_scaling_total[0]}\")\n", - "print(f\"alpha_total (across all, first coord): {alpha_total[0]}\")\n", - "\n", - "for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities):\n", - " er = scaling_factor * luminosity**scaling_exponent\n", - " print(\n", - " f\" scaling galaxy @ {tuple(centre)}: \"\n", - " f\"einstein_radius = {scaling_factor:.2f} * {luminosity:.3f} ** {scaling_exponent:.1f} = {er:.4f}\"\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Manual Ray-Tracing__\n", - "\n", - "The source-plane grid is the image-plane grid minus the total deflection. We compute it by hand and compare to\n", - "the `Tracer`-produced grid; they should be identical to within floating-point precision." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_source_manual = masked_grid - alpha_total\n", - "\n", - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=masked_grid)\n", - "grid_source_tracer = traced_grid_list[1]\n", - "\n", - "print(f\"\\nsource-plane grid (first coord, manual): {grid_source_manual[0]}\")\n", - "print(f\"source-plane grid (first coord, tracer): {grid_source_tracer[0]}\")\n", - "\n", - "assert np.allclose(np.asarray(grid_source_manual), np.asarray(grid_source_tracer))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Image__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image_unconvolved = tracer.image_2d_from(grid=masked_grid)\n", - "\n", - "aplt.plot_array(\n", - " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "What `image_2d_from` does internally for our group-scale three-tier lens:\n", - "\n", - " 1. Computes `alpha_lens(theta) = sum_i alpha_main_lens_i + sum_j alpha_extra_individual_j + sum_k alpha_extra_scaling_k`.\n", - " Each `alpha_extra_scaling_k` is the deflection of a profile whose `einstein_radius` was derived from\n", - " `scaling_factor * luminosity_k ** scaling_exponent`.\n", - " 2. Ray-traces the image-plane grid to obtain `grid_source = grid - alpha_lens`.\n", - " 3. Evaluates the source MGE at `grid_source`, producing its image-plane contribution.\n", - "\n", - "For a single-galaxy lens there is just one profile contributing to step 1; for a group with M main lens\n", - "galaxies, N individually-modelled extras, and K scaling-tier extras, there are `M + N + K` contributions. The\n", - "model gains `M * (mass parameters per main lens) + N` free Einstein-radius parameters plus 2 shared scaling\n", - "parameters \u2014 independent of K.\n", - "\n", - "__Likelihood__\n", - "\n", - "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model image\n", - "above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", - "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`. The MGE source's per-Gaussian\n", - "intensities are solved for via the linear-algebra step documented in\n", - "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py`.\n", - "\n", - "We delegate the remaining steps to `FitImaging`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(\n", - " f\"\\nLog likelihood of the manual group scaling-relation fit: {fit.log_likelihood}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The group-scale scaling-relation `log_likelihood` differs from a single-main-lens scaling case in exactly one\n", - "place: the lens-plane deflection sums over multiple main lens galaxies in addition to the two extras tiers.\n", - "Every other step (ray-tracing, source-plane evaluation, PSF convolution, chi-squared, noise normalization,\n", - "linear algebra) is shared with the standard imaging likelihood.\n", - "\n", - "This three-tier API is the production pattern for group-scale strong lenses with many foreground galaxies. It\n", - "lets the lens model stay tractable even when 10s or 100s of scaling-tier galaxies sit on the relation \u2014 the\n", - "lens-plane deflection is still a simple per-galaxy sum, but the scaling-tier contributions are parameterized\n", - "through luminosity rather than per-galaxy free parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Group Scaling Relation__\n", + "\n", + "This script describes the additional steps required to compute the `log_likelihood` for a group-scale strong\n", + "lens whose foreground galaxy population is split across three tiers \u2014 main lens galaxies (modelled via the\n", + "group `lens_dict` API), individually-modelled extras (each with its own free `einstein_radius`), and\n", + "scaling-tier extras (whose Einstein radii are derived from a shared reference-anchored relation\n", + "`einstein_radius = einstein_radius_ref * (luminosity / luminosity_ref) ** 0.5`, the Lenstool convention).\n", + "\n", + "This script does NOT repeat the steps shared with single-plane lensing (mask, image-plane grid, PSF convolution,\n", + "chi-squared, noise normalization, linear-algebra solver for MGE source intensities). It documents only the part\n", + "of the likelihood function which is specific to the group three-tier API: the multi-tier lens-plane deflection\n", + "composition.\n", + "\n", + "__Prerequisites__\n", + "\n", + " - `autolens_workspace/scripts/imaging/features/scaling_relation/likelihood_function.py` \u2014 the single-main-lens\n", + " walkthrough. The three-tier version below generalises that example across the `lens_dict` API.\n", + " - `autolens_workspace/scripts/group/start_here.py` \u2014 the group `lens_dict` API, including how\n", + " `main_lens_centres.json` is loaded.\n", + " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 canonical single-plane walkthrough.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a `Basis`\n", + " of linear Gaussians is solved for via linear algebra.\n", + " - `autolens_workspace/scripts/group/features/scaling_relation/modeling.py` \u2014 search-based three-tier version.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order (see above).\n", + "- **Dataset & Mask.**\n", + "- **Centres + Luminosities:** Load each tier's data.\n", + "- **Galaxies:** lens_dict + extras + scaling-tier + MGE source.\n", + "- **Three-Tier Deflection:** Per-tier deflection sums plus the scaling-relation evaluation.\n", + "- **Manual Ray-Tracing:** Hand-compute the source-plane grid and confirm it matches `Tracer`.\n", + "- **Source-Plane Image:** Source MGE evaluated at the ray-traced grid.\n", + "- **Fit Check.**\n", + "- **Wrap Up.**\n", + "\n", + "__What Changes For A Group Three-Tier Scaling Relation__\n", + "\n", + "For a single-galaxy lens with one mass profile, the lens-plane deflection is just one profile evaluated at\n", + "each image-plane coordinate:\n", + "\n", + " alpha_lens(theta) = alpha_total(theta ; Isothermal parameters)\n", + "\n", + "For a group-scale lens with the three-tier API, the lens-plane deflection is the SUM across every main lens\n", + "galaxy AND every extras / scaling-tier galaxy:\n", + "\n", + " alpha_lens(theta) = sum_i alpha_main_lens_i(theta)\n", + " + sum_j alpha_extra_individual_j(theta)\n", + " + sum_k alpha_extra_scaling_k(theta)\n", + "\n", + " where alpha_extra_scaling_k is the deflection of a mass profile whose\n", + " einstein_radius_k = einstein_radius_ref * (luminosity_k / luminosity_ref) ** 0.5.\n", + "\n", + "The model gains exactly 1 free parameter (`einstein_radius_ref`) regardless of how many galaxies\n", + "sit on the scaling tier. Every other step of the likelihood (PSF convolution, chi-squared, noise normalization,\n", + "MGE linear-algebra solver) is unchanged." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the group scaling_relation dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"scaling_relation\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/group/features/scaling_relation/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 8.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres + Luminosities__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "scaling_table = al.galaxy_table_from_csv(\n", + " file_path=dataset_path / \"scaling_galaxies.csv\"\n", + ")\n", + "scaling_galaxies_centres = scaling_table.centres\n", + "scaling_galaxies_luminosities = scaling_table.luminosities" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "Three populations participate in the ray-tracing:\n", + "\n", + " - `lens_dict` (z=0.5): one `IsothermalSph` mass per main lens centre \u2014 here, just one with `einstein_radius=4.0`.\n", + " - `individual_extras` (z=0.5): two `IsothermalSph` masses with simulator-true Einstein radii 0.8 and 1.0.\n", + " - `scaling_extras` (z=0.5): two `IsothermalSph` masses with Einstein radii from the scaling relation.\n", + " - `source` (z=1.0): a small MGE light component (10 linear Gaussians)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 10\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "main_lens_einstein_radii = [4.0]\n", + "\n", + "lens_dict = {}\n", + "for i, (centre, er) in enumerate(zip(main_lens_centres, main_lens_einstein_radii)):\n", + " lens_dict[f\"lens_{i}\"] = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=tuple(centre), einstein_radius=er),\n", + " )\n", + "\n", + "individual_extras_einstein_radii = [0.8, 1.0]\n", + "individual_extras = [\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=tuple(centre), einstein_radius=er),\n", + " )\n", + " for centre, er in zip(extra_galaxies_centres, individual_extras_einstein_radii)\n", + "]\n", + "\n", + "einstein_radius_ref = 0.135\n", + "scaling_exponent = 0.5\n", + "luminosity_ref = max(scaling_galaxies_luminosities)\n", + "\n", + "scaling_extras = []\n", + "for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities):\n", + " einstein_radius = (\n", + " einstein_radius_ref * (luminosity / luminosity_ref) ** scaling_exponent\n", + " )\n", + " scaling_extras.append(\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(\n", + " centre=tuple(centre), einstein_radius=einstein_radius\n", + " ),\n", + " )\n", + " )\n", + "\n", + "source = al.Galaxy(redshift=1.0, bulge=build_source_basis(centre=(0.0, 0.1)))\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=list(lens_dict.values()) + individual_extras + scaling_extras + [source]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Three-Tier Deflection__\n", + "\n", + "The `Tracer.traced_grid_2d_list_from(...)` call performs the standard image-plane \u2192 source-plane ray-tracing.\n", + "Internally it queries every mass profile across all three tiers and sums their deflections.\n", + "\n", + "To make the decomposition concrete we re-compute the same source-plane grid by hand. Each profile exposes its\n", + "own `deflections_yx_2d_from`; the SUM of per-galaxy contributions across the three tiers is what the tracer\n", + "applies internally:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_grid = dataset.grid\n", + "\n", + "alpha_main_per_lens = [\n", + " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in lens_dict.values()\n", + "]\n", + "alpha_individual = [\n", + " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in individual_extras\n", + "]\n", + "alpha_scaling = [\n", + " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in scaling_extras\n", + "]\n", + "\n", + "alpha_main_total = sum(alpha_main_per_lens)\n", + "alpha_individual_total = sum(alpha_individual)\n", + "alpha_scaling_total = sum(alpha_scaling)\n", + "\n", + "alpha_total = alpha_main_total + alpha_individual_total + alpha_scaling_total\n", + "\n", + "print(f\"alpha_main_lens (tier sum, first coord) : {alpha_main_total[0]}\")\n", + "print(f\"alpha_individual (tier sum, first coord) : {alpha_individual_total[0]}\")\n", + "print(f\"alpha_scaling (tier sum, first coord) : {alpha_scaling_total[0]}\")\n", + "print(f\"alpha_total (across all, first coord): {alpha_total[0]}\")\n", + "\n", + "for centre, luminosity in zip(scaling_galaxies_centres, scaling_galaxies_luminosities):\n", + " er = einstein_radius_ref * (luminosity / luminosity_ref) ** scaling_exponent\n", + " print(\n", + " f\" scaling galaxy @ {tuple(centre)}: \"\n", + " f\"einstein_radius = {einstein_radius_ref:.3f} * ({luminosity:.3f} / {luminosity_ref:.3f}) ** {scaling_exponent:.1f} = {er:.4f}\"\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Manual Ray-Tracing__\n", + "\n", + "The source-plane grid is the image-plane grid minus the total deflection. We compute it by hand and compare to\n", + "the `Tracer`-produced grid; they should be identical to within floating-point precision." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_source_manual = masked_grid - alpha_total\n", + "\n", + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=masked_grid)\n", + "grid_source_tracer = traced_grid_list[1]\n", + "\n", + "print(f\"\\nsource-plane grid (first coord, manual): {grid_source_manual[0]}\")\n", + "print(f\"source-plane grid (first coord, tracer): {grid_source_tracer[0]}\")\n", + "\n", + "assert np.allclose(np.asarray(grid_source_manual), np.asarray(grid_source_tracer))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Image__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image_unconvolved = tracer.image_2d_from(grid=masked_grid)\n", + "\n", + "aplt.plot_array(\n", + " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "What `image_2d_from` does internally for our group-scale three-tier lens:\n", + "\n", + " 1. Computes `alpha_lens(theta) = sum_i alpha_main_lens_i + sum_j alpha_extra_individual_j + sum_k alpha_extra_scaling_k`.\n", + " Each `alpha_extra_scaling_k` is the deflection of a profile whose `einstein_radius` was derived from\n", + " `einstein_radius_ref * (luminosity_k / luminosity_ref) ** 0.5`.\n", + " 2. Ray-traces the image-plane grid to obtain `grid_source = grid - alpha_lens`.\n", + " 3. Evaluates the source MGE at `grid_source`, producing its image-plane contribution.\n", + "\n", + "For a single-galaxy lens there is just one profile contributing to step 1; for a group with M main lens\n", + "galaxies, N individually-modelled extras, and K scaling-tier extras, there are `M + N + K` contributions. The\n", + "model gains `M * (mass parameters per main lens) + N` free Einstein-radius parameters plus 2 shared scaling\n", + "parameters \u2014 independent of K.\n", + "\n", + "__Likelihood__\n", + "\n", + "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model image\n", + "above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", + "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`. The MGE source's per-Gaussian\n", + "intensities are solved for via the linear-algebra step documented in\n", + "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py`.\n", + "\n", + "We delegate the remaining steps to `FitImaging`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(\n", + " f\"\\nLog likelihood of the manual group scaling-relation fit: {fit.log_likelihood}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The group-scale scaling-relation `log_likelihood` differs from a single-main-lens scaling case in exactly one\n", + "place: the lens-plane deflection sums over multiple main lens galaxies in addition to the two extras tiers.\n", + "Every other step (ray-tracing, source-plane evaluation, PSF convolution, chi-squared, noise normalization,\n", + "linear algebra) is shared with the standard imaging likelihood.\n", + "\n", + "This three-tier API is the production pattern for group-scale strong lenses with many foreground galaxies. It\n", + "lets the lens model stay tractable even when 10s or 100s of scaling-tier galaxies sit on the relation \u2014 the\n", + "lens-plane deflection is still a simple per-galaxy sum, but the scaling-tier contributions are parameterized\n", + "through luminosity rather than per-galaxy free parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/scaling_relation/modeling.ipynb b/notebooks/group/features/scaling_relation/modeling.ipynb index 74fc27d34..10069bedc 100644 --- a/notebooks/group/features/scaling_relation/modeling.ipynb +++ b/notebooks/group/features/scaling_relation/modeling.ipynb @@ -1,570 +1,621 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features (Group): Scaling Relations\n", - "============================================\n", - "\n", - "Group-scale strong lenses can have many galaxies in the foreground beyond the primary lens. As the number grows,\n", - "modelling each galaxy individually becomes impractical: a system with 10 companions would gain 10 extra Einstein-radius\n", - "free parameters, and the data is rarely informative enough to constrain them all.\n", - "\n", - "This example demonstrates the **three-tier modeling API** used by the production group pipelines, in which foreground\n", - "galaxies are split into three distinct populations:\n", - "\n", - " - **Main lens galaxies** (`main_lens_centres.json`): the primary lens(es). Modelled with an MGE bulge + free\n", - " `Isothermal` mass + `ExternalShear` (on `lens_0` only). These dominate the lensing.\n", - "\n", - " - **Extra galaxies** (`extra_galaxies_centres.json`): nearby companion galaxies modelled individually, each with its\n", - " own MGE bulge and bounded Einstein radius. Their light is fit and their mass is fit but constrained to a sensible\n", - " range. Use this tier for the brighter / closer companions that contribute non-trivially to the lensing on their own.\n", - "\n", - " - **Scaling galaxies** (`scaling_galaxies_centres.json`): further-out, fainter companions whose Einstein radii are\n", - " tied together via a shared scaling relation:\n", - "\n", - " einstein_radius = scaling_factor * (luminosity ** scaling_exponent)\n", - "\n", - " The free parameters are `scaling_factor` and `scaling_exponent` only \u2014 adding more scaling galaxies does not grow\n", - " the model. Use this tier for the long tail of fainter companions.\n", - "\n", - "Splitting galaxies across these three tiers is the standard pattern in production group fits (see\n", - "`z_projects/euclid_group/scripts/group.py`). It gives the lensing-significant galaxies the model flexibility they need\n", - "while keeping the model tractable as the number of foreground galaxies grows.\n", - "\n", - "For the simpler imaging-only counterpart \u2014 one main lens + one tier of extras on a scaling relation \u2014 see\n", - "`autolens_workspace/scripts/imaging/features/scaling_relation/modeling.py`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Three-Tier API:** Why split foreground galaxies into main, extra and scaling tiers.\n", - "- **Centres:** Three JSON files, one per tier, loaded with `al.from_json`.\n", - "- **Luminosities:** The scaling galaxies need a measured luminosity each; in this tutorial we hardcode them.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Main Lens Galaxies:** MGE bulge + free `Isothermal` mass; `ExternalShear` only on `lens_0`.\n", - "- **Extra Galaxies:** MGE bulge with fixed centre + `IsothermalSph` with bounded uniform `einstein_radius`.\n", - "- **Scaling Galaxies:** MGE bulge with fixed centre + `Isothermal` mass via shared scaling relation.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Search and Analysis:** Configure the non-linear search and run the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Three-Tier API__\n", - "\n", - "The three-tier split is the load-bearing idea here. To make the right choice for a given galaxy, ask:\n", - "\n", - " - Is it bright enough that fitting its light independently meaningfully helps the lens model? -> main or extra tier.\n", - " - Does it dominate the lensing? -> main tier.\n", - " - Is it close enough / bright enough to need a free Einstein radius? -> extra tier.\n", - " - Is it part of the long tail of fainter companions, where individually it contributes little but collectively it\n", - " matters? -> scaling tier.\n", - "\n", - "In this example we have one main galaxy, two extras, and two scaling galaxies, but the same code scales naturally to\n", - "many more on each tier \u2014 the JSON centre files and the per-galaxy loops are the only things that grow." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "This example uses its own dataset under `dataset/group/scaling_relation/`, simulated by the paired simulator at\n", - "`scripts/group/features/scaling_relation/simulator.py`. The dataset has three centre JSON files \u2014 one per tier \u2014 so\n", - "we exercise the full three-tier API." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"scaling_relation\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist, run the paired simulator first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/features/scaling_relation/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We use a slightly larger mask radius than `group/modeling.py` (8.5\") to enclose the scaling galaxies, which are placed\n", - "further out from the lens than the extras." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 8.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "Centres for the main and extra tiers come from JSON files (one (y, x) tuple per galaxy each). The scaling tier loads\n", - "its centres AND luminosities from a single CSV via `al.galaxy_table_from_csv` \u2014 see the next section." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "print(f\"Main lens centres: {main_lens_centres}\")\n", - "print(f\"Extra galaxies centres: {extra_galaxies_centres}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling Galaxy Centres + Luminosities__\n", - "\n", - "The scaling relation needs both centres AND a measured luminosity per scaling galaxy. There are two equally-supported\n", - "ways to provide them in PyAutoLens \u2014 both shown below so you can pick whichever fits your workflow.\n", - "\n", - "**Option A \u2014 CSV via `al.galaxy_table_from_csv` (recommended for non-trivial galaxy counts).** The simulator writes a\n", - "`scaling_galaxies.csv` with columns `y, x, luminosity` (and optional `redshift`) alongside the centre JSONs. We load it\n", - "in one call which returns a typed `GalaxyTable` with `.centres` (a `Grid2DIrregular`), `.luminosities`, and (optionally)\n", - "`.redshifts`. This scales naturally to populations of tens or hundreds of galaxies \u2014 the source of truth lives in a\n", - "single editable file.\n", - "\n", - "**Option B \u2014 JSON centres + hardcoded luminosity list (the original API, fine for short, fixed-length tutorials).**\n", - "Load the centres from `scaling_galaxies_centres.json` with `al.from_json` (the same loader used for the main and\n", - "extras tiers above) and define the luminosities as a Python list. Concise and obvious for small populations; awkward\n", - "once you have more than a handful.\n", - "\n", - "In a real analysis the luminosities come from a prior light-only fit. Two production patterns for obtaining them:\n", - "\n", - " - **Standalone light-only fit.** Run a single-stage non-linear search whose model is just MGE bulges for every galaxy\n", - " (no mass, no source). After the fit, compute total luminosity per galaxy from the bulge gaussian parameters:\n", - " `total_luminosity = sum(2 * pi * sigma**2 / axis_ratio * intensity) / pixel_scale**2`. The standalone example at\n", - " `scripts/group/features/scaling_relation/modeling_for_luminosities.py` writes its result as a `scaling_galaxies.csv`\n", - " in the dataset folder, which can then be loaded directly via Option A here.\n", - "\n", - " - **As the `source_lp[0]` stage of a SLaM pipeline.** Every group SLaM script defines a `source_lp_0` function whose\n", - " sole purpose is to fit a light-only MGE model to the main lens, extra galaxies and scaling galaxies in one go.\n", - " Subsequent stages chain from this result to compute luminosities and bound / scale the per-galaxy mass models.\n", - " See:\n", - "\n", - " scripts/group/slam.py (the canonical SLaM pipeline for group lenses)\n", - " scripts/group/features/pixelization/slam.py (pixelization variant)\n", - "\n", - " Search for `source_lp_0(` in either file \u2014 the pattern is identical and is documented in the function's header\n", - " docstring. The luminosity computation lives in the `source_lp_1` function that consumes the `source_lp_result_0`\n", - " result.\n", - "\n", - "We use Option A by default below. The Option B equivalent is shown commented out \u2014 uncomment it (and comment out\n", - "Option A) to switch." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Option A: CSV (recommended)\n", - "scaling_galaxies_table = al.galaxy_table_from_csv(\n", - " file_path=dataset_path / \"scaling_galaxies.csv\"\n", - ")\n", - "scaling_galaxies_centres = scaling_galaxies_table.centres\n", - "scaling_galaxies_luminosity_list = scaling_galaxies_table.luminosities\n", - "\n", - "# Option B: JSON centres + hardcoded luminosities (uncomment to use instead)\n", - "# scaling_galaxies_centres = al.from_json(\n", - "# file_path=dataset_path / \"scaling_galaxies_centres.json\"\n", - "# )\n", - "# scaling_galaxies_luminosity_list = [0.45, 0.45]\n", - "# assert len(scaling_galaxies_luminosity_list) == len(list(scaling_galaxies_centres)), (\n", - "# \"Number of scaling-galaxy luminosities must match the number of scaling-galaxy centres.\"\n", - "# )\n", - "\n", - "print(f\"Scaling galaxies centres: {scaling_galaxies_centres}\")\n", - "print(f\"Scaling galaxies luminosities: {scaling_galaxies_luminosity_list}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "One MGE bulge + free `Isothermal` mass per main lens; `ExternalShear` only on `lens_0`. Mirrors `group/modeling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "Each modelled with its own MGE bulge (fixed centre) + `IsothermalSph` mass with a bounded uniform `einstein_radius`.\n", - "Each adds 1 free Einstein-radius parameter to the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=tuple(centre)\n", - " )\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = tuple(centre)\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=1.5)\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling Galaxies__\n", - "\n", - "The scaling-relation tier. The two relation parameters are defined ONCE outside the loop \u2014 every scaling galaxy's\n", - "mass is a function of these same two parameters plus its own (fixed) luminosity.\n", - "\n", - "Adding more scaling galaxies (e.g. by lengthening the centres + luminosity lists) does not add any free parameters\n", - "to the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "scaling_exponent = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - "scaling_galaxies_list = []\n", - "\n", - "for scaling_galaxy_centre, scaling_galaxy_luminosity in zip(\n", - " scaling_galaxies_centres, scaling_galaxies_luminosity_list\n", - "):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_fixed=tuple(scaling_galaxy_centre),\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = tuple(scaling_galaxy_centre)\n", - " mass.einstein_radius = scaling_factor * scaling_galaxy_luminosity**scaling_exponent\n", - "\n", - " scaling_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " scaling_galaxies_list.append(scaling_galaxy)\n", - "\n", - "scaling_galaxies = af.Collection(scaling_galaxies_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "Each tier sits in its own top-level collection. This makes `model.info` and `result.info` easy to read \u2014 main lenses\n", - "appear under `galaxies`, individually-modelled companions under `extra_galaxies`, and the scaling-relation tier under\n", - "`scaling_galaxies`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=scaling_galaxies,\n", - ")\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling at every galaxy centre \u2014 main, extras and scaling alike." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = (\n", - " list(main_lens_centres)\n", - " + list(extra_galaxies_centres)\n", - " + list(scaling_galaxies_centres)\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\",\n", - " name=\"scaling_relation\",\n", - " unique_tag=dataset_name,\n", - " n_live=200,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "`result.info` shows all three tiers separately. The recovered `scaling_factor` and `scaling_exponent` should be\n", - "close to the truth values used by the simulator (0.3 and 1.0, given the simulator's chosen luminosities and Einstein\n", - "radii)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This example showed the full three-tier modeling API. The same structure scales naturally to systems with many more\n", - "foreground galaxies \u2014 the only thing that grows is the JSON centre files. The relation can also be applied to other\n", - "mass profiles or other measured quantities by swapping the `Isothermal` for any other `MassProfile` or the luminosity\n", - "for stellar mass / velocity dispersion.\n", - "\n", - "Related examples:\n", - "\n", - " - `autolens_workspace/scripts/imaging/features/scaling_relation/modeling.py` \u2014 the imaging-only counterpart, with one\n", - " `extra_galaxies` collection containing both individually-modelled and relational galaxies.\n", - " - `autolens_workspace/scripts/group/features/scaling_relation/modeling_for_luminosities.py` \u2014 the standalone\n", - " light-only fit that produces the `scaling_galaxies_luminosity_list` used here.\n", - " - `autolens_workspace/scripts/group/slam.py` and friends \u2014 the SLaM pipeline equivalent (`source_lp[0]` stage)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features (Group): Scaling Relations\n", + "============================================\n", + "\n", + "Group-scale strong lenses can have many galaxies in the foreground beyond the primary lens. As the number grows,\n", + "modelling each galaxy individually becomes impractical: a system with 10 companions would gain 10 extra Einstein-radius\n", + "free parameters, and the data is rarely informative enough to constrain them all.\n", + "\n", + "This example demonstrates the **three-tier modeling API** used by the production group pipelines, in which foreground\n", + "galaxies are split into three distinct populations:\n", + "\n", + " - **Main lens galaxies** (`main_lens_centres.json`): the primary lens(es). Modelled with an MGE bulge + free\n", + " `Isothermal` mass + `ExternalShear` (on `lens_0` only). These dominate the lensing.\n", + "\n", + " - **Extra galaxies** (`extra_galaxies_centres.json`): nearby companion galaxies modelled individually, each with its\n", + " own MGE bulge and bounded Einstein radius. Their light is fit and their mass is fit but constrained to a sensible\n", + " range. Use this tier for the brighter / closer companions that contribute non-trivially to the lensing on their own.\n", + "\n", + " - **Scaling galaxies** (`scaling_galaxies_centres.json`): further-out, fainter companions whose Einstein radii are\n", + " tied together via a shared scaling relation:\n", + "\n", + " einstein_radius = einstein_radius_ref * (luminosity / luminosity_ref) ** 0.5\n", + "\n", + " anchored to a *reference galaxy* (the brightest scaling-tier member), with the exponent fixed at the\n", + " Faber-Jackson value of 0.5 \u2014 the convention used by Lenstool and standard in published group- and\n", + " cluster-scale analyses. The only free parameter is `einstein_radius_ref` \u2014 adding more scaling galaxies\n", + " does not grow the model. Use this tier for the long tail of fainter companions.\n", + "\n", + "Splitting galaxies across these three tiers is the standard pattern in production group fits (see\n", + "`z_projects/euclid_group/scripts/group.py`). It gives the lensing-significant galaxies the model flexibility they need\n", + "while keeping the model tractable as the number of foreground galaxies grows.\n", + "\n", + "For the simpler imaging-only counterpart \u2014 one main lens + one tier of extras on a scaling relation \u2014 see\n", + "`autolens_workspace/scripts/imaging/features/scaling_relation/modeling.py`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Three-Tier API:** Why split foreground galaxies into main, extra and scaling tiers.\n", + "- **Centres:** Three JSON files, one per tier, loaded with `al.from_json`.\n", + "- **Luminosities:** The scaling galaxies need a measured luminosity each; in this tutorial we hardcode them.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Main Lens Galaxies:** MGE bulge + free `Isothermal` mass; `ExternalShear` only on `lens_0`.\n", + "- **Extra Galaxies:** MGE bulge with fixed centre + `IsothermalSph` with bounded uniform `einstein_radius`.\n", + "- **Scaling Galaxies:** MGE bulge with fixed centre + `Isothermal` mass via shared scaling relation.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Search and Analysis:** Configure the non-linear search and run the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Three-Tier API__\n", + "\n", + "The three-tier split is the load-bearing idea here. To make the right choice for a given galaxy, ask:\n", + "\n", + " - Is it bright enough that fitting its light independently meaningfully helps the lens model? -> main or extra tier.\n", + " - Does it dominate the lensing? -> main tier.\n", + " - Is it close enough / bright enough to need a free Einstein radius? -> extra tier.\n", + " - Is it part of the long tail of fainter companions, where individually it contributes little but collectively it\n", + " matters? -> scaling tier.\n", + "\n", + "In this example we have one main galaxy, two extras, and two scaling galaxies, but the same code scales naturally to\n", + "many more on each tier \u2014 the JSON centre files and the per-galaxy loops are the only things that grow." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "This example uses its own dataset under `dataset/group/scaling_relation/`, simulated by the paired simulator at\n", + "`scripts/group/features/scaling_relation/simulator.py`. The dataset has three centre JSON files \u2014 one per tier \u2014 so\n", + "we exercise the full three-tier API." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"scaling_relation\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist, run the paired simulator first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/features/scaling_relation/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We use a slightly larger mask radius than `group/modeling.py` (8.5\") to enclose the scaling galaxies, which are placed\n", + "further out from the lens than the extras." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 8.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "Centres for the main and extra tiers come from JSON files (one (y, x) tuple per galaxy each). The scaling tier loads\n", + "its centres AND luminosities from a single CSV via `al.galaxy_table_from_csv` \u2014 see the next section." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "print(f\"Main lens centres: {main_lens_centres}\")\n", + "print(f\"Extra galaxies centres: {extra_galaxies_centres}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling Galaxy Centres + Luminosities__\n", + "\n", + "The scaling relation needs both centres AND a measured luminosity per scaling galaxy. There are two equally-supported\n", + "ways to provide them in PyAutoLens \u2014 both shown below so you can pick whichever fits your workflow.\n", + "\n", + "**Option A \u2014 CSV via `al.galaxy_table_from_csv` (recommended for non-trivial galaxy counts).** The simulator writes a\n", + "`scaling_galaxies.csv` with columns `y, x, luminosity` (and optional `redshift`) alongside the centre JSONs. We load it\n", + "in one call which returns a typed `GalaxyTable` with `.centres` (a `Grid2DIrregular`), `.luminosities`, and (optionally)\n", + "`.redshifts`. This scales naturally to populations of tens or hundreds of galaxies \u2014 the source of truth lives in a\n", + "single editable file.\n", + "\n", + "**Option B \u2014 JSON centres + hardcoded luminosity list (the original API, fine for short, fixed-length tutorials).**\n", + "Load the centres from `scaling_galaxies_centres.json` with `al.from_json` (the same loader used for the main and\n", + "extras tiers above) and define the luminosities as a Python list. Concise and obvious for small populations; awkward\n", + "once you have more than a handful.\n", + "\n", + "In a real analysis the luminosities come from a prior light-only fit. Two production patterns for obtaining them:\n", + "\n", + " - **Standalone light-only fit.** Run a single-stage non-linear search whose model is just MGE bulges for every galaxy\n", + " (no mass, no source). After the fit, compute total luminosity per galaxy from the bulge gaussian parameters:\n", + " `total_luminosity = sum(2 * pi * sigma**2 / axis_ratio * intensity) / pixel_scale**2`. The standalone example at\n", + " `scripts/group/features/scaling_relation/modeling_for_luminosities.py` writes its result as a `scaling_galaxies.csv`\n", + " in the dataset folder, which can then be loaded directly via Option A here.\n", + "\n", + " - **As the `source_lp[0]` stage of a SLaM pipeline.** Every group SLaM script defines a `source_lp_0` function whose\n", + " sole purpose is to fit a light-only MGE model to the main lens, extra galaxies and scaling galaxies in one go.\n", + " Subsequent stages chain from this result to compute luminosities and bound / scale the per-galaxy mass models.\n", + " See:\n", + "\n", + " scripts/group/slam.py (the canonical SLaM pipeline for group lenses)\n", + " scripts/group/features/pixelization/slam.py (pixelization variant)\n", + "\n", + " Search for `source_lp_0(` in either file \u2014 the pattern is identical and is documented in the function's header\n", + " docstring. The luminosity computation lives in the `source_lp_1` function that consumes the `source_lp_result_0`\n", + " result.\n", + "\n", + "We use Option A by default below. The Option B equivalent is shown commented out \u2014 uncomment it (and comment out\n", + "Option A) to switch." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Option A: CSV (recommended)\n", + "scaling_galaxies_table = al.galaxy_table_from_csv(\n", + " file_path=dataset_path / \"scaling_galaxies.csv\"\n", + ")\n", + "scaling_galaxies_centres = scaling_galaxies_table.centres\n", + "scaling_galaxies_luminosity_list = scaling_galaxies_table.luminosities\n", + "\n", + "# Option B: JSON centres + hardcoded luminosities (uncomment to use instead)\n", + "# scaling_galaxies_centres = al.from_json(\n", + "# file_path=dataset_path / \"scaling_galaxies_centres.json\"\n", + "# )\n", + "# scaling_galaxies_luminosity_list = [0.45, 0.45]\n", + "# assert len(scaling_galaxies_luminosity_list) == len(list(scaling_galaxies_centres)), (\n", + "# \"Number of scaling-galaxy luminosities must match the number of scaling-galaxy centres.\"\n", + "# )\n", + "\n", + "print(f\"Scaling galaxies centres: {scaling_galaxies_centres}\")\n", + "print(f\"Scaling galaxies luminosities: {scaling_galaxies_luminosity_list}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "One MGE bulge + free `Isothermal` mass per main lens; `ExternalShear` only on `lens_0`. Mirrors `group/modeling.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "Each modelled with its own MGE bulge (fixed centre) + `IsothermalSph` mass with a bounded uniform `einstein_radius`.\n", + "Each adds 1 free Einstein-radius parameter to the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=tuple(centre)\n", + " )\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = tuple(centre)\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=1.5)\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling Galaxies__\n", + "\n", + "The scaling-relation tier, in the reference-anchored convention used by Lenstool and essentially every published\n", + "group- and cluster-scale analysis (Limousin et al. 2005; Eliasdottir et al. 2007; Bergamini et al. 2019). The\n", + "normalization ``einstein_radius_ref`` is the Einstein radius of the *brightest* scaling member \u2014 a physically\n", + "interpretable quantity with an easy-to-motivate prior range \u2014 and it is defined ONCE outside the loop: every\n", + "scaling galaxy's mass derives from it via its luminosity ratio to the reference. The exponent is *fixed* at the\n", + "Faber-Jackson value (einstein_radius \u221d sigma\u00b2 and sigma \u221d L^(1/4) give einstein_radius \u221d L^(1/2)) rather than\n", + "fitted, avoiding the normalization-slope degeneracy. Only luminosity ratios enter, so the luminosity units are\n", + "irrelevant; magnitude catalogues convert via ``L / L_ref = 10 ** (0.4 * (m_ref - m))``.\n", + "\n", + "The dPIE-profile cluster-scale analogue \u2014 which also scales the truncation radius (``rs \u221d L^0.5``, mirroring\n", + "Lenstool's r_cut scaling) \u2014 is ``scripts/cluster/modeling.py``. To free the exponent as a systematics test,\n", + "replace the fixed value with e.g. ``af.UniformPrior(lower_limit=0.0, upper_limit=1.0)``.\n", + "\n", + "Adding more scaling galaxies (e.g. by lengthening the centres + luminosity lists) does not add any free parameters\n", + "to the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "einstein_radius_ref = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "scaling_exponent = 0.5\n", + "\n", + "luminosity_ref = max(scaling_galaxies_luminosity_list)\n", + "\n", + "scaling_galaxies_list = []\n", + "\n", + "for scaling_galaxy_centre, scaling_galaxy_luminosity in zip(\n", + " scaling_galaxies_centres, scaling_galaxies_luminosity_list\n", + "):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_fixed=tuple(scaling_galaxy_centre),\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = tuple(scaling_galaxy_centre)\n", + " luminosity_ratio = scaling_galaxy_luminosity / luminosity_ref\n", + " mass.einstein_radius = einstein_radius_ref * luminosity_ratio**scaling_exponent\n", + "\n", + " scaling_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " scaling_galaxies_list.append(scaling_galaxy)\n", + "\n", + "scaling_galaxies = af.Collection(scaling_galaxies_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "Each tier sits in its own top-level collection. This makes `model.info` and `result.info` easy to read \u2014 main lenses\n", + "appear under `galaxies`, individually-modelled companions under `extra_galaxies`, and the scaling-relation tier under\n", + "`scaling_galaxies`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=scaling_galaxies,\n", + ")\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling at every galaxy centre \u2014 main, extras and scaling alike." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = (\n", + " list(main_lens_centres)\n", + " + list(extra_galaxies_centres)\n", + " + list(scaling_galaxies_centres)\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\",\n", + " name=\"scaling_relation\",\n", + " unique_tag=dataset_name,\n", + " n_live=200,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "`result.info` shows all three tiers separately. The recovered `einstein_radius_ref` should be close to the truth\n", + "value used by the simulator (0.135, the Einstein radius of the brightest scaling member)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This example showed the full three-tier modeling API. The same structure scales naturally to systems with many more\n", + "foreground galaxies \u2014 the only thing that grows is the JSON centre files. The relation can also be applied to other\n", + "mass profiles or other measured quantities by swapping the `Isothermal` for any other `MassProfile` or the luminosity\n", + "for stellar mass / velocity dispersion.\n", + "\n", + "Related examples:\n", + "\n", + " - `autolens_workspace/scripts/imaging/features/scaling_relation/modeling.py` \u2014 the imaging-only counterpart, with one\n", + " `extra_galaxies` collection containing both individually-modelled and relational galaxies.\n", + " - `autolens_workspace/scripts/group/features/scaling_relation/modeling_for_luminosities.py` \u2014 the standalone\n", + " light-only fit that produces the `scaling_galaxies_luminosity_list` used here.\n", + " - `autolens_workspace/scripts/group/slam.py` and friends \u2014 the SLaM pipeline equivalent (`source_lp[0]` stage)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/scaling_relation/modeling_for_luminosities.ipynb b/notebooks/group/features/scaling_relation/modeling_for_luminosities.ipynb index 2fecbf759..6ae72af08 100644 --- a/notebooks/group/features/scaling_relation/modeling_for_luminosities.ipynb +++ b/notebooks/group/features/scaling_relation/modeling_for_luminosities.ipynb @@ -1,462 +1,499 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features (Group): Fit Galaxy Luminosities for a Scaling Relation\n", - "=========================================================================\n", - "\n", - "The scaling-relation modeling examples (`modeling.py` in this directory and\n", - "`scripts/imaging/features/scaling_relation/modeling.py`) need a measured **luminosity** for every galaxy that sits on\n", - "the relation:\n", - "\n", - " einstein_radius = scaling_factor * (luminosity ** scaling_exponent)\n", - "\n", - "Those tutorials hardcode the luminosity list for readability. In a production fit the luminosities have to be measured\n", - "from the data itself. This example shows the standard standalone way to do that:\n", - "\n", - " 1. Load the dataset and all three centre JSON files (main, extras, scaling).\n", - " 2. Build a model with **light only** \u2014 MGE bulges for every galaxy, no mass profiles, no source.\n", - " 3. Fit that model with a single non-linear search.\n", - " 4. Compute the total luminosity per galaxy from the fitted MGE bulge gaussians.\n", - " 5. Print the luminosities so they can be pasted into the scaling-relation modeling script (or write them to JSON for\n", - " re-use).\n", - "\n", - "Why a separate light-only fit? Because the lensing model is dominated by the **mass** of the galaxies, but the scaling\n", - "relation needs the **light** of each galaxy as an *input*. Fitting the light first decouples the two, gives the relation\n", - "a stable per-galaxy measurement, and avoids degeneracies between the relation parameters and individual luminosities.\n", - "\n", - "The same idea is implemented inside the production SLaM pipelines as the `source_lp_0` stage. See\n", - "`scripts/group/slam.py` and `scripts/group/features/pixelization/slam.py` for the chained-pipeline equivalent. This\n", - "script is the standalone, single-stage version of that step \u2014 useful for users who don't want to commit to a full SLaM\n", - "pipeline just to obtain luminosities.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask.\n", - "- **Centres:** Load all three centre JSON files.\n", - "- **Light-only Model:** MGE bulge per galaxy, no mass, no source.\n", - "- **Search and Analysis:** Configure the non-linear search and run the model-fit.\n", - "- **Compute Luminosities:** Extract `total_luminosity` per galaxy from the fitted MGE bulge gaussians.\n", - "- **Output:** Print and save the luminosities for the scaling-relation modeling script.\n", - "- **Wrap Up:** What to do next.\n", - "\n", - "__Light-only Model Sizing__\n", - "\n", - "The MGE bulge for the main lens uses 30 gaussians across 2 bases (matching the production SLaM `source_lp_0` stage).\n", - "The extras and scaling galaxies use 10 gaussians each in 1 basis \u2014 the smaller companions don't need the extra\n", - "flexibility and shrinking the basis keeps the model dimensionality low.\n", - "\n", - "__Output Galaxy Order__\n", - "\n", - "After the fit, the tracer's galaxies appear in the order:\n", - "\n", - " main_lenses[0..n_main-1], extras[0..n_extra-1], scaling[0..n_scaling-1]\n", - "\n", - "Use this ordering when indexing `tracer.galaxies` to compute per-galaxy luminosities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"scaling_relation\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)\n", - "\n", - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/features/scaling_relation/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "pixel_scale = float(dataset.pixel_scales[0])\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 8.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "scaling_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"scaling_galaxies_centres.json\"\n", - ")\n", - "\n", - "n_main = len(list(main_lens_centres))\n", - "n_extra = len(list(extra_galaxies_centres))\n", - "n_scaling = len(list(scaling_galaxies_centres))\n", - "\n", - "print(f\"Main lens galaxies: {n_main}\")\n", - "print(f\"Extra galaxies: {n_extra}\")\n", - "print(f\"Scaling galaxies: {n_scaling}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Light-only Model__\n", - "\n", - "Each galaxy gets an MGE bulge and nothing else \u2014 no mass, no source. The main lens uses a richer 30-gaussian / 2-basis\n", - "MGE because it dominates the light; companions use 10 gaussians in 1 basis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=False,\n", - " centre=tuple(centre),\n", - " centre_sigma=0.1,\n", - " )\n", - " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, redshift=0.5, bulge=bulge)\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_prior_is_uniform=True,\n", - " centre=tuple(centre),\n", - " ell_comps_prior_is_uniform=True,\n", - " )\n", - " extra_galaxies_list.append(af.Model(al.Galaxy, redshift=0.5, bulge=bulge))\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "scaling_galaxies_list = []\n", - "\n", - "for centre in scaling_galaxies_centres:\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_prior_is_uniform=True,\n", - " centre=tuple(centre),\n", - " ell_comps_prior_is_uniform=True,\n", - " )\n", - " scaling_galaxies_list.append(af.Model(al.Galaxy, redshift=0.5, bulge=bulge))\n", - "\n", - "scaling_galaxies = af.Collection(scaling_galaxies_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "No source galaxy: this is purely a light-only fit. Keep the three populations in their own collections so the post-fit\n", - "tracer ordering is predictable: main lenses first, then extras, then scaling galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=scaling_galaxies,\n", - ")\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = (\n", - " list(main_lens_centres)\n", - " + list(extra_galaxies_centres)\n", - " + list(scaling_galaxies_centres)\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "A single Nautilus search. With no mass and no source, the parameter space is much smaller than a full lens fit so we\n", - "can use a smaller `n_live` than `modeling.py` and converge faster." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\") / \"features\" / \"scaling_relation\",\n", - " name=\"modeling_for_luminosities\",\n", - " unique_tag=dataset_name,\n", - " n_live=100 + 30 * (n_main + n_extra + n_scaling),\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Compute Luminosities__\n", - "\n", - "The fit's `tracer_linear_light_profiles_to_light_profiles` property converts the linear-inversion-solved intensities\n", - "back to standard `LightProfile`s with concrete `intensity` values. This is the canonical way to read MGE intensities\n", - "out of a result.\n", - "\n", - "For an MGE basis composed of 2D Gaussians, the per-gaussian luminosity is:\n", - "\n", - " L_g = 2 * pi * sigma^2 / axis_ratio * intensity\n", - "\n", - "The total luminosity of the galaxy is the sum of `L_g` over all gaussians in its `bulge.profile_list`, divided by\n", - "`pixel_scale**2` to convert from per-pixel to per-arcsec^2.\n", - "\n", - "This is the same formula used by the SLaM pipelines in `source_lp_1` (see `scripts/group/slam.py`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit_tracer = (\n", - " result.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", - ")\n", - "\n", - "\n", - "def total_luminosity_from(galaxy):\n", - " return (\n", - " sum(\n", - " 2.0 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", - " for g in galaxy.bulge.profile_list\n", - " )\n", - " / pixel_scale**2\n", - " )\n", - "\n", - "\n", - "main_luminosities = [\n", - " total_luminosity_from(fit_tracer.galaxies[i]) for i in range(n_main)\n", - "]\n", - "extra_luminosities = [\n", - " total_luminosity_from(fit_tracer.galaxies[n_main + i]) for i in range(n_extra)\n", - "]\n", - "scaling_luminosities = [\n", - " total_luminosity_from(fit_tracer.galaxies[n_main + n_extra + i])\n", - " for i in range(n_scaling)\n", - "]\n", - "\n", - "print(\"Main lens luminosities:\", main_luminosities)\n", - "print(\"Extra galaxy luminosities:\", extra_luminosities)\n", - "print(\"Scaling galaxy luminosities:\", scaling_luminosities)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Write the centres + luminosities to a `scaling_galaxies.csv` next to the centre JSONs. This is the same file the\n", - "scaling-relation modeling script consumes via `al.galaxy_table_from_csv`, so the result of this fit can be chained\n", - "into the next step with no manual copy/paste." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "csv_path = dataset_path / \"scaling_galaxies.csv\"\n", - "\n", - "al.galaxy_table_to_csv(\n", - " centres=list(scaling_galaxies_centres),\n", - " luminosities=[float(l) for l in scaling_luminosities],\n", - " file_path=csv_path,\n", - ")\n", - "\n", - "print(f\"Wrote scaling-galaxy centres + luminosities to: {csv_path}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The CSV at `dataset_path / \"scaling_galaxies.csv\"` is the canonical chain-point. The downstream modeling script\n", - "loads it directly:\n", - "\n", - " table = al.galaxy_table_from_csv(file_path=dataset_path / \"scaling_galaxies.csv\")\n", - " scaling_galaxies_centres = table.centres\n", - " scaling_galaxies_luminosity_list = table.luminosities\n", - "\n", - "Re-running this script overwrites the CSV in place, so iterating on the light fit and re-running the lens fit is\n", - "just two `python ...` invocations.\n", - "\n", - "For the chained-pipeline alternative \u2014 where the light fit is the `source_lp[0]` stage of a SLaM run instead of a\n", - "standalone search \u2014 see `scripts/group/slam.py` (`__SOURCE LP PIPELINE \u2014 stage 0__`) and\n", - "`scripts/group/features/pixelization/slam.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features (Group): Fit Galaxy Luminosities for a Scaling Relation\n", + "=========================================================================\n", + "\n", + "The scaling-relation modeling examples (`modeling.py` in this directory and\n", + "`scripts/imaging/features/scaling_relation/modeling.py`) need a measured **luminosity** for every galaxy that sits on\n", + "the relation:\n", + "\n", + " einstein_radius = einstein_radius_ref * (luminosity / luminosity_ref) ** 0.5\n", + "\n", + "Those tutorials hardcode the luminosity list for readability. In a production fit the luminosities have to be measured\n", + "from the data itself. This example shows the standard standalone way to do that:\n", + "\n", + " 1. Load the dataset and all three centre JSON files (main, extras, scaling).\n", + " 2. Build a model with **light only** \u2014 MGE bulges for every galaxy, no mass profiles, no source.\n", + " 3. Fit that model with a single non-linear search.\n", + " 4. Compute the total luminosity per galaxy from the fitted MGE bulge gaussians.\n", + " 5. Print the luminosities so they can be pasted into the scaling-relation modeling script (or write them to JSON for\n", + " re-use).\n", + "\n", + "Why a separate light-only fit? Because the lensing model is dominated by the **mass** of the galaxies, but the scaling\n", + "relation needs the **light** of each galaxy as an *input*. Fitting the light first decouples the two, gives the relation\n", + "a stable per-galaxy measurement, and avoids degeneracies between the relation parameters and individual luminosities.\n", + "\n", + "The same idea is implemented inside the production SLaM pipelines as the `source_lp_0` stage. See\n", + "`scripts/group/slam.py` and `scripts/group/features/pixelization/slam.py` for the chained-pipeline equivalent. This\n", + "script is the standalone, single-stage version of that step \u2014 useful for users who don't want to commit to a full SLaM\n", + "pipeline just to obtain luminosities.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask.\n", + "- **Centres:** Load all three centre JSON files.\n", + "- **Light-only Model:** MGE bulge per galaxy, no mass, no source.\n", + "- **Search and Analysis:** Configure the non-linear search and run the model-fit.\n", + "- **Compute Luminosities:** Extract `total_luminosity` per galaxy from the fitted MGE bulge gaussians.\n", + "- **Output:** Print and save the luminosities for the scaling-relation modeling script.\n", + "- **Wrap Up:** What to do next.\n", + "\n", + "__Light-only Model Sizing__\n", + "\n", + "The MGE bulge for the main lens uses 30 gaussians across 2 bases (matching the production SLaM `source_lp_0` stage).\n", + "The extras and scaling galaxies use 10 gaussians each in 1 basis \u2014 the smaller companions don't need the extra\n", + "flexibility and shrinking the basis keeps the model dimensionality low.\n", + "\n", + "__Output Galaxy Order__\n", + "\n", + "After the fit, the tracer's galaxies appear in the order:\n", + "\n", + " main_lenses[0..n_main-1], extras[0..n_extra-1], scaling[0..n_scaling-1]\n", + "\n", + "Use this ordering when indexing `tracer.galaxies` to compute per-galaxy luminosities." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"scaling_relation\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/features/scaling_relation/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "pixel_scale = float(dataset.pixel_scales[0])\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 8.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "scaling_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"scaling_galaxies_centres.json\"\n", + ")\n", + "\n", + "n_main = len(list(main_lens_centres))\n", + "n_extra = len(list(extra_galaxies_centres))\n", + "n_scaling = len(list(scaling_galaxies_centres))\n", + "\n", + "print(f\"Main lens galaxies: {n_main}\")\n", + "print(f\"Extra galaxies: {n_extra}\")\n", + "print(f\"Scaling galaxies: {n_scaling}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Light-only Model__\n", + "\n", + "Each galaxy gets an MGE bulge and nothing else \u2014 no mass, no source. The main lens uses a richer 30-gaussian / 2-basis\n", + "MGE because it dominates the light; companions use 10 gaussians in 1 basis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=False,\n", + " centre=tuple(centre),\n", + " centre_sigma=0.1,\n", + " )\n", + " lens_dict[f\"lens_{i}\"] = af.Model(al.Galaxy, redshift=0.5, bulge=bulge)\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_prior_is_uniform=True,\n", + " centre=tuple(centre),\n", + " ell_comps_prior_is_uniform=True,\n", + " )\n", + " extra_galaxies_list.append(af.Model(al.Galaxy, redshift=0.5, bulge=bulge))\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "scaling_galaxies_list = []\n", + "\n", + "for centre in scaling_galaxies_centres:\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_prior_is_uniform=True,\n", + " centre=tuple(centre),\n", + " ell_comps_prior_is_uniform=True,\n", + " )\n", + " scaling_galaxies_list.append(af.Model(al.Galaxy, redshift=0.5, bulge=bulge))\n", + "\n", + "scaling_galaxies = af.Collection(scaling_galaxies_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "No source galaxy: this is purely a light-only fit. Keep the three populations in their own collections so the post-fit\n", + "tracer ordering is predictable: main lenses first, then extras, then scaling galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=scaling_galaxies,\n", + ")\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = (\n", + " list(main_lens_centres)\n", + " + list(extra_galaxies_centres)\n", + " + list(scaling_galaxies_centres)\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "A single Nautilus search. With no mass and no source, the parameter space is much smaller than a full lens fit so we\n", + "can use a smaller `n_live` than `modeling.py` and converge faster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\") / \"features\" / \"scaling_relation\",\n", + " name=\"modeling_for_luminosities\",\n", + " unique_tag=dataset_name,\n", + " n_live=100 + 30 * (n_main + n_extra + n_scaling),\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Compute Luminosities__\n", + "\n", + "The fit's `tracer_linear_light_profiles_to_light_profiles` property converts the linear-inversion-solved intensities\n", + "back to standard `LightProfile`s with concrete `intensity` values. This is the canonical way to read MGE intensities\n", + "out of a result.\n", + "\n", + "For an MGE basis composed of 2D Gaussians, the per-gaussian luminosity is:\n", + "\n", + " L_g = 2 * pi * sigma^2 / axis_ratio * intensity\n", + "\n", + "The total luminosity of the galaxy is the sum of `L_g` over all gaussians in its `bulge.profile_list`, divided by\n", + "`pixel_scale**2` to convert from per-pixel to per-arcsec^2.\n", + "\n", + "This is the same formula used by the SLaM pipelines in `source_lp_1` (see `scripts/group/slam.py`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_tracer = (\n", + " result.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", + ")\n", + "\n", + "\n", + "def total_luminosity_from(galaxy):\n", + " return (\n", + " sum(\n", + " 2.0 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", + " for g in galaxy.bulge.profile_list\n", + " )\n", + " / pixel_scale**2\n", + " )\n", + "\n", + "\n", + "main_luminosities = [\n", + " total_luminosity_from(fit_tracer.galaxies[i]) for i in range(n_main)\n", + "]\n", + "extra_luminosities = [\n", + " total_luminosity_from(fit_tracer.galaxies[n_main + i]) for i in range(n_extra)\n", + "]\n", + "scaling_luminosities = [\n", + " total_luminosity_from(fit_tracer.galaxies[n_main + n_extra + i])\n", + " for i in range(n_scaling)\n", + "]\n", + "\n", + "print(\"Main lens luminosities:\", main_luminosities)\n", + "print(\"Extra galaxy luminosities:\", extra_luminosities)\n", + "print(\"Scaling galaxy luminosities:\", scaling_luminosities)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Write the centres + luminosities to a `scaling_galaxies.csv` next to the centre JSONs. This is the same file the\n", + "scaling-relation modeling script consumes via `al.galaxy_table_from_csv`, so the result of this fit can be chained\n", + "into the next step with no manual copy/paste." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "csv_path = dataset_path / \"scaling_galaxies.csv\"\n", + "\n", + "al.galaxy_table_to_csv(\n", + " centres=list(scaling_galaxies_centres),\n", + " luminosities=[float(l) for l in scaling_luminosities],\n", + " file_path=csv_path,\n", + ")\n", + "\n", + "print(f\"Wrote scaling-galaxy centres + luminosities to: {csv_path}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The CSV at `dataset_path / \"scaling_galaxies.csv\"` is the canonical chain-point. The downstream modeling script\n", + "loads it directly:\n", + "\n", + " table = al.galaxy_table_from_csv(file_path=dataset_path / \"scaling_galaxies.csv\")\n", + " scaling_galaxies_centres = table.centres\n", + " scaling_galaxies_luminosity_list = table.luminosities\n", + "\n", + "Re-running this script overwrites the CSV in place, so iterating on the light fit and re-running the lens fit is\n", + "just two `python ...` invocations.\n", + "\n", + "For the chained-pipeline alternative \u2014 where the light fit is the `source_lp[0]` stage of a SLaM run instead of a\n", + "standalone search \u2014 see `scripts/group/slam.py` (`__SOURCE LP PIPELINE \u2014 stage 0__`) and\n", + "`scripts/group/features/pixelization/slam.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/features/scaling_relation/simulator.ipynb b/notebooks/group/features/scaling_relation/simulator.ipynb index e7449c1c8..e76b43e98 100644 --- a/notebooks/group/features/scaling_relation/simulator.ipynb +++ b/notebooks/group/features/scaling_relation/simulator.ipynb @@ -1,504 +1,543 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Group Scaling Relation\n", - "=================================\n", - "\n", - "This script simulates a group-scale strong lens with three populations of galaxies in the foreground, designed to\n", - "exercise the three-tier modeling API used by `group/features/scaling_relation/modeling.py`:\n", - "\n", - " - One **main lens galaxy** at the origin, which dominates the light and mass of the system.\n", - " - Two **extra galaxies** offset from the lens, modelled individually in the fit (one Einstein radius per galaxy).\n", - " - Two **scaling galaxies** further out, modelled via a shared luminosity-mass scaling relation.\n", - "\n", - "Each population's centres are saved to a separate JSON file (`main_lens_centres.json`, `extra_galaxies_centres.json`,\n", - "`scaling_galaxies_centres.json`) so the modeling script can load them directly.\n", - "\n", - "This dataset is independent of `dataset/group/simple` so the existing group examples are unaffected." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"scaling_relation\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "A 250x250 0.1\"/pixel grid wide enough to contain all five galaxies plus the lensed source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(250, 250),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Three populations:\n", - "\n", - " - `main_lens_centres`: primary lens at the origin.\n", - " - `extra_galaxies_centres`: closer companions modelled individually in the fit.\n", - " - `scaling_galaxies_centres`: further-out, fainter companions modelled via a shared scaling relation.\n", - "\n", - "The scaling galaxies are deliberately placed further out and given fainter light profiles below \u2014 this matches the\n", - "typical observational scenario in which the scaling-relation tier is reserved for galaxies that contribute only a small\n", - "amount of lensing individually but matter collectively." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [(0.0, 0.0)]\n", - "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]\n", - "scaling_galaxies_centres = [(6.5, 0.0), (-1.0, 7.0)]\n", - "\n", - "all_galaxy_centres = (\n", - " main_lens_centres + extra_galaxies_centres + scaling_galaxies_centres\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling is applied at every galaxy centre (main + extras + scaling) for accurate light evaluation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_galaxy_centres,\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A simple Gaussian PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")\n", - "\n", - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxy__\n", - "\n", - "A bright spherical Sersic + Isothermal mass at the origin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "main_lens_galaxies = [lens_0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "Two companion galaxies modelled with their own light + mass. These are the brighter, closer-in tier; in the modeling\n", - "script they receive their own free `einstein_radius` parameter each." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling Galaxies__\n", - "\n", - "Two further-out, fainter companions whose true Einstein radii are consistent with the relation\n", - "``einstein_radius = 0.3 * luminosity ** 1.0`` (luminosities ~0.45 -> Einstein radii ~0.135). Modelling these\n", - "individually would add 2 free parameters; on a scaling relation they add zero (the 2 relation parameters are shared\n", - "across all scaling galaxies, so adding more does not grow the model)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scaling_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(6.5, 0.0), intensity=0.45, effective_radius=0.6, sersic_index=2.5\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(6.5, 0.0), einstein_radius=0.135),\n", - ")\n", - "\n", - "scaling_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-1.0, 7.0), intensity=0.45, effective_radius=0.6, sersic_index=2.5\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-1.0, 7.0), einstein_radius=0.135),\n", - ")\n", - "\n", - "scaling_galaxies = [scaling_galaxy_0, scaling_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Tracer order: main lenses, extras, scaling galaxies, source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=main_lens_galaxies + extra_galaxies + scaling_galaxies + [source_galaxy]\n", - ")\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the truth tracer for later inspection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centre JSON Files__\n", - "\n", - "One JSON per population, loaded individually by the modeling script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(scaling_galaxies_centres),\n", - " file_path=Path(dataset_path, \"scaling_galaxies_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Population CSVs__\n", - "\n", - "The modeling script loads luminosities (and centres) for both the extras and the scaling tier from CSVs written here.\n", - "The simulator knows the truth values of the per-galaxy luminosities so we write them out alongside the centre JSONs.\n", - "\n", - "The CSV schema is `y, x, luminosity, redshift?` \u2014 see `al.galaxy_table_from_csv` /\n", - "`al.galaxy_table_to_csv` (`autogalaxy/galaxy/galaxy_table.py`). Centre JSONs above are kept for backward compatibility;\n", - "new consumers should prefer the CSVs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_luminosities = [0.9, 0.9]\n", - "scaling_galaxies_luminosities = [0.45, 0.45]\n", - "\n", - "al.galaxy_table_to_csv(\n", - " centres=extra_galaxies_centres,\n", - " luminosities=extra_galaxies_luminosities,\n", - " file_path=Path(dataset_path, \"extra_galaxies.csv\"),\n", - ")\n", - "\n", - "al.galaxy_table_to_csv(\n", - " centres=scaling_galaxies_centres,\n", - " luminosities=scaling_galaxies_luminosities,\n", - " file_path=Path(dataset_path, \"scaling_galaxies.csv\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Solve for the lensed source positions; written for the SLaM-style pipelines that consume this dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1),\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.01,\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=positions,\n", - " file_path=dataset_path / \"positions.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Group Scaling Relation\n", + "=================================\n", + "\n", + "This script simulates a group-scale strong lens with three populations of galaxies in the foreground, designed to\n", + "exercise the three-tier modeling API used by `group/features/scaling_relation/modeling.py`:\n", + "\n", + " - One **main lens galaxy** at the origin, which dominates the light and mass of the system.\n", + " - Two **extra galaxies** offset from the lens, modelled individually in the fit (one Einstein radius per galaxy).\n", + " - Two **scaling galaxies** further out, modelled via a shared luminosity-mass scaling relation.\n", + "\n", + "Each population's centres are saved to a separate JSON file (`main_lens_centres.json`, `extra_galaxies_centres.json`,\n", + "`scaling_galaxies_centres.json`) so the modeling script can load them directly.\n", + "\n", + "This dataset is independent of `dataset/group/simple` so the existing group examples are unaffected." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"scaling_relation\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "A 250x250 0.1\"/pixel grid wide enough to contain all five galaxies plus the lensed source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(250, 250),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Three populations:\n", + "\n", + " - `main_lens_centres`: primary lens at the origin.\n", + " - `extra_galaxies_centres`: closer companions modelled individually in the fit.\n", + " - `scaling_galaxies_centres`: further-out, fainter companions modelled via a shared scaling relation.\n", + "\n", + "The scaling galaxies are deliberately placed further out and given fainter light profiles below \u2014 this matches the\n", + "typical observational scenario in which the scaling-relation tier is reserved for galaxies that contribute only a small\n", + "amount of lensing individually but matter collectively." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [(0.0, 0.0)]\n", + "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]\n", + "scaling_galaxies_centres = [(6.5, 0.0), (-1.0, 7.0)]\n", + "\n", + "all_galaxy_centres = (\n", + " main_lens_centres + extra_galaxies_centres + scaling_galaxies_centres\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling is applied at every galaxy centre (main + extras + scaling) for accurate light evaluation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_galaxy_centres,\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A simple Gaussian PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")\n", + "\n", + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxy__\n", + "\n", + "A bright spherical Sersic + Isothermal mass at the origin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "main_lens_galaxies = [lens_0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "Two companion galaxies modelled with their own light + mass. These are the brighter, closer-in tier; in the modeling\n", + "script they receive their own free `einstein_radius` parameter each." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling Galaxies__\n", + "\n", + "Two further-out, fainter companions whose true Einstein radii are consistent with the reference-anchored\n", + "relation ``einstein_radius = einstein_radius_ref * (luminosity / luminosity_ref) ** 0.5`` with\n", + "``einstein_radius_ref = 0.135`` anchored to the brightest scaling member (``luminosity_ref = 0.45``; both members\n", + "share that luminosity here, so their radii are equal). Modelling these individually would add 2 free parameters;\n", + "on a scaling relation they add zero (the single free normalization is shared across all scaling galaxies, so\n", + "adding more does not grow the model)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scaling_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(6.5, 0.0), intensity=0.45, effective_radius=0.6, sersic_index=2.5\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(6.5, 0.0), einstein_radius=0.135),\n", + ")\n", + "\n", + "scaling_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-1.0, 7.0), intensity=0.45, effective_radius=0.6, sersic_index=2.5\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-1.0, 7.0), einstein_radius=0.135),\n", + ")\n", + "\n", + "scaling_galaxies = [scaling_galaxy_0, scaling_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Tracer order: main lenses, extras, scaling galaxies, source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=main_lens_galaxies + extra_galaxies + scaling_galaxies + [source_galaxy]\n", + ")\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the truth tracer for later inspection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centre JSON Files__\n", + "\n", + "One JSON per population, loaded individually by the modeling script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(scaling_galaxies_centres),\n", + " file_path=Path(dataset_path, \"scaling_galaxies_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Population CSVs__\n", + "\n", + "The modeling script loads luminosities (and centres) for both the extras and the scaling tier from CSVs written here.\n", + "The simulator knows the truth values of the per-galaxy luminosities so we write them out alongside the centre JSONs.\n", + "\n", + "The CSV schema is `y, x, luminosity, redshift?` \u2014 see `al.galaxy_table_from_csv` /\n", + "`al.galaxy_table_to_csv` (`autogalaxy/galaxy/galaxy_table.py`). Centre JSONs above are kept for backward compatibility;\n", + "new consumers should prefer the CSVs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_luminosities = [0.9, 0.9]\n", + "scaling_galaxies_luminosities = [0.45, 0.45]\n", + "\n", + "al.galaxy_table_to_csv(\n", + " centres=extra_galaxies_centres,\n", + " luminosities=extra_galaxies_luminosities,\n", + " file_path=Path(dataset_path, \"extra_galaxies.csv\"),\n", + ")\n", + "\n", + "al.galaxy_table_to_csv(\n", + " centres=scaling_galaxies_centres,\n", + " luminosities=scaling_galaxies_luminosities,\n", + " file_path=Path(dataset_path, \"scaling_galaxies.csv\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Solve for the lensed source positions; written for the SLaM-style pipelines that consume this dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1),\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.01,\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=positions,\n", + " file_path=dataset_path / \"positions.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/fit.ipynb b/notebooks/group/fit.ipynb index cc1f1479a..49ac725e3 100644 --- a/notebooks/group/fit.ipynb +++ b/notebooks/group/fit.ipynb @@ -1,778 +1,815 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fits: Group\n", - "===========\n", - "\n", - "This guide shows how to fit data using the `FitImaging` object for group-scale strong lenses, including visualizing\n", - "and interpreting its results.\n", - "\n", - "A group-scale lens differs from a galaxy-scale lens in that there are multiple lens galaxies contributing to the\n", - "lensing. In this example, there is a single main lens galaxy and two extra galaxies nearby whose mass contributes\n", - "significantly to the ray-tracing and must therefore be included in the model.\n", - "\n", - "References\n", - "----------\n", - "\n", - "This example uses functionality described fully in other examples in the `guides` package:\n", - "\n", - "- `guides/plot`: Using the plotting API (`aplt.plot_array`, `aplt.subplot_fit_imaging`, etc.) to visualize figures.\n", - "- `guides/units`: The source code unit conventions (e.g. arc seconds for distances and how to convert to physical units).\n", - "- `guides/data_structures`: The bespoke data structures used to store 1D and 2d arrays.\n", - "\n", - "__Contents__\n", - "\n", - "- **Loading Data:** We begin by loading the group-scale strong lens dataset `simple` from .fits files, which is the.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Galaxy Centres:** For group-scale lenses we load the centres of the main lens galaxies and extra galaxies from JSON.\n", - "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", - "- **Bad Fit:** A bad lens model will show features in the residual-map and chi-squared map.\n", - "- **Fit Quantities:** The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", - "- **Figures of Merit:** There are single valued floats which quantify the goodness of fit.\n", - "- **Plane Quantities:** The `FitImaging` object has specific quantities which break down each image of each plane.\n", - "- **Unmasked Quantities:** All of the quantities above are computed using the mask which was used to fit the data.\n", - "- **Pixel Counting:** An alternative way to quantify residuals like the lens light residuals is pixel counting.\n", - "- **Outputting Results:** You may wish to output certain results to .fits files for later inspection.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "\n", - "__JAX__\n", - "\n", - "Same JAX story as `scripts/imaging/fit.py`: `FitImaging` runs on either\n", - "backend. For the standard analysis-driven path see `start_here.py` /\n", - "`modeling.py`. For JIT-ing library methods directly see\n", - "`scripts/guides/lens_calc.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Data__\n", - "\n", - "We begin by loading the group-scale strong lens dataset `simple` from .fits files, which is the dataset\n", - "we will use to demonstrate fitting.\n", - "\n", - "This dataset was simulated using the `group/simulator` example, read through that to have a better\n", - "understanding of how the data this example fits was generated.\n", - "\n", - "The group-scale dataset has a larger field of view than a typical galaxy-scale lens, because it includes\n", - "emission from multiple lens galaxies and a more extended lensing configuration." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `aplt.subplot_imaging_dataset` contains a subplot which plots all the key properties of the dataset simultaneously.\n", - "\n", - "This includes the observed image data, RMS noise map, Point Spread Function and other information." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We now mask the data, so that regions where there is no signal (e.g. the edges) are omitted from the fit.\n", - "\n", - "We use a ``Mask2D`` object, which for this example is a 7.5\" circular mask. This is larger than a typical\n", - "galaxy-scale lens mask because the group-scale lens has emission spread over a wider area due to the\n", - "multiple lens galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now combine the imaging dataset with the mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now plot the image with the mask applied, where the image automatically zooms around the mask to make the lensed\n", - "source appear bigger." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=dataset.data, title=\"Image Data With Mask Applied\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The mask is also used to compute a `Grid2D`, where the (y,x) arc-second coordinates are only computed in unmasked\n", - "pixels within the masks' circle.\n", - "\n", - "As shown in the previous overview example, this grid will be used to perform lensing calculations when fitting the\n", - "data below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=dataset.grid, title=\"Grid2D of Masked Dataset\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "For group-scale lenses we load the centres of the main lens galaxies and extra galaxies from JSON files. These\n", - "centres are used during modeling to fix or constrain the positions of the galaxies.\n", - "\n", - "The main lens galaxy is at (0.0, 0.0) and the two extra galaxies are at (3.5, 2.5) and (-4.4, -5.0)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "print(f\"Main lens centres: {main_lens_centres}\")\n", - "print(f\"Extra galaxies centres: {extra_galaxies_centres}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fitting__\n", - "\n", - "Following the previous overview example, we can make a tracer from a collection of light profiles, mass profiles\n", - "and galaxies.\n", - "\n", - "The combination of light and mass profiles below is the same as those used to generate the simulated\n", - "dataset we loaded above.\n", - "\n", - "For a group-scale lens, we have multiple lens galaxies: a main lens galaxy and extra galaxies. The fit\n", - "handles all of these galaxies simultaneously, computing the combined deflection field from all mass\n", - "profiles to ray-trace the source galaxy light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Because the tracer's light and mass profiles are the same used to make the dataset, its image is nearly the same as the\n", - "observed image.\n", - "\n", - "However, the tracer's image does appear different to the data, in that its ring appears a bit thinner. This is\n", - "because its image has not been blurred with the telescope optics PSF, which the data has.\n", - "\n", - "[For those not familiar with Astronomy data, the PSF describes how the observed emission of the galaxy is blurred by\n", - "the telescope optics when it is observed. It mimicks this blurring effect via a 2D convolution operation]." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now use a `FitImaging` object to fit this tracer to the dataset.\n", - "\n", - "The fit creates a `model_image` which we fit the data with, which includes performing the step of blurring the tracer`s\n", - "image with the imaging dataset's PSF. We can see this by comparing the tracer`s image (which isn't PSF convolved) and\n", - "the fit`s model image (which is).\n", - "\n", - "For a group-scale lens, the model image includes contributions from all lens galaxies (main and extra) as well as\n", - "the lensed source galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.plot_array(array=fit.model_data, title=\"Model Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit does a lot more than just blur the tracer's image with the PSF, it also creates the following:\n", - "\n", - " - The `residual_map`: The `model_image` subtracted from the observed dataset`s `data`.\n", - " - The `normalized_residual_map`: The `residual_map `divided by the observed dataset's `noise_map`.\n", - " - The `chi_squared_map`: The `normalized_residual_map` squared.\n", - "\n", - "For a good lens model where the model image and tracer are representative of the strong lens system the\n", - "residuals, normalized residuals and chi-squareds are minimized:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", - "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", - "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi Squared Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A subplot can be plotted which contains all of the above quantities, as well as other information contained in the\n", - "tracer such as the source-plane image, a zoom in of the source-plane and a normalized residual map where the colorbar\n", - "goes from 1.0 sigma to -1.0 sigma, to highlight regions where the fit is poor." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit also provides us with a ``log_likelihood``, a single value quantifying how good the tracer fitted the dataset.\n", - "\n", - "Lens modeling, described in the next overview example, effectively tries to maximize this log likelihood value." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Bad Fit__\n", - "\n", - "A bad lens model will show features in the residual-map and chi-squared map.\n", - "\n", - "We can produce such an image by creating a tracer with different lens and source galaxies. In the example below, we\n", - "change the centre of the main lens galaxy's mass from (0.0, 0.0) to (0.2, 0.2), which leads to residuals appearing\n", - "in the fit. For a group-scale lens, even a small offset in the main lens mass centre can produce significant\n", - "residuals because the main lens dominates the total deflection field." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.2, 0.2), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A new fit using this tracer shows residuals, normalized residuals and chi-squared which are non-zero." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We also note that its likelihood decreases." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit Quantities__\n", - "\n", - "The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", - "\n", - "There is a `model_image`, which is the image-plane image of the tracer we inspected in the previous tutorial\n", - "blurred with the imaging data's PSF.\n", - "\n", - "This is the image that is fitted to the data in order to compute the log likelihood and therefore quantify the\n", - "goodness-of-fit.\n", - "\n", - "If you are unclear on what `slim` means, refer to the section `Data Structure` at the top of this example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_data.slim)\n", - "\n", - "# The native property provides quantities in 2D NumPy Arrays.\n", - "# print(fit.model_data.native)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There are numerous ndarrays showing the goodness of fit:\n", - "\n", - " - `residual_map`: Residuals = (Data - Model_Data).\n", - " - `normalized_residual_map`: Normalized_Residual = (Data - Model_Data) / Noise\n", - " - `chi_squared_map`: Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.residual_map.slim)\n", - "print(fit.normalized_residual_map.slim)\n", - "print(fit.chi_squared_map.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Figures of Merit__\n", - "\n", - "There are single valued floats which quantify the goodness of fit:\n", - "\n", - " - `chi_squared`: The sum of the `chi_squared_map`.\n", - "\n", - " - `noise_normalization`: The normalizing noise term in the likelihood function\n", - " where [Noise_Term] = sum(log(2*pi*[Noise]**2.0)).\n", - "\n", - " - `log_likelihood`: The log likelihood value of the fit where [LogLikelihood] = -0.5*[Chi_Squared_Term + Noise_Term]." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.chi_squared)\n", - "print(fit.noise_normalization)\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Plane Quantities__\n", - "\n", - "The `FitImaging` object has specific quantities which break down each image of each plane:\n", - "\n", - " - `model_images_of_planes_list`: Model-images of each individual plane, which for a group-scale lens includes the\n", - " model images of the main lens galaxy, each extra galaxy and the lensed source galaxy. All images are convolved\n", - " with the imaging's PSF.\n", - "\n", - " - `subtracted_images_of_planes_list`: Subtracted images of each individual plane, which are the data's image with\n", - " all other plane's model-images subtracted. For example, the first subtracted image has the source galaxy's and\n", - " extra galaxies' model images subtracted, leaving only the main lens galaxy's emission. This is especially useful\n", - " for group-scale lenses where isolating the light contribution of each galaxy is important.\n", - "\n", - "For group-scale lenses, there are more galaxies contributing to each plane compared to galaxy-scale lenses.\n", - "All lens galaxies (main and extra) are at the same redshift and therefore in the same plane, while the\n", - "source galaxy is in a separate background plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_images_of_planes_list[0].slim)\n", - "print(fit.model_images_of_planes_list[1].slim)\n", - "\n", - "print(fit.subtracted_images_of_planes_list[0].slim)\n", - "print(fit.subtracted_images_of_planes_list[1].slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Unmasked Quantities__\n", - "\n", - "All of the quantities above are computed using the mask which was used to fit the data.\n", - "\n", - "The `FitImaging` can also compute the unmasked blurred image of each plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.unmasked_blurred_image.native)\n", - "print(fit.unmasked_blurred_image_of_planes_list[0].native)\n", - "print(fit.unmasked_blurred_image_of_planes_list[1].native)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We can use the `Mask2D` object to mask regions of one of the fit's maps and estimate quantities of it.\n", - "\n", - "Below, we estimate the average absolute normalized residuals within a 1.0\" circular mask, which would inform us of\n", - "how accurate the lens light subtraction of a model fit is and if it leaves any significant residuals.\n", - "\n", - "For group-scale lenses, this is particularly useful for evaluating how well each individual galaxy's light\n", - "has been subtracted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(\n", - " shape_native=fit.dataset.shape_native,\n", - " pixel_scales=fit.dataset.pixel_scales,\n", - " radius=1.0,\n", - ")\n", - "\n", - "normalized_residuals = fit.normalized_residual_map.apply_mask(mask=mask)\n", - "\n", - "print(np.mean(np.abs(normalized_residuals.slim)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixel Counting__\n", - "\n", - "An alternative way to quantify residuals like the lens light residuals is pixel counting. For example, we could sum\n", - "up the number of pixels whose chi-squared values are above 10 which indicates a poor fit to the data.\n", - "\n", - "Whereas computing the mean above the average level of residuals, pixel counting informs us how spatially large the\n", - "residuals extend." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(\n", - " shape_native=fit.dataset.shape_native,\n", - " pixel_scales=fit.dataset.pixel_scales,\n", - " radius=1.0,\n", - ")\n", - "\n", - "chi_squared_map = fit.chi_squared_map.apply_mask(mask=mask)\n", - "\n", - "print(np.sum(chi_squared_map > 10.0))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Outputting Results__\n", - "\n", - "You may wish to output certain results to .fits files for later inspection.\n", - "\n", - "For example, one could output the lens light subtracted image of the lensed source galaxy to a .fits file such that\n", - "we could fit this source-only image again with an independent pipeline. For group-scale lenses, this subtracted\n", - "image has the light of all lens galaxies (main and extra) removed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_subtracted_image = fit.subtracted_images_of_planes_list[1]\n", - "aplt.fits_array(\n", - " array=lens_subtracted_image,\n", - " file_path=dataset_path / \"lens_subtracted_data.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fits: Group\n", + "===========\n", + "\n", + "This guide shows how to fit data using the `FitImaging` object for group-scale strong lenses, including visualizing\n", + "and interpreting its results.\n", + "\n", + "A group-scale lens differs from a galaxy-scale lens in that there are multiple lens galaxies contributing to the\n", + "lensing. In this example, there is a single main lens galaxy and two extra galaxies nearby whose mass contributes\n", + "significantly to the ray-tracing and must therefore be included in the model.\n", + "\n", + "References\n", + "----------\n", + "\n", + "This example uses functionality described fully in other examples in the `guides` package:\n", + "\n", + "- `guides/plot`: Using the plotting API (`aplt.plot_array`, `aplt.subplot_fit_imaging`, etc.) to visualize figures.\n", + "- `guides/units`: The source code unit conventions (e.g. arc seconds for distances and how to convert to physical units).\n", + "- `guides/data_structures`: The bespoke data structures used to store 1D and 2d arrays.\n", + "\n", + "__Contents__\n", + "\n", + "- **Loading Data:** We begin by loading the group-scale strong lens dataset `simple` from .fits files, which is the.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Galaxy Centres:** For group-scale lenses we load the centres of the main lens galaxies and extra galaxies from JSON.\n", + "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", + "- **Bad Fit:** A bad lens model will show features in the residual-map and chi-squared map.\n", + "- **Fit Quantities:** The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", + "- **Figures of Merit:** There are single valued floats which quantify the goodness of fit.\n", + "- **Plane Quantities:** The `FitImaging` object has specific quantities which break down each image of each plane.\n", + "- **Unmasked Quantities:** All of the quantities above are computed using the mask which was used to fit the data.\n", + "- **Pixel Counting:** An alternative way to quantify residuals like the lens light residuals is pixel counting.\n", + "- **Outputting Results:** You may wish to output certain results to .fits files for later inspection.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "\n", + "__JAX__\n", + "\n", + "Same JAX story as `scripts/imaging/fit.py`: `FitImaging` runs on either\n", + "backend. For the standard analysis-driven path see `start_here.py` /\n", + "`modeling.py`. For JIT-ing library methods directly see\n", + "`scripts/guides/lens_calc.py`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Data__\n", + "\n", + "We begin by loading the group-scale strong lens dataset `simple` from .fits files, which is the dataset\n", + "we will use to demonstrate fitting.\n", + "\n", + "This dataset was simulated using the `group/simulator` example, read through that to have a better\n", + "understanding of how the data this example fits was generated.\n", + "\n", + "The group-scale dataset has a larger field of view than a typical galaxy-scale lens, because it includes\n", + "emission from multiple lens galaxies and a more extended lensing configuration." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `aplt.subplot_imaging_dataset` contains a subplot which plots all the key properties of the dataset simultaneously.\n", + "\n", + "This includes the observed image data, RMS noise map, Point Spread Function and other information." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We now mask the data, so that regions where there is no signal (e.g. the edges) are omitted from the fit.\n", + "\n", + "We use a ``Mask2D`` object, which for this example is a 7.5\" circular mask. This is larger than a typical\n", + "galaxy-scale lens mask because the group-scale lens has emission spread over a wider area due to the\n", + "multiple lens galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now combine the imaging dataset with the mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now plot the image with the mask applied, where the image automatically zooms around the mask to make the lensed\n", + "source appear bigger." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=dataset.data, title=\"Image Data With Mask Applied\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The mask is also used to compute a `Grid2D`, where the (y,x) arc-second coordinates are only computed in unmasked\n", + "pixels within the masks' circle.\n", + "\n", + "As shown in the previous overview example, this grid will be used to perform lensing calculations when fitting the\n", + "data below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=dataset.grid, title=\"Grid2D of Masked Dataset\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "For group-scale lenses we load the centres of the main lens galaxies and extra galaxies from JSON files. These\n", + "centres are used during modeling to fix or constrain the positions of the galaxies.\n", + "\n", + "The main lens galaxy is at (0.0, 0.0) and the two extra galaxies are at (3.5, 2.5) and (-4.4, -5.0)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "print(f\"Main lens centres: {main_lens_centres}\")\n", + "print(f\"Extra galaxies centres: {extra_galaxies_centres}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fitting__\n", + "\n", + "Following the previous overview example, we can make a tracer from a collection of light profiles, mass profiles\n", + "and galaxies.\n", + "\n", + "The combination of light and mass profiles below is the same as those used to generate the simulated\n", + "dataset we loaded above.\n", + "\n", + "For a group-scale lens, we have multiple lens galaxies: a main lens galaxy and extra galaxies. The fit\n", + "handles all of these galaxies simultaneously, computing the combined deflection field from all mass\n", + "profiles to ray-trace the source galaxy light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Because the tracer's light and mass profiles are the same used to make the dataset, its image is nearly the same as the\n", + "observed image.\n", + "\n", + "However, the tracer's image does appear different to the data, in that its ring appears a bit thinner. This is\n", + "because its image has not been blurred with the telescope optics PSF, which the data has.\n", + "\n", + "[For those not familiar with Astronomy data, the PSF describes how the observed emission of the galaxy is blurred by\n", + "the telescope optics when it is observed. It mimicks this blurring effect via a 2D convolution operation]." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now use a `FitImaging` object to fit this tracer to the dataset.\n", + "\n", + "The fit creates a `model_image` which we fit the data with, which includes performing the step of blurring the tracer`s\n", + "image with the imaging dataset's PSF. We can see this by comparing the tracer`s image (which isn't PSF convolved) and\n", + "the fit`s model image (which is).\n", + "\n", + "For a group-scale lens, the model image includes contributions from all lens galaxies (main and extra) as well as\n", + "the lensed source galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.plot_array(array=fit.model_data, title=\"Model Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit does a lot more than just blur the tracer's image with the PSF, it also creates the following:\n", + "\n", + " - The `residual_map`: The `model_image` subtracted from the observed dataset`s `data`.\n", + " - The `normalized_residual_map`: The `residual_map `divided by the observed dataset's `noise_map`.\n", + " - The `chi_squared_map`: The `normalized_residual_map` squared.\n", + "\n", + "For a good lens model where the model image and tracer are representative of the strong lens system the\n", + "residuals, normalized residuals and chi-squareds are minimized:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", + "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", + "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi Squared Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A subplot can be plotted which contains all of the above quantities, as well as other information contained in the\n", + "tracer such as the source-plane image, a zoom in of the source-plane and a normalized residual map where the colorbar\n", + "goes from 1.0 sigma to -1.0 sigma, to highlight regions where the fit is poor." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit also provides us with a ``log_likelihood``, a single value quantifying how good the tracer fitted the dataset.\n", + "\n", + "Lens modeling, described in the next overview example, effectively tries to maximize this log likelihood value." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Bad Fit__\n", + "\n", + "A bad lens model will show features in the residual-map and chi-squared map.\n", + "\n", + "We can produce such an image by creating a tracer with different lens and source galaxies. In the example below, we\n", + "change the centre of the main lens galaxy's mass from (0.0, 0.0) to (0.2, 0.2), which leads to residuals appearing\n", + "in the fit. For a group-scale lens, even a small offset in the main lens mass centre can produce significant\n", + "residuals because the main lens dominates the total deflection field." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.2, 0.2), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A new fit using this tracer shows residuals, normalized residuals and chi-squared which are non-zero." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We also note that its likelihood decreases." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit Quantities__\n", + "\n", + "The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", + "\n", + "There is a `model_image`, which is the image-plane image of the tracer we inspected in the previous tutorial\n", + "blurred with the imaging data's PSF.\n", + "\n", + "This is the image that is fitted to the data in order to compute the log likelihood and therefore quantify the\n", + "goodness-of-fit.\n", + "\n", + "If you are unclear on what `slim` means, refer to the section `Data Structure` at the top of this example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_data.slim)\n", + "\n", + "# The native property provides quantities in 2D NumPy Arrays.\n", + "# print(fit.model_data.native)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There are numerous ndarrays showing the goodness of fit:\n", + "\n", + " - `residual_map`: Residuals = (Data - Model_Data).\n", + " - `normalized_residual_map`: Normalized_Residual = (Data - Model_Data) / Noise\n", + " - `chi_squared_map`: Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.residual_map.slim)\n", + "print(fit.normalized_residual_map.slim)\n", + "print(fit.chi_squared_map.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Figures of Merit__\n", + "\n", + "There are single valued floats which quantify the goodness of fit:\n", + "\n", + " - `chi_squared`: The sum of the `chi_squared_map`.\n", + "\n", + " - `noise_normalization`: The normalizing noise term in the likelihood function\n", + " where [Noise_Term] = sum(log(2*pi*[Noise]**2.0)).\n", + "\n", + " - `log_likelihood`: The log likelihood value of the fit where [LogLikelihood] = -0.5*[Chi_Squared_Term + Noise_Term]." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.chi_squared)\n", + "print(fit.noise_normalization)\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Plane Quantities__\n", + "\n", + "The `FitImaging` object has specific quantities which break down each image of each plane:\n", + "\n", + " - `model_images_of_planes_list`: Model-images of each individual plane, which for a group-scale lens includes the\n", + " model images of the main lens galaxy, each extra galaxy and the lensed source galaxy. All images are convolved\n", + " with the imaging's PSF.\n", + "\n", + " - `subtracted_images_of_planes_list`: Subtracted images of each individual plane, which are the data's image with\n", + " all other plane's model-images subtracted. For example, the first subtracted image has the source galaxy's and\n", + " extra galaxies' model images subtracted, leaving only the main lens galaxy's emission. This is especially useful\n", + " for group-scale lenses where isolating the light contribution of each galaxy is important.\n", + "\n", + "For group-scale lenses, there are more galaxies contributing to each plane compared to galaxy-scale lenses.\n", + "All lens galaxies (main and extra) are at the same redshift and therefore in the same plane, while the\n", + "source galaxy is in a separate background plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_images_of_planes_list[0].slim)\n", + "print(fit.model_images_of_planes_list[1].slim)\n", + "\n", + "print(fit.subtracted_images_of_planes_list[0].slim)\n", + "print(fit.subtracted_images_of_planes_list[1].slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Unmasked Quantities__\n", + "\n", + "All of the quantities above are computed using the mask which was used to fit the data.\n", + "\n", + "The `FitImaging` can also compute the unmasked blurred image of each plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.unmasked_blurred_image.native)\n", + "print(fit.unmasked_blurred_image_of_planes_list[0].native)\n", + "print(fit.unmasked_blurred_image_of_planes_list[1].native)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We can use the `Mask2D` object to mask regions of one of the fit's maps and estimate quantities of it.\n", + "\n", + "Below, we estimate the average absolute normalized residuals within a 1.0\" circular mask, which would inform us of\n", + "how accurate the lens light subtraction of a model fit is and if it leaves any significant residuals.\n", + "\n", + "For group-scale lenses, this is particularly useful for evaluating how well each individual galaxy's light\n", + "has been subtracted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=fit.dataset.shape_native,\n", + " pixel_scales=fit.dataset.pixel_scales,\n", + " radius=1.0,\n", + ")\n", + "\n", + "normalized_residuals = fit.normalized_residual_map.apply_mask(mask=mask)\n", + "\n", + "print(np.mean(np.abs(normalized_residuals.slim)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixel Counting__\n", + "\n", + "An alternative way to quantify residuals like the lens light residuals is pixel counting. For example, we could sum\n", + "up the number of pixels whose chi-squared values are above 10 which indicates a poor fit to the data.\n", + "\n", + "Whereas computing the mean above the average level of residuals, pixel counting informs us how spatially large the\n", + "residuals extend." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=fit.dataset.shape_native,\n", + " pixel_scales=fit.dataset.pixel_scales,\n", + " radius=1.0,\n", + ")\n", + "\n", + "chi_squared_map = fit.chi_squared_map.apply_mask(mask=mask)\n", + "\n", + "print(np.sum(chi_squared_map > 10.0))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Outputting Results__\n", + "\n", + "You may wish to output certain results to .fits files for later inspection.\n", + "\n", + "For example, one could output the lens light subtracted image of the lensed source galaxy to a .fits file such that\n", + "we could fit this source-only image again with an independent pipeline. For group-scale lenses, this subtracted\n", + "image has the light of all lens galaxies (main and extra) removed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_subtracted_image = fit.subtracted_images_of_planes_list[1]\n", + "aplt.fits_array(\n", + " array=lens_subtracted_image,\n", + " file_path=dataset_path / \"lens_subtracted_data.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/likelihood_function.ipynb b/notebooks/group/likelihood_function.ipynb index 99615f7b7..2ad1aad18 100644 --- a/notebooks/group/likelihood_function.ipynb +++ b/notebooks/group/likelihood_function.ipynb @@ -1,788 +1,825 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Group__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data of\n", - "a group-scale strong lens, where there are multiple lens galaxies whose mass profiles all contribute to the\n", - "ray-tracing of the source galaxy.\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers using, so that readers can understand the likelihood\n", - " function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - " - To illustrate how group-scale lensing differs from galaxy-scale lensing: multiple mass profiles from\n", - " multiple galaxies all contribute to the deflection angles, meaning the total deflection is the sum of\n", - " deflections from every galaxy in the group.\n", - "\n", - "Accompanying this script is the imaging `likelihood_function.py` which provides the same step-by-step guide\n", - "for a single lens galaxy. This script extends that to the group scale.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Masked Image Grid:** To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", - "- **Main Lens Galaxy:** The main lens galaxy is at the centre of the group.\n", - "- **Extra Galaxies:** The two extra galaxies are companion galaxies near the main lens.\n", - "- **Source Galaxy Light Profile:** The source galaxy is fitted using an analytic light profile, in this example a cored elliptical.\n", - "- **Lens Light:** Compute a 2D image of each lens galaxy's light and sum them together.\n", - "- **Lens Galaxy Mass:** We next consider the mass profiles of all galaxies in the group.\n", - "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\\\theta$ from the.\n", - "- **Source Image:** We pass the traced grid and blurring grid of coordinates to the source galaxy to evaluate its 2D.\n", - "- **Convolution:** Convolve the 2D image of the lens galaxies and source above with the PSF in real-space (as opposed.\n", - "- **Likelihood Function:** We now quantify the goodness-of-fit of our group-scale lens model.\n", - "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", - "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", - "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the lens model, by combining the two terms.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autoarray as aa\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "In order to perform a likelihood evaluation, we first load a dataset.\n", - "\n", - "This example fits a simulated group-scale strong lens where the imaging resolution is 0.1 arcsecond-per-pixel\n", - "resolution. The group consists of one main lens galaxy and two extra companion galaxies whose mass contributes\n", - "significantly to the ray-tracing." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"group\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This guide uses in-built visualization tools for plotting.\n", - "\n", - "For example, using the `aplt.subplot_imaging_dataset` the imaging dataset we perform a likelihood evaluation on is plotted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "The likelihood is only evaluated using image pixels contained within a 2D mask, which we choose before performing\n", - "lens modeling.\n", - "\n", - "Below, we define a 2D circular mask with a 7.5\" radius. This is larger than the mask used for galaxy-scale lenses\n", - "because group-scale systems have lensed images that extend over a wider area due to the combined mass of multiple\n", - "galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When we plot the masked imaging, only the circular masked region is shown." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=masked_dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling evaluates a light profile using multiple samples of its intensity per image-pixel.\n", - "\n", - "For simplicity, we disable over sampling in this guide by setting `sub_size=1`.\n", - "\n", - "a full description of over sampling and how to use it is given in `autolens_workspace/*/guides/over_sampling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masked Image Grid__\n", - "\n", - "To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", - "\n", - "These are given by `masked_dataset.grids.lp`, which we can plot and see is a uniform grid of (y,x) Cartesian\n", - "coordinates which have had the 7.5\" circular mask applied.\n", - "\n", - "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", - "used to perform ray-tracing and evaluate a light profile the intensity of the profile at the centre of each\n", - "image-pixel is computed, making it straight forward to compute the light profile's image to the image data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=masked_dataset.grids.lp, title=\"\")\n", - "\n", - "print(\n", - " f\"(y,x) coordinates of first ten unmasked image-pixels {masked_dataset.grid[0:9]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy Light (Setup)__\n", - "\n", - "To perform a likelihood evaluation we now compose our lens model.\n", - "\n", - "For a group-scale lens, there are multiple galaxies whose light and mass must be modeled. We define each galaxy\n", - "individually.\n", - "\n", - "A light profile is defined by its intensity $I (\\eta_{\\rm l}) $, for example the Sersic profile:\n", - "\n", - "$I_{\\rm Ser} (\\eta_{\\rm l}) = I \\exp \\bigg\\{ -k \\bigg[ \\bigg( \\frac{\\eta}{R} \\bigg)^{\\frac{1}{n}} - 1 \\bigg] \\bigg\\}$\n", - "\n", - "Where:\n", - "\n", - " - $\\eta$ are the elliptical coordinates of the masked image-grid.\n", - " - $I$ is the `intensity`, which controls the overall brightness of the Sersic profile.\n", - " - $n$ is the ``sersic_index``, which via $k$ controls the steepness of the inner profile.\n", - " - $R$ is the `effective_radius`, which defines the arc-second radius of a circle containing half the light.\n", - "\n", - "__Main Lens Galaxy__\n", - "\n", - "The main lens galaxy is at the centre of the group. It has a spherical Sersic light profile and a spherical\n", - "isothermal mass profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "The two extra galaxies are companion galaxies near the main lens. They each have their own light and mass profiles.\n", - "\n", - "Getting the mass of these extra galaxies right is crucial: their mass profiles contribute to the total deflection\n", - "angles, and if they are wrong the ray-traced source-plane coordinates will be incorrect, leading to a poor fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Light Profile__\n", - "\n", - "The source galaxy is fitted using an analytic light profile, in this example a cored elliptical Sersic." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Using the masked 2D grid defined above, we can calculate and plot images of each galaxy's light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_2d_lens = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", - "\n", - "aplt.plot_array(array=image_2d_lens, title=\"Main Lens Galaxy Image\")\n", - "\n", - "image_2d_extra_0 = extra_galaxy_0.image_2d_from(grid=masked_dataset.grid)\n", - "\n", - "aplt.plot_array(array=image_2d_extra_0, title=\"Extra Galaxy 0 Image\")\n", - "\n", - "image_2d_extra_1 = extra_galaxy_1.image_2d_from(grid=masked_dataset.grid)\n", - "\n", - "aplt.plot_array(array=image_2d_extra_1, title=\"Extra Galaxy 1 Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Light__\n", - "\n", - "Compute a 2D image of each lens galaxy's light and sum them together.\n", - "\n", - "For group-scale lenses, the total lens light is the sum of the images of ALL lens galaxies (main + extra). This is\n", - "a key difference from galaxy-scale lensing where there is typically only one lens galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", - "extra_0_image_2d = extra_galaxy_0.image_2d_from(grid=masked_dataset.grid)\n", - "extra_1_image_2d = extra_galaxy_1.image_2d_from(grid=masked_dataset.grid)\n", - "\n", - "total_lens_image_2d = lens_image_2d + extra_0_image_2d + extra_1_image_2d\n", - "\n", - "aplt.plot_array(array=total_lens_image_2d, title=\"Total Lens Light (All Galaxies)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To convolve the lens light with the imaging data's PSF, we need the `blurring_image`. This represents all flux\n", - "values not within the mask, which are close enough to it that their flux blurs into the mask after PSF convolution.\n", - "\n", - "We compute blurring images for ALL lens galaxies and sum them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n", - "extra_0_blurring_image_2d = extra_galaxy_0.image_2d_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")\n", - "extra_1_blurring_image_2d = extra_galaxy_1.image_2d_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")\n", - "\n", - "total_lens_blurring_image_2d = (\n", - " lens_blurring_image_2d + extra_0_blurring_image_2d + extra_1_blurring_image_2d\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy Mass__\n", - "\n", - "We next consider the mass profiles of all galaxies in the group.\n", - "\n", - "A mass profile is defined by its convergence $\\kappa (\\eta)$, which is related to\n", - "the surface density of the mass distribution as\n", - "\n", - "$\\kappa(\\eta)=\\frac{\\Sigma(\\eta)}{\\Sigma_\\mathrm{crit}},$\n", - "\n", - "where\n", - "\n", - "$\\Sigma_\\mathrm{crit}=\\frac{{\\rm c}^2}{4{\\rm \\pi} {\\rm G}}\\frac{D_{\\rm s}}{D_{\\rm l} D_{\\rm ls}},$\n", - "\n", - "For the isothermal profile used by all galaxies in this group:\n", - "\n", - "$\\kappa(\\eta) = \\frac{1.0}{1 + q} \\bigg( \\frac{\\theta_{\\rm E}}{\\eta} \\bigg)$\n", - "\n", - "Where $\\theta_{\\rm E}$ is the `einstein_radius`.\n", - "\n", - "From each mass profile we can compute its deflection angles, which describe how due to gravitational lensing\n", - "image-pixels are ray-traced to the source plane:\n", - "\n", - "$\\\\vec{{\\\\alpha}}_{\\\\rm x,y} (\\\\vec{x}) = \\\\frac{1}{\\\\pi} \\\\int \\\\frac{\\\\vec{x} - \\\\vec{x'}}{\\\\left | \\\\vec{x} - \\\\vec{x'} \\\\right |^2} \\\\kappa(\\\\vec{x'}) d\\\\vec{x'} \\\\, ,$" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_lens = lens_galaxy.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_extra_0 = extra_galaxy_0.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_extra_1 = extra_galaxy_1.deflections_yx_2d_from(grid=masked_dataset.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\\\theta$ from the image-plane to its (y,x)\n", - "source-plane coordinate $\\\\beta$ using the summed deflection angles $\\\\alpha$ of ALL mass profiles:\n", - "\n", - " $\\\\beta = \\\\theta - \\\\alpha(\\\\theta)$\n", - "\n", - "For group-scale lensing, the total deflection angle $\\\\alpha$ is the sum of deflection angles from ALL galaxies:\n", - "\n", - " $\\\\alpha_{\\\\rm total} = \\\\alpha_{\\\\rm lens} + \\\\alpha_{\\\\rm extra\\\\_0} + \\\\alpha_{\\\\rm extra\\\\_1}$\n", - "\n", - "This is the fundamental reason why getting the mass of extra galaxies right matters: each galaxy's mass profile\n", - "contributes to the total deflection, and errors in any of them lead to incorrect source-plane coordinates.\n", - "\n", - "The `Tracer` object handles this automatically by including all galaxies when computing ray-traced coordinates." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", - "\n", - "aplt.plot_grid(grid=traced_grid, title=\"Source Plane Grid (Traced)\")\n", - "\n", - "traced_blurring_grid = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")[-1]\n", - "\n", - "aplt.plot_grid(grid=traced_blurring_grid, title=\"Source Plane Blurring Grid (Traced)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Image__\n", - "\n", - "We pass the traced grid and blurring grid of coordinates to the source galaxy to evaluate its 2D image.\n", - "\n", - "This step is identical to galaxy-scale lensing -- the source galaxy's light profile is evaluated on the\n", - "source-plane grid. The difference is that the source-plane grid was computed using deflection angles from\n", - "multiple galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_image_2d = source_galaxy.image_2d_from(grid=traced_grid)\n", - "\n", - "source_blurring_image_2d = source_galaxy.image_2d_from(grid=traced_blurring_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens + Source Light Addition__\n", - "\n", - "We add the total lens galaxy light (from ALL galaxies) and the source galaxy image together, to create an overall\n", - "image of the group-scale strong lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = total_lens_image_2d + source_image_2d\n", - "\n", - "aplt.plot_array(array=image, title=\"Total Image (All Galaxies + Source)\")\n", - "\n", - "blurring_image_2d = total_lens_blurring_image_2d + source_blurring_image_2d" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Convolution__\n", - "\n", - "Convolve the 2D image of the lens galaxies and source above with the PSF in real-space (as opposed to via an FFT)\n", - "using a `Kernal2D`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", - " image=image, blurring_image=blurring_image_2d\n", - ")\n", - "\n", - "aplt.plot_array(array=convolved_image_2d, title=\"Convolved Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our group-scale lens model.\n", - "\n", - "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", - "\n", - "The likelihood function for parametric lens modeling consists of two terms:\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "We now explain what each of these terms mean.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", - "\n", - " - `model_data` = `convolved_image_2d`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "The chi-squared therefore quantifies if our fit to the data is accurate or not.\n", - "\n", - "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image\n", - "for, corresponding to a fit with a lower likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image = convolved_image_2d\n", - "\n", - "residual_map = masked_dataset.data - model_image\n", - "normalized_residual_map = residual_map / masked_dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(chi_squared)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", - "\n", - "aplt.plot_array(array=chi_squared_map, title=\"Chi-Squared Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", - "\n", - "The final term in the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", - "of the log of every noise-map value squared.\n", - "\n", - "Given the `noise_map` is fixed, this term does not change during the lens modeling process and has no impact on the\n", - "model we infer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "We can now, finally, compute the `log_likelihood` of the lens model, by combining the two terms computed above using\n", - "the likelihood function defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(figure_of_merit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This step-by-step process to perform a likelihood function evaluation is what is performed in the `FitImaging` object.\n", - "\n", - "The `FitImaging` object handles all of the steps above automatically: it sums the light from all galaxies, computes\n", - "the total deflection angles from all mass profiles, ray-traces the grid to the source plane, evaluates the source\n", - "light, convolves with the PSF, and computes the log likelihood.\n", - "\n", - "For group-scale lenses, the key advantage of the `FitImaging` and `Tracer` objects is that they automatically handle\n", - "the summation of light and mass contributions from an arbitrary number of galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)\n", - "fit_figure_of_merit = fit.figure_of_merit\n", - "print(fit_figure_of_merit)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `Nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the group-scale parametric likelihood function.\n", - "\n", - "The key differences from galaxy-scale lensing are:\n", - "\n", - " - Multiple lens galaxies (main + extra) each contribute light profiles whose images are summed together.\n", - " - Multiple mass profiles from ALL galaxies contribute to the deflection angles, and the total deflection\n", - " is the sum of deflections from every galaxy.\n", - " - Getting the mass of extra galaxies right is important because their deflection angles affect the\n", - " source-plane coordinates and therefore the quality of the source reconstruction.\n", - " - The `FitImaging` and `Tracer` objects handle all of this automatically for an arbitrary number of galaxies.\n", - "\n", - "__JAX__\n", - "\n", - "Same JAX recipe as `scripts/imaging/likelihood_function.py`: wrap the\n", - "hand-rolled construction in `@jax.jit`, register tracer classes once via\n", - "`autolens.jax.register_tracer_classes(tracer)`, validate via\n", - "`Fitness._vmap`. The group dataset adds nothing JAX-specific beyond the\n", - "imaging case \u2014 the extra galaxies are just more `Galaxy` instances and\n", - "the registration walker handles them uniformly.\n", - "\n", - "See `scripts/imaging/likelihood_function.py` `__JAX__` for the full\n", - "recipe." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Group__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data of\n", + "a group-scale strong lens, where there are multiple lens galaxies whose mass profiles all contribute to the\n", + "ray-tracing of the source galaxy.\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers using, so that readers can understand the likelihood\n", + " function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + " - To illustrate how group-scale lensing differs from galaxy-scale lensing: multiple mass profiles from\n", + " multiple galaxies all contribute to the deflection angles, meaning the total deflection is the sum of\n", + " deflections from every galaxy in the group.\n", + "\n", + "Accompanying this script is the imaging `likelihood_function.py` which provides the same step-by-step guide\n", + "for a single lens galaxy. This script extends that to the group scale.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Masked Image Grid:** To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", + "- **Main Lens Galaxy:** The main lens galaxy is at the centre of the group.\n", + "- **Extra Galaxies:** The two extra galaxies are companion galaxies near the main lens.\n", + "- **Source Galaxy Light Profile:** The source galaxy is fitted using an analytic light profile, in this example a cored elliptical.\n", + "- **Lens Light:** Compute a 2D image of each lens galaxy's light and sum them together.\n", + "- **Lens Galaxy Mass:** We next consider the mass profiles of all galaxies in the group.\n", + "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\\\theta$ from the.\n", + "- **Source Image:** We pass the traced grid and blurring grid of coordinates to the source galaxy to evaluate its 2D.\n", + "- **Convolution:** Convolve the 2D image of the lens galaxies and source above with the PSF in real-space (as opposed.\n", + "- **Likelihood Function:** We now quantify the goodness-of-fit of our group-scale lens model.\n", + "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", + "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", + "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the lens model, by combining the two terms.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autoarray as aa\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "In order to perform a likelihood evaluation, we first load a dataset.\n", + "\n", + "This example fits a simulated group-scale strong lens where the imaging resolution is 0.1 arcsecond-per-pixel\n", + "resolution. The group consists of one main lens galaxy and two extra companion galaxies whose mass contributes\n", + "significantly to the ray-tracing." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"group\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This guide uses in-built visualization tools for plotting.\n", + "\n", + "For example, using the `aplt.subplot_imaging_dataset` the imaging dataset we perform a likelihood evaluation on is plotted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "The likelihood is only evaluated using image pixels contained within a 2D mask, which we choose before performing\n", + "lens modeling.\n", + "\n", + "Below, we define a 2D circular mask with a 7.5\" radius. This is larger than the mask used for galaxy-scale lenses\n", + "because group-scale systems have lensed images that extend over a wider area due to the combined mass of multiple\n", + "galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When we plot the masked imaging, only the circular masked region is shown." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=masked_dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling evaluates a light profile using multiple samples of its intensity per image-pixel.\n", + "\n", + "For simplicity, we disable over sampling in this guide by setting `sub_size=1`.\n", + "\n", + "a full description of over sampling and how to use it is given in `autolens_workspace/*/guides/over_sampling.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masked Image Grid__\n", + "\n", + "To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", + "\n", + "These are given by `masked_dataset.grids.lp`, which we can plot and see is a uniform grid of (y,x) Cartesian\n", + "coordinates which have had the 7.5\" circular mask applied.\n", + "\n", + "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", + "used to perform ray-tracing and evaluate a light profile the intensity of the profile at the centre of each\n", + "image-pixel is computed, making it straight forward to compute the light profile's image to the image data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=masked_dataset.grids.lp, title=\"\")\n", + "\n", + "print(\n", + " f\"(y,x) coordinates of first ten unmasked image-pixels {masked_dataset.grid[0:9]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy Light (Setup)__\n", + "\n", + "To perform a likelihood evaluation we now compose our lens model.\n", + "\n", + "For a group-scale lens, there are multiple galaxies whose light and mass must be modeled. We define each galaxy\n", + "individually.\n", + "\n", + "A light profile is defined by its intensity $I (\\eta_{\\rm l}) $, for example the Sersic profile:\n", + "\n", + "$I_{\\rm Ser} (\\eta_{\\rm l}) = I \\exp \\bigg\\{ -k \\bigg[ \\bigg( \\frac{\\eta}{R} \\bigg)^{\\frac{1}{n}} - 1 \\bigg] \\bigg\\}$\n", + "\n", + "Where:\n", + "\n", + " - $\\eta$ are the elliptical coordinates of the masked image-grid.\n", + " - $I$ is the `intensity`, which controls the overall brightness of the Sersic profile.\n", + " - $n$ is the ``sersic_index``, which via $k$ controls the steepness of the inner profile.\n", + " - $R$ is the `effective_radius`, which defines the arc-second radius of a circle containing half the light.\n", + "\n", + "__Main Lens Galaxy__\n", + "\n", + "The main lens galaxy is at the centre of the group. It has a spherical Sersic light profile and a spherical\n", + "isothermal mass profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "The two extra galaxies are companion galaxies near the main lens. They each have their own light and mass profiles.\n", + "\n", + "Getting the mass of these extra galaxies right is crucial: their mass profiles contribute to the total deflection\n", + "angles, and if they are wrong the ray-traced source-plane coordinates will be incorrect, leading to a poor fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Light Profile__\n", + "\n", + "The source galaxy is fitted using an analytic light profile, in this example a cored elliptical Sersic." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Using the masked 2D grid defined above, we can calculate and plot images of each galaxy's light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_2d_lens = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", + "\n", + "aplt.plot_array(array=image_2d_lens, title=\"Main Lens Galaxy Image\")\n", + "\n", + "image_2d_extra_0 = extra_galaxy_0.image_2d_from(grid=masked_dataset.grid)\n", + "\n", + "aplt.plot_array(array=image_2d_extra_0, title=\"Extra Galaxy 0 Image\")\n", + "\n", + "image_2d_extra_1 = extra_galaxy_1.image_2d_from(grid=masked_dataset.grid)\n", + "\n", + "aplt.plot_array(array=image_2d_extra_1, title=\"Extra Galaxy 1 Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Light__\n", + "\n", + "Compute a 2D image of each lens galaxy's light and sum them together.\n", + "\n", + "For group-scale lenses, the total lens light is the sum of the images of ALL lens galaxies (main + extra). This is\n", + "a key difference from galaxy-scale lensing where there is typically only one lens galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", + "extra_0_image_2d = extra_galaxy_0.image_2d_from(grid=masked_dataset.grid)\n", + "extra_1_image_2d = extra_galaxy_1.image_2d_from(grid=masked_dataset.grid)\n", + "\n", + "total_lens_image_2d = lens_image_2d + extra_0_image_2d + extra_1_image_2d\n", + "\n", + "aplt.plot_array(array=total_lens_image_2d, title=\"Total Lens Light (All Galaxies)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To convolve the lens light with the imaging data's PSF, we need the `blurring_image`. This represents all flux\n", + "values not within the mask, which are close enough to it that their flux blurs into the mask after PSF convolution.\n", + "\n", + "We compute blurring images for ALL lens galaxies and sum them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n", + "extra_0_blurring_image_2d = extra_galaxy_0.image_2d_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")\n", + "extra_1_blurring_image_2d = extra_galaxy_1.image_2d_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")\n", + "\n", + "total_lens_blurring_image_2d = (\n", + " lens_blurring_image_2d + extra_0_blurring_image_2d + extra_1_blurring_image_2d\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy Mass__\n", + "\n", + "We next consider the mass profiles of all galaxies in the group.\n", + "\n", + "A mass profile is defined by its convergence $\\kappa (\\eta)$, which is related to\n", + "the surface density of the mass distribution as\n", + "\n", + "$\\kappa(\\eta)=\\frac{\\Sigma(\\eta)}{\\Sigma_\\mathrm{crit}},$\n", + "\n", + "where\n", + "\n", + "$\\Sigma_\\mathrm{crit}=\\frac{{\\rm c}^2}{4{\\rm \\pi} {\\rm G}}\\frac{D_{\\rm s}}{D_{\\rm l} D_{\\rm ls}},$\n", + "\n", + "For the isothermal profile used by all galaxies in this group:\n", + "\n", + "$\\kappa(\\eta) = \\frac{1.0}{1 + q} \\bigg( \\frac{\\theta_{\\rm E}}{\\eta} \\bigg)$\n", + "\n", + "Where $\\theta_{\\rm E}$ is the `einstein_radius`.\n", + "\n", + "From each mass profile we can compute its deflection angles, which describe how due to gravitational lensing\n", + "image-pixels are ray-traced to the source plane:\n", + "\n", + "$\\\\vec{{\\\\alpha}}_{\\\\rm x,y} (\\\\vec{x}) = \\\\frac{1}{\\\\pi} \\\\int \\\\frac{\\\\vec{x} - \\\\vec{x'}}{\\\\left | \\\\vec{x} - \\\\vec{x'} \\\\right |^2} \\\\kappa(\\\\vec{x'}) d\\\\vec{x'} \\\\, ,$" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_lens = lens_galaxy.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_extra_0 = extra_galaxy_0.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_extra_1 = extra_galaxy_1.deflections_yx_2d_from(grid=masked_dataset.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\\\theta$ from the image-plane to its (y,x)\n", + "source-plane coordinate $\\\\beta$ using the summed deflection angles $\\\\alpha$ of ALL mass profiles:\n", + "\n", + " $\\\\beta = \\\\theta - \\\\alpha(\\\\theta)$\n", + "\n", + "For group-scale lensing, the total deflection angle $\\\\alpha$ is the sum of deflection angles from ALL galaxies:\n", + "\n", + " $\\\\alpha_{\\\\rm total} = \\\\alpha_{\\\\rm lens} + \\\\alpha_{\\\\rm extra\\\\_0} + \\\\alpha_{\\\\rm extra\\\\_1}$\n", + "\n", + "This is the fundamental reason why getting the mass of extra galaxies right matters: each galaxy's mass profile\n", + "contributes to the total deflection, and errors in any of them lead to incorrect source-plane coordinates.\n", + "\n", + "The `Tracer` object handles this automatically by including all galaxies when computing ray-traced coordinates." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", + "\n", + "aplt.plot_grid(grid=traced_grid, title=\"Source Plane Grid (Traced)\")\n", + "\n", + "traced_blurring_grid = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")[-1]\n", + "\n", + "aplt.plot_grid(grid=traced_blurring_grid, title=\"Source Plane Blurring Grid (Traced)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Image__\n", + "\n", + "We pass the traced grid and blurring grid of coordinates to the source galaxy to evaluate its 2D image.\n", + "\n", + "This step is identical to galaxy-scale lensing -- the source galaxy's light profile is evaluated on the\n", + "source-plane grid. The difference is that the source-plane grid was computed using deflection angles from\n", + "multiple galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_image_2d = source_galaxy.image_2d_from(grid=traced_grid)\n", + "\n", + "source_blurring_image_2d = source_galaxy.image_2d_from(grid=traced_blurring_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens + Source Light Addition__\n", + "\n", + "We add the total lens galaxy light (from ALL galaxies) and the source galaxy image together, to create an overall\n", + "image of the group-scale strong lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = total_lens_image_2d + source_image_2d\n", + "\n", + "aplt.plot_array(array=image, title=\"Total Image (All Galaxies + Source)\")\n", + "\n", + "blurring_image_2d = total_lens_blurring_image_2d + source_blurring_image_2d" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Convolution__\n", + "\n", + "Convolve the 2D image of the lens galaxies and source above with the PSF in real-space (as opposed to via an FFT)\n", + "using a `Kernal2D`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", + " image=image, blurring_image=blurring_image_2d\n", + ")\n", + "\n", + "aplt.plot_array(array=convolved_image_2d, title=\"Convolved Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our group-scale lens model.\n", + "\n", + "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", + "\n", + "The likelihood function for parametric lens modeling consists of two terms:\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "We now explain what each of these terms mean.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", + "\n", + " - `model_data` = `convolved_image_2d`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "The chi-squared therefore quantifies if our fit to the data is accurate or not.\n", + "\n", + "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image\n", + "for, corresponding to a fit with a lower likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image = convolved_image_2d\n", + "\n", + "residual_map = masked_dataset.data - model_image\n", + "normalized_residual_map = residual_map / masked_dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(chi_squared)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", + "\n", + "aplt.plot_array(array=chi_squared_map, title=\"Chi-Squared Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", + "\n", + "The final term in the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", + "of the log of every noise-map value squared.\n", + "\n", + "Given the `noise_map` is fixed, this term does not change during the lens modeling process and has no impact on the\n", + "model we infer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "We can now, finally, compute the `log_likelihood` of the lens model, by combining the two terms computed above using\n", + "the likelihood function defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(figure_of_merit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This step-by-step process to perform a likelihood function evaluation is what is performed in the `FitImaging` object.\n", + "\n", + "The `FitImaging` object handles all of the steps above automatically: it sums the light from all galaxies, computes\n", + "the total deflection angles from all mass profiles, ray-traces the grid to the source plane, evaluates the source\n", + "light, convolves with the PSF, and computes the log likelihood.\n", + "\n", + "For group-scale lenses, the key advantage of the `FitImaging` and `Tracer` objects is that they automatically handle\n", + "the summation of light and mass contributions from an arbitrary number of galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)\n", + "fit_figure_of_merit = fit.figure_of_merit\n", + "print(fit_figure_of_merit)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `Nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the group-scale parametric likelihood function.\n", + "\n", + "The key differences from galaxy-scale lensing are:\n", + "\n", + " - Multiple lens galaxies (main + extra) each contribute light profiles whose images are summed together.\n", + " - Multiple mass profiles from ALL galaxies contribute to the deflection angles, and the total deflection\n", + " is the sum of deflections from every galaxy.\n", + " - Getting the mass of extra galaxies right is important because their deflection angles affect the\n", + " source-plane coordinates and therefore the quality of the source reconstruction.\n", + " - The `FitImaging` and `Tracer` objects handle all of this automatically for an arbitrary number of galaxies.\n", + "\n", + "__JAX__\n", + "\n", + "Same JAX recipe as `scripts/imaging/likelihood_function.py`: wrap the\n", + "hand-rolled construction in `@jax.jit`, register tracer classes once via\n", + "`autolens.jax.register_tracer_classes(tracer)`, validate via\n", + "`Fitness._vmap`. The group dataset adds nothing JAX-specific beyond the\n", + "imaging case \u2014 the extra galaxies are just more `Galaxy` instances and\n", + "the registration walker handles them uniformly.\n", + "\n", + "See `scripts/imaging/likelihood_function.py` `__JAX__` for the full\n", + "recipe." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/modeling.ipynb b/notebooks/group/modeling.ipynb index a29602b51..c4c750425 100644 --- a/notebooks/group/modeling.ipynb +++ b/notebooks/group/modeling.ipynb @@ -1,1051 +1,1088 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Group: Modeling\n", - "===============\n", - "\n", - "This script models an example strong lens on the 'group' scale, which typically have one or more \"main\" lens galaxies\n", - "and smaller extra galaxies nearby, whose light may blur with the source light and whose mass contributes significantly\n", - "to the ray-tracing, meaning both are therefore included in the strong lens model.\n", - "\n", - "This example uses a list-based model composition API, where:\n", - "\n", - " - Main lens galaxies are built in a loop over centres loaded from a JSON file and stored in the model as\n", - " `lens_0`, `lens_1`, etc. Only the first main lens galaxy (`lens_0`) carries an `ExternalShear`.\n", - "\n", - " - Extra galaxies are built in a loop over centres loaded from a separate JSON file and stored in an\n", - " `extra_galaxies` collection. Their mass centres are fixed to the observed centres of light and their\n", - " Einstein radii are given a uniform prior.\n", - "\n", - "This list-based approach scales naturally to systems with many main lens galaxies and many extra galaxies.\n", - "The centres are loaded from JSON files (`main_lens_centres.json` and `extra_galaxies_centres.json`) rather\n", - "than being hardcoded, so the same script works for different datasets without code changes.\n", - "\n", - "__Contents__\n", - "\n", - "- **Scaling Relations:** This example models the mass of each galaxy individually, which means the number of dimensions of.\n", - "- **Example:** This script fits an `Imaging` dataset of a 'group-scale' strong lens where.\n", - "- **Simulation:** Overview of how the simulated dataset was generated.\n", - "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Main Galaxies and Extra Galaxies:** For a group-scale lens, we designate there to be two types of lens galaxies in the system.\n", - "- **Centres:** The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files in the.\n", - "- **Redshifts:** In this example all line of sight galaxies are at the same redshift as the lens galaxy, meaning.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Coordinates:** Coordinate system assumptions for the model-fit.\n", - "- **Improved Lens Model:** The previous model used Sersic light profiles for the lens, source and extra galaxies.\n", - "- **Linear Light Profiles:** The MGE model below uses a **linear light profile** for the bulge and disk via the ``lp_linear``.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Unique Identifier:** In the path above, the `unique_identifier` appears as a collection of characters, where this.\n", - "- **Iterations Per Update:** Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and.\n", - "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", - "- **VRAM Use:** When running AutoLens with JAX on a GPU, the analysis must fit within the GPU's available VRAM.\n", - "- **Run Times:** Profiling the expected run time of the model-fit.\n", - "- **Output Folder Layout:** Description of the structure of the `output` folder where results are written.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Features:** The examples in the `autolens_workspace/*/imaging/features` package illustrate other lens modeling.\n", - "- **HowToLens:** This `start_here.py` script, and the features examples above, do not explain many details of how.\n", - "- **Modeling Customization:** The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative.\n", - "\n", - "__Scaling Relations__\n", - "\n", - "This example models the mass of each galaxy individually, which means the number of dimensions of the model increases\n", - "as we model group scale lenses with more galaxies. This can lead to a model that is slow to fit and poorly constrained.\n", - "There may also not be enough information in the data to constrain every galaxy's mass.\n", - "\n", - "A common approach to overcome this is to put many of the extra galaxies a scaling relation, where the mass of the\n", - "galaxies are related to their light via a observationally motivated scaling relation. This means that as more\n", - "galaxies are included in the lens model, the dimensionality of the model does not increase. Furthermore, their\n", - "luminosities act as priors on their masses, which helps ensure the model is well constrained.\n", - "\n", - "Lens modeling using scaling relations is fully support and described in the `features/scaling_relation.ipynb` example.\n", - "If your group has many extra galaxies (e.g. more than 5) you probably want to read this example once you are confident\n", - "with this one.\n", - "\n", - "__Example__\n", - "\n", - "This script fits an `Imaging` dataset of a 'group-scale' strong lens where\n", - "\n", - " - There is a main lens galaxy whose lens galaxy's light is an MGE.\n", - " - There is a main lens galaxy whose total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - There are two extra lens galaxies whose light models are `SersicSph` profiles and total mass distributions\n", - " are `IsothermalSph` models.\n", - " - The source galaxy's light is an MGE.\n", - "\n", - "__Simulation__\n", - "\n", - "This script fits a simulated `Imaging` dataset of a strong lens, which is produced in the\n", - "script `autolens_workspace/*/imaging/simulator.py`\n", - "\n", - "__Data Preparation__\n", - "\n", - "The `Imaging` dataset fitted in this example confirms to a number of standard that make it suitable to be fitted in\n", - "**PyAutoLens**.\n", - "\n", - "If you are intending to fit your own strong lens data, you will need to ensure it conforms to these standards, which are\n", - "described in the script `autolens_workspace/*/imaging/data_preparation/start_here.ipynb`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens group dataset `simple`, which is the dataset we will use to perform lens modeling.\n", - "\n", - "This is loaded via .fits files, which is a data format used by astronomers to store images.\n", - "\n", - "The `pixel_scales` define the arc-second to pixel conversion factor of the image, which for the dataset we are using\n", - "is 0.1\" / pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", \"group\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use an `aplt.subplot_imaging_dataset` the plot the data, including:\n", - "\n", - " - `data`: The image of the strong lens.\n", - " - `noise_map`: The noise-map of the image, which quantifies the noise in every pixel as their RMS values.\n", - " - `psf`: The point spread function of the image, which describes the blurring of the image by the telescope optics.\n", - " - `signal_to_noise_map`: Quantifies the signal-to-noise in every pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "The model-fit requires a 2D mask defining the regions of the image we fit the lens model to the data.\n", - "\n", - "We create a 7.5 arcsecond circular mask and apply it to the `Imaging` object that the lens model fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we plot the masked data, the mask removes the exterior regions of the image where there is no emission from the\n", - "lens and lensed source galaxies.\n", - "\n", - "The mask used to fit the data can be customized, as described in\n", - "the script `autolens_workspace/*/guides/modeling/customize`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Galaxies and Extra Galaxies__\n", - "\n", - "For a group-scale lens, we designate there to be two types of lens galaxies in the system:\n", - "\n", - " - `main_galaxies`: The main lens galaxies which likely make up the majority of light and mass in the lens system.\n", - " These are modeled individually and stored as `lens_0`, `lens_1`, etc. in the model's `galaxies` collection.\n", - " Their centres are loaded from the `main_lens_centres.json` file. Only the first main lens galaxy (`lens_0`)\n", - " carries an `ExternalShear`.\n", - "\n", - " - `extra_galaxies`: The extra galaxies which are nearby the lens system and contribute to the lensing of the source\n", - " galaxy. These are modeled with a more restrictive model, for example with their centres fixed to the observed\n", - " centre of light and their mass distributions modeled using a scaling relation. These are grouped into a single\n", - " `extra_galaxies` collection. Their centres are loaded from the `extra_galaxies_centres.json` file.\n", - "\n", - "In this simple example group scale lens, there is one main lens galaxy and two extra galaxies.\n", - "\n", - "__Centres__\n", - "\n", - "The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files in the dataset\n", - "directory. This makes the script reusable across different datasets without hardcoding centre values.\n", - "\n", - "For the main lens galaxies, the centres are loaded from `main_lens_centres.json` (e.g. `[(0.0, 0.0)]`).\n", - "For the extra galaxies, the centres are loaded from `extra_galaxies_centres.json` (e.g. `[(3.5, 2.5), (-4.4, -5.0)]`).\n", - "\n", - "If the centres of the extra galaxies are treated as free parameters, there are too many\n", - "parameters and the model may not be fitted accurately.\n", - "\n", - "For group-scale lenses we therefore manually specify the centres of the extra galaxies, which are fixed to the observed\n", - "centres of light of the galaxies.\n", - "\n", - "In a real analysis, one must determine the centres of the galaxies before modeling them, which can be done as follows:\n", - "\n", - " - Use the GUI tool in the `data_preparation/point_source/gui/extra_galaxies_centres.py` script to determine the centres\n", - " of the extra galaxies.\n", - "\n", - " - Use image processing software like Source Extractor (https://sextractor.readthedocs.io/en/latest/).\n", - "\n", - " - Fit every galaxy individually with a light profile (e.g. an `Sersic`).\n", - "\n", - "__Redshifts__\n", - "\n", - "In this example all line of sight galaxies are at the same redshift as the lens galaxy, meaning multi-plane lensing\n", - "is not used.\n", - "\n", - "If you have redshift information on the line of sight galaxies and some of their redshifts are different to the lens\n", - "galaxy, you can easily extend this example below to perform multi-plane lensing.\n", - "\n", - "You would simply define a `redshift_list` and use this to set up the extra `Galaxy` redshifts.\n", - "\n", - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The main lens galaxy's light is a `Sersic` light profile [7 parameters].\n", - "\n", - " - The main lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - There are two extra lens galaxies with linear `SersicSph` light and `IsothermalSph` total mass distributions, with\n", - " centres fixed to the observed centres of light [8 parameters].\n", - "\n", - " - The source galaxy's light is a point `SersicCore` [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=28.\n", - "\n", - "__Model Composition (List-Based API)__\n", - "\n", - "The API below for composing a lens model uses the list-based approach, where main lens galaxies are built in a loop\n", - "and stored as `lens_0`, `lens_1`, etc. Extra galaxies are similarly built in a loop and collected into an\n", - "`extra_galaxies` collection.\n", - "\n", - "This list-based API scales naturally: adding more main lens galaxies or extra galaxies simply means adding more\n", - "entries to the respective JSON files. The model composition code does not need to change.\n", - "\n", - "A full description of model composition is provided by the model cookbook:\n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", - "\n", - "__Coordinates__\n", - "\n", - "The model fitting default settings assume that the lens galaxy centre is near the coordinates (0.0\", 0.0\").\n", - "\n", - "If for your dataset the lens is not centred at (0.0\", 0.0\"), we recommend that you either:\n", - "\n", - " - Reduce your data so that the centre is (`autolens_workspace/*/data_preparation`).\n", - " - Manually override the lens model priors (`autolens_workspace/*/guides/modeling/customize`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Load centres from JSON files:\n", - "\n", - "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", - "extra_galaxies_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = af.Model(al.lp.Sersic)\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light\n", - "\n", - " bulge = af.Model(al.lp.SersicSph)\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "This shows the group scale model, with separate entries for each main lens galaxy (e.g. `lens_0`), the source galaxy\n", - "and the extra galaxies collection.\n", - "\n", - "The `info` below may not display optimally on your computer screen, for example the whitespace between parameter\n", - "names on the left and parameter priors on the right may lead them to appear across multiple lines. This is a\n", - "common issue in Jupyter notebooks.\n", - "\n", - "The`info_whitespace_length` parameter in the file `config/general.yaml` in the [output] section can be changed to\n", - "increase or decrease the amount of whitespace (The Jupyter notebook kernel will need to be reset for this change to\n", - "appear in a notebook)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Improved Lens Model__\n", - "\n", - "The previous model used Sersic light profiles for the lens, source and extra galaxies. This makes the model API concise,\n", - "readable, and easy to follow.\n", - "\n", - "However, single Sersic profiles perform poorly for most strong lenses. Symmetric profiles (e.g. elliptical Sersics)\n", - "typically leave significant residuals because they cannot capture the irregular and asymmetric morphology of real\n", - "galaxies (e.g. isophotal twists, radially varying ellipticity).\n", - "\n", - "For the extra galaxies, each Sersic also introduces 5 non-linear parameters, which means that as we add more extra\n", - "galaxies the model becomes increasingly complex, making it difficult to fit accurately and efficiently.\n", - "\n", - "This example therefore uses a lens model that combines two features, described in detail elsewhere (but a brief\n", - "overview is provided below):\n", - "\n", - "- **Linear light profiles** (see ``autolens_workspace/*/imaging/features/linear_light_profiles``)\n", - "- **Multi-Gaussian Expansion (MGE) light profiles** (see ``autolens_workspace/*/imaging/features/multi_gaussian_expansion``)\n", - "\n", - "These features avoid wasted effort trying to fit Sersic profiles to complex data, which is likely to fail unless the\n", - "lens is extremely simple. This does mean the model composition is more complex and as a user its a steeper learning\n", - "curve to understand the API, but its worth it for the improved accuracy and speed of lens modeling.\n", - "\n", - "__Multi-Gaussian Expansion (MGE)__\n", - "\n", - "A Multi-Gaussian Expansion (MGE) decomposes the lens and source light into ~50-100 Gaussians with varying ellipticities\n", - "and sizes. An MGE captures irregular features far more effectively than Sersic profiles, leading to more accurate lens m\n", - "odels.\n", - "\n", - "Remarkably, modeling with MGEs is also significantly faster than using Sersics: they remain efficient in JAX (on CPU\n", - "or GPU), require fewer non-linear parameters despite their flexibility, and yield simpler parameter spaces that\n", - "sample in far fewer iterations.\n", - "\n", - "The MGE is extremely important for group-scale lenses. Every time we add an extra galaxy, the MGE does not add\n", - "any extra non-linear parameters, unlike light profiles like Sersics. This means we can model the light of many\n", - "extra galaxies, ensuring the lens light model is accurate, without making the model slow to fit or poorly constrained.\n", - "\n", - "__Linear Light Profiles__\n", - "\n", - "The MGE model below uses a **linear light profile** for the bulge and disk via the ``lp_linear`` API, instead of the\n", - "standard ``lp`` light profiles used above.\n", - "\n", - "A linear light profile solves for the *intensity* of each component via a linear inversion, rather than treating it as\n", - "a free parameter. This reduces the dimensionality of the non-linear parameter space: a model with ~80 Gaussians\n", - "does not introduce ~80 additional free parameters.\n", - "\n", - "Linear light profiles therefore improve speed and accuracy, and they are used by default in all modeling example.\n", - "\n", - "__List-Based MGE Model__\n", - "\n", - "The improved model below uses the same list-based API as the simple model above, but replaces the Sersic light\n", - "profiles with MGE models created via `al.model_util.mge_model_from`. The main lens galaxies use 20 Gaussians\n", - "with uniform centre priors, while extra galaxies use 10 Gaussians with centres fixed to the observed positions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Main Lens Galaxies:\n", - "\n", - "lens_dict = {}\n", - "\n", - "for i, centre in enumerate(main_lens_centres):\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", - " )\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", - "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of\n", - "over sampling to ensure the lensed images are evaluated accurately.\n", - "\n", - "For a new user, the details of over-sampling are not important, therefore just be aware that calculations either:\n", - "\n", - " (i) use adaptive over sampling for the foregorund lens's light, which ensures high accuracy across.\n", - " (ii) use cored light profiles for the background source galaxy, where the core ensures low levels of over-sampling\n", - " produce numerically accurate but fast to compute results.\n", - "\n", - "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) is performed to ensure the lens\n", - "calculations are accurate across the full field of the group.\n", - "\n", - "Once you are more experienced, you should read up on over-sampling in more detail via\n", - "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The imaging subplot updates the bottom two panels to reflect the update to over sampling, which now uses a higher\n", - "values in the centre.\n", - "\n", - "Whilst you may not yet understand the details of over-sampling, you can at least track it visually in the plots\n", - "and later learnt more about it in the `over_sampling.ipynb` guide." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "This shows the group scale model, with separate entries for each main lens galaxy, the source galaxy and the\n", - "extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The lens model is fitted to the data using a non-linear search.\n", - "\n", - "All examples in the autolens workspace use the nested sampling algorithm\n", - "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has revealed gives the most\n", - "accurate and efficient modeling results.\n", - "\n", - "Nautilus has one main setting that trades-off accuracy and computational run-time, the number of `live_points`.\n", - "A higher number of live points gives a more accurate result, but increases the run-time. A lower value give\n", - "less reliable lens modeling (e.g. the fit may infer a local maxima), but is faster.\n", - "\n", - "The suitable value depends on the model complexity whereby models with more parameters require more live points.\n", - "The default value of 200 is sufficient for the vast majority of common lens models. Lower values often given reliable\n", - "results though, and speed up the run-times. In this example, given the model is quite simple (N=21 parameters), we\n", - "reduce the number of live points to 100 to speed up the run-time.\n", - "\n", - "__Unique Identifier__\n", - "\n", - "In the path above, the `unique_identifier` appears as a collection of characters, where this identifier is generated\n", - "based on the model, search and dataset that are used in the fit.\n", - "\n", - "An identical combination of model and search generates the same identifier, meaning that rerunning the script will use\n", - "the existing results to resume the model-fit. In contrast, if you change the model or search, a new unique identifier\n", - "will be generated, ensuring that the model-fit results are output into a separate folder.\n", - "\n", - "We additionally want the unique identifier to be specific to the dataset fitted, so that if we fit different datasets\n", - "with the same model and search results are output to a different folder. We achieve this below by passing\n", - "the `dataset_name` to the search's `unique_tag`.\n", - "\n", - "__Iterations Per Update__\n", - "\n", - "Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and its best fit\n", - "image to the Jupyter Notebook display and to hard-disk.\n", - "\n", - "This process takes around ~10 seconds, so we don't want it to happen too often so as to slow down the overall\n", - "fit, but we also want it to happen frequently enough that we can track the progress.\n", - "\n", - "The value of 10000 below means this output happens every few minutes on GPU and every ~10 minutes on CPU, a good balance.\n", - "\n", - "__Live Visual Update__\n", - "\n", - "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", - "live display surface:\n", - "\n", - "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", - " watch the fit converge without leaving your terminal.\n", - "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", - " refreshes in place every `iterations_per_quick_update`.\n", - "\n", - "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", - "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"group\"), # The path where results and output are stored.\n", - " name=\"modeling\", # The name of the fit and folder results are output to.\n", - " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", - " n_live=150, # The number of Nautilus \"live\" points, increase for more complex models.\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=10000, # Every N iterations the max likelihood model is visualized in the Jupter Notebook and output to hard-disk.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We next create an `AnalysisImaging` object, which can be given many inputs customizing how the lens model is\n", - "fitted to the data (in this example they are omitted for simplicity).\n", - "\n", - "Internally, this object defines the `log_likelihood_function` used by the non-linear search to fit the model to\n", - "the `Imaging` dataset.\n", - "\n", - "It is not vital that you as a user understand the details of how the `log_likelihood_function` fits a lens model to\n", - "data, but interested readers can find a step-by-step guide of the likelihood\n", - "function at ``autolens_workspace/*/imaging/log_likelihood_function`\n", - "\n", - "__JAX__\n", - "\n", - "`AnalysisImaging` defaults to `use_jax=True`. For group-scale fits the\n", - "JAX speedup is substantial \u2014 the multi-galaxy deflection sum dominates\n", - "runtime on CPU but vectorises cleanly on GPU. Search driver wraps the\n", - "likelihood in `jax.vmap(jax.jit(...))`. Force NumPy with `use_jax=False`\n", - "(or `PYAUTO_DISABLE_JAX=1`) when debugging." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM Use__\n", - "\n", - "When running AutoLens with JAX on a GPU, the analysis must fit within the GPU's\n", - "available VRAM. If insufficient VRAM is available, the analysis will fail with an\n", - "out-of-memory error, typically during JIT compilation or the first likelihood call.\n", - "\n", - "Two factors dictate the VRAM usage of an analysis:\n", - "\n", - "- The number of arrays and other data structures JAX must store in VRAM to fit the model\n", - " to the data in the likelihood function. This is dictated by the model complexity and dataset size.\n", - " For a MGE model its relatively low, but for other models (e.g. pixelized sources) it can be much higher.\n", - "\n", - "- The `batch_size` sets how many likelihood evaluations are performed simultaneously.\n", - " Increasing the batch size increases VRAM usage but can reduce overall run time,\n", - " while decreasing it lowers VRAM usage at the cost of slower execution.\n", - "\n", - "Before running an analysis, users should check that the estimated VRAM usage for the\n", - "chosen batch size is comfortably below their GPU's total VRAM.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits.\n", - "\n", - "For a MGE model with the low resolution dataset fitted in this example VRAM use is relatively quite high (~1.8GB),\n", - "illustrating how group scale modeling is already quite VRAM intensive. For more complex models (e.g. pixelized sources)\n", - "and higher resolution datasets it can be much higher (> 1GB going beyond 10GB)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "Lens modeling can be a computationally expensive process. When fitting complex models to high resolution datasets\n", - "run times can be of order hours, days, weeks or even months.\n", - "\n", - "Run times are dictated by two factors:\n", - "\n", - " - The log likelihood evaluation time: the time it takes for a single `instance` of the lens model to be fitted to\n", - " the dataset such that a log likelihood is returned.\n", - "\n", - " - The number of iterations (e.g. log likelihood evaluations) performed by the non-linear search: more complex lens\n", - " models require more iterations to converge to a solution.\n", - "\n", - "For this analysis, the log likelihood evaluation time is < 0.005 seconds on GPU, < 0.05 seconds on CPU, which is\n", - "extremely fast for group-scale lens modeling.\n", - "\n", - "To estimate the expected overall run time of the model-fit we multiply the log likelihood evaluation time by an\n", - "estimate of the number of iterations the non-linear search will perform, which is around 20000 to 40000 for this model.\n", - "\n", - "GPU run times are around 15 minutes, CPU run times are around 45 minutes.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results).\n", - "\n", - "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce\n", - "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output Folder Layout__\n", - "\n", - "Now the fit is running you should checkout the `autolens_workspace/output` folder. This is where results are\n", - "written to hard-disk in human-readable formats \u2014 `.json`, `.csv`, `.fits`, `.png` and plain text.\n", - "\n", - "As the fit progresses, results are written on the fly using the highest likelihood model found by the\n", - "non-linear search so far. This means you can inspect the model-fit as it runs, without waiting for the\n", - "non-linear search to terminate.\n", - "\n", - "Each completed fit lives at a path like::\n", - "\n", - " output/group//modeling//\n", - " files/ <- JSON + CSV: loadable Python objects\n", - " tracer.json <- max log likelihood Tracer (all lens galaxies + source)\n", - " model.json <- fitted af.Collection model\n", - " samples.csv <- full Nautilus samples\n", - " samples_summary.json <- max log likelihood parameter values + errors\n", - " samples_info.json <- metadata about the samples\n", - " search.json <- non-linear search configuration\n", - " settings.json <- search settings\n", - " cosmology.json <- cosmology used for the fit\n", - " covariance.csv <- parameter covariance matrix\n", - " image/ <- FITS + PNG: imaging products\n", - " dataset.fits <- data, noise-map and PSF\n", - " fit.fits <- model image, residuals, chi-squared map\n", - " tracer.fits <- tracer image-plane images per galaxy\n", - " source_plane_images.fits <- source plane reconstructions\n", - " model_galaxy_images.fits <- per-galaxy model images (lens_0, lens_1, ..., source)\n", - " galaxy_images.fits <- per-galaxy images\n", - " dataset.png, fit.png, tracer.png <- visualisations\n", - " model.info <- human-readable model summary\n", - " model.results <- human-readable fit summary\n", - " search.summary <- search run summary\n", - " search_internal/ <- internal files used to resume / visualise the search\n", - " metadata <- run metadata\n", - "\n", - "The `` is a 32-character identifier derived from the model, search and dataset, so re-running the\n", - "same configuration resumes from the existing fit automatically.\n", - "\n", - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", - "\n", - "The result contains entries for each main lens galaxy (e.g. `lens_0`), the source galaxy and the extra galaxies.\n", - "\n", - "[Above, we discussed that the `info_whitespace_length` parameter in the config files could b changed to make\n", - "the `model.info` attribute display optimally on your computer. This attribute also controls the whitespace of the\n", - "`result.info` attribute.]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Result` object also contains:\n", - "\n", - " - The model corresponding to the maximum log likelihood solution in parameter space.\n", - " - The corresponding maximum log likelihood `Tracer` and `FitImaging` objects.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It also contains information on the posterior as estimated by the non-linear search (in this example `Nautilus`).\n", - "\n", - "Below, we make a corner plot of the \"Probability Density Function\" of every parameter in the model-fit.\n", - "\n", - "The plot is labeled with short hand parameter names (e.g. `sersic_index` is mapped to the short hand\n", - "parameter `n`). These mappings ate specified in the `config/notation.yaml` file and can be customized by users.\n", - "\n", - "The superscripts of labels correspond to the name each component was given in the model (e.g. for the `Isothermal`\n", - "mass its name `mass` defined when making the `Model` above is used)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This script gives a concise overview of the basic modeling API, fitting one the simplest lens models possible.\n", - "\n", - "Lets now consider what features you should read about to improve your lens modeling, especially if you are aiming\n", - "to fit more complex models to your data.\n", - "\n", - "This is especially important for group scale modeling, in order to reduce the complexity of the model.\n", - "\n", - "__Features__\n", - "\n", - "The examples in the `autolens_workspace/*/imaging/features` package illustrate other lens modeling features.\n", - "\n", - "We recommend you checkout the following features, because they make lens modeling in general more reliable and\n", - "efficient (you will therefore benefit from using these features irrespective of the quality of your data and\n", - "scientific topic of study).\n", - "\n", - "We recommend you now checkout the following features:\n", - "\n", - "- ``scaling_relation``: This feature allows you to model the light and mass of the extra galaxies using a scaling relation.\n", - "- ``linear_light_profiles``: The model light profiles use linear algebra to solve for their intensity, reducing model complexity.\n", - "- ``multi_gaussian_expansion``: The lens (or source) light is modeled as ~25-100 Gaussian basis functions\n", - "- ``pixelization``: The source is reconstructed using an adaptive rectangular or Delaunay mesh\n", - "- ``no_lens_light``: The foreground lens's light is not present in the data and thus omitted from the model.\n", - "\n", - "For group scale modeling, the multi Gaussian expansion is particularly important, as this can dramatically reduce the\n", - "dimensionality of the model and improve the accuracy of the fit for both the lens and source galaxies.\n", - "\n", - "It is also recommended you read through the `imaging` package, to get a complete picture of how point-source\n", - "modeling works.\n", - "\n", - "__Data Preparation__\n", - "\n", - "If you are looking to fit your own CCD imaging data of a strong lens, checkout\n", - "the `autolens_workspace/*/imaging/data_preparation/start_here.ipynb` script for an overview of how data should be\n", - "prepared before being modeled.\n", - "\n", - "__HowToLens__\n", - "\n", - "This `start_here.py` script, and the features examples above, do not explain many details of how lens modeling is\n", - "performed, for example:\n", - "\n", - " - How does PyAutoLens perform ray-tracing and lensing calculations in order to fit a lens model?\n", - " - How is a lens model fitted to data? What quantifies the goodness of fit (e.g. how is a log likelihood computed?).\n", - " - How does Nautilus find the highest likelihood lens models? What exactly is a \"non-linear search\"?\n", - "\n", - "You do not need to be able to answer these questions in order to fit lens models with PyAutoLens and do science.\n", - "However, having a deeper understanding of how it all works is both interesting and will benefit you as a scientist\n", - "\n", - "This deeper insight is offered by the **HowToLens** Jupyter notebook lectures, which live in their own repository:\n", - "https://github.com/PyAutoLabs/HowToLens.\n", - "\n", - "I recommend that you check them out if you are interested in more details!\n", - "\n", - "__Modeling Customization__\n", - "\n", - "The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative non-linear searches,\n", - "other than Nautilus, that can be used to fit lens models.\n", - "\n", - "They also provide details on how to customize the model-fit, for example the priors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Group: Modeling\n", + "===============\n", + "\n", + "This script models an example strong lens on the 'group' scale, which typically have one or more \"main\" lens galaxies\n", + "and smaller extra galaxies nearby, whose light may blur with the source light and whose mass contributes significantly\n", + "to the ray-tracing, meaning both are therefore included in the strong lens model.\n", + "\n", + "This example uses a list-based model composition API, where:\n", + "\n", + " - Main lens galaxies are built in a loop over centres loaded from a JSON file and stored in the model as\n", + " `lens_0`, `lens_1`, etc. Only the first main lens galaxy (`lens_0`) carries an `ExternalShear`.\n", + "\n", + " - Extra galaxies are built in a loop over centres loaded from a separate JSON file and stored in an\n", + " `extra_galaxies` collection. Their mass centres are fixed to the observed centres of light and their\n", + " Einstein radii are given a uniform prior.\n", + "\n", + "This list-based approach scales naturally to systems with many main lens galaxies and many extra galaxies.\n", + "The centres are loaded from JSON files (`main_lens_centres.json` and `extra_galaxies_centres.json`) rather\n", + "than being hardcoded, so the same script works for different datasets without code changes.\n", + "\n", + "__Contents__\n", + "\n", + "- **Scaling Relations:** This example models the mass of each galaxy individually, which means the number of dimensions of.\n", + "- **Example:** This script fits an `Imaging` dataset of a 'group-scale' strong lens where.\n", + "- **Simulation:** Overview of how the simulated dataset was generated.\n", + "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Main Galaxies and Extra Galaxies:** For a group-scale lens, we designate there to be two types of lens galaxies in the system.\n", + "- **Centres:** The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files in the.\n", + "- **Redshifts:** In this example all line of sight galaxies are at the same redshift as the lens galaxy, meaning.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Coordinates:** Coordinate system assumptions for the model-fit.\n", + "- **Improved Lens Model:** The previous model used Sersic light profiles for the lens, source and extra galaxies.\n", + "- **Linear Light Profiles:** The MGE model below uses a **linear light profile** for the bulge and disk via the ``lp_linear``.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Unique Identifier:** In the path above, the `unique_identifier` appears as a collection of characters, where this.\n", + "- **Iterations Per Update:** Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and.\n", + "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", + "- **VRAM Use:** When running AutoLens with JAX on a GPU, the analysis must fit within the GPU's available VRAM.\n", + "- **Run Times:** Profiling the expected run time of the model-fit.\n", + "- **Output Folder Layout:** Description of the structure of the `output` folder where results are written.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Features:** The examples in the `autolens_workspace/*/imaging/features` package illustrate other lens modeling.\n", + "- **HowToLens:** This `start_here.py` script, and the features examples above, do not explain many details of how.\n", + "- **Modeling Customization:** The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative.\n", + "\n", + "__Scaling Relations__\n", + "\n", + "This example models the mass of each galaxy individually, which means the number of dimensions of the model increases\n", + "as we model group scale lenses with more galaxies. This can lead to a model that is slow to fit and poorly constrained.\n", + "There may also not be enough information in the data to constrain every galaxy's mass.\n", + "\n", + "A common approach to overcome this is to put many of the extra galaxies a scaling relation, where the mass of the\n", + "galaxies are related to their light via a observationally motivated scaling relation. This means that as more\n", + "galaxies are included in the lens model, the dimensionality of the model does not increase. Furthermore, their\n", + "luminosities act as priors on their masses, which helps ensure the model is well constrained.\n", + "\n", + "Lens modeling using scaling relations is fully support and described in the `features/scaling_relation.ipynb` example.\n", + "If your group has many extra galaxies (e.g. more than 5) you probably want to read this example once you are confident\n", + "with this one.\n", + "\n", + "__Example__\n", + "\n", + "This script fits an `Imaging` dataset of a 'group-scale' strong lens where\n", + "\n", + " - There is a main lens galaxy whose lens galaxy's light is an MGE.\n", + " - There is a main lens galaxy whose total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - There are two extra lens galaxies whose light models are `SersicSph` profiles and total mass distributions\n", + " are `IsothermalSph` models.\n", + " - The source galaxy's light is an MGE.\n", + "\n", + "__Simulation__\n", + "\n", + "This script fits a simulated `Imaging` dataset of a strong lens, which is produced in the\n", + "script `autolens_workspace/*/imaging/simulator.py`\n", + "\n", + "__Data Preparation__\n", + "\n", + "The `Imaging` dataset fitted in this example confirms to a number of standard that make it suitable to be fitted in\n", + "**PyAutoLens**.\n", + "\n", + "If you are intending to fit your own strong lens data, you will need to ensure it conforms to these standards, which are\n", + "described in the script `autolens_workspace/*/imaging/data_preparation/start_here.ipynb`." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens group dataset `simple`, which is the dataset we will use to perform lens modeling.\n", + "\n", + "This is loaded via .fits files, which is a data format used by astronomers to store images.\n", + "\n", + "The `pixel_scales` define the arc-second to pixel conversion factor of the image, which for the dataset we are using\n", + "is 0.1\" / pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", \"group\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use an `aplt.subplot_imaging_dataset` the plot the data, including:\n", + "\n", + " - `data`: The image of the strong lens.\n", + " - `noise_map`: The noise-map of the image, which quantifies the noise in every pixel as their RMS values.\n", + " - `psf`: The point spread function of the image, which describes the blurring of the image by the telescope optics.\n", + " - `signal_to_noise_map`: Quantifies the signal-to-noise in every pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "The model-fit requires a 2D mask defining the regions of the image we fit the lens model to the data.\n", + "\n", + "We create a 7.5 arcsecond circular mask and apply it to the `Imaging` object that the lens model fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we plot the masked data, the mask removes the exterior regions of the image where there is no emission from the\n", + "lens and lensed source galaxies.\n", + "\n", + "The mask used to fit the data can be customized, as described in\n", + "the script `autolens_workspace/*/guides/modeling/customize`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Galaxies and Extra Galaxies__\n", + "\n", + "For a group-scale lens, we designate there to be two types of lens galaxies in the system:\n", + "\n", + " - `main_galaxies`: The main lens galaxies which likely make up the majority of light and mass in the lens system.\n", + " These are modeled individually and stored as `lens_0`, `lens_1`, etc. in the model's `galaxies` collection.\n", + " Their centres are loaded from the `main_lens_centres.json` file. Only the first main lens galaxy (`lens_0`)\n", + " carries an `ExternalShear`.\n", + "\n", + " - `extra_galaxies`: The extra galaxies which are nearby the lens system and contribute to the lensing of the source\n", + " galaxy. These are modeled with a more restrictive model, for example with their centres fixed to the observed\n", + " centre of light and their mass distributions modeled using a scaling relation. These are grouped into a single\n", + " `extra_galaxies` collection. Their centres are loaded from the `extra_galaxies_centres.json` file.\n", + "\n", + "In this simple example group scale lens, there is one main lens galaxy and two extra galaxies.\n", + "\n", + "__Centres__\n", + "\n", + "The centres of both the main lens galaxies and the extra galaxies are loaded from JSON files in the dataset\n", + "directory. This makes the script reusable across different datasets without hardcoding centre values.\n", + "\n", + "For the main lens galaxies, the centres are loaded from `main_lens_centres.json` (e.g. `[(0.0, 0.0)]`).\n", + "For the extra galaxies, the centres are loaded from `extra_galaxies_centres.json` (e.g. `[(3.5, 2.5), (-4.4, -5.0)]`).\n", + "\n", + "If the centres of the extra galaxies are treated as free parameters, there are too many\n", + "parameters and the model may not be fitted accurately.\n", + "\n", + "For group-scale lenses we therefore manually specify the centres of the extra galaxies, which are fixed to the observed\n", + "centres of light of the galaxies.\n", + "\n", + "In a real analysis, one must determine the centres of the galaxies before modeling them, which can be done as follows:\n", + "\n", + " - Use the GUI tool in the `data_preparation/point_source/gui/extra_galaxies_centres.py` script to determine the centres\n", + " of the extra galaxies.\n", + "\n", + " - Use image processing software like Source Extractor (https://sextractor.readthedocs.io/en/latest/).\n", + "\n", + " - Fit every galaxy individually with a light profile (e.g. an `Sersic`).\n", + "\n", + "__Redshifts__\n", + "\n", + "In this example all line of sight galaxies are at the same redshift as the lens galaxy, meaning multi-plane lensing\n", + "is not used.\n", + "\n", + "If you have redshift information on the line of sight galaxies and some of their redshifts are different to the lens\n", + "galaxy, you can easily extend this example below to perform multi-plane lensing.\n", + "\n", + "You would simply define a `redshift_list` and use this to set up the extra `Galaxy` redshifts.\n", + "\n", + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The main lens galaxy's light is a `Sersic` light profile [7 parameters].\n", + "\n", + " - The main lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - There are two extra lens galaxies with linear `SersicSph` light and `IsothermalSph` total mass distributions, with\n", + " centres fixed to the observed centres of light [8 parameters].\n", + "\n", + " - The source galaxy's light is a point `SersicCore` [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=28.\n", + "\n", + "__Model Composition (List-Based API)__\n", + "\n", + "The API below for composing a lens model uses the list-based approach, where main lens galaxies are built in a loop\n", + "and stored as `lens_0`, `lens_1`, etc. Extra galaxies are similarly built in a loop and collected into an\n", + "`extra_galaxies` collection.\n", + "\n", + "This list-based API scales naturally: adding more main lens galaxies or extra galaxies simply means adding more\n", + "entries to the respective JSON files. The model composition code does not need to change.\n", + "\n", + "A full description of model composition is provided by the model cookbook:\n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", + "\n", + "__Coordinates__\n", + "\n", + "The model fitting default settings assume that the lens galaxy centre is near the coordinates (0.0\", 0.0\").\n", + "\n", + "If for your dataset the lens is not centred at (0.0\", 0.0\"), we recommend that you either:\n", + "\n", + " - Reduce your data so that the centre is (`autolens_workspace/*/data_preparation`).\n", + " - Manually override the lens model priors (`autolens_workspace/*/guides/modeling/customize`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Load centres from JSON files:\n", + "\n", + "main_lens_centres = al.from_json(file_path=dataset_path / \"main_lens_centres.json\")\n", + "extra_galaxies_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = af.Model(al.lp.Sersic)\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light\n", + "\n", + " bulge = af.Model(al.lp.SersicSph)\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "This shows the group scale model, with separate entries for each main lens galaxy (e.g. `lens_0`), the source galaxy\n", + "and the extra galaxies collection.\n", + "\n", + "The `info` below may not display optimally on your computer screen, for example the whitespace between parameter\n", + "names on the left and parameter priors on the right may lead them to appear across multiple lines. This is a\n", + "common issue in Jupyter notebooks.\n", + "\n", + "The`info_whitespace_length` parameter in the file `config/general.yaml` in the [output] section can be changed to\n", + "increase or decrease the amount of whitespace (The Jupyter notebook kernel will need to be reset for this change to\n", + "appear in a notebook)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Improved Lens Model__\n", + "\n", + "The previous model used Sersic light profiles for the lens, source and extra galaxies. This makes the model API concise,\n", + "readable, and easy to follow.\n", + "\n", + "However, single Sersic profiles perform poorly for most strong lenses. Symmetric profiles (e.g. elliptical Sersics)\n", + "typically leave significant residuals because they cannot capture the irregular and asymmetric morphology of real\n", + "galaxies (e.g. isophotal twists, radially varying ellipticity).\n", + "\n", + "For the extra galaxies, each Sersic also introduces 5 non-linear parameters, which means that as we add more extra\n", + "galaxies the model becomes increasingly complex, making it difficult to fit accurately and efficiently.\n", + "\n", + "This example therefore uses a lens model that combines two features, described in detail elsewhere (but a brief\n", + "overview is provided below):\n", + "\n", + "- **Linear light profiles** (see ``autolens_workspace/*/imaging/features/linear_light_profiles``)\n", + "- **Multi-Gaussian Expansion (MGE) light profiles** (see ``autolens_workspace/*/imaging/features/multi_gaussian_expansion``)\n", + "\n", + "These features avoid wasted effort trying to fit Sersic profiles to complex data, which is likely to fail unless the\n", + "lens is extremely simple. This does mean the model composition is more complex and as a user its a steeper learning\n", + "curve to understand the API, but its worth it for the improved accuracy and speed of lens modeling.\n", + "\n", + "__Multi-Gaussian Expansion (MGE)__\n", + "\n", + "A Multi-Gaussian Expansion (MGE) decomposes the lens and source light into ~50-100 Gaussians with varying ellipticities\n", + "and sizes. An MGE captures irregular features far more effectively than Sersic profiles, leading to more accurate lens m\n", + "odels.\n", + "\n", + "Remarkably, modeling with MGEs is also significantly faster than using Sersics: they remain efficient in JAX (on CPU\n", + "or GPU), require fewer non-linear parameters despite their flexibility, and yield simpler parameter spaces that\n", + "sample in far fewer iterations.\n", + "\n", + "The MGE is extremely important for group-scale lenses. Every time we add an extra galaxy, the MGE does not add\n", + "any extra non-linear parameters, unlike light profiles like Sersics. This means we can model the light of many\n", + "extra galaxies, ensuring the lens light model is accurate, without making the model slow to fit or poorly constrained.\n", + "\n", + "__Linear Light Profiles__\n", + "\n", + "The MGE model below uses a **linear light profile** for the bulge and disk via the ``lp_linear`` API, instead of the\n", + "standard ``lp`` light profiles used above.\n", + "\n", + "A linear light profile solves for the *intensity* of each component via a linear inversion, rather than treating it as\n", + "a free parameter. This reduces the dimensionality of the non-linear parameter space: a model with ~80 Gaussians\n", + "does not introduce ~80 additional free parameters.\n", + "\n", + "Linear light profiles therefore improve speed and accuracy, and they are used by default in all modeling example.\n", + "\n", + "__List-Based MGE Model__\n", + "\n", + "The improved model below uses the same list-based API as the simple model above, but replaces the Sersic light\n", + "profiles with MGE models created via `al.model_util.mge_model_from`. The main lens galaxies use 20 Gaussians\n", + "with uniform centre priors, while extra galaxies use 10 Gaussians with centres fixed to the observed positions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Main Lens Galaxies:\n", + "\n", + "lens_dict = {}\n", + "\n", + "for i, centre in enumerate(main_lens_centres):\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=centre\n", + " )\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", + "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of\n", + "over sampling to ensure the lensed images are evaluated accurately.\n", + "\n", + "For a new user, the details of over-sampling are not important, therefore just be aware that calculations either:\n", + "\n", + " (i) use adaptive over sampling for the foregorund lens's light, which ensures high accuracy across.\n", + " (ii) use cored light profiles for the background source galaxy, where the core ensures low levels of over-sampling\n", + " produce numerically accurate but fast to compute results.\n", + "\n", + "Over sampling at each galaxy centre (both main lens galaxies and extra galaxies) is performed to ensure the lens\n", + "calculations are accurate across the full field of the group.\n", + "\n", + "Once you are more experienced, you should read up on over-sampling in more detail via\n", + "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = list(main_lens_centres) + list(extra_galaxies_centres)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The imaging subplot updates the bottom two panels to reflect the update to over sampling, which now uses a higher\n", + "values in the centre.\n", + "\n", + "Whilst you may not yet understand the details of over-sampling, you can at least track it visually in the plots\n", + "and later learnt more about it in the `over_sampling.ipynb` guide." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "This shows the group scale model, with separate entries for each main lens galaxy, the source galaxy and the\n", + "extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The lens model is fitted to the data using a non-linear search.\n", + "\n", + "All examples in the autolens workspace use the nested sampling algorithm\n", + "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has revealed gives the most\n", + "accurate and efficient modeling results.\n", + "\n", + "Nautilus has one main setting that trades-off accuracy and computational run-time, the number of `live_points`.\n", + "A higher number of live points gives a more accurate result, but increases the run-time. A lower value give\n", + "less reliable lens modeling (e.g. the fit may infer a local maxima), but is faster.\n", + "\n", + "The suitable value depends on the model complexity whereby models with more parameters require more live points.\n", + "The default value of 200 is sufficient for the vast majority of common lens models. Lower values often given reliable\n", + "results though, and speed up the run-times. In this example, given the model is quite simple (N=21 parameters), we\n", + "reduce the number of live points to 100 to speed up the run-time.\n", + "\n", + "__Unique Identifier__\n", + "\n", + "In the path above, the `unique_identifier` appears as a collection of characters, where this identifier is generated\n", + "based on the model, search and dataset that are used in the fit.\n", + "\n", + "An identical combination of model and search generates the same identifier, meaning that rerunning the script will use\n", + "the existing results to resume the model-fit. In contrast, if you change the model or search, a new unique identifier\n", + "will be generated, ensuring that the model-fit results are output into a separate folder.\n", + "\n", + "We additionally want the unique identifier to be specific to the dataset fitted, so that if we fit different datasets\n", + "with the same model and search results are output to a different folder. We achieve this below by passing\n", + "the `dataset_name` to the search's `unique_tag`.\n", + "\n", + "__Iterations Per Update__\n", + "\n", + "Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and its best fit\n", + "image to the Jupyter Notebook display and to hard-disk.\n", + "\n", + "This process takes around ~10 seconds, so we don't want it to happen too often so as to slow down the overall\n", + "fit, but we also want it to happen frequently enough that we can track the progress.\n", + "\n", + "The value of 10000 below means this output happens every few minutes on GPU and every ~10 minutes on CPU, a good balance.\n", + "\n", + "__Live Visual Update__\n", + "\n", + "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", + "live display surface:\n", + "\n", + "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", + " watch the fit converge without leaving your terminal.\n", + "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", + " refreshes in place every `iterations_per_quick_update`.\n", + "\n", + "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", + "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"group\"), # The path where results and output are stored.\n", + " name=\"modeling\", # The name of the fit and folder results are output to.\n", + " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", + " n_live=150, # The number of Nautilus \"live\" points, increase for more complex models.\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=10000, # Every N iterations the max likelihood model is visualized in the Jupter Notebook and output to hard-disk.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We next create an `AnalysisImaging` object, which can be given many inputs customizing how the lens model is\n", + "fitted to the data (in this example they are omitted for simplicity).\n", + "\n", + "Internally, this object defines the `log_likelihood_function` used by the non-linear search to fit the model to\n", + "the `Imaging` dataset.\n", + "\n", + "It is not vital that you as a user understand the details of how the `log_likelihood_function` fits a lens model to\n", + "data, but interested readers can find a step-by-step guide of the likelihood\n", + "function at ``autolens_workspace/*/imaging/log_likelihood_function`\n", + "\n", + "__JAX__\n", + "\n", + "`AnalysisImaging` defaults to `use_jax=True`. For group-scale fits the\n", + "JAX speedup is substantial \u2014 the multi-galaxy deflection sum dominates\n", + "runtime on CPU but vectorises cleanly on GPU. Search driver wraps the\n", + "likelihood in `jax.vmap(jax.jit(...))`. Force NumPy with `use_jax=False`\n", + "(or `PYAUTO_DISABLE_JAX=1`) when debugging." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM Use__\n", + "\n", + "When running AutoLens with JAX on a GPU, the analysis must fit within the GPU's\n", + "available VRAM. If insufficient VRAM is available, the analysis will fail with an\n", + "out-of-memory error, typically during JIT compilation or the first likelihood call.\n", + "\n", + "Two factors dictate the VRAM usage of an analysis:\n", + "\n", + "- The number of arrays and other data structures JAX must store in VRAM to fit the model\n", + " to the data in the likelihood function. This is dictated by the model complexity and dataset size.\n", + " For a MGE model its relatively low, but for other models (e.g. pixelized sources) it can be much higher.\n", + "\n", + "- The `batch_size` sets how many likelihood evaluations are performed simultaneously.\n", + " Increasing the batch size increases VRAM usage but can reduce overall run time,\n", + " while decreasing it lowers VRAM usage at the cost of slower execution.\n", + "\n", + "Before running an analysis, users should check that the estimated VRAM usage for the\n", + "chosen batch size is comfortably below their GPU's total VRAM.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits.\n", + "\n", + "For a MGE model with the low resolution dataset fitted in this example VRAM use is relatively quite high (~1.8GB),\n", + "illustrating how group scale modeling is already quite VRAM intensive. For more complex models (e.g. pixelized sources)\n", + "and higher resolution datasets it can be much higher (> 1GB going beyond 10GB)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "Lens modeling can be a computationally expensive process. When fitting complex models to high resolution datasets\n", + "run times can be of order hours, days, weeks or even months.\n", + "\n", + "Run times are dictated by two factors:\n", + "\n", + " - The log likelihood evaluation time: the time it takes for a single `instance` of the lens model to be fitted to\n", + " the dataset such that a log likelihood is returned.\n", + "\n", + " - The number of iterations (e.g. log likelihood evaluations) performed by the non-linear search: more complex lens\n", + " models require more iterations to converge to a solution.\n", + "\n", + "For this analysis, the log likelihood evaluation time is < 0.005 seconds on GPU, < 0.05 seconds on CPU, which is\n", + "extremely fast for group-scale lens modeling.\n", + "\n", + "To estimate the expected overall run time of the model-fit we multiply the log likelihood evaluation time by an\n", + "estimate of the number of iterations the non-linear search will perform, which is around 20000 to 40000 for this model.\n", + "\n", + "GPU run times are around 15 minutes, CPU run times are around 45 minutes.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results).\n", + "\n", + "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce\n", + "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output Folder Layout__\n", + "\n", + "Now the fit is running you should checkout the `autolens_workspace/output` folder. This is where results are\n", + "written to hard-disk in human-readable formats \u2014 `.json`, `.csv`, `.fits`, `.png` and plain text.\n", + "\n", + "As the fit progresses, results are written on the fly using the highest likelihood model found by the\n", + "non-linear search so far. This means you can inspect the model-fit as it runs, without waiting for the\n", + "non-linear search to terminate.\n", + "\n", + "Each completed fit lives at a path like::\n", + "\n", + " output/group//modeling//\n", + " files/ <- JSON + CSV: loadable Python objects\n", + " tracer.json <- max log likelihood Tracer (all lens galaxies + source)\n", + " model.json <- fitted af.Collection model\n", + " samples.csv <- full Nautilus samples\n", + " samples_summary.json <- max log likelihood parameter values + errors\n", + " samples_info.json <- metadata about the samples\n", + " search.json <- non-linear search configuration\n", + " settings.json <- search settings\n", + " cosmology.json <- cosmology used for the fit\n", + " covariance.csv <- parameter covariance matrix\n", + " image/ <- FITS + PNG: imaging products\n", + " dataset.fits <- data, noise-map and PSF\n", + " fit.fits <- model image, residuals, chi-squared map\n", + " tracer.fits <- tracer image-plane images per galaxy\n", + " source_plane_images.fits <- source plane reconstructions\n", + " model_galaxy_images.fits <- per-galaxy model images (lens_0, lens_1, ..., source)\n", + " galaxy_images.fits <- per-galaxy images\n", + " dataset.png, fit.png, tracer.png <- visualisations\n", + " model.info <- human-readable model summary\n", + " model.results <- human-readable fit summary\n", + " search.summary <- search run summary\n", + " search_internal/ <- internal files used to resume / visualise the search\n", + " metadata <- run metadata\n", + "\n", + "The `` is a 32-character identifier derived from the model, search and dataset, so re-running the\n", + "same configuration resumes from the existing fit automatically.\n", + "\n", + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", + "\n", + "The result contains entries for each main lens galaxy (e.g. `lens_0`), the source galaxy and the extra galaxies.\n", + "\n", + "[Above, we discussed that the `info_whitespace_length` parameter in the config files could b changed to make\n", + "the `model.info` attribute display optimally on your computer. This attribute also controls the whitespace of the\n", + "`result.info` attribute.]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Result` object also contains:\n", + "\n", + " - The model corresponding to the maximum log likelihood solution in parameter space.\n", + " - The corresponding maximum log likelihood `Tracer` and `FitImaging` objects.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It also contains information on the posterior as estimated by the non-linear search (in this example `Nautilus`).\n", + "\n", + "Below, we make a corner plot of the \"Probability Density Function\" of every parameter in the model-fit.\n", + "\n", + "The plot is labeled with short hand parameter names (e.g. `sersic_index` is mapped to the short hand\n", + "parameter `n`). These mappings ate specified in the `config/notation.yaml` file and can be customized by users.\n", + "\n", + "The superscripts of labels correspond to the name each component was given in the model (e.g. for the `Isothermal`\n", + "mass its name `mass` defined when making the `Model` above is used)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This script gives a concise overview of the basic modeling API, fitting one the simplest lens models possible.\n", + "\n", + "Lets now consider what features you should read about to improve your lens modeling, especially if you are aiming\n", + "to fit more complex models to your data.\n", + "\n", + "This is especially important for group scale modeling, in order to reduce the complexity of the model.\n", + "\n", + "__Features__\n", + "\n", + "The examples in the `autolens_workspace/*/imaging/features` package illustrate other lens modeling features.\n", + "\n", + "We recommend you checkout the following features, because they make lens modeling in general more reliable and\n", + "efficient (you will therefore benefit from using these features irrespective of the quality of your data and\n", + "scientific topic of study).\n", + "\n", + "We recommend you now checkout the following features:\n", + "\n", + "- ``scaling_relation``: This feature allows you to model the light and mass of the extra galaxies using a scaling relation.\n", + "- ``linear_light_profiles``: The model light profiles use linear algebra to solve for their intensity, reducing model complexity.\n", + "- ``multi_gaussian_expansion``: The lens (or source) light is modeled as ~25-100 Gaussian basis functions\n", + "- ``pixelization``: The source is reconstructed using an adaptive rectangular or Delaunay mesh\n", + "- ``no_lens_light``: The foreground lens's light is not present in the data and thus omitted from the model.\n", + "\n", + "For group scale modeling, the multi Gaussian expansion is particularly important, as this can dramatically reduce the\n", + "dimensionality of the model and improve the accuracy of the fit for both the lens and source galaxies.\n", + "\n", + "It is also recommended you read through the `imaging` package, to get a complete picture of how point-source\n", + "modeling works.\n", + "\n", + "__Data Preparation__\n", + "\n", + "If you are looking to fit your own CCD imaging data of a strong lens, checkout\n", + "the `autolens_workspace/*/imaging/data_preparation/start_here.ipynb` script for an overview of how data should be\n", + "prepared before being modeled.\n", + "\n", + "__HowToLens__\n", + "\n", + "This `start_here.py` script, and the features examples above, do not explain many details of how lens modeling is\n", + "performed, for example:\n", + "\n", + " - How does PyAutoLens perform ray-tracing and lensing calculations in order to fit a lens model?\n", + " - How is a lens model fitted to data? What quantifies the goodness of fit (e.g. how is a log likelihood computed?).\n", + " - How does Nautilus find the highest likelihood lens models? What exactly is a \"non-linear search\"?\n", + "\n", + "You do not need to be able to answer these questions in order to fit lens models with PyAutoLens and do science.\n", + "However, having a deeper understanding of how it all works is both interesting and will benefit you as a scientist\n", + "\n", + "This deeper insight is offered by the **HowToLens** Jupyter notebook lectures, which live in their own repository:\n", + "https://github.com/PyAutoLabs/HowToLens.\n", + "\n", + "I recommend that you check them out if you are interested in more details!\n", + "\n", + "__Modeling Customization__\n", + "\n", + "The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative non-linear searches,\n", + "other than Nautilus, that can be used to fit lens models.\n", + "\n", + "They also provide details on how to customize the model-fit, for example the priors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/simulator.ipynb b/notebooks/group/simulator.ipynb index 4ac790659..846bf9364 100644 --- a/notebooks/group/simulator.ipynb +++ b/notebooks/group/simulator.ipynb @@ -1,576 +1,613 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Group\n", - "================\n", - "\n", - "This script simulates an example strong lens on the 'group' scale, where there is a single primary lens galaxy\n", - "and two smaller extra galaxies nearby, whose mass contributes significantly to the ray-tracing and is therefore\n", - "included in the strong lens model.\n", - "\n", - "This script simulates `Imaging` of a 'group-scale' strong lens where:\n", - "\n", - " - The group consists of one main lens galaxy and two extra galaxies whose light distributions are `SersicSph`\n", - " profiles and total mass distributions are `IsothermalSph` profiles.\n", - " - A single source galaxy is observed whose `LightProfile` is a `SersicCore`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Main Lens Galaxies vs Extra Galaxies:** For group-scale lens modeling, galaxies are organized into two categories.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Grid:** Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and.\n", - "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Main Lens Galaxies:** The main lens galaxy is at the origin (0.0, 0.0).\n", - "- **Extra Galaxies:** The two extra galaxies are companion galaxies near the lens system.\n", - "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", - "- **Ray Tracing:** Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging`.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "- **Centre JSON Files:** Save the centres of the main lens galaxies and extra galaxies as JSON files.\n", - "\n", - "__Main Lens Galaxies vs Extra Galaxies__\n", - "\n", - "For group-scale lens modeling, galaxies are organized into two categories:\n", - "\n", - " - `main_lens_galaxies`: The primary lens galaxies that dominate the light and mass of the system. These are\n", - " modeled individually with unique parametric light and mass profiles.\n", - "\n", - " - `extra_galaxies`: Companion galaxies near the lens system that contribute to lensing but are modeled with\n", - " more restrictive assumptions (e.g. fixed centres, scaling relations).\n", - "\n", - "Centres for each category are saved to separate JSON files (`main_lens_centres.json` and\n", - "`extra_galaxies_centres.json`) so that the modeling scripts can load them directly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", - "\n", - " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/image.fits`.\n", - " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.fits`.\n", - " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/psf.fits`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"group\"\n", - "dataset_name = \"simple\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output.\n", - "\n", - "In this example, this is: `/autolens_workspace/dataset/group/simple`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and therefore simulated\n", - "on, via the inputs:\n", - "\n", - " - `shape_native`: The (y_pixels, x_pixels) 2D shape of the grid defining the shape of the data that is simulated.\n", - " - `pixel_scales`: The arc-second to pixel conversion factor of the grid and data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(250, 250),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Define the centres of the main lens galaxies and extra galaxies. These are used for over-sampling and are also\n", - "output to JSON files so that the modeling scripts can load them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = [(0.0, 0.0)]\n", - "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", - "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of\n", - "over sampling to ensure the lensed images are evaluated accurately.\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", - "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "An adaptive oversampling scheme is implemented, evaluating the central regions at (0.0\", 0.0\") of the light profile at a\n", - "resolution of 32x32, transitioning to 8x8 in intermediate areas, and 2x2 in the outskirts. This ensures precise and\n", - "accurate image simulation while focusing computational resources on the bright regions that demand higher oversampling.\n", - "\n", - "This adaptive over sampling is also applied at the centre of every other galaxy in the group.\n", - "\n", - "An adaptive oversampling grid cannot be defined for the lensed source because its light appears in different regions of\n", - "the image plane for each dataset. For this reason, most workspace examples utilize cored light profiles for the\n", - "source galaxy. Cored light profiles change gradually in their central regions, allowing accurate evaluation without\n", - "requiring oversampling.\n", - "\n", - "Once you are more experienced, you should read up on over-sampling in more detail via\n", - "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=main_lens_centres + extra_galaxies_centres,\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the `Imaging` dataset we first create a simulator, which defines the exposure time, background sky,\n", - "noise levels and psf of the dataset that is simulated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Main Lens Galaxies__\n", - "\n", - "The main lens galaxy is at the origin (0.0, 0.0). It has a spherical Sersic light profile and an isothermal\n", - "mass profile.\n", - "\n", - "In the list-based API used by the group modeling scripts, main lens galaxies are stored in a list called\n", - "`main_lens_galaxies`, where each galaxy is referred to as `lens_0`, `lens_1`, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "main_lens_galaxies = [lens_0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "The two extra galaxies are companion galaxies near the lens system. They have spherical Sersic light profiles\n", - "and isothermal mass profiles, with centres offset from the origin.\n", - "\n", - "In the list-based API, extra galaxies are stored in a list called `extra_galaxies`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy__\n", - "\n", - "The source galaxy whose lensed images we simulate. It uses a cored Sersic profile so that adaptive over-sampling\n", - "is not required for the source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset.\n", - "\n", - "The tracer combines main lens galaxies, extra galaxies and the source galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets plot the simulated `Imaging` dataset before we output it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future.\n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centre JSON Files__\n", - "\n", - "Save the centres of the main lens galaxies and extra galaxies as JSON files. These are loaded by the group\n", - "modeling scripts to set up the lens model (e.g. fixing centres of extra galaxies, defining scaling relations)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(main_lens_centres),\n", - " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Solve for the lensed positions of the source galaxy, which are used as input for the group\n", - "modeling scripts (e.g. SLaM pipeline) to help the non-linear search converge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1),\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.01,\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=positions,\n", - " file_path=dataset_path / \"positions.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished.\n", - "\n", - "__JAX Variant__\n", - "\n", - "Same pattern as `scripts/imaging/simulator.py` `__JAX Variant__`:\n", - "instantiate `al.SimulatorImaging(use_jax=True)` and wrap\n", - "`via_tracer_from` in `@jax.jit`. The simulator handles pytree\n", - "registration internally.\n", - "\n", - "See `scripts/imaging/simulator.py` for the runnable variant block." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Group\n", + "================\n", + "\n", + "This script simulates an example strong lens on the 'group' scale, where there is a single primary lens galaxy\n", + "and two smaller extra galaxies nearby, whose mass contributes significantly to the ray-tracing and is therefore\n", + "included in the strong lens model.\n", + "\n", + "This script simulates `Imaging` of a 'group-scale' strong lens where:\n", + "\n", + " - The group consists of one main lens galaxy and two extra galaxies whose light distributions are `SersicSph`\n", + " profiles and total mass distributions are `IsothermalSph` profiles.\n", + " - A single source galaxy is observed whose `LightProfile` is a `SersicCore`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Main Lens Galaxies vs Extra Galaxies:** For group-scale lens modeling, galaxies are organized into two categories.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Grid:** Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and.\n", + "- **Galaxy Centres:** Define the centres of the main lens galaxies and extra galaxies.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Main Lens Galaxies:** The main lens galaxy is at the origin (0.0, 0.0).\n", + "- **Extra Galaxies:** The two extra galaxies are companion galaxies near the lens system.\n", + "- **Source Galaxy:** The source galaxy whose lensed images we simulate.\n", + "- **Ray Tracing:** Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging`.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "- **Centre JSON Files:** Save the centres of the main lens galaxies and extra galaxies as JSON files.\n", + "\n", + "__Main Lens Galaxies vs Extra Galaxies__\n", + "\n", + "For group-scale lens modeling, galaxies are organized into two categories:\n", + "\n", + " - `main_lens_galaxies`: The primary lens galaxies that dominate the light and mass of the system. These are\n", + " modeled individually with unique parametric light and mass profiles.\n", + "\n", + " - `extra_galaxies`: Companion galaxies near the lens system that contribute to lensing but are modeled with\n", + " more restrictive assumptions (e.g. fixed centres, scaling relations).\n", + "\n", + "Centres for each category are saved to separate JSON files (`main_lens_centres.json` and\n", + "`extra_galaxies_centres.json`) so that the modeling scripts can load them directly." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", + "\n", + " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/image.fits`.\n", + " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.fits`.\n", + " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/psf.fits`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"group\"\n", + "dataset_name = \"simple\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output.\n", + "\n", + "In this example, this is: `/autolens_workspace/dataset/group/simple`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and therefore simulated\n", + "on, via the inputs:\n", + "\n", + " - `shape_native`: The (y_pixels, x_pixels) 2D shape of the grid defining the shape of the data that is simulated.\n", + " - `pixel_scales`: The arc-second to pixel conversion factor of the grid and data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(250, 250),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Define the centres of the main lens galaxies and extra galaxies. These are used for over-sampling and are also\n", + "output to JSON files so that the modeling scripts can load them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = [(0.0, 0.0)]\n", + "extra_galaxies_centres = [(3.5, 2.5), (-4.4, -5.0)]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", + "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of\n", + "over sampling to ensure the lensed images are evaluated accurately.\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", + "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "An adaptive oversampling scheme is implemented, evaluating the central regions at (0.0\", 0.0\") of the light profile at a\n", + "resolution of 32x32, transitioning to 8x8 in intermediate areas, and 2x2 in the outskirts. This ensures precise and\n", + "accurate image simulation while focusing computational resources on the bright regions that demand higher oversampling.\n", + "\n", + "This adaptive over sampling is also applied at the centre of every other galaxy in the group.\n", + "\n", + "An adaptive oversampling grid cannot be defined for the lensed source because its light appears in different regions of\n", + "the image plane for each dataset. For this reason, most workspace examples utilize cored light profiles for the\n", + "source galaxy. Cored light profiles change gradually in their central regions, allowing accurate evaluation without\n", + "requiring oversampling.\n", + "\n", + "Once you are more experienced, you should read up on over-sampling in more detail via\n", + "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=main_lens_centres + extra_galaxies_centres,\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the `Imaging` dataset we first create a simulator, which defines the exposure time, background sky,\n", + "noise levels and psf of the dataset that is simulated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Main Lens Galaxies__\n", + "\n", + "The main lens galaxy is at the origin (0.0, 0.0). It has a spherical Sersic light profile and an isothermal\n", + "mass profile.\n", + "\n", + "In the list-based API used by the group modeling scripts, main lens galaxies are stored in a list called\n", + "`main_lens_galaxies`, where each galaxy is referred to as `lens_0`, `lens_1`, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=2.0, sersic_index=4.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "main_lens_galaxies = [lens_0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "The two extra galaxies are companion galaxies near the lens system. They have spherical Sersic light profiles\n", + "and isothermal mass profiles, with centres offset from the origin.\n", + "\n", + "In the list-based API, extra galaxies are stored in a list called `extra_galaxies`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-4.4, -5.0), intensity=0.9, effective_radius=0.8, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "extra_galaxies = [extra_galaxy_0, extra_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy__\n", + "\n", + "The source galaxy whose lensed images we simulate. It uses a cored Sersic profile so that adaptive over-sampling\n", + "is not required for the source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Use all galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset.\n", + "\n", + "The tracer combines main lens galaxies, extra galaxies and the source galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=main_lens_galaxies + extra_galaxies + [source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets plot the simulated `Imaging` dataset before we output it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future.\n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centre JSON Files__\n", + "\n", + "Save the centres of the main lens galaxies and extra galaxies as JSON files. These are loaded by the group\n", + "modeling scripts to set up the lens model (e.g. fixing centres of extra galaxies, defining scaling relations)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(main_lens_centres),\n", + " file_path=Path(dataset_path, \"main_lens_centres.json\"),\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Solve for the lensed positions of the source galaxy, which are used as input for the group\n", + "modeling scripts (e.g. SLaM pipeline) to help the non-linear search converge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.1),\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.01,\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=positions,\n", + " file_path=dataset_path / \"positions.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished.\n", + "\n", + "__JAX Variant__\n", + "\n", + "Same pattern as `scripts/imaging/simulator.py` `__JAX Variant__`:\n", + "instantiate `al.SimulatorImaging(use_jax=True)` and wrap\n", + "`via_tracer_from` in `@jax.jit`. The simulator handles pytree\n", + "registration internally.\n", + "\n", + "See `scripts/imaging/simulator.py` for the runnable variant block." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/slam.ipynb b/notebooks/group/slam.ipynb index e191f7753..3af681224 100644 --- a/notebooks/group/slam.ipynb +++ b/notebooks/group/slam.ipynb @@ -1,1213 +1,1250 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM (Source, Light and Mass): Group Scale\n", - "==========================================\n", - "\n", - "This script uses the SLaM pipelines to fit a group-scale strong lens, including extra galaxies and\n", - "scaling galaxies surrounding the main lens whose light and mass are both modeled.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **Extra Galaxies and Scaling Galaxies:** This group-scale SLaM pipeline handles two distinct categories of companion galaxy, which differ in.\n", - "- **This Script:** Using a SOURCE LP PIPELINE (two searches), SOURCE PIX PIPELINE (two searches), LIGHT LP PIPELINE.\n", - "- **SOURCE LP PIPELINE 0:** Not present in `slam_start_here.py`.\n", - "- **SOURCE LP PIPELINE 1:** Equivalent to `source_lp` in `slam_start_here.py`, except lens light is fixed from `source_lp[0]`.\n", - "- **SOURCE PIX PIPELINE 1:** Equivalent to `source_pix_1` in `slam_start_here.py`, except a Hilbert image mesh is used instead.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `source_pix_1` above, except the adapt data for the Hilbert image mesh is capped at a.\n", - "- **LIGHT LP PIPELINE:** Identical to `light_lp` in `slam_start_here.py`, except extra galaxies receive a fresh free MGE.\n", - "- **MASS TOTAL PIPELINE:** Identical to `mass_total` in `slam_start_here.py`, except extra galaxies receive a new.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Galaxy Centres:** main_lens_centres.json \u2014 required; determines the number of main lenses.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "- **Group** (`group/modeling`):\n", - " How we model group-scale strong lenses in PyAutoLens, including how we include extra galaxies in\n", - " the lens model.\n", - "\n", - "__Extra Galaxies and Scaling Galaxies__\n", - "\n", - "This group-scale SLaM pipeline handles two distinct categories of companion galaxy, which differ in\n", - "how their masses are parameterized:\n", - "\n", - "**Extra Galaxies**\n", - "\n", - "Extra galaxies are a small number of nearby companions whose light and mass are modeled individually.\n", - "In each pipeline stage, they receive a free MGE light profile and an `Isothermal` mass profile whose\n", - "Einstein radius prior is bounded by a value derived from the galaxy's luminosity:\n", - "\n", - " upper_limit = min(5 * 0.5 * total_luminosity^0.6, 5.0)\n", - "\n", - "This luminosity-informed bound prevents unphysically large mass assignments while keeping the mass\n", - "free per galaxy. The extra-galaxy models are stored in `model.extra_galaxies` (a `Collection`).\n", - "\n", - "**Scaling Galaxies**\n", - "\n", - "Scaling galaxies are a larger ensemble of companions whose masses are constrained through a shared\n", - "luminosity-to-mass scaling relation rather than being individually free. They each carry a free MGE\n", - "light profile, but their Einstein radii follow:\n", - "\n", - " einstein_radius = scaling_factor * total_luminosity^scaling_relation\n", - "\n", - "where `scaling_factor` and `scaling_relation` are two shared free parameters whose priors are\n", - "`UniformPrior(0, 0.5)` and `UniformPrior(0, 2)` respectively. This reduces the number of mass\n", - "parameters considerably when many companion galaxies are present. Scaling-galaxy models are stored\n", - "in `model.scaling_galaxies`.\n", - "\n", - "The choice between these two categories is determined by which JSON file each galaxy's centre\n", - "appears in (`extra_galaxies_centres.json` vs `scaling_galaxies_centres.json`).\n", - "\n", - "**Comparison with the Galaxy-Scale SLaM Pipeline**\n", - "\n", - "The galaxy-scale extra-galaxies SLaM pipeline (`features/extra_galaxies/slam`) models each\n", - "companion galaxy with a fully free `IsothermalSph` mass \u2014 it does not use luminosity bounds or a\n", - "shared scaling relation. The group-scale pipeline introduced here is appropriate when:\n", - "\n", - "- there are more companion galaxies than can be modeled independently, or\n", - "- prior physical knowledge motivates a luminosity-to-mass scaling.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE (two searches), SOURCE PIX PIPELINE (two searches), LIGHT LP PIPELINE\n", - "and TOTAL MASS PIPELINE this SLaM modeling script fits `Imaging` data of a group-scale strong lens\n", - "where in the final model:\n", - "\n", - " - Each main lens galaxy has a free MGE bulge and a `PowerLaw` total mass.\n", - " - Each extra galaxy has a free MGE bulge and a luminosity-bounded `Isothermal` mass.\n", - " - Each scaling galaxy has a free MGE bulge and a mass set by a shared scaling relation.\n", - " - The source galaxy's light is a Delaunay `Pixelization` with `AdaptSplit` regularization.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "\n", - "def _load_centres(path):\n", - " \"\"\"Load a centres JSON file, returning an empty list if the file is absent.\"\"\"\n", - " try:\n", - " return al.Grid2DIrregular(al.from_json(file_path=path))\n", - " except FileNotFoundError:\n", - " return al.Grid2DIrregular([])\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 0__\n", - "\n", - "Not present in `slam_start_here.py`. Fits light only \u2014 no mass, no source \u2014 for every galaxy\n", - "simultaneously, giving the next search clean fixed light models to build on.\n", - "\n", - "Fits multiple main-lens galaxies (`lens_0`, `lens_1`, ...) under `galaxies`, extra galaxies\n", - "under `extra_galaxies`, and scaling galaxies under `scaling_galaxies`. `n_live` scales with\n", - "the total number of galaxies across all three categories." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_0(\n", - " dataset,\n", - " settings_search,\n", - " main_lens_centres,\n", - " extra_lens_centres,\n", - " scaling_lens_centres,\n", - " mask_radius,\n", - " redshift_lens,\n", - " n_batch=50,\n", - "):\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " # --- main lens light models (one per centre, light only) ---\n", - " lens_dict = {}\n", - " for i, centre in enumerate(main_lens_centres):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=False,\n", - " centre=(centre[0], centre[1]),\n", - " centre_sigma=0.1,\n", - " )\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=bulge, disk=None, point=None\n", - " )\n", - "\n", - " # --- extra lens galaxy light models ---\n", - " extra_light_models = []\n", - " for centre in extra_lens_centres:\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_prior_is_uniform=True,\n", - " centre=(centre[0], centre[1]),\n", - " ell_comps_prior_is_uniform=True,\n", - " )\n", - " extra_light_models.append(\n", - " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", - "\n", - " # --- scaling galaxy light models ---\n", - " scaling_light_models = []\n", - " for centre in scaling_lens_centres:\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_prior_is_uniform=True,\n", - " centre=(centre[0], centre[1]),\n", - " ell_comps_prior_is_uniform=True,\n", - " )\n", - " scaling_light_models.append(\n", - " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", - " )\n", - "\n", - " scaling_galaxies = (\n", - " af.Collection(scaling_light_models) if scaling_light_models else None\n", - " )\n", - "\n", - " n_extra = len(extra_galaxies) if extra_galaxies is not None else 0\n", - " n_scaling = len(scaling_galaxies) if scaling_galaxies is not None else 0\n", - " n_live = 100 + 30 * len(lens_dict) + 30 * n_extra + 30 * n_scaling\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=scaling_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[0]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " n_like_max=1000000,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 1__\n", - "\n", - "Equivalent to `source_lp` in `slam_start_here.py`, except lens light is fixed from\n", - "`source_lp[0]` rather than free, and mass and source are introduced here for the first time.\n", - "\n", - "Multiple main-lens galaxies each get an `Isothermal` mass; only `lens_0` carries an\n", - "`ExternalShear`. Extra-galaxy Einstein radii are bounded by a luminosity-derived prior\n", - "(`min(5 * 0.5 * L^0.6, 5.0)`). Scaling galaxies share two free parameters,\n", - "`scaling_factor` and `scaling_relation`, so their masses follow\n", - "`einstein_radius = scaling_factor * luminosity^scaling_relation`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_1(\n", - " dataset,\n", - " settings_search,\n", - " source_lp_result_0,\n", - " positions,\n", - " pixel_scale,\n", - " redshift_lens,\n", - " redshift_source,\n", - " source_mge_radius,\n", - " n_batch=50,\n", - "):\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_main = sum(\n", - " 1 for k in vars(source_lp_result_0.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - " n_extra = (\n", - " len(list(source_lp_result_0.instance.extra_galaxies))\n", - " if source_lp_result_0.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - " n_scaling = (\n", - " len(list(source_lp_result_0.instance.scaling_galaxies))\n", - " if source_lp_result_0.instance.scaling_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " tracer = (\n", - " source_lp_result_0.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", - " )\n", - "\n", - " # Source MGE centred on primary lens bulge from stage 0.\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=source_mge_radius,\n", - " total_gaussians=30,\n", - " centre=source_lp_result_0.instance.galaxies.lens_0.bulge.centre,\n", - " centre_prior_is_uniform=False,\n", - " centre_sigma=0.6,\n", - " )\n", - "\n", - " # --- main lens full models (light fixed from stage 0, mass + shear free) ---\n", - " # Only lens_0 carries the ExternalShear; one shear per group system.\n", - " lens_dict = {}\n", - " for i in range(n_main):\n", - " lp0_lens = getattr(source_lp_result_0.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = lp0_lens.bulge.centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=5.0)\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lp0_lens.bulge,\n", - " disk=lp0_lens.disk,\n", - " point=lp0_lens.point,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", - " )\n", - "\n", - " # --- extra lens galaxy models (light fixed, mass bounded by luminosity) ---\n", - " # Tracer order: [lens_0..lens_{n_main-1}, extra_0..extra_{n_extra-1}, scaling_0..]\n", - " extra_mass_models = []\n", - " for i in range(n_extra):\n", - " lp0_extra = source_lp_result_0.instance.extra_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = lp0_extra.bulge.centre\n", - " mass.ell_comps = lp0_extra.bulge.ell_comps\n", - "\n", - " luminosity_per_gaussian_list = [\n", - " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", - " for g in tracer.galaxies[n_main + i].bulge.profile_list\n", - " ]\n", - " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", - " mass.einstein_radius = af.UniformPrior(\n", - " lower_limit=0.0,\n", - " upper_limit=min(5 * 0.5 * total_luminosity**0.6, 5.0),\n", - " )\n", - "\n", - " extra_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=lp0_extra.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", - "\n", - " # --- scaling lens galaxy models (light fixed, shared luminosity scaling relation) ---\n", - " scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - " scaling_relation = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - " scaling_mass_models = []\n", - " for i in range(n_scaling):\n", - " lp0_scaling = source_lp_result_0.instance.scaling_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = lp0_scaling.bulge.centre\n", - " mass.ell_comps = lp0_scaling.bulge.ell_comps\n", - "\n", - " luminosity_per_gaussian_list = [\n", - " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", - " for g in tracer.galaxies[n_main + n_extra + i].bulge.profile_list\n", - " ]\n", - " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", - " mass.einstein_radius = scaling_factor * total_luminosity**scaling_relation\n", - "\n", - " scaling_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=lp0_scaling.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " scaling_galaxies = (\n", - " af.Collection(scaling_mass_models) if scaling_mass_models else None\n", - " )\n", - "\n", - " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=scaling_galaxies,\n", - " )\n", - "\n", - " n_extra_model = len(extra_galaxies) if extra_galaxies is not None else 0\n", - " n_scaling_model = len(scaling_galaxies) if scaling_galaxies is not None else 0\n", - " n_live = 150 + 30 * n_main + 30 * n_extra_model + 30 * n_scaling_model\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " n_like_max=200000,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Equivalent to `source_pix_1` in `slam_start_here.py`, except a Hilbert image mesh is used\n", - "instead of `RectangularAdaptDensity`, with edge-point padding of 30 pixels. The Hilbert pixel\n", - "count is set by `al.model_util.hilbert_pixels_from_pixel_scale`.\n", - "\n", - "Pixelization over-sampling is signal-adaptive: pixels above the S/N threshold use sub-size 4,\n", - "the rest sub-size 2. The re-sampled dataset and adapt_images are returned alongside the result.\n", - "Extra and scaling galaxy models are carried forward as free `model` parameters, not fixed instances." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " dataset,\n", - " mask,\n", - " settings_search,\n", - " source_lp_result_1,\n", - " over_sample_size,\n", - " pixel_scale,\n", - " mask_radius,\n", - " positions,\n", - " n_batch=20,\n", - "):\n", - " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", - " edge_pixels_total = 30\n", - " signal_to_noise_threshold = 3.0\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result_1\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Hilbert(\n", - " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", - " )\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=mask,\n", - " adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"],\n", - " )\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " over_sample_size_pixelization = np.where(\n", - " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", - " 4,\n", - " 2,\n", - " )\n", - " over_sample_size_pixelization = al.Array2D(\n", - " values=over_sample_size_pixelization, mask=mask\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=over_sample_size_pixelization,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result_1.positions_likelihood_from(\n", - " factor=2.0, positions=positions, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_lp_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lp_lens_instance = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", - " lp_lens_model = getattr(source_lp_result_1.model.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=lp_lens_model.mass,\n", - " mass_result=lp_lens_model.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lp_lens_instance.redshift,\n", - " bulge=lp_lens_instance.bulge,\n", - " disk=lp_lens_instance.disk,\n", - " point=lp_lens_instance.point,\n", - " mass=mass,\n", - " shear=lp_lens_model.shear,\n", - " )\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=af.Model(al.reg.AdaptSplit),\n", - " ),\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=source_lp_result_1.model.extra_galaxies,\n", - " scaling_galaxies=source_lp_result_1.model.scaling_galaxies,\n", - " )\n", - "\n", - " n_live = 150 + 50 * (n_lenses - 1)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - " return result, dataset, adapt_images\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `source_pix_1` above, except the adapt data for the Hilbert image mesh is capped\n", - "at a S/N threshold of 3.0 to prevent over-concentration of source pixels on the brightest peak,\n", - "and extra and scaling galaxy models are fixed as instances from `source_pix[1]`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " dataset,\n", - " mask,\n", - " settings_search,\n", - " source_lp_result_1,\n", - " source_pix_result_1,\n", - " over_sample_size,\n", - " pixel_scale,\n", - " mask_radius,\n", - " n_batch=20,\n", - "):\n", - " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", - " edge_pixels_total = 30\n", - " signal_to_noise_threshold = 3.0\n", - " signal_to_noise_threshold_image_mesh = 3.0\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_data_snr_max = galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", - " adapt_data_snr_max[adapt_data_snr_max > signal_to_noise_threshold_image_mesh] = (\n", - " signal_to_noise_threshold_image_mesh\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Hilbert(\n", - " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", - " )\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=mask, adapt_data=adapt_data_snr_max\n", - " )\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " over_sample_size_pixelization = np.where(\n", - " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", - " 4,\n", - " 2,\n", - " )\n", - " over_sample_size_pixelization = al.Array2D(\n", - " values=over_sample_size_pixelization, mask=mask\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size,\n", - " over_sample_size_pixelization=over_sample_size_pixelization,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lp_lens_instance = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", - " pix1_lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lp_lens_instance.redshift,\n", - " bulge=lp_lens_instance.bulge,\n", - " disk=lp_lens_instance.disk,\n", - " point=lp_lens_instance.point,\n", - " mass=pix1_lens_instance.mass,\n", - " shear=pix1_lens_instance.shear,\n", - " )\n", - "\n", - " source = af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=af.Model(al.reg.AdaptSplit),\n", - " ),\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", - " scaling_galaxies=source_pix_result_1.instance.scaling_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - " return result, dataset, adapt_images\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `light_lp` in `slam_start_here.py`, except extra galaxies receive a fresh free\n", - "MGE bulge (centred on the `source_pix[1]` mass centre) with mass fixed from `source_pix[1]`,\n", - "and scaling galaxies are fully fixed from `source_pix[2]`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " dataset,\n", - " settings_search,\n", - " source_lp_result_0,\n", - " source_pix_result_1,\n", - " source_pix_result_2,\n", - " adapt_images,\n", - " mask_radius,\n", - " redshift_lens,\n", - " n_batch=20,\n", - "):\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " n_lenses = sum(\n", - " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - " n_extra = (\n", - " len(list(source_pix_result_1.instance.extra_galaxies))\n", - " if source_pix_result_1.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " # --- main lens light models (MGE centred on stage-0 bulge centre) ---\n", - " lens_bulge_list = []\n", - " for i in range(n_lenses):\n", - " lp0_lens = getattr(source_lp_result_0.instance.galaxies, f\"lens_{i}\")\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " centre=lp0_lens.bulge.centre,\n", - " )\n", - " lens_bulge_list.append(bulge)\n", - "\n", - " # --- extra lens galaxy light models (free MGE, mass fixed from source_pix[1]) ---\n", - " extra_light_models = []\n", - " for i in range(n_extra):\n", - " pix1_extra = source_pix_result_1.instance.extra_galaxies[i]\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_prior_is_uniform=True,\n", - " centre=pix1_extra.mass.centre,\n", - " )\n", - " extra_light_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=pix1_extra.mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_pix_result_2, source_is_model=False\n", - " )\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lens_instance.redshift,\n", - " bulge=lens_bulge_list[i],\n", - " disk=None,\n", - " point=None,\n", - " mass=lens_instance.mass,\n", - " shear=lens_instance.shear,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=source_pix_result_2.instance.scaling_galaxies,\n", - " )\n", - "\n", - " n_live = 300 + 100 * (n_lenses - 1)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `mass_total` in `slam_start_here.py`, except extra galaxies receive a new\n", - "luminosity-bounded `Isothermal` mass (using `light[1]` luminosities) and scaling galaxies\n", - "receive a new shared luminosity scaling relation, both paired with their fixed `light[1]` bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " dataset,\n", - " settings_search,\n", - " source_pix_result_1,\n", - " source_pix_result_2,\n", - " light_result,\n", - " adapt_images,\n", - " positions,\n", - " pixel_scale,\n", - " redshift_lens,\n", - " n_batch=20,\n", - "):\n", - " # Total mass model for each main lens galaxy.\n", - " n_lenses = sum(\n", - " 1 for k in vars(light_result.instance.galaxies) if k.startswith(\"lens_\")\n", - " )\n", - " n_extra = (\n", - " len(list(light_result.instance.extra_galaxies))\n", - " if light_result.instance.extra_galaxies is not None\n", - " else 0\n", - " )\n", - " n_scaling = (\n", - " len(list(light_result.instance.scaling_galaxies))\n", - " if light_result.instance.scaling_galaxies is not None\n", - " else 0\n", - " )\n", - "\n", - " # --- extra galaxies: fixed light, free mass ---\n", - " extra_mass_models = []\n", - " for i in range(n_extra):\n", - " light_extra = light_result.instance.extra_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = light_extra.bulge.centre\n", - " mass.ell_comps = light_extra.bulge.ell_comps\n", - "\n", - " luminosity_per_gaussian_list = [\n", - " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", - " for g in light_extra.bulge.profile_list\n", - " ]\n", - " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", - " mass.einstein_radius = af.UniformPrior(\n", - " lower_limit=0.0,\n", - " upper_limit=min(5 * 0.5 * total_luminosity**0.6, 5.0),\n", - " )\n", - "\n", - " extra_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=light_extra.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", - "\n", - " # --- scaling galaxies: fixed light, free shared scaling relation ---\n", - " scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - " scaling_relation = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - " scaling_mass_models = []\n", - " for i in range(n_scaling):\n", - " light_scaling = light_result.instance.scaling_galaxies[i]\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = light_scaling.bulge.centre\n", - " mass.ell_comps = light_scaling.bulge.ell_comps\n", - "\n", - " luminosity_per_gaussian_list = [\n", - " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", - " for g in light_scaling.bulge.profile_list\n", - " ]\n", - " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", - " mass.einstein_radius = scaling_factor * total_luminosity**scaling_relation\n", - "\n", - " scaling_mass_models.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=light_scaling.bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " scaling_galaxies = (\n", - " af.Collection(scaling_mass_models) if scaling_mass_models else None\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " light_result.positions_likelihood_from(\n", - " factor=3.0, positions=positions, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " lens_dict = {}\n", - " for i in range(n_lenses):\n", - " lens_model = getattr(source_pix_result_1.model.galaxies, f\"lens_{i}\")\n", - " light_lens_instance = getattr(light_result.instance.galaxies, f\"lens_{i}\")\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.PowerLaw),\n", - " mass_result=lens_model.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " lens_dict[f\"lens_{i}\"] = af.Model(\n", - " al.Galaxy,\n", - " redshift=lens_model.redshift,\n", - " bulge=light_lens_instance.bulge,\n", - " disk=light_lens_instance.disk,\n", - " point=light_lens_instance.point,\n", - " mass=mass,\n", - " shear=lens_model.shear,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(**lens_dict, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - " scaling_galaxies=scaling_galaxies,\n", - " )\n", - "\n", - " n_live = 200 + 100 * (n_lenses - 1)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=n_live,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "pixel_scale = 0.1\n", - "mask_radius = 6.0\n", - "mask_centre = (0.0, 0.0)\n", - "redshift_lens = 0.5\n", - "redshift_source = 1.0\n", - "source_mge_radius = 1.0\n", - "n_batch = 20\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=pixel_scale,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "main_lens_centres.json \u2014 required; determines the number of main lenses.\n", - "extra_galaxies_centres.json \u2014 optional; empty list if absent.\n", - "scaling_galaxies_centres.json \u2014 optional; empty list if absent.\n", - "\n", - "All three files contain a list of [y, x] arcsecond coordinates." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "main_lens_centres = _load_centres(dataset_path / \"main_lens_centres.json\")\n", - "extra_lens_centres = _load_centres(dataset_path / \"extra_galaxies_centres.json\")\n", - "scaling_lens_centres = _load_centres(dataset_path / \"scaling_galaxies_centres.json\")\n", - "\n", - "all_galaxy_centres = al.Grid2DIrregular(\n", - " main_lens_centres.in_list\n", - " + extra_lens_centres.in_list\n", - " + scaling_lens_centres.in_list\n", - ")\n", - "\n", - "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " centre=mask_centre,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=list(all_galaxy_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"group\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result_0 = source_lp_0(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " main_lens_centres=main_lens_centres,\n", - " extra_lens_centres=extra_lens_centres,\n", - " scaling_lens_centres=scaling_lens_centres,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - ")\n", - "\n", - "source_lp_result_1 = source_lp_1(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " source_lp_result_0=source_lp_result_0,\n", - " positions=positions,\n", - " pixel_scale=pixel_scale,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - " source_mge_radius=source_mge_radius,\n", - ")\n", - "\n", - "source_pix_result_1, dataset, adapt_images = source_pix_1(\n", - " dataset=dataset,\n", - " mask=mask,\n", - " settings_search=settings_search,\n", - " source_lp_result_1=source_lp_result_1,\n", - " over_sample_size=over_sample_size,\n", - " pixel_scale=pixel_scale,\n", - " mask_radius=mask_radius,\n", - " positions=positions,\n", - " n_batch=n_batch,\n", - ")\n", - "\n", - "source_pix_result_2, dataset, adapt_images = source_pix_2(\n", - " dataset=dataset,\n", - " mask=mask,\n", - " settings_search=settings_search,\n", - " source_lp_result_1=source_lp_result_1,\n", - " source_pix_result_1=source_pix_result_1,\n", - " over_sample_size=over_sample_size,\n", - " pixel_scale=pixel_scale,\n", - " mask_radius=mask_radius,\n", - " n_batch=n_batch,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " source_lp_result_0=source_lp_result_0,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " adapt_images=adapt_images,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " n_batch=n_batch,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " dataset=dataset,\n", - " settings_search=settings_search,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " light_result=light_result,\n", - " adapt_images=adapt_images,\n", - " positions=positions,\n", - " pixel_scale=pixel_scale,\n", - " redshift_lens=redshift_lens,\n", - " n_batch=n_batch,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM (Source, Light and Mass): Group Scale\n", + "==========================================\n", + "\n", + "This script uses the SLaM pipelines to fit a group-scale strong lens, including extra galaxies and\n", + "scaling galaxies surrounding the main lens whose light and mass are both modeled.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **Extra Galaxies and Scaling Galaxies:** This group-scale SLaM pipeline handles two distinct categories of companion galaxy, which differ in.\n", + "- **This Script:** Using a SOURCE LP PIPELINE (two searches), SOURCE PIX PIPELINE (two searches), LIGHT LP PIPELINE.\n", + "- **SOURCE LP PIPELINE 0:** Not present in `slam_start_here.py`.\n", + "- **SOURCE LP PIPELINE 1:** Equivalent to `source_lp` in `slam_start_here.py`, except lens light is fixed from `source_lp[0]`.\n", + "- **SOURCE PIX PIPELINE 1:** Equivalent to `source_pix_1` in `slam_start_here.py`, except a Hilbert image mesh is used instead.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `source_pix_1` above, except the adapt data for the Hilbert image mesh is capped at a.\n", + "- **LIGHT LP PIPELINE:** Identical to `light_lp` in `slam_start_here.py`, except extra galaxies receive a fresh free MGE.\n", + "- **MASS TOTAL PIPELINE:** Identical to `mass_total` in `slam_start_here.py`, except extra galaxies receive a new.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Galaxy Centres:** main_lens_centres.json \u2014 required; determines the number of main lenses.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "- **Group** (`group/modeling`):\n", + " How we model group-scale strong lenses in PyAutoLens, including how we include extra galaxies in\n", + " the lens model.\n", + "\n", + "__Extra Galaxies and Scaling Galaxies__\n", + "\n", + "This group-scale SLaM pipeline handles two distinct categories of companion galaxy, which differ in\n", + "how their masses are parameterized:\n", + "\n", + "**Extra Galaxies**\n", + "\n", + "Extra galaxies are a small number of nearby companions whose light and mass are modeled individually.\n", + "In each pipeline stage, they receive a free MGE light profile and an `Isothermal` mass profile whose\n", + "Einstein radius prior is bounded by a value derived from the galaxy's luminosity:\n", + "\n", + " upper_limit = min(5 * 0.5 * total_luminosity^0.6, 5.0)\n", + "\n", + "This luminosity-informed bound prevents unphysically large mass assignments while keeping the mass\n", + "free per galaxy. The extra-galaxy models are stored in `model.extra_galaxies` (a `Collection`).\n", + "\n", + "**Scaling Galaxies**\n", + "\n", + "Scaling galaxies are a larger ensemble of companions whose masses are constrained through a shared\n", + "luminosity-to-mass scaling relation rather than being individually free. They each carry a free MGE\n", + "light profile, but their Einstein radii follow:\n", + "\n", + " einstein_radius = scaling_factor * total_luminosity^scaling_relation\n", + "\n", + "where `scaling_factor` and `scaling_relation` are two shared free parameters whose priors are\n", + "`UniformPrior(0, 0.5)` and `UniformPrior(0, 2)` respectively. This reduces the number of mass\n", + "parameters considerably when many companion galaxies are present. Scaling-galaxy models are stored\n", + "in `model.scaling_galaxies`.\n", + "\n", + "The choice between these two categories is determined by which JSON file each galaxy's centre\n", + "appears in (`extra_galaxies_centres.json` vs `scaling_galaxies_centres.json`).\n", + "\n", + "**Comparison with the Galaxy-Scale SLaM Pipeline**\n", + "\n", + "The galaxy-scale extra-galaxies SLaM pipeline (`features/extra_galaxies/slam`) models each\n", + "companion galaxy with a fully free `IsothermalSph` mass \u2014 it does not use luminosity bounds or a\n", + "shared scaling relation. The group-scale pipeline introduced here is appropriate when:\n", + "\n", + "- there are more companion galaxies than can be modeled independently, or\n", + "- prior physical knowledge motivates a luminosity-to-mass scaling.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE (two searches), SOURCE PIX PIPELINE (two searches), LIGHT LP PIPELINE\n", + "and TOTAL MASS PIPELINE this SLaM modeling script fits `Imaging` data of a group-scale strong lens\n", + "where in the final model:\n", + "\n", + " - Each main lens galaxy has a free MGE bulge and a `PowerLaw` total mass.\n", + " - Each extra galaxy has a free MGE bulge and a luminosity-bounded `Isothermal` mass.\n", + " - Each scaling galaxy has a free MGE bulge and a mass set by a shared scaling relation.\n", + " - The source galaxy's light is a Delaunay `Pixelization` with `AdaptSplit` regularization.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "\n", + "def _load_centres(path):\n", + " \"\"\"Load a centres JSON file, returning an empty list if the file is absent.\"\"\"\n", + " try:\n", + " return al.Grid2DIrregular(al.from_json(file_path=path))\n", + " except FileNotFoundError:\n", + " return al.Grid2DIrregular([])\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 0__\n", + "\n", + "Not present in `slam_start_here.py`. Fits light only \u2014 no mass, no source \u2014 for every galaxy\n", + "simultaneously, giving the next search clean fixed light models to build on.\n", + "\n", + "Fits multiple main-lens galaxies (`lens_0`, `lens_1`, ...) under `galaxies`, extra galaxies\n", + "under `extra_galaxies`, and scaling galaxies under `scaling_galaxies`. `n_live` scales with\n", + "the total number of galaxies across all three categories." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_0(\n", + " dataset,\n", + " settings_search,\n", + " main_lens_centres,\n", + " extra_lens_centres,\n", + " scaling_lens_centres,\n", + " mask_radius,\n", + " redshift_lens,\n", + " n_batch=50,\n", + "):\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " # --- main lens light models (one per centre, light only) ---\n", + " lens_dict = {}\n", + " for i, centre in enumerate(main_lens_centres):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=False,\n", + " centre=(centre[0], centre[1]),\n", + " centre_sigma=0.1,\n", + " )\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=bulge, disk=None, point=None\n", + " )\n", + "\n", + " # --- extra lens galaxy light models ---\n", + " extra_light_models = []\n", + " for centre in extra_lens_centres:\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_prior_is_uniform=True,\n", + " centre=(centre[0], centre[1]),\n", + " ell_comps_prior_is_uniform=True,\n", + " )\n", + " extra_light_models.append(\n", + " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", + "\n", + " # --- scaling galaxy light models ---\n", + " scaling_light_models = []\n", + " for centre in scaling_lens_centres:\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_prior_is_uniform=True,\n", + " centre=(centre[0], centre[1]),\n", + " ell_comps_prior_is_uniform=True,\n", + " )\n", + " scaling_light_models.append(\n", + " af.Model(al.Galaxy, redshift=redshift_lens, bulge=bulge)\n", + " )\n", + "\n", + " scaling_galaxies = (\n", + " af.Collection(scaling_light_models) if scaling_light_models else None\n", + " )\n", + "\n", + " n_extra = len(extra_galaxies) if extra_galaxies is not None else 0\n", + " n_scaling = len(scaling_galaxies) if scaling_galaxies is not None else 0\n", + " n_live = 100 + 30 * len(lens_dict) + 30 * n_extra + 30 * n_scaling\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=scaling_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[0]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " n_like_max=1000000,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 1__\n", + "\n", + "Equivalent to `source_lp` in `slam_start_here.py`, except lens light is fixed from\n", + "`source_lp[0]` rather than free, and mass and source are introduced here for the first time.\n", + "\n", + "Multiple main-lens galaxies each get an `Isothermal` mass; only `lens_0` carries an\n", + "`ExternalShear`. Extra-galaxy Einstein radii are bounded by a luminosity-derived prior\n", + "(`min(5 * 0.5 * L^0.6, 5.0)`). Scaling galaxies share two free parameters,\n", + "`scaling_factor` and `scaling_relation`, so their masses follow\n", + "`einstein_radius = scaling_factor * luminosity^scaling_relation`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_1(\n", + " dataset,\n", + " settings_search,\n", + " source_lp_result_0,\n", + " positions,\n", + " pixel_scale,\n", + " redshift_lens,\n", + " redshift_source,\n", + " source_mge_radius,\n", + " n_batch=50,\n", + "):\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_main = sum(\n", + " 1 for k in vars(source_lp_result_0.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + " n_extra = (\n", + " len(list(source_lp_result_0.instance.extra_galaxies))\n", + " if source_lp_result_0.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + " n_scaling = (\n", + " len(list(source_lp_result_0.instance.scaling_galaxies))\n", + " if source_lp_result_0.instance.scaling_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " tracer = (\n", + " source_lp_result_0.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles\n", + " )\n", + "\n", + " # Source MGE centred on primary lens bulge from stage 0.\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=source_mge_radius,\n", + " total_gaussians=30,\n", + " centre=source_lp_result_0.instance.galaxies.lens_0.bulge.centre,\n", + " centre_prior_is_uniform=False,\n", + " centre_sigma=0.6,\n", + " )\n", + "\n", + " # --- main lens full models (light fixed from stage 0, mass + shear free) ---\n", + " # Only lens_0 carries the ExternalShear; one shear per group system.\n", + " lens_dict = {}\n", + " for i in range(n_main):\n", + " lp0_lens = getattr(source_lp_result_0.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = lp0_lens.bulge.centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=5.0)\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lp0_lens.bulge,\n", + " disk=lp0_lens.disk,\n", + " point=lp0_lens.point,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear) if i == 0 else None,\n", + " )\n", + "\n", + " # --- extra lens galaxy models (light fixed, mass bounded by luminosity) ---\n", + " # Tracer order: [lens_0..lens_{n_main-1}, extra_0..extra_{n_extra-1}, scaling_0..]\n", + " extra_mass_models = []\n", + " for i in range(n_extra):\n", + " lp0_extra = source_lp_result_0.instance.extra_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = lp0_extra.bulge.centre\n", + " mass.ell_comps = lp0_extra.bulge.ell_comps\n", + "\n", + " luminosity_per_gaussian_list = [\n", + " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", + " for g in tracer.galaxies[n_main + i].bulge.profile_list\n", + " ]\n", + " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", + " mass.einstein_radius = af.UniformPrior(\n", + " lower_limit=0.0,\n", + " upper_limit=min(5 * 0.5 * total_luminosity**0.6, 5.0),\n", + " )\n", + "\n", + " extra_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=lp0_extra.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", + "\n", + " # --- scaling lens galaxy models (light fixed, shared luminosity scaling relation) ---\n", + " scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + " scaling_relation = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + " scaling_mass_models = []\n", + " for i in range(n_scaling):\n", + " lp0_scaling = source_lp_result_0.instance.scaling_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = lp0_scaling.bulge.centre\n", + " mass.ell_comps = lp0_scaling.bulge.ell_comps\n", + "\n", + " luminosity_per_gaussian_list = [\n", + " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", + " for g in tracer.galaxies[n_main + n_extra + i].bulge.profile_list\n", + " ]\n", + " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", + " mass.einstein_radius = scaling_factor * total_luminosity**scaling_relation\n", + "\n", + " scaling_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=lp0_scaling.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " scaling_galaxies = (\n", + " af.Collection(scaling_mass_models) if scaling_mass_models else None\n", + " )\n", + "\n", + " source = af.Model(al.Galaxy, redshift=redshift_source, bulge=source_bulge)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=scaling_galaxies,\n", + " )\n", + "\n", + " n_extra_model = len(extra_galaxies) if extra_galaxies is not None else 0\n", + " n_scaling_model = len(scaling_galaxies) if scaling_galaxies is not None else 0\n", + " n_live = 150 + 30 * n_main + 30 * n_extra_model + 30 * n_scaling_model\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " n_like_max=200000,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Equivalent to `source_pix_1` in `slam_start_here.py`, except a Hilbert image mesh is used\n", + "instead of `RectangularAdaptDensity`, with edge-point padding of 30 pixels. The Hilbert pixel\n", + "count is set by `al.model_util.hilbert_pixels_from_pixel_scale`.\n", + "\n", + "Pixelization over-sampling is signal-adaptive: pixels above the S/N threshold use sub-size 4,\n", + "the rest sub-size 2. The re-sampled dataset and adapt_images are returned alongside the result.\n", + "Extra and scaling galaxy models are carried forward as free `model` parameters, not fixed instances." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " dataset,\n", + " mask,\n", + " settings_search,\n", + " source_lp_result_1,\n", + " over_sample_size,\n", + " pixel_scale,\n", + " mask_radius,\n", + " positions,\n", + " n_batch=20,\n", + "):\n", + " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", + " edge_pixels_total = 30\n", + " signal_to_noise_threshold = 3.0\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result_1\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Hilbert(\n", + " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", + " )\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=mask,\n", + " adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"],\n", + " )\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " over_sample_size_pixelization = np.where(\n", + " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", + " 4,\n", + " 2,\n", + " )\n", + " over_sample_size_pixelization = al.Array2D(\n", + " values=over_sample_size_pixelization, mask=mask\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=over_sample_size_pixelization,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result_1.positions_likelihood_from(\n", + " factor=2.0, positions=positions, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_lp_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lp_lens_instance = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", + " lp_lens_model = getattr(source_lp_result_1.model.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=lp_lens_model.mass,\n", + " mass_result=lp_lens_model.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lp_lens_instance.redshift,\n", + " bulge=lp_lens_instance.bulge,\n", + " disk=lp_lens_instance.disk,\n", + " point=lp_lens_instance.point,\n", + " mass=mass,\n", + " shear=lp_lens_model.shear,\n", + " )\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=af.Model(al.reg.AdaptSplit),\n", + " ),\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=source_lp_result_1.model.extra_galaxies,\n", + " scaling_galaxies=source_lp_result_1.model.scaling_galaxies,\n", + " )\n", + "\n", + " n_live = 150 + 50 * (n_lenses - 1)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + " return result, dataset, adapt_images\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `source_pix_1` above, except the adapt data for the Hilbert image mesh is capped\n", + "at a S/N threshold of 3.0 to prevent over-concentration of source pixels on the brightest peak,\n", + "and extra and scaling galaxy models are fixed as instances from `source_pix[1]`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " dataset,\n", + " mask,\n", + " settings_search,\n", + " source_lp_result_1,\n", + " source_pix_result_1,\n", + " over_sample_size,\n", + " pixel_scale,\n", + " mask_radius,\n", + " n_batch=20,\n", + "):\n", + " hilbert_pixels = al.model_util.hilbert_pixels_from_pixel_scale(pixel_scale)\n", + " edge_pixels_total = 30\n", + " signal_to_noise_threshold = 3.0\n", + " signal_to_noise_threshold_image_mesh = 3.0\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_data_snr_max = galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", + " adapt_data_snr_max[adapt_data_snr_max > signal_to_noise_threshold_image_mesh] = (\n", + " signal_to_noise_threshold_image_mesh\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Hilbert(\n", + " pixels=hilbert_pixels, weight_power=3.5, weight_floor=0.01\n", + " )\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=mask, adapt_data=adapt_data_snr_max\n", + " )\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " over_sample_size_pixelization = np.where(\n", + " galaxy_image_name_dict[\"('galaxies', 'source')\"] > signal_to_noise_threshold,\n", + " 4,\n", + " 2,\n", + " )\n", + " over_sample_size_pixelization = al.Array2D(\n", + " values=over_sample_size_pixelization, mask=mask\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size,\n", + " over_sample_size_pixelization=over_sample_size_pixelization,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lp_lens_instance = getattr(source_lp_result_1.instance.galaxies, f\"lens_{i}\")\n", + " pix1_lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lp_lens_instance.redshift,\n", + " bulge=lp_lens_instance.bulge,\n", + " disk=lp_lens_instance.disk,\n", + " point=lp_lens_instance.point,\n", + " mass=pix1_lens_instance.mass,\n", + " shear=pix1_lens_instance.shear,\n", + " )\n", + "\n", + " source = af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=hilbert_pixels, zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=af.Model(al.reg.AdaptSplit),\n", + " ),\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", + " scaling_galaxies=source_pix_result_1.instance.scaling_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " result = search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + " return result, dataset, adapt_images\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `light_lp` in `slam_start_here.py`, except extra galaxies receive a fresh free\n", + "MGE bulge (centred on the `source_pix[1]` mass centre) with mass fixed from `source_pix[1]`,\n", + "and scaling galaxies are fully fixed from `source_pix[2]`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " dataset,\n", + " settings_search,\n", + " source_lp_result_0,\n", + " source_pix_result_1,\n", + " source_pix_result_2,\n", + " adapt_images,\n", + " mask_radius,\n", + " redshift_lens,\n", + " n_batch=20,\n", + "):\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " n_lenses = sum(\n", + " 1 for k in vars(source_pix_result_1.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + " n_extra = (\n", + " len(list(source_pix_result_1.instance.extra_galaxies))\n", + " if source_pix_result_1.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " # --- main lens light models (MGE centred on stage-0 bulge centre) ---\n", + " lens_bulge_list = []\n", + " for i in range(n_lenses):\n", + " lp0_lens = getattr(source_lp_result_0.instance.galaxies, f\"lens_{i}\")\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " centre=lp0_lens.bulge.centre,\n", + " )\n", + " lens_bulge_list.append(bulge)\n", + "\n", + " # --- extra lens galaxy light models (free MGE, mass fixed from source_pix[1]) ---\n", + " extra_light_models = []\n", + " for i in range(n_extra):\n", + " pix1_extra = source_pix_result_1.instance.extra_galaxies[i]\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_prior_is_uniform=True,\n", + " centre=pix1_extra.mass.centre,\n", + " )\n", + " extra_light_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=bulge, mass=pix1_extra.mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_light_models) if extra_light_models else None\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_pix_result_2, source_is_model=False\n", + " )\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lens_instance = getattr(source_pix_result_1.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lens_instance.redshift,\n", + " bulge=lens_bulge_list[i],\n", + " disk=None,\n", + " point=None,\n", + " mass=lens_instance.mass,\n", + " shear=lens_instance.shear,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=source_pix_result_2.instance.scaling_galaxies,\n", + " )\n", + "\n", + " n_live = 300 + 100 * (n_lenses - 1)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `mass_total` in `slam_start_here.py`, except extra galaxies receive a new\n", + "luminosity-bounded `Isothermal` mass (using `light[1]` luminosities) and scaling galaxies\n", + "receive a new shared luminosity scaling relation, both paired with their fixed `light[1]` bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " dataset,\n", + " settings_search,\n", + " source_pix_result_1,\n", + " source_pix_result_2,\n", + " light_result,\n", + " adapt_images,\n", + " positions,\n", + " pixel_scale,\n", + " redshift_lens,\n", + " n_batch=20,\n", + "):\n", + " # Total mass model for each main lens galaxy.\n", + " n_lenses = sum(\n", + " 1 for k in vars(light_result.instance.galaxies) if k.startswith(\"lens_\")\n", + " )\n", + " n_extra = (\n", + " len(list(light_result.instance.extra_galaxies))\n", + " if light_result.instance.extra_galaxies is not None\n", + " else 0\n", + " )\n", + " n_scaling = (\n", + " len(list(light_result.instance.scaling_galaxies))\n", + " if light_result.instance.scaling_galaxies is not None\n", + " else 0\n", + " )\n", + "\n", + " # --- extra galaxies: fixed light, free mass ---\n", + " extra_mass_models = []\n", + " for i in range(n_extra):\n", + " light_extra = light_result.instance.extra_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = light_extra.bulge.centre\n", + " mass.ell_comps = light_extra.bulge.ell_comps\n", + "\n", + " luminosity_per_gaussian_list = [\n", + " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", + " for g in light_extra.bulge.profile_list\n", + " ]\n", + " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", + " mass.einstein_radius = af.UniformPrior(\n", + " lower_limit=0.0,\n", + " upper_limit=min(5 * 0.5 * total_luminosity**0.6, 5.0),\n", + " )\n", + "\n", + " extra_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=light_extra.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_mass_models) if extra_mass_models else None\n", + "\n", + " # --- scaling galaxies: fixed light, free shared scaling relation ---\n", + " scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + " scaling_relation = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + " scaling_mass_models = []\n", + " for i in range(n_scaling):\n", + " light_scaling = light_result.instance.scaling_galaxies[i]\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = light_scaling.bulge.centre\n", + " mass.ell_comps = light_scaling.bulge.ell_comps\n", + "\n", + " luminosity_per_gaussian_list = [\n", + " 2 * np.pi * g.sigma**2 / g.axis_ratio() * g.intensity\n", + " for g in light_scaling.bulge.profile_list\n", + " ]\n", + " total_luminosity = np.sum(luminosity_per_gaussian_list) / pixel_scale**2\n", + " mass.einstein_radius = scaling_factor * total_luminosity**scaling_relation\n", + "\n", + " scaling_mass_models.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=light_scaling.bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " scaling_galaxies = (\n", + " af.Collection(scaling_mass_models) if scaling_mass_models else None\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " light_result.positions_likelihood_from(\n", + " factor=3.0, positions=positions, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " lens_dict = {}\n", + " for i in range(n_lenses):\n", + " lens_model = getattr(source_pix_result_1.model.galaxies, f\"lens_{i}\")\n", + " light_lens_instance = getattr(light_result.instance.galaxies, f\"lens_{i}\")\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.PowerLaw),\n", + " mass_result=lens_model.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " lens_dict[f\"lens_{i}\"] = af.Model(\n", + " al.Galaxy,\n", + " redshift=lens_model.redshift,\n", + " bulge=light_lens_instance.bulge,\n", + " disk=light_lens_instance.disk,\n", + " point=light_lens_instance.point,\n", + " mass=mass,\n", + " shear=lens_model.shear,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(**lens_dict, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + " scaling_galaxies=scaling_galaxies,\n", + " )\n", + "\n", + " n_live = 200 + 100 * (n_lenses - 1)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=n_live,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "pixel_scale = 0.1\n", + "mask_radius = 6.0\n", + "mask_centre = (0.0, 0.0)\n", + "redshift_lens = 0.5\n", + "redshift_source = 1.0\n", + "source_mge_radius = 1.0\n", + "n_batch = 20\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=pixel_scale,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "main_lens_centres.json \u2014 required; determines the number of main lenses.\n", + "extra_galaxies_centres.json \u2014 optional; empty list if absent.\n", + "scaling_galaxies_centres.json \u2014 optional; empty list if absent.\n", + "\n", + "All three files contain a list of [y, x] arcsecond coordinates." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "main_lens_centres = _load_centres(dataset_path / \"main_lens_centres.json\")\n", + "extra_lens_centres = _load_centres(dataset_path / \"extra_galaxies_centres.json\")\n", + "scaling_lens_centres = _load_centres(dataset_path / \"scaling_galaxies_centres.json\")\n", + "\n", + "all_galaxy_centres = al.Grid2DIrregular(\n", + " main_lens_centres.in_list\n", + " + extra_lens_centres.in_list\n", + " + scaling_lens_centres.in_list\n", + ")\n", + "\n", + "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " centre=mask_centre,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=list(all_galaxy_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"group\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result_0 = source_lp_0(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " main_lens_centres=main_lens_centres,\n", + " extra_lens_centres=extra_lens_centres,\n", + " scaling_lens_centres=scaling_lens_centres,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + ")\n", + "\n", + "source_lp_result_1 = source_lp_1(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " source_lp_result_0=source_lp_result_0,\n", + " positions=positions,\n", + " pixel_scale=pixel_scale,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + " source_mge_radius=source_mge_radius,\n", + ")\n", + "\n", + "source_pix_result_1, dataset, adapt_images = source_pix_1(\n", + " dataset=dataset,\n", + " mask=mask,\n", + " settings_search=settings_search,\n", + " source_lp_result_1=source_lp_result_1,\n", + " over_sample_size=over_sample_size,\n", + " pixel_scale=pixel_scale,\n", + " mask_radius=mask_radius,\n", + " positions=positions,\n", + " n_batch=n_batch,\n", + ")\n", + "\n", + "source_pix_result_2, dataset, adapt_images = source_pix_2(\n", + " dataset=dataset,\n", + " mask=mask,\n", + " settings_search=settings_search,\n", + " source_lp_result_1=source_lp_result_1,\n", + " source_pix_result_1=source_pix_result_1,\n", + " over_sample_size=over_sample_size,\n", + " pixel_scale=pixel_scale,\n", + " mask_radius=mask_radius,\n", + " n_batch=n_batch,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " source_lp_result_0=source_lp_result_0,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " adapt_images=adapt_images,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " n_batch=n_batch,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " dataset=dataset,\n", + " settings_search=settings_search,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " light_result=light_result,\n", + " adapt_images=adapt_images,\n", + " positions=positions,\n", + " pixel_scale=pixel_scale,\n", + " redshift_lens=redshift_lens,\n", + " n_batch=n_batch,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/group/source_science.ipynb b/notebooks/group/source_science.ipynb index f89a318c1..e6f9b17d5 100644 --- a/notebooks/group/source_science.ipynb +++ b/notebooks/group/source_science.ipynb @@ -1,502 +1,539 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Source Science: Group\n", - "=====================\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", - "size of the source.\n", - "\n", - "This example shows how to perform these calculations for a group-scale lens, where multiple galaxies at the same\n", - "redshift contribute to the lensing potential. The key difference from the single-galaxy imaging case is that ALL\n", - "mass profiles in the group must be included when computing ray-tracing and magnification.\n", - "\n", - "__Contents__\n", - "\n", - "- **Simulated Dataset:** We load and plot the `simple` group example dataset, which is simulated imaging of a group-scale.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Source Values:** Source science calculations for real lenses are performed using the best-fitting model inferred.\n", - "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", - "- **Source Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", - "- **Impact of Extra Galaxies:** For group-scale lenses, the magnification is determined by ALL mass profiles in the group: the main.\n", - "- **Tracer:** Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to.\n", - "- **Parametric Source Models:** If your lens modeling uses a parametric source model (e.g." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulated Dataset__\n", - "\n", - "We load and plot the `simple` group example dataset, which is simulated imaging of a group-scale strong lens\n", - "that we will use to demonstrate source science calculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/group/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We apply a 7.5 arcsecond circular mask and apply it to the `Imaging` object.\n", - "\n", - "A larger mask radius is needed for group-scale lenses because the lensed images are spread over a wider area\n", - "due to the extended mass distribution.\n", - "\n", - "Source science calculations are typically performed on masked datasets to ensure only the lensed source is used\n", - "in the calculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 7.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Values__\n", - "\n", - "Source science calculations for real lenses are performed using the best-fitting model inferred from a dataset,\n", - "and this example demonstrates how to use this below.\n", - "\n", - "However, for simplicity, we demonstrate these calculations using the source model used to simulate the dataset,\n", - "which we refer to as the \"true\" source model. When analysing real strong lenses, a true underlying model is not known,\n", - "but for simulated datasets it is.\n", - "\n", - "For a group-scale lens, we must include ALL galaxies in the group that contribute to the lensing potential. This\n", - "is the main lens galaxy plus the extra galaxies in the group. Omitting any mass profile would produce incorrect\n", - "ray-tracing and therefore incorrect magnification estimates.\n", - "\n", - "The `tracer` below corresponds to the same tracer used to simulate the `simple` group dataset, and therefore\n", - "represents the true model. We also include the 2D grid of (y,x) coordinates which simulate the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", - ")\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.4,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the image of the tracer, we confirm it looks identical to the simulated dataset but does not have\n", - "CCD imaging features such as noise or blurring from a PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Flux__\n", - "\n", - "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see\n", - "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", - "\n", - "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", - "we compute the image of the light profile, and then sum the image.\n", - "\n", - "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", - "we will assume everything is in electrons per second (`e- s^-1`), which is typical for Hubble Space Telescope imaging data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Source Galaxy's Intensity {source_galaxy.bulge.intensity} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The total flux, in units of `e- s^-1` , is computed by summing the image of the light profile over all pixels.\n", - "\n", - "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", - "and high enough resolution that it captures all the light from the light profile.\n", - "\n", - "Note that we are using the source galaxy's true light profile, which corresponds to its emission in the source-plane.\n", - "For real datasets, we have to infer this via lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels\n", - "\n", - "print(f\"Total Source Flux: {total_flux} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Below, we will compare how this true source flux compares to the inferred source fluxes we compute using different\n", - "source modeling techniques (e.g. parametric and pixelized source models). Converting the flux to magnitudes or\n", - "other quantities used for tasks like SED fitting is described in the `autolens_workspace/*/guides/units/flux` example.\n", - "\n", - "__Source Magnification__\n", - "\n", - "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and\n", - "total surface brightness in the source-plane.\n", - "\n", - "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area.\n", - "We therefore explicitly mention how area folds into the calculation below.\n", - "\n", - "To ensure the magnification is stable and that we resolve all source emission in both the image-plane and source-plane\n", - "we use a very high resolution grid, higher than we used to compute the total flux above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We repeat our calculation of the source's total flux in the source-plane using this higher resolution grid, note\n", - "that we do not take the area into account, the reason for this is explained below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_source_plane_flux = np.sum(image) # in units e- s^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now need the total flux of the lensed source in the image-plane, that is how much flux we measure after\n", - "gravitational lensing.\n", - "\n", - "To calculate this, we first ray-trace the grid above from the image-plane to the source-plane using the tracer\n", - "and then pass it to the source galaxy's light profile to compute the lensed image.\n", - "\n", - "For group-scale lenses, the ray-tracing is performed through ALL mass profiles in the group. The tracer\n", - "automatically handles this multi-galaxy ray-tracing, combining the deflections from the main lens galaxy\n", - "and all extra galaxies at the same redshift." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "source_plane_grid = traced_grid_list[-1]\n", - "\n", - "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(\n", - " lensed_source_image\n", - ") # in units e- s^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now take the ratio of the total image-plane flux to source-plane flux to estimate the magnification.\n", - "\n", - "Because both fluxes were computed on grids with the same total area and area per pixel, we do not need to\n", - "explicitly account for area in this calculation. This is because the area terms cancel out when taking the ratio.\n", - "Were the grid areas different, we would need to include area terms in the calculation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Magnification (all group galaxies): {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Impact of Extra Galaxies__\n", - "\n", - "For group-scale lenses, the magnification is determined by ALL mass profiles in the group: the main lens galaxy\n", - "plus the extra galaxies. Omitting the extra galaxy masses would give an incorrect magnification estimate.\n", - "\n", - "Below, we demonstrate this by computing the magnification using only the main lens galaxy and comparing it\n", - "to the full group magnification computed above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_main_only = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "traced_grid_list_main_only = tracer_main_only.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "source_plane_grid_main_only = traced_grid_list_main_only[-1]\n", - "\n", - "lensed_source_image_main_only = source_galaxy.bulge.image_2d_from(\n", - " grid=source_plane_grid_main_only\n", - ")\n", - "\n", - "total_image_plane_flux_main_only = np.sum(lensed_source_image_main_only)\n", - "\n", - "source_magnification_main_only = (\n", - " total_image_plane_flux_main_only / total_source_plane_flux\n", - ")\n", - "\n", - "print(f\"Source Magnification (main lens only): {source_magnification_main_only}\")\n", - "print(\n", - " f\"Magnification difference when omitting extra galaxies: \"\n", - " f\"{source_magnification - source_magnification_main_only:.4f}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The magnification values differ, showing that omitting the extra galaxies from the group leads to an incorrect\n", - "estimate. For accurate source science with group-scale lenses, all contributing mass profiles must be included.\n", - "\n", - "__Tracer__\n", - "\n", - "Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to compute\n", - "source science calculations for real datasets.\n", - "\n", - "The code below shows how using a tracer, composed of the full group of lens galaxies and source galaxies, we can\n", - "compute the source flux and magnification. It reproduces the calculations above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "image_plane_grid = traced_grid_list[0]\n", - "source_plane_grid = traced_grid_list[-1]\n", - "\n", - "lensed_source_image = tracer.planes[-1].image_2d_from(grid=source_plane_grid)\n", - "source_plane_image = tracer.planes[-1].image_2d_from(grid=image_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(lensed_source_image)\n", - "total_source_plane_flux = np.sum(source_plane_image)\n", - "\n", - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Plane Total Flux via Tracer: {total_source_plane_flux} e- s^-1\")\n", - "print(f\"Source Magnification via Tracer: {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parametric Source Models__\n", - "\n", - "If your lens modeling uses a parametric source model (e.g. Sersic, Multi Gaussian Expansion), the only object\n", - "you need to perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", - "\n", - "Alternatively, as done above, you can manually set up a tracer using the lens and source galaxies inferred\n", - "by lens modeling.\n", - "\n", - "For group-scale lenses, ensure that the tracer includes all group member galaxies with their mass profiles,\n", - "as each contributes to the ray-tracing and therefore to the magnification of the source.\n", - "\n", - "Therefore, you may now wish to go to your results, extract the `max_log_likelihood_tracer`, and use it to compute\n", - "the source flux and magnification as shown above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Source Science: Group\n", + "=====================\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", + "size of the source.\n", + "\n", + "This example shows how to perform these calculations for a group-scale lens, where multiple galaxies at the same\n", + "redshift contribute to the lensing potential. The key difference from the single-galaxy imaging case is that ALL\n", + "mass profiles in the group must be included when computing ray-tracing and magnification.\n", + "\n", + "__Contents__\n", + "\n", + "- **Simulated Dataset:** We load and plot the `simple` group example dataset, which is simulated imaging of a group-scale.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Source Values:** Source science calculations for real lenses are performed using the best-fitting model inferred.\n", + "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", + "- **Source Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", + "- **Impact of Extra Galaxies:** For group-scale lenses, the magnification is determined by ALL mass profiles in the group: the main.\n", + "- **Tracer:** Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to.\n", + "- **Parametric Source Models:** If your lens modeling uses a parametric source model (e.g." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulated Dataset__\n", + "\n", + "We load and plot the `simple` group example dataset, which is simulated imaging of a group-scale strong lens\n", + "that we will use to demonstrate source science calculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"group\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/group/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We apply a 7.5 arcsecond circular mask and apply it to the `Imaging` object.\n", + "\n", + "A larger mask radius is needed for group-scale lenses because the lensed images are spread over a wider area\n", + "due to the extended mass distribution.\n", + "\n", + "Source science calculations are typically performed on masked datasets to ensure only the lensed source is used\n", + "in the calculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 7.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Values__\n", + "\n", + "Source science calculations for real lenses are performed using the best-fitting model inferred from a dataset,\n", + "and this example demonstrates how to use this below.\n", + "\n", + "However, for simplicity, we demonstrate these calculations using the source model used to simulate the dataset,\n", + "which we refer to as the \"true\" source model. When analysing real strong lenses, a true underlying model is not known,\n", + "but for simulated datasets it is.\n", + "\n", + "For a group-scale lens, we must include ALL galaxies in the group that contribute to the lensing potential. This\n", + "is the main lens galaxy plus the extra galaxies in the group. Omitting any mass profile would produce incorrect\n", + "ray-tracing and therefore incorrect magnification estimates.\n", + "\n", + "The `tracer` below corresponds to the same tracer used to simulate the `simple` group dataset, and therefore\n", + "represents the true model. We also include the 2D grid of (y,x) coordinates which simulate the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=4.0),\n", + ")\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.8),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=(-4.4, -5.0), einstein_radius=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.4,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the image of the tracer, we confirm it looks identical to the simulated dataset but does not have\n", + "CCD imaging features such as noise or blurring from a PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Flux__\n", + "\n", + "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see\n", + "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", + "\n", + "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", + "we compute the image of the light profile, and then sum the image.\n", + "\n", + "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", + "we will assume everything is in electrons per second (`e- s^-1`), which is typical for Hubble Space Telescope imaging data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Source Galaxy's Intensity {source_galaxy.bulge.intensity} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The total flux, in units of `e- s^-1` , is computed by summing the image of the light profile over all pixels.\n", + "\n", + "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", + "and high enough resolution that it captures all the light from the light profile.\n", + "\n", + "Note that we are using the source galaxy's true light profile, which corresponds to its emission in the source-plane.\n", + "For real datasets, we have to infer this via lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels\n", + "\n", + "print(f\"Total Source Flux: {total_flux} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below, we will compare how this true source flux compares to the inferred source fluxes we compute using different\n", + "source modeling techniques (e.g. parametric and pixelized source models). Converting the flux to magnitudes or\n", + "other quantities used for tasks like SED fitting is described in the `autolens_workspace/*/guides/units/flux` example.\n", + "\n", + "__Source Magnification__\n", + "\n", + "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and\n", + "total surface brightness in the source-plane.\n", + "\n", + "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area.\n", + "We therefore explicitly mention how area folds into the calculation below.\n", + "\n", + "To ensure the magnification is stable and that we resolve all source emission in both the image-plane and source-plane\n", + "we use a very high resolution grid, higher than we used to compute the total flux above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We repeat our calculation of the source's total flux in the source-plane using this higher resolution grid, note\n", + "that we do not take the area into account, the reason for this is explained below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_source_plane_flux = np.sum(image) # in units e- s^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now need the total flux of the lensed source in the image-plane, that is how much flux we measure after\n", + "gravitational lensing.\n", + "\n", + "To calculate this, we first ray-trace the grid above from the image-plane to the source-plane using the tracer\n", + "and then pass it to the source galaxy's light profile to compute the lensed image.\n", + "\n", + "For group-scale lenses, the ray-tracing is performed through ALL mass profiles in the group. The tracer\n", + "automatically handles this multi-galaxy ray-tracing, combining the deflections from the main lens galaxy\n", + "and all extra galaxies at the same redshift." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "source_plane_grid = traced_grid_list[-1]\n", + "\n", + "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(\n", + " lensed_source_image\n", + ") # in units e- s^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now take the ratio of the total image-plane flux to source-plane flux to estimate the magnification.\n", + "\n", + "Because both fluxes were computed on grids with the same total area and area per pixel, we do not need to\n", + "explicitly account for area in this calculation. This is because the area terms cancel out when taking the ratio.\n", + "Were the grid areas different, we would need to include area terms in the calculation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Magnification (all group galaxies): {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Impact of Extra Galaxies__\n", + "\n", + "For group-scale lenses, the magnification is determined by ALL mass profiles in the group: the main lens galaxy\n", + "plus the extra galaxies. Omitting the extra galaxy masses would give an incorrect magnification estimate.\n", + "\n", + "Below, we demonstrate this by computing the magnification using only the main lens galaxy and comparing it\n", + "to the full group magnification computed above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_main_only = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "traced_grid_list_main_only = tracer_main_only.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "source_plane_grid_main_only = traced_grid_list_main_only[-1]\n", + "\n", + "lensed_source_image_main_only = source_galaxy.bulge.image_2d_from(\n", + " grid=source_plane_grid_main_only\n", + ")\n", + "\n", + "total_image_plane_flux_main_only = np.sum(lensed_source_image_main_only)\n", + "\n", + "source_magnification_main_only = (\n", + " total_image_plane_flux_main_only / total_source_plane_flux\n", + ")\n", + "\n", + "print(f\"Source Magnification (main lens only): {source_magnification_main_only}\")\n", + "print(\n", + " f\"Magnification difference when omitting extra galaxies: \"\n", + " f\"{source_magnification - source_magnification_main_only:.4f}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The magnification values differ, showing that omitting the extra galaxies from the group leads to an incorrect\n", + "estimate. For accurate source science with group-scale lenses, all contributing mass profiles must be included.\n", + "\n", + "__Tracer__\n", + "\n", + "Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to compute\n", + "source science calculations for real datasets.\n", + "\n", + "The code below shows how using a tracer, composed of the full group of lens galaxies and source galaxies, we can\n", + "compute the source flux and magnification. It reproduces the calculations above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "image_plane_grid = traced_grid_list[0]\n", + "source_plane_grid = traced_grid_list[-1]\n", + "\n", + "lensed_source_image = tracer.planes[-1].image_2d_from(grid=source_plane_grid)\n", + "source_plane_image = tracer.planes[-1].image_2d_from(grid=image_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(lensed_source_image)\n", + "total_source_plane_flux = np.sum(source_plane_image)\n", + "\n", + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Plane Total Flux via Tracer: {total_source_plane_flux} e- s^-1\")\n", + "print(f\"Source Magnification via Tracer: {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parametric Source Models__\n", + "\n", + "If your lens modeling uses a parametric source model (e.g. Sersic, Multi Gaussian Expansion), the only object\n", + "you need to perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", + "\n", + "Alternatively, as done above, you can manually set up a tracer using the lens and source galaxies inferred\n", + "by lens modeling.\n", + "\n", + "For group-scale lenses, ensure that the tracer includes all group member galaxies with their mass profiles,\n", + "as each contributes to the ray-tracing and therefore to the magnification of the source.\n", + "\n", + "Therefore, you may now wish to go to your results, extract the `max_log_likelihood_tracer`, and use it to compute\n", + "the source flux and magnification as shown above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/advanced/add_a_profile.ipynb b/notebooks/guides/advanced/add_a_profile.ipynb index af1eb0ecf..940f7c8e9 100644 --- a/notebooks/guides/advanced/add_a_profile.ipynb +++ b/notebooks/guides/advanced/add_a_profile.ipynb @@ -1,826 +1,863 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Misc: Add A Profile\n", - "===================\n", - "\n", - "**PyAutoLens** supports a wide range of mass and light profiles for modelling the\n", - "lens and source galaxies in strong gravitational lensing systems.\n", - "\n", - "For some science cases, it may be necessary to define custom profiles. This could\n", - "involve implementing a new profile that is not currently supported by\n", - "**PyAutoLens**, or introducing a new parameterization of an existing profile. Both\n", - "of these possibilities are covered in this example tutorial.\n", - "\n", - "We begin by explaining how to add a new _mass profile_, as this introduces the core\n", - "concepts required for defining custom profiles in **PyAutoLens**. These concepts\n", - "are then applied to show how custom _light profiles_ can be implemented.\n", - "\n", - "__Contents__\n", - "\n", - "- **Source Code:** This example includes direct links to the source code of the classes used to define mass and light.\n", - "- **Example Mass Profile:** The mass profiles available in **PyAutoLens** are located in its parent package, **PyAutoGalaxy**.\n", - "- **Inheritance Structure:** Let us next consider the inheritance structure of the ``Isothermal`` profile, defined by the class.\n", - "- **Data Structure Decorators:** Different grids can be input into each mass profile function (e.g.\n", - "- **Transform Decorator:** Overview of transform decorator for this example.\n", - "- **Lens Modeling:** **PyAutoLens** assumes that all input parameters of a mass profile (for example, those listed in.\n", - "- **Lens Modeling Configs:** In most **PyAutoLens** examples, you will notice we compose models without manually specifying.\n", - "- **Deflections:** We are therefore ready to implement a mass profile, and the best place to start is the.\n", - "- **Spherical Template:** radial_grid removal of ell_comps.\n", - "- **Physical Profiles:** Show how to wrap existing profiles with physical units?\n", - "- **Light Profiles:** Pretty much the same but need to add text.\n", - "\n", - "__Source Code__\n", - "\n", - "This example includes direct links to the source code of the classes used to define\n", - "mass and light profiles, allowing you to see exactly how they are implemented.\n", - "\n", - "The tutorial is fully standalone and, by the end, should enable you to implement a\n", - "custom profile without needing to dive deeply into the **PyAutoLens** codebase.\n", - "That said, we still recommend exploring the source code to better understand how\n", - "everything fits together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Example Mass Profile__\n", - "\n", - "The mass profiles available in **PyAutoLens** are located in its parent package, **PyAutoGalaxy**: \n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy\n", - "\n", - "All light and mass profiles are found in the following python package:\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/tree/main/autogalaxy/profiles\n", - "\n", - "Mass profiles are in the following package:\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/tree/main/autogalaxy/profiles/mass\n", - "\n", - "Lets look at an example mass profile. We'll use the `Isothermal` profile, which is located in the `total` package\n", - "because it represents a total (stars + dark matter) mass distribution:\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/tree/main/autogalaxy/profiles/mass/total\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/profiles/mass/total/isothermal.py\n", - "\n", - "For simplicity, a shortened version of the `Isothermal` profile is shown below. \n", - "\n", - "This has docstrings updated to focus on the key aspects of implementing a new profiles and simplifies the \n", - "inheritance structure of the profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from typing import Tuple\n", - "\n", - "import autoarray as aa\n", - "\n", - "from autogalaxy.profiles.mass.total.isothermal import psi_from\n", - "from autogalaxy.profiles.mass.abstract.abstract import MassProfile\n", - "\n", - "\n", - "class Isothermal(MassProfile):\n", - " def __init__(\n", - " self,\n", - " centre: Tuple[float, float] = (0.0, 0.0),\n", - " ell_comps: Tuple[float, float] = (0.0, 0.0),\n", - " einstein_radius: float = 1.0,\n", - " ):\n", - " \"\"\"\n", - " Represents an elliptical isothermal density distribution, which is equivalent to the elliptical power-law\n", - " density distribution for the value slope = 2.0.\n", - "\n", - " Parameters\n", - " ----------\n", - " centre\n", - " The (y,x) arc-second coordinates of the profile centre.\n", - " ell_comps\n", - " The first and second ellipticity components of the elliptical coordinate system.\n", - " einstein_radius\n", - " The arc-second Einstein radius.\n", - " \"\"\"\n", - "\n", - " super().__init__(\n", - " centre=centre,\n", - " ell_comps=ell_comps,\n", - " )\n", - "\n", - " self.einstein_radius = einstein_radius\n", - " self.slope = 2.0\n", - "\n", - " @aa.grid_dec.to_vector_yx\n", - " @aa.grid_dec.transform(rotate_back=True)\n", - " def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", - " \"\"\"\n", - " Calculate the deflection angles on a grid of (y,x) arc-second coordinates.\n", - "\n", - " The input grid of (y,x) coordinates are transformed to a coordinate system centred on the profile centre with\n", - " and rotated based on the position angle defined from its `ell_comps` (this is described fully below).\n", - "\n", - " Because this method computes deflection components using the rotated grid coordinates (i.e. the\n", - " components are expressed in the profile's frame), ``rotate_back=True`` is set so the decorator\n", - " automatically rotates them back to the observer frame.\n", - "\n", - " The numerical backend can be selected via the ``xp`` argument, allowing this\n", - " method to be used with both NumPy and JAX (e.g. inside ``jax.jit``-compiled\n", - " code). This is described fully later in this example.\n", - "\n", - " Parameters\n", - " ----------\n", - " grid\n", - " The grid of (y,x) arc-second coordinates the deflection angles are computed on.\n", - " xp\n", - " The numerical backend to use, either `numpy` or `jax.numpy`.\n", - " \"\"\"\n", - "\n", - " factor = (\n", - " 2.0\n", - " * self.einstein_radius_rescaled(xp)\n", - " * self.axis_ratio(xp)\n", - " / xp.sqrt(1 - self.axis_ratio(xp) ** 2)\n", - " )\n", - "\n", - " psi = psi_from(\n", - " grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=0.0, xp=xp\n", - " )\n", - "\n", - " deflection_y = xp.arctanh(\n", - " xp.divide(\n", - " xp.multiply(xp.sqrt(1 - self.axis_ratio(xp) ** 2), grid.array[:, 0]),\n", - " psi,\n", - " )\n", - " )\n", - " deflection_x = xp.arctan(\n", - " xp.divide(\n", - " xp.multiply(xp.sqrt(1 - self.axis_ratio(xp) ** 2), grid.array[:, 1]),\n", - " psi,\n", - " )\n", - " )\n", - " return xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__JAX, Numpy and xp__\n", - "\n", - "Throughout this tutorial, and in the profile functions above, you will see functions and methods that accept an\n", - "argument called ``xp``. The default input value above is ``xp=np``, which sets it to the standard NumPy library,\n", - "imported using the statement ``import numpy as np``. This means all arithmetic operations use NumPy in a way\n", - "you are likely familiar with.\n", - "\n", - "However, to enable mass profile calculations to run on a GPU, a library called JAX is used which mirrors the\n", - "NumPy API. Conventionally, this is imported as ``jnp`` using the statement ``import jax.numpy as jnp``. When the\n", - "source code is running in JAX mode, the input ``xp`` to the functions above will be the ``jnp`` library instead\n", - "of ``np``. This is why mass profiles support the ``xp`` argument and API: they need to be able to run using\n", - "either NumPy or JAX.\n", - "\n", - "The PyAutoLens source code runs in pure NumPy by default, where ``xp`` is always set to ``np``. This only\n", - "changes if you manually call a function passing ``xp=jnp``, or when certain high-level objects, such as the\n", - "``Analysis`` class, are used. These objects automatically set ``xp=jnp`` when a likelihood is evaluated for\n", - "lens modelling.\n", - "\n", - "Your final mass profile should therefore use the ``xp`` API throughout, ensuring compatibility with both NumPy\n", - "and JAX and allowing it to work seamlessly with the PyAutoLens source code. You may find it easier to first\n", - "write your functions in pure NumPy (which you are likely most familiar with), and then convert them to use the\n", - "``xp`` API and test them with JAX afterwards. While using ``xp`` makes the API slightly more verbose, it is a\n", - "small price to pay for the significant speed-ups available when running JAX on a GPU.\n", - "\n", - "__Inheritance Structure__\n", - "\n", - "Let us next consider the inheritance structure of the ``Isothermal`` profile,\n", - "defined by the class declaration::\n", - "\n", - " class Isothermal(MassProfile):\n", - "\n", - "In Python, inheritance means that a class can reuse and extend the behaviour of\n", - "another class. By inheriting from ``MassProfile``, the ``Isothermal`` profile\n", - "automatically has access to all methods and attributes defined in\n", - "``MassProfile``. This allows ``Isothermal`` to make use of shared functionality\n", - "(such as common calculations and interfaces) without reimplementing it.\n", - "\n", - "The key mechanism used to enable this inheritance is the ``super`` function. For\n", - "example, in the ``Isothermal`` initializer we see::\n", - "\n", - " super().__init__(\n", - " centre=centre,\n", - " ell_comps=ell_comps,\n", - " )\n", - "\n", - "This line calls the ``__init__`` method of the parent ``MassProfile`` class,\n", - "ensuring that all required base-class setup is performed before adding any\n", - "``Isothermal``-specific behaviour.\n", - "\n", - "It is important to emphasize that you do not need to fully understand the full\n", - "inheritance structure of the **PyAutoLens** profiles or the layout of the source\n", - "code to define your own custom profiles. This discussion is included simply to\n", - "highlight that all calculations involving mass and light profiles are built on\n", - "a set of abstract base classes, which your custom profiles will automatically\n", - "inherit from.\n", - "\n", - "__Inheritance (MassProfile)__\n", - "\n", - "All mass profiles in **PyAutoLens** inherit from the `MassProfile` abstract base class, which is located in the\n", - "following package:\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/profiles/mass/abstract/abstract.py\n", - "\n", - "This contains functions which are useful for any mass profile, which your custom mass profile will inherit.\n", - "\n", - "For example, it includes the function `mass_angular_within_circle_from`, which computes the mass of the profile\n", - "within an input circle of radius `radius`.\n", - "\n", - "__Inheritance (GeometryProfile)__\n", - "\n", - "The `MassProfile` class inherits from the `GeometryProfile` abstract base class, which is located here:\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/profiles/geometry_profiles.py\n", - "\n", - "This contains functions which are useful for any elliptical (and spherical) profile, which your custom \n", - "mass profile will again inherit (e.g. `radial_grid_from`).\n", - "\n", - "They typically perform coordinate transforms between the profile's elliptical (spherical) coordinate system and the \n", - "input 2D grid of (y,x) coordinates. For example, the function `transformed_to_reference_frame_grid_from` transforms\n", - "the (y,x) coordinates to the profile's elliptical coordinate system.\n", - "\n", - "The convention of this calculation is key for ensuring you implement your custom profile correctly. We illustrate\n", - "it fully below.\n", - "\n", - "__Inheritance (OperateDeflections)__\n", - "\n", - "Nearly all lensing quantities (e.g. `convergence`, `potential`, `magnification`) can be derived from the deflection\n", - "angles of a mass profile. \n", - "\n", - "Mass profiles therefore also inherit from the `OperateDeflections` abstract base class, which contains numerous \n", - "functions for computing these lensing quantities from the deflection angles. This is located here:\n", - "\n", - "https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/operate/deflections.py\n", - "\n", - "This means that once you've implemented a deflections angles calculation for your mass profile, you can compute all\n", - "lensing quantities from it without having to write any additional code!\n", - "\n", - "__Whats Going On With Those Decorators?__\n", - "\n", - "A decorator in Python is a special syntax used to modify or extend the behaviour of a function or method without \n", - "changing its implementation. It is denoted by the `@decorator_name` syntax placed above the function definition,\n", - "with two decorators shown in the example above: `@aa.grid_dec.to_vector_yx` and `@aa.grid_dec.transform`. Lets\n", - "now consider what these do.\n", - "\n", - "__Data Structure Decorators__\n", - "\n", - "Different grids can be input into each mass profile function (e.g. `Grid2D`, `Grid2DIrregular`). Depending on the input \n", - "grid, this changes the structure of the output array. \n", - "\n", - "For example, if a `Grid2D` is input, which is defined on a uniform grid of 2D coordinates, the output deflection angles\n", - "are also defined on a uniform grid and are returned as a `VectorYX2D` object. If a `Grid2DIrregular` is input, \n", - "which is defined on an irregular grid of 2D coordinates, the output deflection angles are also defined on an irregular \n", - "grid and are returned as a `VectorYX2DIrregular` object.\n", - "\n", - "The `@aa.grid_dec.to_vector_yx` decorator handles this structure conversion for vector quantities, such that the output\n", - "vector structure matches the input grid structure.\n", - "\n", - "The function `deflections_yx_2d_from` returns 2D vectors, but other mass profile methods, like `convergence_2d_from` and\n", - "`potential_2d_from`, return scalar quantities. These methods use the `@aa.grid_dec.to_array` decorator, which behaves\n", - "analogously to the `@aa.grid_dec.to_vector_yx` decorator but for scalar quantities (e.g. for an input `Grid2D`, the output\n", - "is an `Array2D` object, for an input `Grid2DIrregular`, the output is an `ArrayIrregular` object).\n", - "\n", - "For your custom mass profile, you basically just need to copy and paste these decorators above your mass profile \n", - "functions and not worry about them any further.\n", - " \n", - "__Transform Decorator__\n", - " \n", - "The second decorator is the `@aa.grid_dec.transform` decorator. This one we will have a closer look at, as it\n", - "will influence how you implement your mass profile functions.\n", - "\n", - "The `transform` decorator is used to transform the input grid of (y,x) coordinates to the mass profile's elliptical \n", - "coordinate system. It does this by calling the function `transform_grid_2d_to_reference_frame`, which I have\n", - "provided below for convenience:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def transform_grid_2d_to_reference_frame(\n", - " grid_2d: np.ndarray, centre: Tuple[float, float], angle: float, xp=np\n", - ") -> np.ndarray:\n", - " \"\"\"\n", - " Transform a 2D grid of (y,x) coordinates to a new reference frame.\n", - "\n", - " This transformation includes:\n", - "\n", - " 1) A translation to a new (y,x) centre value, by subtracting the centre from every coordinate on the grid.\n", - " 2) A rotation of the grid around this new centre, which is performed clockwise from an input angle.\n", - "\n", - " Parameters\n", - " ----------\n", - " grid\n", - " The 2d grid of (y, x) coordinates which are transformed to a new reference frame.\n", - " \"\"\"\n", - "\n", - " shifted_grid_2d = grid_2d - xp.array(centre)\n", - "\n", - " radius = xp.sqrt(xp.sum(xp.square(shifted_grid_2d), axis=1))\n", - " theta_coordinate_to_profile = xp.arctan2(\n", - " shifted_grid_2d[:, 0], shifted_grid_2d[:, 1]\n", - " ) - xp.radians(angle)\n", - "\n", - " return xp.vstack(\n", - " [\n", - " radius * xp.sin(theta_coordinate_to_profile),\n", - " radius * xp.cos(theta_coordinate_to_profile),\n", - " ]\n", - " ).T\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A simple example of this function is shifting a grid to a mass profile's centre (which simply subtracts the centre\n", - "coordinates from every coordinate on the grid):" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2DIrregular(values=[(0.0, 0.0), (1.0, 1.0), (2.0, 2.0)])\n", - "\n", - "mass_profile_centre = (0.5, 0.5)\n", - "\n", - "transformed_grid = transform_grid_2d_to_reference_frame(\n", - " grid_2d=grid, centre=mass_profile_centre, angle=0.0\n", - ")\n", - "\n", - "print(f\"Grid Coordinates Before: {grid}\")\n", - "print(f\"Grid Coordinates After: {transformed_grid}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `angle` input is the rotation angle of the mass profile's ellipse counter-clockwise from the positive x-axis.\n", - "\n", - "It is computed from the `ell_comps` of the mass profile, which are the elliptical components of the mass profile's\n", - "ellipse." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def axis_ratio_and_angle_from(\n", - " ell_comps: Tuple[float, float], xp=np\n", - ") -> Tuple[float, float]:\n", - " \"\"\"\n", - " Returns the axis-ratio and position angle in degrees (-45 < angle < 135.0) from input elliptical components e1\n", - " and e2 of a light or mass profile.\n", - "\n", - " The elliptical components of a light or mass profile are given by:\n", - "\n", - " elliptical_component_y = ell_comps[0] = (1-axis_ratio)/(1+axis_ratio) * sin(2 * angle)\n", - " elliptical_component_x = ell_comps[1] = (1-axis_ratio)/(1+axis_ratio) * cos(2 * angle)\n", - "\n", - " The axis-ratio and angle are therefore given by:\n", - "\n", - " axis_ratio = (1 - fac) / (1 + fac)\n", - " angle = 0.5 * arctan(ell_comps[0] / ell_comps[1])\n", - "\n", - " where `fac = sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2).\n", - "\n", - " This function returns the axis-ratio and angle in degrees.\n", - "\n", - " An additional check is performed which requires the angle is between -45 and 135 degrees. This ensures that\n", - " for certain values of `ell_comps` the angle does not jump from one boundary to another (e.g. without this check\n", - " certain values of `ell_comps` return -1.0 degrees and others 179.0 degrees). This ensures that when error\n", - " estimates are computed from samples of a lens model via marginalization, the calculation is not biased by the\n", - " angle jumping between these two values.\n", - "\n", - " Parameters\n", - " ----------\n", - " ell_comps\n", - " The elliptical components of the light or mass profile which are converted to an angle.\n", - " \"\"\"\n", - " angle = xp.arctan2(ell_comps[0], ell_comps[1]) / 2\n", - " angle *= 180.0 / xp.pi\n", - "\n", - " angle = xp.where(angle < -45, angle + 180, angle)\n", - "\n", - " fac = xp.sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2)\n", - " if xp.__name__.startswith(\"jax\"):\n", - " import jax\n", - "\n", - " fac = jax.lax.min(fac, 0.999)\n", - " else: # NumPy\n", - " fac = np.minimum(fac, 0.999)\n", - "\n", - " axis_ratio = (1 - fac) / (1 + fac)\n", - " return axis_ratio, angle\n", - "\n", - "\n", - "mass_profile_ell_comps = (0.5, 0.5)\n", - "mass_profile_angle = axis_ratio_and_angle_from(ell_comps=mass_profile_ell_comps)[1]\n", - "\n", - "print(f\"\\nMass Profile Angle (degrees) {mass_profile_angle}\")\n", - "\n", - "transformed_grid = transform_grid_2d_to_reference_frame(\n", - " grid_2d=grid, centre=mass_profile_centre, angle=mass_profile_angle\n", - ")\n", - "\n", - "print(f\"Grid Coordinates Before: {grid}\")\n", - "print(f\"Grid Coordinates After: {transformed_grid}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `@aa.grid_dec.transform` packages all the above calculations up and uses the mass profile `centre` and `ell_comps` \n", - "to perform them before your function is called.\n", - "\n", - "The class below demonstrates this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autogalaxy.profiles.geometry_profiles import EllProfile\n", - "\n", - "\n", - "class ExampleMass(EllProfile):\n", - " def __init__(\n", - " self,\n", - " centre: Tuple[float, float] = (0.0, 0.0),\n", - " ell_comps: Tuple[float, float] = (0.0, 0.0),\n", - " ):\n", - " super().__init__(centre=centre, ell_comps=ell_comps)\n", - "\n", - " @aa.grid_dec.transform\n", - " def deflections_yx_2d_from(self, grid: aa.Grid2D, xp=np, **kwargs):\n", - " print(\n", - " f\"\\n Grid In Deflections After Transform \"\n", - " f\"Which is Same As Transformed Grid Above: {grid}\"\n", - " )\n", - "\n", - "\n", - "mass = ExampleMass(centre=(0.5, 0.5), ell_comps=(0.5, 0.5))\n", - "mass.deflections_yx_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Do I Rotate Back?__\n", - "\n", - "Whether the returned result needs back-rotation depends on what kind of quantity the function returns\n", - "and which coordinate frame its components are expressed in.\n", - "\n", - "**Scalars** (convergence, potential): frame-invariant \u2014 no back-rotation needed.\n", - "\n", - "**Vectors** (deflection angles): it depends on how the function computes them. The ``@aa.grid_dec.transform``\n", - "decorator transforms the input grid into the profile's rotated reference frame. If the function then computes\n", - "deflection components ``(alpha_y, alpha_x)`` using that rotated grid \u2014 as the ``Isothermal`` example above does \u2014\n", - "those components are expressed in the profile's frame. Since the ray-tracing code expects observer-frame\n", - "components, they must be rotated back. Setting ``rotate_back=True`` handles this automatically by calling\n", - "``self.rotated_grid_from_reference_frame_from(...)`` on the result.\n", - "\n", - "However, back-rotation is **not** always needed for vectors. If a function computes a scalar quantity in the\n", - "profile frame (e.g. a radial deflection magnitude) and then reconstructs the Cartesian vector using\n", - "observer-frame geometry, the result is already in the observer frame. In that case, ``rotate_back`` should\n", - "remain ``False``.\n", - "\n", - "The rule is: look at which coordinate basis the returned components are expressed in. If they use the rotated\n", - "basis, set ``rotate_back=True``. If they are already observer-frame components, leave it as ``False``.\n", - "\n", - "**Spin-2 quantities** (shear): these transform under a coordinate rotation by twice the profile angle (the\n", - "spin-2 transformation law). This is not handled by ``rotate_back`` \u2014 shear methods must apply the ``2 * angle``\n", - "rotation manually via ``self.rotated_grid_from_reference_frame_from(grid=..., angle=self.angle(xp) * 2)``.\n", - "\n", - "__Lens Modeling__\n", - "\n", - "**PyAutoLens** assumes that all input parameters of a mass profile (for example,\n", - "those listed in its ``__init__`` constructor) are free parameters that can be\n", - "fitted during lens modelling using a non-linear search.\n", - "\n", - "If a parameter in the ``__init__`` constructor is a float (e.g. the\n", - "``einstein_radius`` of the ``Isothermal`` profile), it is treated as a single\n", - "free parameter. If a parameter is a tuple of floats (e.g. the ``centre`` of the\n", - "``Isothermal`` profile), each element of the tuple is treated as a separate free\n", - "parameter.\n", - "\n", - "We demonstrate this behaviour using a simple example mass profile defined\n", - "below. We compose it as a model using ``af.Model`` and print its ``info``,\n", - "which summarizes its free parameters and shows that no priors have yet been\n", - "assigned." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import autofit as af\n", - "from typing import Tuple\n", - "\n", - "\n", - "class LensModelExample:\n", - " def __init__(\n", - " self, # <-- **PyAutoLens** assumes these input parameters are free.\n", - " centre: Tuple[float, float] = (\n", - " 0.0,\n", - " 0.0,\n", - " ), # <-- Two free parameters because this is a tuple.\n", - " ell_comps: Tuple[float, float] = (\n", - " 0.0,\n", - " 0.0,\n", - " ), # <-- Also two free parameters.\n", - " einstein_radius: float = 1.0, # <-- A single free parameter.\n", - " your_parameter_here: float = 2.0, # <-- Add any custom parameters you need.\n", - " ):\n", - " pass\n", - "\n", - "\n", - "lens_model_example = af.Model(LensModelExample)\n", - "\n", - "print(lens_model_example.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For this example model, we can manually assign priors to its parameters as shown\n", - "below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_model_example.centre.centre_0 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0)\n", - "lens_model_example.centre.centre_1 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0)\n", - "lens_model_example.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - "lens_model_example.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - "lens_model_example.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=3.0)\n", - "lens_model_example.your_parameter_here = af.UniformPrior(\n", - " lower_limit=0.0, upper_limit=5.0\n", - ")\n", - "\n", - "print(lens_model_example.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The exact same API applies to the ``Isothermal`` class defined above, which has\n", - "three ``__init__`` parameters: ``centre``, ``ell_comps``, and\n", - "``einstein_radius``.\n", - "\n", - "When composed as a model, the ``Isothermal`` profile therefore has five free\n", - "parameters in total." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = af.Model(Isothermal)\n", - "\n", - "print(mass.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As before, we must manually assign priors to these parameters for lens modelling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass.centre.centre_0 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0)\n", - "mass.centre.centre_1 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0)\n", - "mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - "mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - "mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=3.0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Provided you follow this same convention when defining your own mass profile,\n", - "it will fully support lens modelling in **PyAutoLens** without requiring any\n", - "additional code.\n", - "\n", - "__Lens Modeling Configs__\n", - "\n", - "In most **PyAutoLens** examples, you will notice we compose models without manually specifying priors. This is because\n", - "**PyAutoLens** uses configs to set up priors on all of the model's parameters. \n", - "\n", - "These configs are stored in the `config/priors/mass//.yaml`, for \n", - "example `config/priors/mass/total/isothermal.yaml`.\n", - "\n", - "If you add your mass profile to the **PyAutoLens** source code you can add a config file for it to this folder and\n", - "**PyAutoLens** will automatically use it to set up the priors on your mass profile.\n", - "\n", - "You should also add your mass profile and its parameters to the `config/notation.yaml` file, so that **PyAutoLens**\n", - "knows how to label your mass profile in plots.\n", - "\n", - "__Deflections__\n", - "\n", - "We are therefore ready to implement a mass profile, and the best place to start is the `deflections_yx_2d_from` function.\n", - "\n", - "In fact, this is the only function you need to implement in order for lens modeling to work. This is because pretty much\n", - "all lensing calculations can be computed from the deflection angles.\n", - "\n", - "However, we recommend you also implement analytic functions for the `convergence` and `potential` of your mass profile.\n", - "They are often used for separate calculations outside of lens modeling and are commonly visualized in plots.\n", - "\n", - "The template below is a good starting point for your mass profile and explains what functions you need to implement\n", - "and what are optional." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class TemplateMass(EllProfile):\n", - " def __init__(\n", - " self,\n", - " centre: Tuple[float, float] = (0.0, 0.0),\n", - " ell_comps: Tuple[float, float] = (0.0, 0.0),\n", - " # Your parameters here.\n", - " ):\n", - " super().__init__(centre=centre, ell_comps=ell_comps)\n", - "\n", - " # Note that for a Spherical profile, which does not have an `ell_comps` parameter,\n", - " # you can remove it from the __init__ constructor and pass (0.0, 0.0) below, e.g.\n", - "\n", - " # super().__init__(centre=centre, ell_comps=(0.0, 0.0))\n", - "\n", - " @aa.grid_dec.to_vector_yx\n", - " @aa.grid_dec.transform(rotate_back=True)\n", - " def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", - " \"\"\"\n", - " REQUIRED: The function is key for all lensing calculations and must be implemented.\n", - "\n", - " Set ``rotate_back=True`` if your function computes deflection components using the rotated\n", - " grid coordinates (i.e. the components are expressed in the profile's frame). The decorator\n", - " will rotate them back to the observer frame automatically. If your function reconstructs\n", - " observer-frame components from scalar quantities, leave ``rotate_back=False``.\n", - " \"\"\"\n", - " pass\n", - "\n", - " @aa.grid_dec.to_array\n", - " @aa.grid_dec.transform\n", - " def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", - " \"\"\"\n", - " RECOMMENDED: The convergence is used for visualization and inspecting properties of the mass profile.\n", - " \"\"\"\n", - " pass\n", - "\n", - " @aa.grid_dec.to_array\n", - " @aa.grid_dec.transform\n", - " def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", - " \"\"\"\n", - " RECOMMENDED: The gravitational potential is used for visualization and inspecting properties of the mass\n", - " profile.\n", - " \"\"\"\n", - " pass\n", - "\n", - " def convergence_func(self, grid_radius: float) -> float:\n", - " \"\"\"\n", - " Optional: A 1D function which returns the convergence at a given 1D coordinate (e.g. radius). This is used\n", - " for computing integrated mass quantities.\n", - " \"\"\"\n", - " pass\n", - "\n", - " @staticmethod\n", - " def potential_func(u, y, x, axis_ratio, slope, core_radius):\n", - " _eta_u = np.sqrt((u * ((x**2) + (y**2 / (1 - (1 - axis_ratio**2) * u)))))\n", - " return (\n", - " (_eta_u / u)\n", - " * ((3.0 - slope) * _eta_u) ** -1.0\n", - " * _eta_u ** (3.0 - slope)\n", - " / ((1 - (1 - axis_ratio**2) * u) ** 0.5)\n", - " )\n", - "\n", - " @aa.grid_dec.to_vector_yx\n", - " @aa.grid_dec.transform\n", - " def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", - " \"\"\"\n", - " OPTIONAL: Shears are used for weak lensing calculations and inspection properties of the mass profile.\n", - "\n", - " Shears can reliably be calculated via methods inherited from the `OperateDeflections` class. Providing an\n", - " analytic calculation here can speed this up and provide more accurate results.\n", - " \"\"\"\n", - " pass\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Spherical Template__\n", - "\n", - "radial_grid\n", - "removal of ell_comps\n", - "\n", - "__Physical Profiles__\n", - "\n", - "Show how to wrap existing profiles with physical units?\n", - "\n", - "__Light Profiles__\n", - "\n", - "Pretty much the same but need to add text." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Misc: Add A Profile\n", + "===================\n", + "\n", + "**PyAutoLens** supports a wide range of mass and light profiles for modelling the\n", + "lens and source galaxies in strong gravitational lensing systems.\n", + "\n", + "For some science cases, it may be necessary to define custom profiles. This could\n", + "involve implementing a new profile that is not currently supported by\n", + "**PyAutoLens**, or introducing a new parameterization of an existing profile. Both\n", + "of these possibilities are covered in this example tutorial.\n", + "\n", + "We begin by explaining how to add a new _mass profile_, as this introduces the core\n", + "concepts required for defining custom profiles in **PyAutoLens**. These concepts\n", + "are then applied to show how custom _light profiles_ can be implemented.\n", + "\n", + "__Contents__\n", + "\n", + "- **Source Code:** This example includes direct links to the source code of the classes used to define mass and light.\n", + "- **Example Mass Profile:** The mass profiles available in **PyAutoLens** are located in its parent package, **PyAutoGalaxy**.\n", + "- **Inheritance Structure:** Let us next consider the inheritance structure of the ``Isothermal`` profile, defined by the class.\n", + "- **Data Structure Decorators:** Different grids can be input into each mass profile function (e.g.\n", + "- **Transform Decorator:** Overview of transform decorator for this example.\n", + "- **Lens Modeling:** **PyAutoLens** assumes that all input parameters of a mass profile (for example, those listed in.\n", + "- **Lens Modeling Configs:** In most **PyAutoLens** examples, you will notice we compose models without manually specifying.\n", + "- **Deflections:** We are therefore ready to implement a mass profile, and the best place to start is the.\n", + "- **Spherical Template:** radial_grid removal of ell_comps.\n", + "- **Physical Profiles:** Show how to wrap existing profiles with physical units?\n", + "- **Light Profiles:** Pretty much the same but need to add text.\n", + "\n", + "__Source Code__\n", + "\n", + "This example includes direct links to the source code of the classes used to define\n", + "mass and light profiles, allowing you to see exactly how they are implemented.\n", + "\n", + "The tutorial is fully standalone and, by the end, should enable you to implement a\n", + "custom profile without needing to dive deeply into the **PyAutoLens** codebase.\n", + "That said, we still recommend exploring the source code to better understand how\n", + "everything fits together." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Example Mass Profile__\n", + "\n", + "The mass profiles available in **PyAutoLens** are located in its parent package, **PyAutoGalaxy**: \n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy\n", + "\n", + "All light and mass profiles are found in the following python package:\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/tree/main/autogalaxy/profiles\n", + "\n", + "Mass profiles are in the following package:\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/tree/main/autogalaxy/profiles/mass\n", + "\n", + "Lets look at an example mass profile. We'll use the `Isothermal` profile, which is located in the `total` package\n", + "because it represents a total (stars + dark matter) mass distribution:\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/tree/main/autogalaxy/profiles/mass/total\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/profiles/mass/total/isothermal.py\n", + "\n", + "For simplicity, a shortened version of the `Isothermal` profile is shown below. \n", + "\n", + "This has docstrings updated to focus on the key aspects of implementing a new profiles and simplifies the \n", + "inheritance structure of the profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from typing import Tuple\n", + "\n", + "import autoarray as aa\n", + "\n", + "from autogalaxy.profiles.mass.total.isothermal import psi_from\n", + "from autogalaxy.profiles.mass.abstract.abstract import MassProfile\n", + "\n", + "\n", + "class Isothermal(MassProfile):\n", + " def __init__(\n", + " self,\n", + " centre: Tuple[float, float] = (0.0, 0.0),\n", + " ell_comps: Tuple[float, float] = (0.0, 0.0),\n", + " einstein_radius: float = 1.0,\n", + " ):\n", + " \"\"\"\n", + " Represents an elliptical isothermal density distribution, which is equivalent to the elliptical power-law\n", + " density distribution for the value slope = 2.0.\n", + "\n", + " Parameters\n", + " ----------\n", + " centre\n", + " The (y,x) arc-second coordinates of the profile centre.\n", + " ell_comps\n", + " The first and second ellipticity components of the elliptical coordinate system.\n", + " einstein_radius\n", + " The arc-second Einstein radius.\n", + " \"\"\"\n", + "\n", + " super().__init__(\n", + " centre=centre,\n", + " ell_comps=ell_comps,\n", + " )\n", + "\n", + " self.einstein_radius = einstein_radius\n", + " self.slope = 2.0\n", + "\n", + " @aa.grid_dec.to_vector_yx\n", + " @aa.grid_dec.transform(rotate_back=True)\n", + " def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", + " \"\"\"\n", + " Calculate the deflection angles on a grid of (y,x) arc-second coordinates.\n", + "\n", + " The input grid of (y,x) coordinates are transformed to a coordinate system centred on the profile centre with\n", + " and rotated based on the position angle defined from its `ell_comps` (this is described fully below).\n", + "\n", + " Because this method computes deflection components using the rotated grid coordinates (i.e. the\n", + " components are expressed in the profile's frame), ``rotate_back=True`` is set so the decorator\n", + " automatically rotates them back to the observer frame.\n", + "\n", + " The numerical backend can be selected via the ``xp`` argument, allowing this\n", + " method to be used with both NumPy and JAX (e.g. inside ``jax.jit``-compiled\n", + " code). This is described fully later in this example.\n", + "\n", + " Parameters\n", + " ----------\n", + " grid\n", + " The grid of (y,x) arc-second coordinates the deflection angles are computed on.\n", + " xp\n", + " The numerical backend to use, either `numpy` or `jax.numpy`.\n", + " \"\"\"\n", + "\n", + " factor = (\n", + " 2.0\n", + " * self.einstein_radius_rescaled(xp)\n", + " * self.axis_ratio(xp)\n", + " / xp.sqrt(1 - self.axis_ratio(xp) ** 2)\n", + " )\n", + "\n", + " psi = psi_from(\n", + " grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=0.0, xp=xp\n", + " )\n", + "\n", + " deflection_y = xp.arctanh(\n", + " xp.divide(\n", + " xp.multiply(xp.sqrt(1 - self.axis_ratio(xp) ** 2), grid.array[:, 0]),\n", + " psi,\n", + " )\n", + " )\n", + " deflection_x = xp.arctan(\n", + " xp.divide(\n", + " xp.multiply(xp.sqrt(1 - self.axis_ratio(xp) ** 2), grid.array[:, 1]),\n", + " psi,\n", + " )\n", + " )\n", + " return xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__JAX, Numpy and xp__\n", + "\n", + "Throughout this tutorial, and in the profile functions above, you will see functions and methods that accept an\n", + "argument called ``xp``. The default input value above is ``xp=np``, which sets it to the standard NumPy library,\n", + "imported using the statement ``import numpy as np``. This means all arithmetic operations use NumPy in a way\n", + "you are likely familiar with.\n", + "\n", + "However, to enable mass profile calculations to run on a GPU, a library called JAX is used which mirrors the\n", + "NumPy API. Conventionally, this is imported as ``jnp`` using the statement ``import jax.numpy as jnp``. When the\n", + "source code is running in JAX mode, the input ``xp`` to the functions above will be the ``jnp`` library instead\n", + "of ``np``. This is why mass profiles support the ``xp`` argument and API: they need to be able to run using\n", + "either NumPy or JAX.\n", + "\n", + "The PyAutoLens source code runs in pure NumPy by default, where ``xp`` is always set to ``np``. This only\n", + "changes if you manually call a function passing ``xp=jnp``, or when certain high-level objects, such as the\n", + "``Analysis`` class, are used. These objects automatically set ``xp=jnp`` when a likelihood is evaluated for\n", + "lens modelling.\n", + "\n", + "Your final mass profile should therefore use the ``xp`` API throughout, ensuring compatibility with both NumPy\n", + "and JAX and allowing it to work seamlessly with the PyAutoLens source code. You may find it easier to first\n", + "write your functions in pure NumPy (which you are likely most familiar with), and then convert them to use the\n", + "``xp`` API and test them with JAX afterwards. While using ``xp`` makes the API slightly more verbose, it is a\n", + "small price to pay for the significant speed-ups available when running JAX on a GPU.\n", + "\n", + "__Inheritance Structure__\n", + "\n", + "Let us next consider the inheritance structure of the ``Isothermal`` profile,\n", + "defined by the class declaration::\n", + "\n", + " class Isothermal(MassProfile):\n", + "\n", + "In Python, inheritance means that a class can reuse and extend the behaviour of\n", + "another class. By inheriting from ``MassProfile``, the ``Isothermal`` profile\n", + "automatically has access to all methods and attributes defined in\n", + "``MassProfile``. This allows ``Isothermal`` to make use of shared functionality\n", + "(such as common calculations and interfaces) without reimplementing it.\n", + "\n", + "The key mechanism used to enable this inheritance is the ``super`` function. For\n", + "example, in the ``Isothermal`` initializer we see::\n", + "\n", + " super().__init__(\n", + " centre=centre,\n", + " ell_comps=ell_comps,\n", + " )\n", + "\n", + "This line calls the ``__init__`` method of the parent ``MassProfile`` class,\n", + "ensuring that all required base-class setup is performed before adding any\n", + "``Isothermal``-specific behaviour.\n", + "\n", + "It is important to emphasize that you do not need to fully understand the full\n", + "inheritance structure of the **PyAutoLens** profiles or the layout of the source\n", + "code to define your own custom profiles. This discussion is included simply to\n", + "highlight that all calculations involving mass and light profiles are built on\n", + "a set of abstract base classes, which your custom profiles will automatically\n", + "inherit from.\n", + "\n", + "__Inheritance (MassProfile)__\n", + "\n", + "All mass profiles in **PyAutoLens** inherit from the `MassProfile` abstract base class, which is located in the\n", + "following package:\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/profiles/mass/abstract/abstract.py\n", + "\n", + "This contains functions which are useful for any mass profile, which your custom mass profile will inherit.\n", + "\n", + "For example, it includes the function `mass_angular_within_circle_from`, which computes the mass of the profile\n", + "within an input circle of radius `radius`.\n", + "\n", + "__Inheritance (GeometryProfile)__\n", + "\n", + "The `MassProfile` class inherits from the `GeometryProfile` abstract base class, which is located here:\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/profiles/geometry_profiles.py\n", + "\n", + "This contains functions which are useful for any elliptical (and spherical) profile, which your custom \n", + "mass profile will again inherit (e.g. `radial_grid_from`).\n", + "\n", + "They typically perform coordinate transforms between the profile's elliptical (spherical) coordinate system and the \n", + "input 2D grid of (y,x) coordinates. For example, the function `transformed_to_reference_frame_grid_from` transforms\n", + "the (y,x) coordinates to the profile's elliptical coordinate system.\n", + "\n", + "The convention of this calculation is key for ensuring you implement your custom profile correctly. We illustrate\n", + "it fully below.\n", + "\n", + "__Inheritance (OperateDeflections)__\n", + "\n", + "Nearly all lensing quantities (e.g. `convergence`, `potential`, `magnification`) can be derived from the deflection\n", + "angles of a mass profile. \n", + "\n", + "Mass profiles therefore also inherit from the `OperateDeflections` abstract base class, which contains numerous \n", + "functions for computing these lensing quantities from the deflection angles. This is located here:\n", + "\n", + "https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/operate/deflections.py\n", + "\n", + "This means that once you've implemented a deflections angles calculation for your mass profile, you can compute all\n", + "lensing quantities from it without having to write any additional code!\n", + "\n", + "__Whats Going On With Those Decorators?__\n", + "\n", + "A decorator in Python is a special syntax used to modify or extend the behaviour of a function or method without \n", + "changing its implementation. It is denoted by the `@decorator_name` syntax placed above the function definition,\n", + "with two decorators shown in the example above: `@aa.grid_dec.to_vector_yx` and `@aa.grid_dec.transform`. Lets\n", + "now consider what these do.\n", + "\n", + "__Data Structure Decorators__\n", + "\n", + "Different grids can be input into each mass profile function (e.g. `Grid2D`, `Grid2DIrregular`). Depending on the input \n", + "grid, this changes the structure of the output array. \n", + "\n", + "For example, if a `Grid2D` is input, which is defined on a uniform grid of 2D coordinates, the output deflection angles\n", + "are also defined on a uniform grid and are returned as a `VectorYX2D` object. If a `Grid2DIrregular` is input, \n", + "which is defined on an irregular grid of 2D coordinates, the output deflection angles are also defined on an irregular \n", + "grid and are returned as a `VectorYX2DIrregular` object.\n", + "\n", + "The `@aa.grid_dec.to_vector_yx` decorator handles this structure conversion for vector quantities, such that the output\n", + "vector structure matches the input grid structure.\n", + "\n", + "The function `deflections_yx_2d_from` returns 2D vectors, but other mass profile methods, like `convergence_2d_from` and\n", + "`potential_2d_from`, return scalar quantities. These methods use the `@aa.grid_dec.to_array` decorator, which behaves\n", + "analogously to the `@aa.grid_dec.to_vector_yx` decorator but for scalar quantities (e.g. for an input `Grid2D`, the output\n", + "is an `Array2D` object, for an input `Grid2DIrregular`, the output is an `ArrayIrregular` object).\n", + "\n", + "For your custom mass profile, you basically just need to copy and paste these decorators above your mass profile \n", + "functions and not worry about them any further.\n", + " \n", + "__Transform Decorator__\n", + " \n", + "The second decorator is the `@aa.grid_dec.transform` decorator. This one we will have a closer look at, as it\n", + "will influence how you implement your mass profile functions.\n", + "\n", + "The `transform` decorator is used to transform the input grid of (y,x) coordinates to the mass profile's elliptical \n", + "coordinate system. It does this by calling the function `transform_grid_2d_to_reference_frame`, which I have\n", + "provided below for convenience:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def transform_grid_2d_to_reference_frame(\n", + " grid_2d: np.ndarray, centre: Tuple[float, float], angle: float, xp=np\n", + ") -> np.ndarray:\n", + " \"\"\"\n", + " Transform a 2D grid of (y,x) coordinates to a new reference frame.\n", + "\n", + " This transformation includes:\n", + "\n", + " 1) A translation to a new (y,x) centre value, by subtracting the centre from every coordinate on the grid.\n", + " 2) A rotation of the grid around this new centre, which is performed clockwise from an input angle.\n", + "\n", + " Parameters\n", + " ----------\n", + " grid\n", + " The 2d grid of (y, x) coordinates which are transformed to a new reference frame.\n", + " \"\"\"\n", + "\n", + " shifted_grid_2d = grid_2d - xp.array(centre)\n", + "\n", + " radius = xp.sqrt(xp.sum(xp.square(shifted_grid_2d), axis=1))\n", + " theta_coordinate_to_profile = xp.arctan2(\n", + " shifted_grid_2d[:, 0], shifted_grid_2d[:, 1]\n", + " ) - xp.radians(angle)\n", + "\n", + " return xp.vstack(\n", + " [\n", + " radius * xp.sin(theta_coordinate_to_profile),\n", + " radius * xp.cos(theta_coordinate_to_profile),\n", + " ]\n", + " ).T\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A simple example of this function is shifting a grid to a mass profile's centre (which simply subtracts the centre\n", + "coordinates from every coordinate on the grid):" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2DIrregular(values=[(0.0, 0.0), (1.0, 1.0), (2.0, 2.0)])\n", + "\n", + "mass_profile_centre = (0.5, 0.5)\n", + "\n", + "transformed_grid = transform_grid_2d_to_reference_frame(\n", + " grid_2d=grid, centre=mass_profile_centre, angle=0.0\n", + ")\n", + "\n", + "print(f\"Grid Coordinates Before: {grid}\")\n", + "print(f\"Grid Coordinates After: {transformed_grid}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `angle` input is the rotation angle of the mass profile's ellipse counter-clockwise from the positive x-axis.\n", + "\n", + "It is computed from the `ell_comps` of the mass profile, which are the elliptical components of the mass profile's\n", + "ellipse." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def axis_ratio_and_angle_from(\n", + " ell_comps: Tuple[float, float], xp=np\n", + ") -> Tuple[float, float]:\n", + " \"\"\"\n", + " Returns the axis-ratio and position angle in degrees (-45 < angle < 135.0) from input elliptical components e1\n", + " and e2 of a light or mass profile.\n", + "\n", + " The elliptical components of a light or mass profile are given by:\n", + "\n", + " elliptical_component_y = ell_comps[0] = (1-axis_ratio)/(1+axis_ratio) * sin(2 * angle)\n", + " elliptical_component_x = ell_comps[1] = (1-axis_ratio)/(1+axis_ratio) * cos(2 * angle)\n", + "\n", + " The axis-ratio and angle are therefore given by:\n", + "\n", + " axis_ratio = (1 - fac) / (1 + fac)\n", + " angle = 0.5 * arctan(ell_comps[0] / ell_comps[1])\n", + "\n", + " where `fac = sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2).\n", + "\n", + " This function returns the axis-ratio and angle in degrees.\n", + "\n", + " An additional check is performed which requires the angle is between -45 and 135 degrees. This ensures that\n", + " for certain values of `ell_comps` the angle does not jump from one boundary to another (e.g. without this check\n", + " certain values of `ell_comps` return -1.0 degrees and others 179.0 degrees). This ensures that when error\n", + " estimates are computed from samples of a lens model via marginalization, the calculation is not biased by the\n", + " angle jumping between these two values.\n", + "\n", + " Parameters\n", + " ----------\n", + " ell_comps\n", + " The elliptical components of the light or mass profile which are converted to an angle.\n", + " \"\"\"\n", + " angle = xp.arctan2(ell_comps[0], ell_comps[1]) / 2\n", + " angle *= 180.0 / xp.pi\n", + "\n", + " angle = xp.where(angle < -45, angle + 180, angle)\n", + "\n", + " fac = xp.sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2)\n", + " if xp.__name__.startswith(\"jax\"):\n", + " import jax\n", + "\n", + " fac = jax.lax.min(fac, 0.999)\n", + " else: # NumPy\n", + " fac = np.minimum(fac, 0.999)\n", + "\n", + " axis_ratio = (1 - fac) / (1 + fac)\n", + " return axis_ratio, angle\n", + "\n", + "\n", + "mass_profile_ell_comps = (0.5, 0.5)\n", + "mass_profile_angle = axis_ratio_and_angle_from(ell_comps=mass_profile_ell_comps)[1]\n", + "\n", + "print(f\"\\nMass Profile Angle (degrees) {mass_profile_angle}\")\n", + "\n", + "transformed_grid = transform_grid_2d_to_reference_frame(\n", + " grid_2d=grid, centre=mass_profile_centre, angle=mass_profile_angle\n", + ")\n", + "\n", + "print(f\"Grid Coordinates Before: {grid}\")\n", + "print(f\"Grid Coordinates After: {transformed_grid}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `@aa.grid_dec.transform` packages all the above calculations up and uses the mass profile `centre` and `ell_comps` \n", + "to perform them before your function is called.\n", + "\n", + "The class below demonstrates this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autogalaxy.profiles.geometry_profiles import EllProfile\n", + "\n", + "\n", + "class ExampleMass(EllProfile):\n", + " def __init__(\n", + " self,\n", + " centre: Tuple[float, float] = (0.0, 0.0),\n", + " ell_comps: Tuple[float, float] = (0.0, 0.0),\n", + " ):\n", + " super().__init__(centre=centre, ell_comps=ell_comps)\n", + "\n", + " @aa.grid_dec.transform\n", + " def deflections_yx_2d_from(self, grid: aa.Grid2D, xp=np, **kwargs):\n", + " print(\n", + " f\"\\n Grid In Deflections After Transform \"\n", + " f\"Which is Same As Transformed Grid Above: {grid}\"\n", + " )\n", + "\n", + "\n", + "mass = ExampleMass(centre=(0.5, 0.5), ell_comps=(0.5, 0.5))\n", + "mass.deflections_yx_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Do I Rotate Back?__\n", + "\n", + "Whether the returned result needs back-rotation depends on what kind of quantity the function returns\n", + "and which coordinate frame its components are expressed in.\n", + "\n", + "**Scalars** (convergence, potential): frame-invariant \u2014 no back-rotation needed.\n", + "\n", + "**Vectors** (deflection angles): it depends on how the function computes them. The ``@aa.grid_dec.transform``\n", + "decorator transforms the input grid into the profile's rotated reference frame. If the function then computes\n", + "deflection components ``(alpha_y, alpha_x)`` using that rotated grid \u2014 as the ``Isothermal`` example above does \u2014\n", + "those components are expressed in the profile's frame. Since the ray-tracing code expects observer-frame\n", + "components, they must be rotated back. Setting ``rotate_back=True`` handles this automatically by calling\n", + "``self.rotated_grid_from_reference_frame_from(...)`` on the result.\n", + "\n", + "However, back-rotation is **not** always needed for vectors. If a function computes a scalar quantity in the\n", + "profile frame (e.g. a radial deflection magnitude) and then reconstructs the Cartesian vector using\n", + "observer-frame geometry, the result is already in the observer frame. In that case, ``rotate_back`` should\n", + "remain ``False``.\n", + "\n", + "The rule is: look at which coordinate basis the returned components are expressed in. If they use the rotated\n", + "basis, set ``rotate_back=True``. If they are already observer-frame components, leave it as ``False``.\n", + "\n", + "**Spin-2 quantities** (shear): these transform under a coordinate rotation by twice the profile angle (the\n", + "spin-2 transformation law). This is not handled by ``rotate_back`` \u2014 shear methods must apply the ``2 * angle``\n", + "rotation manually via ``self.rotated_grid_from_reference_frame_from(grid=..., angle=self.angle(xp) * 2)``.\n", + "\n", + "__Lens Modeling__\n", + "\n", + "**PyAutoLens** assumes that all input parameters of a mass profile (for example,\n", + "those listed in its ``__init__`` constructor) are free parameters that can be\n", + "fitted during lens modelling using a non-linear search.\n", + "\n", + "If a parameter in the ``__init__`` constructor is a float (e.g. the\n", + "``einstein_radius`` of the ``Isothermal`` profile), it is treated as a single\n", + "free parameter. If a parameter is a tuple of floats (e.g. the ``centre`` of the\n", + "``Isothermal`` profile), each element of the tuple is treated as a separate free\n", + "parameter.\n", + "\n", + "We demonstrate this behaviour using a simple example mass profile defined\n", + "below. We compose it as a model using ``af.Model`` and print its ``info``,\n", + "which summarizes its free parameters and shows that no priors have yet been\n", + "assigned." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import autofit as af\n", + "from typing import Tuple\n", + "\n", + "\n", + "class LensModelExample:\n", + " def __init__(\n", + " self, # <-- **PyAutoLens** assumes these input parameters are free.\n", + " centre: Tuple[float, float] = (\n", + " 0.0,\n", + " 0.0,\n", + " ), # <-- Two free parameters because this is a tuple.\n", + " ell_comps: Tuple[float, float] = (\n", + " 0.0,\n", + " 0.0,\n", + " ), # <-- Also two free parameters.\n", + " einstein_radius: float = 1.0, # <-- A single free parameter.\n", + " your_parameter_here: float = 2.0, # <-- Add any custom parameters you need.\n", + " ):\n", + " pass\n", + "\n", + "\n", + "lens_model_example = af.Model(LensModelExample)\n", + "\n", + "print(lens_model_example.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For this example model, we can manually assign priors to its parameters as shown\n", + "below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_model_example.centre.centre_0 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0)\n", + "lens_model_example.centre.centre_1 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0)\n", + "lens_model_example.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + "lens_model_example.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + "lens_model_example.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=3.0)\n", + "lens_model_example.your_parameter_here = af.UniformPrior(\n", + " lower_limit=0.0, upper_limit=5.0\n", + ")\n", + "\n", + "print(lens_model_example.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The exact same API applies to the ``Isothermal`` class defined above, which has\n", + "three ``__init__`` parameters: ``centre``, ``ell_comps``, and\n", + "``einstein_radius``.\n", + "\n", + "When composed as a model, the ``Isothermal`` profile therefore has five free\n", + "parameters in total." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = af.Model(Isothermal)\n", + "\n", + "print(mass.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As before, we must manually assign priors to these parameters for lens modelling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass.centre.centre_0 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0)\n", + "mass.centre.centre_1 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0)\n", + "mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + "mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + "mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=3.0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Provided you follow this same convention when defining your own mass profile,\n", + "it will fully support lens modelling in **PyAutoLens** without requiring any\n", + "additional code.\n", + "\n", + "__Lens Modeling Configs__\n", + "\n", + "In most **PyAutoLens** examples, you will notice we compose models without manually specifying priors. This is because\n", + "**PyAutoLens** uses configs to set up priors on all of the model's parameters. \n", + "\n", + "These configs are stored in the `config/priors/mass//.yaml`, for \n", + "example `config/priors/mass/total/isothermal.yaml`.\n", + "\n", + "If you add your mass profile to the **PyAutoLens** source code you can add a config file for it to this folder and\n", + "**PyAutoLens** will automatically use it to set up the priors on your mass profile.\n", + "\n", + "You should also add your mass profile and its parameters to the `config/notation.yaml` file, so that **PyAutoLens**\n", + "knows how to label your mass profile in plots.\n", + "\n", + "__Deflections__\n", + "\n", + "We are therefore ready to implement a mass profile, and the best place to start is the `deflections_yx_2d_from` function.\n", + "\n", + "In fact, this is the only function you need to implement in order for lens modeling to work. This is because pretty much\n", + "all lensing calculations can be computed from the deflection angles.\n", + "\n", + "However, we recommend you also implement analytic functions for the `convergence` and `potential` of your mass profile.\n", + "They are often used for separate calculations outside of lens modeling and are commonly visualized in plots.\n", + "\n", + "The template below is a good starting point for your mass profile and explains what functions you need to implement\n", + "and what are optional." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class TemplateMass(EllProfile):\n", + " def __init__(\n", + " self,\n", + " centre: Tuple[float, float] = (0.0, 0.0),\n", + " ell_comps: Tuple[float, float] = (0.0, 0.0),\n", + " # Your parameters here.\n", + " ):\n", + " super().__init__(centre=centre, ell_comps=ell_comps)\n", + "\n", + " # Note that for a Spherical profile, which does not have an `ell_comps` parameter,\n", + " # you can remove it from the __init__ constructor and pass (0.0, 0.0) below, e.g.\n", + "\n", + " # super().__init__(centre=centre, ell_comps=(0.0, 0.0))\n", + "\n", + " @aa.grid_dec.to_vector_yx\n", + " @aa.grid_dec.transform(rotate_back=True)\n", + " def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", + " \"\"\"\n", + " REQUIRED: The function is key for all lensing calculations and must be implemented.\n", + "\n", + " Set ``rotate_back=True`` if your function computes deflection components using the rotated\n", + " grid coordinates (i.e. the components are expressed in the profile's frame). The decorator\n", + " will rotate them back to the observer frame automatically. If your function reconstructs\n", + " observer-frame components from scalar quantities, leave ``rotate_back=False``.\n", + " \"\"\"\n", + " pass\n", + "\n", + " @aa.grid_dec.to_array\n", + " @aa.grid_dec.transform\n", + " def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", + " \"\"\"\n", + " RECOMMENDED: The convergence is used for visualization and inspecting properties of the mass profile.\n", + " \"\"\"\n", + " pass\n", + "\n", + " @aa.grid_dec.to_array\n", + " @aa.grid_dec.transform\n", + " def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", + " \"\"\"\n", + " RECOMMENDED: The gravitational potential is used for visualization and inspecting properties of the mass\n", + " profile.\n", + " \"\"\"\n", + " pass\n", + "\n", + " def convergence_func(self, grid_radius: float) -> float:\n", + " \"\"\"\n", + " Optional: A 1D function which returns the convergence at a given 1D coordinate (e.g. radius). This is used\n", + " for computing integrated mass quantities.\n", + " \"\"\"\n", + " pass\n", + "\n", + " @staticmethod\n", + " def potential_func(u, y, x, axis_ratio, slope, core_radius):\n", + " _eta_u = np.sqrt((u * ((x**2) + (y**2 / (1 - (1 - axis_ratio**2) * u)))))\n", + " return (\n", + " (_eta_u / u)\n", + " * ((3.0 - slope) * _eta_u) ** -1.0\n", + " * _eta_u ** (3.0 - slope)\n", + " / ((1 - (1 - axis_ratio**2) * u) ** 0.5)\n", + " )\n", + "\n", + " @aa.grid_dec.to_vector_yx\n", + " @aa.grid_dec.transform\n", + " def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):\n", + " \"\"\"\n", + " OPTIONAL: Shears are used for weak lensing calculations and inspection properties of the mass profile.\n", + "\n", + " Shears can reliably be calculated via methods inherited from the `OperateDeflections` class. Providing an\n", + " analytic calculation here can speed this up and provide more accurate results.\n", + " \"\"\"\n", + " pass\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Spherical Template__\n", + "\n", + "radial_grid\n", + "removal of ell_comps\n", + "\n", + "__Physical Profiles__\n", + "\n", + "Show how to wrap existing profiles with physical units?\n", + "\n", + "__Light Profiles__\n", + "\n", + "Pretty much the same but need to add text." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/advanced/custom_analysis.ipynb b/notebooks/guides/advanced/custom_analysis.ipynb index d087f63ce..dbf7a092e 100644 --- a/notebooks/guides/advanced/custom_analysis.ipynb +++ b/notebooks/guides/advanced/custom_analysis.ipynb @@ -1,815 +1,852 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Misc: Custom Analysis\n", - "=====================\n", - "\n", - "Users familiar with **PyAutoLens** will have seen that `Analysis` classes are used to performed lens modeling of\n", - "different datasets. For example, the `AnalysisImaging` class fits imaging datasets, the `AnalysisInterferometer` class\n", - "fits interferometer datasets.\n", - "\n", - "You may have a dataset which you want to use **PyAutoLenss**'s lensing capabilities to model, but which does not fit\n", - "into one of the standard `Analysis` classes.\n", - "\n", - "A good example (at the time of writing this script) is fitting a weak lensing shear catalogue with a model of the lens\n", - "galaxy's mass. **PyAutoLens** as the lensing capabilities to produce the shears of a mass model, but does not have an\n", - "`Analysis` class to fit these shears to a dataset.\n", - "\n", - "This example demonstrates how you can write your own `Analysis` class to fit a dataset with **PyAutoLens**.\n", - "\n", - "__Contents__\n", - "\n", - "- **PyAutoFit:** The `Analysis` class is the interface between the data and model, whereby a.\n", - "- **Source Code:** This example contains URLs to the locations of the source code of the classes used when creating.\n", - "- **Example Analysis Class:** The `Analysis` classes available in **PyAutoLens** are actually located in both **PyAutoLens** and.\n", - "- **Lens Model:** To illustrate how to write a custom `Analysis` class, we require an example lens model that we will.\n", - "- **Instances:** Instances of the model above can be created, where an input `vector` of parameters is mapped to.\n", - "- **Simple Analysis Example:** For simplicity, a shortened version of an `AnalysisImaging` class is shown below where certain.\n", - "- **Analysis Class Considerations:** Lets quickly think about the design of an `Analysis` class and how this can help us to set up any.\n", - "- **Model Fit:** Perform the model-fit using the search and analysis.\n", - "- **Weak Lensing Example:** Now lets consider how to write our own custom `Analysis` class, for the example of performing a.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **To Do List:** The following `Analysis` cookbook from the **PyAutoFit** readthedocs should help you get started.\n", - "\n", - "__PyAutoFit__\n", - "\n", - "The `Analysis` class is the interface between the data and model, whereby a `log_likelihood_function` is defined\n", - "and called by the non-linear search to fit the model.\n", - "\n", - "You may have performed a similar task yourself, for example by taking a fitting library (e.g. an MCMC method like\n", - "Emcee or nested sampler like Dynesty) and writing a likelihood function that calls it to fit a model to a dataset.\n", - "If you haven't done this, this script will explain how!\n", - "\n", - "**PyAutoLens** uses a library called **PyAutoFit** to set up this interfacebetween the data, model,\n", - "`log_likelihood_function` and non-linear search. **PyAutoFit** is a general purpose library for model fitting,\n", - "and we will see that it has a lot of powerful tools that we can use to customize our `Analysis` class.\n", - "\n", - "You can checkout the **PyAutoFit** readthedocs here:\n", - "\n", - " https://pyautofit.readthedocs.io/en/latest/\n", - "\n", - "The following analysis cookbook provides a concise reference guide to `Analysis` objects, and once you have completed\n", - "this example will be a useful resource for writing your own `Analysis` class:\n", - "\n", - " https://pyautofit.readthedocs.io/en/latest/cookbooks/analysis.html\n", - "\n", - "__Source Code__\n", - "\n", - "This example contains URLs to the locations of the source code of the classes used when creating light and mass\n", - "profiles.\n", - "\n", - "The example itself is standalone and should by the end allow you to implement a custom profile without diving into\n", - "the **PyAutoLens** source code.\n", - "\n", - "We still recommend you take a look to see how things are structured!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Example Analysis Class__\n", - "\n", - "The `Analysis` classes available in **PyAutoLens** are actually located in both **PyAutoLens** and its parent \n", - "package, **PyAutoGalaxy**: \n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy\n", - " https://github.com/PyAutoLabs/PyAutoLens\n", - "\n", - "All classes used for lens modeling are found in the following packages:\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/tree/main/autogalaxy/imaging/model\n", - " https://github.com/PyAutoLabs/PyAutoLens/tree/main/autolens/imaging/model\n", - "\n", - "The `AnalysisImaging` classes are found in the following modules:\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/imaging/model/analysis.py\n", - " https://github.com/PyAutoLabs/PyAutoLens/blob/main/autolens/imaging/model/analysis.py\n", - "\n", - "__Lens Model__\n", - "\n", - "To illustrate how to write a custom `Analysis` class, we require an example lens model that we will use to fit\n", - "the dataset.\n", - "\n", - "We compose a simple lens model with an `IsothermalSph` mass model for the lens and an `Sersic` for the source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - "\n", - "# Source:\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.ExponentialCoreSph)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Instances__\n", - "\n", - "Instances of the model above can be created, where an input `vector` of parameters is mapped to create an instance of \n", - "the Python class of the model.\n", - "\n", - "This is used internally by the `Analysis` class we are about to write, and will be used in \n", - "our `log_likelihood_function`. Therefore, we are quickly highlighting it here.\n", - "\n", - "We first need to know the order of parameters in the model, so we know how to define the input `vector`. This\n", - "information is contained in the models `paths` attribute:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.paths)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We input values for the parameters of our model following the order of paths above.\n", - "\n", - "This creates an `instance` of the lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "instance = model.instance_from_vector(vector=[0.0, 0.0, 1.6, 0.1, 0.1, 0.01, 2.0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This `instance` contains each of the model components we defined above. \n", - "\n", - "The argument names input into each `Collection` define the attribute names of the `instance`. \n", - "\n", - "For example, when composing the `model` above, we used a `Collection` called `galaxies` which had a `lens` and `source` \n", - "attribute. These `lens` and `source` attributes each contained components called `mass` and `bulge` respectively." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Lens Centre = {instance.galaxies.lens.mass.centre}\")\n", - "print(f\"Lens Einstein Radius = {instance.galaxies.lens.mass.einstein_radius}\")\n", - "print(f\"Source Centre = {instance.galaxies.source.bulge.centre}\")\n", - "print(f\"Source Intensity = {instance.galaxies.source.bulge.intensity}\")\n", - "print(f\"Source Effective Radius = {instance.galaxies.source.bulge.effective_radius}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simple Analysis Example__\n", - "\n", - "For simplicity, a shortened version of an `AnalysisImaging` class is shown below where certain functions have been \n", - "edited to make them easy to read and understand. \n", - "\n", - "This has docstrings updated to focus on the key aspects of implementing a new `Analysis` class and simplifies the \n", - "inheritance structure of the profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class AnalysisImaging(af.Analysis):\n", - " def __init__(\n", - " self,\n", - " dataset: al.Imaging,\n", - " cosmology: al.cosmo.LensingCosmology = al.cosmo.Planck15(),\n", - " use_jax: bool = True,\n", - " ):\n", - " \"\"\"\n", - " Fits a lens model to an imaging dataset via a non-linear search.\n", - "\n", - " The `Analysis` class defines the `log_likelihood_function` which fits the model to the dataset and returns the\n", - " log likelihood value defining how well the model fitted the data.\n", - "\n", - " It handles many other tasks, such as visualization, outputting results to hard-disk and storing results in\n", - " a format that can be loaded after the model-fit is complete.\n", - "\n", - " This class is used for model-fits which fit strong lenses composed via a `Tracer` to an imaging dataset.\n", - " Parameters\n", - " ----------\n", - " dataset\n", - " The `Imaging` dataset that the model is fitted to.\n", - " cosmology\n", - " The Cosmology assumed for this analysis.\n", - " \"\"\"\n", - " super().__init__(use_jax=use_jax)\n", - "\n", - " self.dataset = dataset\n", - " self.cosmology = cosmology\n", - "\n", - " def log_likelihood_function(self, instance: af.ModelInstance) -> float:\n", - " \"\"\"\n", - " Given an instance of the model, where the model parameters are set via a non-linear search, fit the model\n", - " instance to the imaging dataset.\n", - "\n", - " This function returns a log likelihood which is used by the non-linear search to guide the model-fit.\n", - "\n", - " For this analysis class, this function performs the following steps:\n", - "\n", - " 1) Extracts all galaxies from the model instance and set up a `Tracer`, which includes ordering the galaxies\n", - " by redshift to set up each `Plane`.\n", - "\n", - " 2) Use the `Tracer` and other attributes to create a `FitImaging` object, which performs steps such as creating\n", - " model images of every galaxy in the tracer, blurring them with the imaging dataset's PSF and computing\n", - " residuals, a chi-squared statistic and the log likelihood.\n", - "\n", - " Parameters\n", - " ----------\n", - " instance\n", - " An instance of the model that is being fitted to the data by this analysis (whose parameters have been set\n", - " via a non-linear search).\n", - "\n", - " Returns\n", - " -------\n", - " float\n", - " The log likelihood indicating how well this model instance fitted the imaging data.\n", - " \"\"\"\n", - "\n", - " \"\"\"\n", - " The `instance` that comes into this method is an instance of the lens model above, which we illustrated\n", - " via print statements how it is structured.\n", - "\n", - " The parameter values are chosen by the non-linear search, based on where it thinks the high likelihood regions \n", - " of parameter space are.\n", - "\n", - " The lines of Python code are commented out below to prevent excessive print statements when we run the\n", - " non-linear search, but feel free to uncomment them and run the search to see the parameters of every instance\n", - " that it fits.\n", - " \"\"\"\n", - "\n", - " # print(f\"Lens Centre = {instance.galaxies.lens.mass.centre}\")\n", - " # print(f\"Lens Einstein Radius = {instance.galaxies.lens.mass.einstein_radius}\")\n", - " # print(f\"Source Centre = {instance.galaxies.source.bulge.centre}\")\n", - " # print(f\"Source Intensity = {instance.galaxies.source.bulge.intensity}\")\n", - " # print(f\"Source Effective Radius = {instance.galaxies.source.bulge.effective_radius}\")\n", - "\n", - " \"\"\"\n", - " You should be familiar with the `Tracer` object, given a list of galaxies it provides all the functionality\n", - " necessary to perform ray-tracing and strong lensing calculations.\n", - " \n", - " One aspect of its design you may not have considered is that the input galaxies can be any size, and it does\n", - " not matter what the galaxies, light or mass profiles are called (e.g. it does not depend on the lens mass\n", - " have the path `galaxies.lens.mass`).\n", - " \n", - " This means that a user can compose a lens model using any combination of light and mass profiles, and the\n", - " `log_likelihood_function` below will still work. You should ensure your `Analysis` class is written generically\n", - " like this. \n", - " \"\"\"\n", - " tracer = al.Tracer(\n", - " galaxies=instance.galaxies,\n", - " cosmology=self.cosmology,\n", - " )\n", - "\n", - " \"\"\"\n", - " You should also be familiar with the `FitImaging` object, which given a tracer and imaging dataset fits the\n", - " tracer's model image to the data, using a chi-squared map to compute the residuals and likelihood.\n", - " \"\"\"\n", - "\n", - " fit = al.FitImaging(dataset=self.dataset, tracer=tracer, xp=self._xp)\n", - "\n", - " \"\"\"\n", - " To get your custom analysis class, running quickly, you may not want to define your own `Fit` class but\n", - " instead just write out manually how the `log_likelihood` is computed. \n", - " \n", - " The commented out code below shows the simplest way to do this, and it is probably suitable for most \n", - " use-cases.\n", - " \n", - " At step-by-step description of what the code is doing is as follows:\n", - " \n", - " 1) Creates an image of the lens and source galaxies from the tracer using its `image_2d_from()` method.\n", - "\n", - " 2) Blurs the tracer`s image with the data's PSF, ensuring the telescope optics are included in the fit. This \n", - " creates what is called the `model_image`.\n", - " \n", - " 3) Computes the difference between this model-image and the observed image, creating the fit`s `residual_map`.\n", - " \n", - " 4) Divides the residual-map by the noise-map, creating the fit`s `normalized_residual_map`.\n", - " \n", - " 5) Squares every value in the normalized residual-map, creating the fit's `chi_squared_map`.\n", - " \n", - " 6) Sums up these chi-squared values and converts them to a `log_likelihood`, which quantifies how good \n", - " this tracer`s fit to the data was (higher log_likelihood = better fit).\n", - " \n", - " Quantities like the `chi_squared_map` and `log_likelihood` are standard quantities used by all model-fitting\n", - " approaches.\n", - " \"\"\"\n", - " # model_data = tracer.blurred_image_2d_from(\n", - " # grid=self.dataset.grid,\n", - " # blurring_grid=self.dataset.grids.blurring,\n", - " # psf=self.dataset.psf,\n", - " # )\n", - "\n", - " # residual_map = self.dataset.data - model_data\n", - " # chi_squared_map = (residual_map / self.dataset.noise_map) ** 2.0\n", - " # chi_squared = sum(chi_squared_map)\n", - " # noise_normalization = np.sum(np.log(2 * np.pi * self.dataset.noise_map**2.0))\n", - " # log_likelihood = -0.5 * (chi_squared + noise_normalization)\n", - "\n", - " \"\"\"\n", - " The `log_likelihood` is returned to the non-linear search, informing it how good a fit this lens model\n", - " was and whether to continue sampling this region of parameter space.\n", - " \"\"\"\n", - "\n", - " return fit.log_likelihood\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis Class Considerations__\n", - "\n", - "Lets quickly think about the design of an `Analysis` class and how this can help us to set up any model-fit we can\n", - "imagine:\n", - "\n", - " - The `__init__` method can be extended to include any data structures needed to perform the analysis. For example, \n", - " the `AnalysisImaging` object in the autolens source code has a `settings` object that customize \n", - " how fits using a `Pixelization` are performed.\n", - " \n", - " - The `log_likelihood_function` can be written in any way that is desired to fit the data. The example above uses\n", - " the `FitImaging` object, but this is not necessary. Furthermore, you could customize this function to assume a \n", - " likelihood function defined by Poisson statistics (the example above assumes Gaussian statistics) or to include\n", - " additional constraints on the model that are specific to your dataset.\n", - "\n", - "__Model Fit__\n", - "\n", - "The standard API for choosing a non-linear search and performing a model-fit can now be used with this `Analysis`\n", - "class." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"custom_analysis\"),\n", - " name=\"strong_lensing_example\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=10000,\n", - ")\n", - "\n", - "# We are using the Analysis class above here!\n", - "\n", - "analysis = AnalysisImaging(dataset=dataset)\n", - "\n", - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Weak Lensing Example__\n", - "\n", - "Now lets consider how to write our own custom `Analysis` class, for the example of performing a weak lensing analysis.\n", - "\n", - "If you are unfamiliar with weak lensing, a brief summary is as follows:\n", - "\n", - " - Weak lensing is the small lensing signal induced into galaxies by lensing due to large-scale structure in the\n", - " universe. \n", - " \n", - "- This signal is much smaller than the strong lensing regime and is often summarised as the small change in the \n", - " ellipticity of a source galaxy's light. \n", - " \n", - "- This change in ellipticity can be measured and is called the `shear`, with the dataset our `Analysis` class will\n", - " fit called a shear catalogue.\n", - "\n", - " - In strong lensing, we typically use the deflection angles of a mass profile to fit the data. For weak lensing\n", - " analysis we compute its shear (via the function `shear_yx_2d_from`) and compare this to the observed shear in the\n", - " shear catalogue data.\n", - "\n", - "__Lens Model__\n", - "\n", - "We first compose our lens model for weak lensing analysis.\n", - "\n", - "This can reuse the **PyAutoLens** API for model composition, but does not require a source galaxy to be included as\n", - "we are simply comparing the mass model shears.`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here is our example `Analysis` class:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class AnalysisShearCatalogue(af.Analysis):\n", - " def __init__(\n", - " self,\n", - " data, # You may wish to group these into a `ShearCatalogue` dataset.\n", - " noise_map,\n", - " grid,\n", - " cosmology: al.cosmo.LensingCosmology = al.cosmo.Planck15(),\n", - " use_jax: bool = True,\n", - " ):\n", - " \"\"\"\n", - " Fits a lens model to a shear catalogue dataset via a non-linear search.\n", - "\n", - " The `Analysis` class defines the `log_likelihood_function` which fits the model to the dataset and returns the\n", - " log likelihood value defining how well the model fitted the data.\n", - "\n", - " It handles many other tasks, such as visualization, outputting results to hard-disk and storing results in\n", - " a format that can be loaded after the model-fit is complete.\n", - "\n", - " This class is used for model-fits which fit strong lenses composed via a `Tracer` to a weak lensing\n", - " shear catalogue dataset.\n", - "\n", - " Parameters\n", - " ----------\n", - " data\n", - " The shear catalogue data.\n", - " noise_map\n", - " An array describing the RMS standard deviation error in each shear measurement point (e.g. the noise-map).\n", - " grid\n", - " The (y,x) coordinates defining where the shears are measured and evaluated (e.g. the locations of the\n", - " galaxies in the shear catalogue).\n", - " cosmology\n", - " The Cosmology assumed for this analysis.\n", - " \"\"\"\n", - " super().__init__(use_jax=use_jax)\n", - "\n", - " self.data = data\n", - " self.noise_map = noise_map\n", - " self.grid = grid\n", - " self.cosmology = cosmology\n", - "\n", - " def log_likelihood_function(self, instance: af.ModelInstance) -> float:\n", - " \"\"\"\n", - " Given an instance of the model, where the model parameters are set via a non-linear search, fit the model\n", - " instance to the imaging dataset.\n", - "\n", - " This function returns a log likelihood which is used by the non-linear search to guide the model-fit.\n", - "\n", - " For this analysis class, this function performs the following steps:\n", - "\n", - " 1) Extracts all galaxies from the model instance and set up a `Tracer`, which includes ordering the galaxies\n", - " by redshift to set up each `Plane`.\n", - "\n", - " 2) Use the `Tracer` to compute the model shear field of the entire strong lensing system.\n", - "\n", - " 3) Compute the shear residuals, a chi-squared statistic and the log likelihood.\n", - "\n", - " Parameters\n", - " ----------\n", - " instance\n", - " An instance of the model that is being fitted to the data by this analysis (whose parameters have been set\n", - " via a non-linear search).\n", - "\n", - " Returns\n", - " -------\n", - " float\n", - " The log likelihood indicating how well this model instance fitted the imaging data.\n", - " \"\"\"\n", - "\n", - " \"\"\"\n", - " For this example, its very easy to compute the model shear field as the `Tracer` object already has this\n", - " functionality built in. \n", - " \"\"\"\n", - " tracer = al.Tracer(\n", - " galaxies=instance.galaxies,\n", - " cosmology=self.cosmology,\n", - " )\n", - "\n", - " model_data = tracer.shear_yx_2d_via_hessian_from(grid=self.grid)\n", - "\n", - " \"\"\"\n", - " We then use this model data and the data itself to compute the residuals, chi-squared and log likelihood.\n", - " \"\"\"\n", - " residual_map = self.data - model_data\n", - " chi_squared_map = (residual_map / self.noise_map) ** 2.0\n", - " chi_squared = sum(chi_squared_map)\n", - " noise_normalization = np.sum(np.log(2 * np.pi * self.noise_map**2.0))\n", - " log_likelihood = -0.5 * (chi_squared + noise_normalization)\n", - "\n", - " return log_likelihood\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "The standard API for choosing a non-linear search and performing a model-fit can now be used with this `Analysis`\n", - "class.\n", - "\n", - "NOTE: Felix can you send me an example shear catalogue so I can get this to run :)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"example_shear_catalogue\"\n", - "dataset_path = Path(\"dataset\") / \"weak_lensing\" / dataset_name\n", - "\n", - "# data = load_shear()\n", - "# noise_map = load_noise_map()\n", - "# grid = load_grid()\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"custom_analysis\"),\n", - " name=\"weak_lensing_example\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=10000,\n", - ")\n", - "\n", - "# We are using the Analysis class above here!\n", - "\n", - "# analysis = AnalysisShearCatalogue(\n", - "# data=data,\n", - "# noise_map=noise_map,\n", - "# grid=grid\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If you are used to using **PyAutoLens**, you'll know that when we run the fit below lots of information about the\n", - "model fit is output to hard-disk (e.g. the best-fit model, error estimates, the model info).\n", - "\n", - "By writing our own `Analysis` class, this is output for free without us having to do anything - pretty cool, huh?\n", - "\n", - "Below, we'll show you how to customize the `Analysis` class even more, to output additional information to hard-disk\n", - "such as visualization and results which you can load elsewhere via the **PyAutoLens** database functionality." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "If you're familiar with **PyAutoLens**'s API, you'll know that the `Result` object returned by the non-linear search\n", - "contains lots of information about the fit. \n", - "\n", - "This includes parameter estimates and errors, details of the non-linear search, etc. \n", - "\n", - "By writing our own `Analysis` class we get all of this information for free, without having to change our code!\n", - "Therefore you should be good to inspect and interpret the results as normal.\n", - "\n", - "The results `info` attribute shows the result in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The result contains the maximum log likelihood instance, which we can use to inspect the result or make plots." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "instance = result.instance\n", - "\n", - "print(f\"Lens Centre = {instance.galaxies.lens.mass.centre}\")\n", - "print(f\"Lens Einstein Radius = {instance.galaxies.lens.mass.einstein_radius}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It also contains information on the posterior as estimated by the non-linear search (in this example `Nautilus`). \n", - "\n", - "Below, we make a corner plot of the \"Probability Density Function\" of every parameter in the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "__To Do List__\n", - "\n", - "The following `Analysis` cookbook from the **PyAutoFit** readthedocs should help you get started customizing your\n", - "own `Analysis` class: \n", - "\n", - "https://pyautofit.readthedocs.io/en/latest/cookbooks/analysis.html\n", - "\n", - "I will extend this guide to include the following in the next few days:\n", - "\n", - " - How to output your own custom visualization.\n", - " - How to extend the `Result` class to include additional information about the model-fit specific to weak lensing \n", - " (e.g. the maximum likelihood shear map).\n", - " - Add methods which output model-specific results to hard-disk in the files folder (e.g. as .json files) to aid in \n", - " the interpretation of results.\n", - " - How to output results to hard-disk in a format that can be loaded into the **PyAutoLens** database.\n", - " " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Misc: Custom Analysis\n", + "=====================\n", + "\n", + "Users familiar with **PyAutoLens** will have seen that `Analysis` classes are used to performed lens modeling of\n", + "different datasets. For example, the `AnalysisImaging` class fits imaging datasets, the `AnalysisInterferometer` class\n", + "fits interferometer datasets.\n", + "\n", + "You may have a dataset which you want to use **PyAutoLenss**'s lensing capabilities to model, but which does not fit\n", + "into one of the standard `Analysis` classes.\n", + "\n", + "A good example (at the time of writing this script) is fitting a weak lensing shear catalogue with a model of the lens\n", + "galaxy's mass. **PyAutoLens** as the lensing capabilities to produce the shears of a mass model, but does not have an\n", + "`Analysis` class to fit these shears to a dataset.\n", + "\n", + "This example demonstrates how you can write your own `Analysis` class to fit a dataset with **PyAutoLens**.\n", + "\n", + "__Contents__\n", + "\n", + "- **PyAutoFit:** The `Analysis` class is the interface between the data and model, whereby a.\n", + "- **Source Code:** This example contains URLs to the locations of the source code of the classes used when creating.\n", + "- **Example Analysis Class:** The `Analysis` classes available in **PyAutoLens** are actually located in both **PyAutoLens** and.\n", + "- **Lens Model:** To illustrate how to write a custom `Analysis` class, we require an example lens model that we will.\n", + "- **Instances:** Instances of the model above can be created, where an input `vector` of parameters is mapped to.\n", + "- **Simple Analysis Example:** For simplicity, a shortened version of an `AnalysisImaging` class is shown below where certain.\n", + "- **Analysis Class Considerations:** Lets quickly think about the design of an `Analysis` class and how this can help us to set up any.\n", + "- **Model Fit:** Perform the model-fit using the search and analysis.\n", + "- **Weak Lensing Example:** Now lets consider how to write our own custom `Analysis` class, for the example of performing a.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **To Do List:** The following `Analysis` cookbook from the **PyAutoFit** readthedocs should help you get started.\n", + "\n", + "__PyAutoFit__\n", + "\n", + "The `Analysis` class is the interface between the data and model, whereby a `log_likelihood_function` is defined\n", + "and called by the non-linear search to fit the model.\n", + "\n", + "You may have performed a similar task yourself, for example by taking a fitting library (e.g. an MCMC method like\n", + "Emcee or nested sampler like Dynesty) and writing a likelihood function that calls it to fit a model to a dataset.\n", + "If you haven't done this, this script will explain how!\n", + "\n", + "**PyAutoLens** uses a library called **PyAutoFit** to set up this interfacebetween the data, model,\n", + "`log_likelihood_function` and non-linear search. **PyAutoFit** is a general purpose library for model fitting,\n", + "and we will see that it has a lot of powerful tools that we can use to customize our `Analysis` class.\n", + "\n", + "You can checkout the **PyAutoFit** readthedocs here:\n", + "\n", + " https://pyautofit.readthedocs.io/en/latest/\n", + "\n", + "The following analysis cookbook provides a concise reference guide to `Analysis` objects, and once you have completed\n", + "this example will be a useful resource for writing your own `Analysis` class:\n", + "\n", + " https://pyautofit.readthedocs.io/en/latest/cookbooks/analysis.html\n", + "\n", + "__Source Code__\n", + "\n", + "This example contains URLs to the locations of the source code of the classes used when creating light and mass\n", + "profiles.\n", + "\n", + "The example itself is standalone and should by the end allow you to implement a custom profile without diving into\n", + "the **PyAutoLens** source code.\n", + "\n", + "We still recommend you take a look to see how things are structured!" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Example Analysis Class__\n", + "\n", + "The `Analysis` classes available in **PyAutoLens** are actually located in both **PyAutoLens** and its parent \n", + "package, **PyAutoGalaxy**: \n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy\n", + " https://github.com/PyAutoLabs/PyAutoLens\n", + "\n", + "All classes used for lens modeling are found in the following packages:\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/tree/main/autogalaxy/imaging/model\n", + " https://github.com/PyAutoLabs/PyAutoLens/tree/main/autolens/imaging/model\n", + "\n", + "The `AnalysisImaging` classes are found in the following modules:\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/imaging/model/analysis.py\n", + " https://github.com/PyAutoLabs/PyAutoLens/blob/main/autolens/imaging/model/analysis.py\n", + "\n", + "__Lens Model__\n", + "\n", + "To illustrate how to write a custom `Analysis` class, we require an example lens model that we will use to fit\n", + "the dataset.\n", + "\n", + "We compose a simple lens model with an `IsothermalSph` mass model for the lens and an `Sersic` for the source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + "# Source:\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.ExponentialCoreSph)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Instances__\n", + "\n", + "Instances of the model above can be created, where an input `vector` of parameters is mapped to create an instance of \n", + "the Python class of the model.\n", + "\n", + "This is used internally by the `Analysis` class we are about to write, and will be used in \n", + "our `log_likelihood_function`. Therefore, we are quickly highlighting it here.\n", + "\n", + "We first need to know the order of parameters in the model, so we know how to define the input `vector`. This\n", + "information is contained in the models `paths` attribute:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.paths)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We input values for the parameters of our model following the order of paths above.\n", + "\n", + "This creates an `instance` of the lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "instance = model.instance_from_vector(vector=[0.0, 0.0, 1.6, 0.1, 0.1, 0.01, 2.0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `instance` contains each of the model components we defined above. \n", + "\n", + "The argument names input into each `Collection` define the attribute names of the `instance`. \n", + "\n", + "For example, when composing the `model` above, we used a `Collection` called `galaxies` which had a `lens` and `source` \n", + "attribute. These `lens` and `source` attributes each contained components called `mass` and `bulge` respectively." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Lens Centre = {instance.galaxies.lens.mass.centre}\")\n", + "print(f\"Lens Einstein Radius = {instance.galaxies.lens.mass.einstein_radius}\")\n", + "print(f\"Source Centre = {instance.galaxies.source.bulge.centre}\")\n", + "print(f\"Source Intensity = {instance.galaxies.source.bulge.intensity}\")\n", + "print(f\"Source Effective Radius = {instance.galaxies.source.bulge.effective_radius}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simple Analysis Example__\n", + "\n", + "For simplicity, a shortened version of an `AnalysisImaging` class is shown below where certain functions have been \n", + "edited to make them easy to read and understand. \n", + "\n", + "This has docstrings updated to focus on the key aspects of implementing a new `Analysis` class and simplifies the \n", + "inheritance structure of the profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class AnalysisImaging(af.Analysis):\n", + " def __init__(\n", + " self,\n", + " dataset: al.Imaging,\n", + " cosmology: al.cosmo.LensingCosmology = al.cosmo.Planck15(),\n", + " use_jax: bool = True,\n", + " ):\n", + " \"\"\"\n", + " Fits a lens model to an imaging dataset via a non-linear search.\n", + "\n", + " The `Analysis` class defines the `log_likelihood_function` which fits the model to the dataset and returns the\n", + " log likelihood value defining how well the model fitted the data.\n", + "\n", + " It handles many other tasks, such as visualization, outputting results to hard-disk and storing results in\n", + " a format that can be loaded after the model-fit is complete.\n", + "\n", + " This class is used for model-fits which fit strong lenses composed via a `Tracer` to an imaging dataset.\n", + " Parameters\n", + " ----------\n", + " dataset\n", + " The `Imaging` dataset that the model is fitted to.\n", + " cosmology\n", + " The Cosmology assumed for this analysis.\n", + " \"\"\"\n", + " super().__init__(use_jax=use_jax)\n", + "\n", + " self.dataset = dataset\n", + " self.cosmology = cosmology\n", + "\n", + " def log_likelihood_function(self, instance: af.ModelInstance) -> float:\n", + " \"\"\"\n", + " Given an instance of the model, where the model parameters are set via a non-linear search, fit the model\n", + " instance to the imaging dataset.\n", + "\n", + " This function returns a log likelihood which is used by the non-linear search to guide the model-fit.\n", + "\n", + " For this analysis class, this function performs the following steps:\n", + "\n", + " 1) Extracts all galaxies from the model instance and set up a `Tracer`, which includes ordering the galaxies\n", + " by redshift to set up each `Plane`.\n", + "\n", + " 2) Use the `Tracer` and other attributes to create a `FitImaging` object, which performs steps such as creating\n", + " model images of every galaxy in the tracer, blurring them with the imaging dataset's PSF and computing\n", + " residuals, a chi-squared statistic and the log likelihood.\n", + "\n", + " Parameters\n", + " ----------\n", + " instance\n", + " An instance of the model that is being fitted to the data by this analysis (whose parameters have been set\n", + " via a non-linear search).\n", + "\n", + " Returns\n", + " -------\n", + " float\n", + " The log likelihood indicating how well this model instance fitted the imaging data.\n", + " \"\"\"\n", + "\n", + " \"\"\"\n", + " The `instance` that comes into this method is an instance of the lens model above, which we illustrated\n", + " via print statements how it is structured.\n", + "\n", + " The parameter values are chosen by the non-linear search, based on where it thinks the high likelihood regions \n", + " of parameter space are.\n", + "\n", + " The lines of Python code are commented out below to prevent excessive print statements when we run the\n", + " non-linear search, but feel free to uncomment them and run the search to see the parameters of every instance\n", + " that it fits.\n", + " \"\"\"\n", + "\n", + " # print(f\"Lens Centre = {instance.galaxies.lens.mass.centre}\")\n", + " # print(f\"Lens Einstein Radius = {instance.galaxies.lens.mass.einstein_radius}\")\n", + " # print(f\"Source Centre = {instance.galaxies.source.bulge.centre}\")\n", + " # print(f\"Source Intensity = {instance.galaxies.source.bulge.intensity}\")\n", + " # print(f\"Source Effective Radius = {instance.galaxies.source.bulge.effective_radius}\")\n", + "\n", + " \"\"\"\n", + " You should be familiar with the `Tracer` object, given a list of galaxies it provides all the functionality\n", + " necessary to perform ray-tracing and strong lensing calculations.\n", + " \n", + " One aspect of its design you may not have considered is that the input galaxies can be any size, and it does\n", + " not matter what the galaxies, light or mass profiles are called (e.g. it does not depend on the lens mass\n", + " have the path `galaxies.lens.mass`).\n", + " \n", + " This means that a user can compose a lens model using any combination of light and mass profiles, and the\n", + " `log_likelihood_function` below will still work. You should ensure your `Analysis` class is written generically\n", + " like this. \n", + " \"\"\"\n", + " tracer = al.Tracer(\n", + " galaxies=instance.galaxies,\n", + " cosmology=self.cosmology,\n", + " )\n", + "\n", + " \"\"\"\n", + " You should also be familiar with the `FitImaging` object, which given a tracer and imaging dataset fits the\n", + " tracer's model image to the data, using a chi-squared map to compute the residuals and likelihood.\n", + " \"\"\"\n", + "\n", + " fit = al.FitImaging(dataset=self.dataset, tracer=tracer, xp=self._xp)\n", + "\n", + " \"\"\"\n", + " To get your custom analysis class, running quickly, you may not want to define your own `Fit` class but\n", + " instead just write out manually how the `log_likelihood` is computed. \n", + " \n", + " The commented out code below shows the simplest way to do this, and it is probably suitable for most \n", + " use-cases.\n", + " \n", + " At step-by-step description of what the code is doing is as follows:\n", + " \n", + " 1) Creates an image of the lens and source galaxies from the tracer using its `image_2d_from()` method.\n", + "\n", + " 2) Blurs the tracer`s image with the data's PSF, ensuring the telescope optics are included in the fit. This \n", + " creates what is called the `model_image`.\n", + " \n", + " 3) Computes the difference between this model-image and the observed image, creating the fit`s `residual_map`.\n", + " \n", + " 4) Divides the residual-map by the noise-map, creating the fit`s `normalized_residual_map`.\n", + " \n", + " 5) Squares every value in the normalized residual-map, creating the fit's `chi_squared_map`.\n", + " \n", + " 6) Sums up these chi-squared values and converts them to a `log_likelihood`, which quantifies how good \n", + " this tracer`s fit to the data was (higher log_likelihood = better fit).\n", + " \n", + " Quantities like the `chi_squared_map` and `log_likelihood` are standard quantities used by all model-fitting\n", + " approaches.\n", + " \"\"\"\n", + " # model_data = tracer.blurred_image_2d_from(\n", + " # grid=self.dataset.grid,\n", + " # blurring_grid=self.dataset.grids.blurring,\n", + " # psf=self.dataset.psf,\n", + " # )\n", + "\n", + " # residual_map = self.dataset.data - model_data\n", + " # chi_squared_map = (residual_map / self.dataset.noise_map) ** 2.0\n", + " # chi_squared = sum(chi_squared_map)\n", + " # noise_normalization = np.sum(np.log(2 * np.pi * self.dataset.noise_map**2.0))\n", + " # log_likelihood = -0.5 * (chi_squared + noise_normalization)\n", + "\n", + " \"\"\"\n", + " The `log_likelihood` is returned to the non-linear search, informing it how good a fit this lens model\n", + " was and whether to continue sampling this region of parameter space.\n", + " \"\"\"\n", + "\n", + " return fit.log_likelihood\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis Class Considerations__\n", + "\n", + "Lets quickly think about the design of an `Analysis` class and how this can help us to set up any model-fit we can\n", + "imagine:\n", + "\n", + " - The `__init__` method can be extended to include any data structures needed to perform the analysis. For example, \n", + " the `AnalysisImaging` object in the autolens source code has a `settings` object that customize \n", + " how fits using a `Pixelization` are performed.\n", + " \n", + " - The `log_likelihood_function` can be written in any way that is desired to fit the data. The example above uses\n", + " the `FitImaging` object, but this is not necessary. Furthermore, you could customize this function to assume a \n", + " likelihood function defined by Poisson statistics (the example above assumes Gaussian statistics) or to include\n", + " additional constraints on the model that are specific to your dataset.\n", + "\n", + "__Model Fit__\n", + "\n", + "The standard API for choosing a non-linear search and performing a model-fit can now be used with this `Analysis`\n", + "class." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"custom_analysis\"),\n", + " name=\"strong_lensing_example\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=10000,\n", + ")\n", + "\n", + "# We are using the Analysis class above here!\n", + "\n", + "analysis = AnalysisImaging(dataset=dataset)\n", + "\n", + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Weak Lensing Example__\n", + "\n", + "Now lets consider how to write our own custom `Analysis` class, for the example of performing a weak lensing analysis.\n", + "\n", + "If you are unfamiliar with weak lensing, a brief summary is as follows:\n", + "\n", + " - Weak lensing is the small lensing signal induced into galaxies by lensing due to large-scale structure in the\n", + " universe. \n", + " \n", + "- This signal is much smaller than the strong lensing regime and is often summarised as the small change in the \n", + " ellipticity of a source galaxy's light. \n", + " \n", + "- This change in ellipticity can be measured and is called the `shear`, with the dataset our `Analysis` class will\n", + " fit called a shear catalogue.\n", + "\n", + " - In strong lensing, we typically use the deflection angles of a mass profile to fit the data. For weak lensing\n", + " analysis we compute its shear (via the function `shear_yx_2d_from`) and compare this to the observed shear in the\n", + " shear catalogue data.\n", + "\n", + "__Lens Model__\n", + "\n", + "We first compose our lens model for weak lensing analysis.\n", + "\n", + "This can reuse the **PyAutoLens** API for model composition, but does not require a source galaxy to be included as\n", + "we are simply comparing the mass model shears.`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here is our example `Analysis` class:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class AnalysisShearCatalogue(af.Analysis):\n", + " def __init__(\n", + " self,\n", + " data, # You may wish to group these into a `ShearCatalogue` dataset.\n", + " noise_map,\n", + " grid,\n", + " cosmology: al.cosmo.LensingCosmology = al.cosmo.Planck15(),\n", + " use_jax: bool = True,\n", + " ):\n", + " \"\"\"\n", + " Fits a lens model to a shear catalogue dataset via a non-linear search.\n", + "\n", + " The `Analysis` class defines the `log_likelihood_function` which fits the model to the dataset and returns the\n", + " log likelihood value defining how well the model fitted the data.\n", + "\n", + " It handles many other tasks, such as visualization, outputting results to hard-disk and storing results in\n", + " a format that can be loaded after the model-fit is complete.\n", + "\n", + " This class is used for model-fits which fit strong lenses composed via a `Tracer` to a weak lensing\n", + " shear catalogue dataset.\n", + "\n", + " Parameters\n", + " ----------\n", + " data\n", + " The shear catalogue data.\n", + " noise_map\n", + " An array describing the RMS standard deviation error in each shear measurement point (e.g. the noise-map).\n", + " grid\n", + " The (y,x) coordinates defining where the shears are measured and evaluated (e.g. the locations of the\n", + " galaxies in the shear catalogue).\n", + " cosmology\n", + " The Cosmology assumed for this analysis.\n", + " \"\"\"\n", + " super().__init__(use_jax=use_jax)\n", + "\n", + " self.data = data\n", + " self.noise_map = noise_map\n", + " self.grid = grid\n", + " self.cosmology = cosmology\n", + "\n", + " def log_likelihood_function(self, instance: af.ModelInstance) -> float:\n", + " \"\"\"\n", + " Given an instance of the model, where the model parameters are set via a non-linear search, fit the model\n", + " instance to the imaging dataset.\n", + "\n", + " This function returns a log likelihood which is used by the non-linear search to guide the model-fit.\n", + "\n", + " For this analysis class, this function performs the following steps:\n", + "\n", + " 1) Extracts all galaxies from the model instance and set up a `Tracer`, which includes ordering the galaxies\n", + " by redshift to set up each `Plane`.\n", + "\n", + " 2) Use the `Tracer` to compute the model shear field of the entire strong lensing system.\n", + "\n", + " 3) Compute the shear residuals, a chi-squared statistic and the log likelihood.\n", + "\n", + " Parameters\n", + " ----------\n", + " instance\n", + " An instance of the model that is being fitted to the data by this analysis (whose parameters have been set\n", + " via a non-linear search).\n", + "\n", + " Returns\n", + " -------\n", + " float\n", + " The log likelihood indicating how well this model instance fitted the imaging data.\n", + " \"\"\"\n", + "\n", + " \"\"\"\n", + " For this example, its very easy to compute the model shear field as the `Tracer` object already has this\n", + " functionality built in. \n", + " \"\"\"\n", + " tracer = al.Tracer(\n", + " galaxies=instance.galaxies,\n", + " cosmology=self.cosmology,\n", + " )\n", + "\n", + " model_data = tracer.shear_yx_2d_via_hessian_from(grid=self.grid)\n", + "\n", + " \"\"\"\n", + " We then use this model data and the data itself to compute the residuals, chi-squared and log likelihood.\n", + " \"\"\"\n", + " residual_map = self.data - model_data\n", + " chi_squared_map = (residual_map / self.noise_map) ** 2.0\n", + " chi_squared = sum(chi_squared_map)\n", + " noise_normalization = np.sum(np.log(2 * np.pi * self.noise_map**2.0))\n", + " log_likelihood = -0.5 * (chi_squared + noise_normalization)\n", + "\n", + " return log_likelihood\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "The standard API for choosing a non-linear search and performing a model-fit can now be used with this `Analysis`\n", + "class.\n", + "\n", + "NOTE: Felix can you send me an example shear catalogue so I can get this to run :)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"example_shear_catalogue\"\n", + "dataset_path = Path(\"dataset\") / \"weak_lensing\" / dataset_name\n", + "\n", + "# data = load_shear()\n", + "# noise_map = load_noise_map()\n", + "# grid = load_grid()\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"custom_analysis\"),\n", + " name=\"weak_lensing_example\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=10000,\n", + ")\n", + "\n", + "# We are using the Analysis class above here!\n", + "\n", + "# analysis = AnalysisShearCatalogue(\n", + "# data=data,\n", + "# noise_map=noise_map,\n", + "# grid=grid\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If you are used to using **PyAutoLens**, you'll know that when we run the fit below lots of information about the\n", + "model fit is output to hard-disk (e.g. the best-fit model, error estimates, the model info).\n", + "\n", + "By writing our own `Analysis` class, this is output for free without us having to do anything - pretty cool, huh?\n", + "\n", + "Below, we'll show you how to customize the `Analysis` class even more, to output additional information to hard-disk\n", + "such as visualization and results which you can load elsewhere via the **PyAutoLens** database functionality." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "If you're familiar with **PyAutoLens**'s API, you'll know that the `Result` object returned by the non-linear search\n", + "contains lots of information about the fit. \n", + "\n", + "This includes parameter estimates and errors, details of the non-linear search, etc. \n", + "\n", + "By writing our own `Analysis` class we get all of this information for free, without having to change our code!\n", + "Therefore you should be good to inspect and interpret the results as normal.\n", + "\n", + "The results `info` attribute shows the result in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The result contains the maximum log likelihood instance, which we can use to inspect the result or make plots." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "instance = result.instance\n", + "\n", + "print(f\"Lens Centre = {instance.galaxies.lens.mass.centre}\")\n", + "print(f\"Lens Einstein Radius = {instance.galaxies.lens.mass.einstein_radius}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It also contains information on the posterior as estimated by the non-linear search (in this example `Nautilus`). \n", + "\n", + "Below, we make a corner plot of the \"Probability Density Function\" of every parameter in the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "__To Do List__\n", + "\n", + "The following `Analysis` cookbook from the **PyAutoFit** readthedocs should help you get started customizing your\n", + "own `Analysis` class: \n", + "\n", + "https://pyautofit.readthedocs.io/en/latest/cookbooks/analysis.html\n", + "\n", + "I will extend this guide to include the following in the next few days:\n", + "\n", + " - How to output your own custom visualization.\n", + " - How to extend the `Result` class to include additional information about the model-fit specific to weak lensing \n", + " (e.g. the maximum likelihood shear map).\n", + " - Add methods which output model-specific results to hard-disk in the files folder (e.g. as .json files) to aid in \n", + " the interpretation of results.\n", + " - How to output results to hard-disk in a format that can be loaded into the **PyAutoLens** database.\n", + " " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/advanced/multi_plane.ipynb b/notebooks/guides/advanced/multi_plane.ipynb index 878659970..da7527143 100644 --- a/notebooks/guides/advanced/multi_plane.ipynb +++ b/notebooks/guides/advanced/multi_plane.ipynb @@ -1,782 +1,819 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Misc: Multi-Plane\n", - "=================\n", - "\n", - "Multi-plane ray-tracing is used when there are more planes than just an image-plane and source-plane. When tracing\n", - "from one plane to another, the redshifts of the different planes must be used to determine scaling factors that are\n", - "applied to the deflection angles.\n", - "\n", - "There are different formalisms for multi-plane ray-tracing, PyAutoLens follows the formalism described in\n", - "this paper: ?.\n", - "\n", - "Examples of multi-plane lensing systems include:\n", - "\n", - " - A standard lens galaxy and source galaxy system, but where there is also a dark matter subhalo whose redshift is\n", - " not at the redshift of the lens galaxy.\n", - "\n", - " - A strong lens system where the deflection due to many dark matter halos down the line-of-sight are included, which\n", - " may be at a large range of different redshifts.\n", - "\n", - " - A galaxy cluster, where the observed different background source galaxies are at a range of different redshifts\n", - " and their deflections due to one another must be included.\n", - "\n", - "__Contents__\n", - "\n", - "- **Example:** To illustrate multi-plane ray-tracing, we first set up a simple lens system, using a `Tracer`.\n", - "- **Ray Tracing:** Multi-plane ray tracing is implemented in the `tracer_util.py` module of the following package.\n", - "- **Profiles With Physical Units:** The above ray-tracing used dimensionless angular units (e.g.\n", - "- **SLACK:** This script was written after discussion on the PyAutoLens Slack channel, where some users modeling.\n", - "\n", - "__Example__\n", - "\n", - "To illustrate multi-plane ray-tracing, we first set up a simple lens system, using a `Tracer` object.\n", - "\n", - "We'll make things simple and assume 3 galaxies at redshifts 0.5, 1.0 and 2.0. We'll use a singular isothermal sphere\n", - "for each galaxy's mass profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from typing import List, Optional, Union\n", - "\n", - "import autoarray as aa\n", - "import autolens as al\n", - "\n", - "lens_0 = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0))\n", - "lens_1 = al.Galaxy(redshift=1.0, mass=al.mp.IsothermalSph(einstein_radius=1.0))\n", - "lens_2 = al.Galaxy(redshift=2.0, mass=al.mp.IsothermalSph(einstein_radius=1.0))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Multi-plane ray tracing is based on the redshifts of the planes that make up the lens system, as opposed to the\n", - "redshifts of the galaxies. \n", - "\n", - "These two things are equivalent, but it means we need to set up the above galaxies as planes in order to perform\n", - "multi-plane ray-tracing." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxies_0 = al.Galaxies(galaxies=[lens_0])\n", - "galaxies_1 = al.Galaxies(galaxies=[lens_1])\n", - "galaxies_2 = al.Galaxies(galaxies=[lens_2])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Multi-plane ray tracing is implemented in the `tracer_util.py` module of the following package:\n", - "\n", - "https://github.com/PyAutoLabs/PyAutoLens/blob/main/autolens/lens/tracer_util.py\n", - "\n", - "It uses the function `traced_grid_2d_list_from`.\n", - "\n", - "Multi-plane ray-tracing also heavily relies on the `scaling_factor_between_redshifts_from` function, which is\n", - "implemented in the `cosmology` package of autolens.\n", - "\n", - "I have copy and pasted both functions below, and put print statements in to show how they works." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def scaling_factor_between_redshifts_from(\n", - " cosmology, redshift_0: float, redshift_1: float, redshift_final: float\n", - ") -> float:\n", - " \"\"\"\n", - " For strong lens systems with more than 2 planes, the deflection angles between different planes must be scaled\n", - " by the angular diameter distances between the planes in order to properly perform multi-plane ray-tracing.\n", - "\n", - " For a system with a first lens galaxy l0 at `redshift_0`, second lens galaxy l1 at `redshift_1` and final\n", - " source galaxy at `redshift_final` this scaling factor is given by:\n", - "\n", - " (D_l0l1 * D_s) / (D_l1* D_l1s)\n", - "\n", - " The critical surface density for lensing, often written as $\\sigma_{cr}$, is given by:\n", - "\n", - " critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l)\n", - "\n", - " D_l0l1 = Angular diameter distance of first lens redshift to second lens redshift.\n", - " D_s = Angular diameter distance of source redshift to earth\n", - " D_l1 = Angular diameter distance of second lens redshift to Earth.\n", - " D_l1s = Angular diameter distance of second lens redshift to source redshift\n", - "\n", - " For systems with more planes this scaling factor is computed multiple times for the different redshift\n", - " combinations and applied recursively when scaling the deflection angles.\n", - "\n", - " Parameters\n", - " ----------\n", - " redshift_0\n", - " The redshift of the first strong lens galaxy.\n", - " redshift_1\n", - " The redshift of the second strong lens galaxy.\n", - " redshift_final\n", - " The redshift of the source galaxy.\n", - " \"\"\"\n", - " angular_diameter_distance_between_redshifts_0_and_1 = (\n", - " cosmology.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_1)\n", - " .to(\"kpc\")\n", - " .value\n", - " )\n", - "\n", - " angular_diameter_distance_to_redshift_final = (\n", - " cosmology.angular_diameter_distance(z=redshift_final).to(\"kpc\").value\n", - " )\n", - "\n", - " angular_diameter_distance_of_redshift_1_to_earth = (\n", - " cosmology.angular_diameter_distance(z=redshift_1).to(\"kpc\").value\n", - " )\n", - "\n", - " angular_diameter_distance_between_redshift_1_and_final = (\n", - " cosmology.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_final)\n", - " .to(\"kpc\")\n", - " .value\n", - " )\n", - "\n", - " return (\n", - " angular_diameter_distance_between_redshifts_0_and_1\n", - " * angular_diameter_distance_to_redshift_final\n", - " ) / (\n", - " angular_diameter_distance_of_redshift_1_to_earth\n", - " * angular_diameter_distance_between_redshift_1_and_final\n", - " )\n", - "\n", - "\n", - "def traced_grid_2d_list_from(\n", - " planes: Union[List[List[al.Galaxy]], List[al.Galaxies]],\n", - " grid: aa.type.Grid2DLike,\n", - " cosmology: al.cosmo.LensingCosmology = al.cosmo.Planck15(),\n", - " plane_index_limit: int = Optional[None],\n", - "):\n", - " \"\"\"\n", - " Returns a ray-traced grid of 2D Cartesian (y,x) coordinates which accounts for multi-plane ray-tracing.\n", - "\n", - " This uses the redshifts and mass profiles of the galaxies contained within the tracer to perform the multi-plane\n", - " ray-tracing calculation.\n", - "\n", - " This function returns a list of 2D (y,x) grids, corresponding to each redshift in the input list of planes. The\n", - " plane redshifts are determined from the redshifts of the galaxies in each plane, whereby there is a unique plane\n", - " at each redshift containing all galaxies at the same redshift.\n", - "\n", - " For example, if the `planes` list contains three lists of galaxies with `redshift`'s z0.5, z=1.0 and z=2.0, the\n", - " returned list of traced grids will contain three entries corresponding to the input grid after ray-tracing to\n", - " redshifts 0.5, 1.0 and 2.0.\n", - "\n", - " An input `AstroPy` cosmology object can change the cosmological model, which is used to compute the scaling\n", - " factors between planes (which are derived from their redshifts and angular diameter distances). It is these\n", - " scaling factors that account for multi-plane ray tracing effects.\n", - "\n", - " The calculation can be terminated early by inputting a `plane_index_limit`. All planes whose integer indexes are\n", - " above this value are omitted from the calculation and not included in the returned list of grids (the size of\n", - " this list is reduced accordingly).\n", - "\n", - " For example, if `planes` has 3 lists of galaxies, but `plane_index_limit=1`, the third plane (corresponding to\n", - " index 2) will not be calculated. The `plane_index_limit` is used to avoid uncessary ray tracing calculations\n", - " of higher redshift planes whose galaxies do not have mass profile (and only have light profiles).\n", - "\n", - " Parameters\n", - " ----------\n", - " galaxies\n", - " The galaxies whose mass profiles are used to perform multi-plane ray-tracing, where the list of galaxies\n", - " has an index for each plane, correspond to each unique redshift in the multi-plane system.\n", - " grid\n", - " The 2D (y, x) coordinates on which multi-plane ray-tracing calculations are performed.\n", - " cosmology\n", - " The cosmology used for ray-tracing from which angular diameter distances between planes are computed.\n", - " plane_index_limit\n", - " The integer index of the last plane which is used to perform ray-tracing, all planes with an index above\n", - " this value are omitted.\n", - "\n", - " Returns\n", - " -------\n", - " traced_grid_list\n", - " A list of 2D (y,x) grids each of which are the input grid ray-traced to a redshift of the input list of planes.\n", - " \"\"\"\n", - "\n", - " traced_grid_list = []\n", - " traced_deflection_list = []\n", - "\n", - " redshift_list = [galaxies[0].redshift for galaxies in planes]\n", - "\n", - " for plane_index, galaxies in enumerate(planes):\n", - " scaled_grid = grid.copy()\n", - "\n", - " if plane_index > 0:\n", - " for previous_plane_index in range(plane_index):\n", - " scaling_factor = cosmology.scaling_factor_between_redshifts_from(\n", - " redshift_0=redshift_list[previous_plane_index],\n", - " redshift_1=galaxies[0].redshift,\n", - " redshift_final=redshift_list[-1],\n", - " )\n", - "\n", - " scaled_deflections = (\n", - " scaling_factor * traced_deflection_list[previous_plane_index]\n", - " )\n", - "\n", - " scaled_grid -= scaled_deflections\n", - "\n", - " traced_grid_list.append(scaled_grid)\n", - "\n", - " if plane_index_limit is not None:\n", - " if plane_index == plane_index_limit:\n", - " return traced_grid_list\n", - "\n", - " deflections_yx_2d = sum(\n", - " map(lambda g: g.deflections_yx_2d_from(grid=scaled_grid), galaxies)\n", - " )\n", - "\n", - " traced_deflection_list.append(deflections_yx_2d)\n", - "\n", - " return traced_grid_list\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Example__\n", - "\n", - "The code below ray-traces a Cartesian coordinate y=1.0\", x=0.0\" to redshift 0.5, 1.0 and 2.0 via multi-plane\n", - "ray-tracing.\n", - "\n", - "The print statements show how the coordinates are transformed as they are ray-traced through each plane and\n", - "therefore how the multi-plane ray-tracing algorithm works." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2DIrregular(values=[(1.0, 0.0)])\n", - "\n", - "traced_grid_2d_list_from(\n", - " planes=[[galaxies_0], [galaxies_1], [galaxies_2]],\n", - " grid=grid,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Profiles With Physical Units__\n", - "\n", - "The above ray-tracing used dimensionless angular units (e.g. the grid was in arc-seconds and mass profile quantities \n", - "like the `einstein_radius` were in arc-seconds).\n", - "\n", - "For certain mass profiles, we define them in physical units (e.g. kpc, solar masses). For example, for the dark matter\n", - "NFW profile called `NFWMCRLudlow` in **PyAutoLens**, it is defined physically with a `mass_at_200` parameter,\n", - "which is the mass in solar masses at which the density profile drops to 200 times the critical density of the Universe.\n", - "\n", - "All internal **PyAutoLens** calculations use dimensionless units, irrespective of whether a mass profile is defined\n", - "in angular dimensionless units of physical units. Therefore, when a physical mass profile is set up, an internal \n", - "conversion is performed which converts its parameters to dimensionless units. This typically requires a cosmology,\n", - "the mass profile redshift and the redshift of the highest redshift plane in the multi-plane system, which are \n", - "often input parameters of physical mass profiles.\n", - "\n", - "For example, when setting up the ``NFWMCRLudlow``'s `mass_at_200`, an internal conversion of this value to the \n", - "dimensionless value used for NFWs, `kappa_s`, is performed. This uses the lens's critical surface mass density, \n", - "`sigma_crit`. This is computed using a cosmology, the NFW redshift and the redshift of the highest redshift \n", - "galaxy (`redshift_source`).\n", - "\n", - "The `scaling_factor` of multi-plane ray-tracing is based on ratios of `sigma_crit` values at different redshifts. For \n", - "NFW profiles in physical units, this can create ambiguity whether the `scaling_factor`'s being applied in multi-plane \n", - "ray-tracing systems are consistent with the `sigma_crit` values used to set up the physical mass profiles values. \n", - "\n", - "**PyAutoLens** uses a convention such that for every physical mass profile in a multi-plane system, their input\n", - "`redshift_source` parameters are the highest redshift plane in the system. When multi-plane ray-tracing algorithm \n", - "computes the `scaling_factor` between planes it correctly scales the lensing parameter (e.g. the `kappa_s` values \n", - "for NFWs), in order to produce the correct deflection angles.\n", - "\n", - "The factor which converts between the physical lens mass and it's lensing strength is sigma_crit. **PyAutoLens**, \n", - "always interprets this with `redshift_source` = `redshift_max_plane`. Therefore, for any profile, if you want the \n", - "projected mass associated with it at some point, you can multiply kappa at that point by sigma_crit(z_profile, z_max).\n", - "\n", - "__SLACK__\n", - "\n", - "This script was written after discussion on the PyAutoLens Slack channel, where some users modeling cluster-scale\n", - "lenses wanted to know how to perform multi-plane ray-tracing in physical units. The following text is the SLACK \n", - "conversion, which if you read in detail should help you fully understand the autolens implementation and details\n", - "of the issue.\n", - "\n", - "\n", - "\n", - "Jack\n", - "\n", - "Hi all - I am working with an undergrad to model cluster-scale lenses with PyAutoLens.\n", - "\n", - "We need to sample in physical units, rather than dimensionless quantities. For example, we want to sample an \n", - "NFW by (log10(M200), c200), rather than (kappa_s, theta_s).\n", - "\n", - "To do so, we will likely be making our own classes along these lines:\n", - "\n", - "class PhysicalNFW(autogalaxy.NFW):\n", - " def __init__(self, center, ell_comps, logM200m, c200m, cosmology):\n", - " kappa_s = ... computing rhos*rs ... / cosmology.critical_density()\n", - " theta_s = ... computing rs ... / cosmology.angular_diamter_distance()\n", - " super(autogalaxy.NFW, self).__init__(\n", - " center=center, ell_comps=ell_comps,\n", - " kappa_s=kappa_s, scale_radius=theta_s\n", - " )\n", - "\n", - "That seems wonderfully nice and simple!\n", - "However, the critical density is ambiguous when we have multiple lens and source planes (our clusters have multiple \n", - "sources at different redshifts). With multiple source planes, which zs should be used to compute the critical density \n", - "to give the right behavior? Is it the redshift of i.e. the next plane after the halo, or the last plane?\n", - "\n", - "\n", - "Jam\n", - "\n", - " However, the critical density is ambiguous when we have multiple lens and source planes (our clusters have multiple \n", - " sources at different redshifts). With multiple source planes, which zs should be used to compute the critical density \n", - " to give the right behavior? Is it the redshift of i.e. the next plane after the halo, or the last plane?\n", - "I have no idea, but I would guess it comes out as a lot of ratios of angular diameter distances. (edited) \n", - "\n", - "\n", - "\n", - "Jack\n", - "\n", - "I think it's an implementation detail of PyAutoLens: when a mass profile defines a bunch of deflection angles, \n", - "which planes does PyAutoLens interpret them as deflections between?\n", - "\n", - "\n", - "Jack\n", - "\n", - "If the mass is in plane i, is it between plane i and i+1?\n", - "\n", - "\n", - "Andrew\n", - "\n", - "So, without being an expert on the internals of PyAutoLens...\n", - "I'm going to use \"critical density\" to \n", - "mean 3H^2/8.pi.G (https://en.wikipedia.org/wiki/Friedmann_equations#Density_parameter), i.e. the 3D density for a \n", - "spatially flat Universe, and \"Sigma_crit\" to mean the critical surface density \n", - "for lensing (https://en.wikipedia.org/wiki/Gravitational_lensing_formalism) (from @Jack's pseudo-code, I'm guessing \n", - "cosmology.critical_density() is Sigma_crit?)\n", - "You will need the critical density at the lens redshift to convert from M200 and c to more physical NFW \n", - "parameters (i.e. rho_0 and R_s here https://en.wikipedia.org/wiki/Navarro%E2%80%93Frenk%E2%80%93White_profile, I \n", - "guess @Jack's rhos, rs).\n", - "\n", - "Projecting this, you get the surface density as a function of position in the image plane (i.e. Sigma(theta))\n", - "One would normally divide by Sigma_crit to get kappa(theta). But it doesn't really make sense to define some \n", - "convergence normalisation (i.e.) kappa_s for autogalaxy's NFW, because with multiple source planes there is no \n", - "one convergence field (because Sigma_crit depends on z_s). So @Jam, do you have a standard way to deal with multiple \n", - "source planes?\n", - "\n", - "That said, the convergences for the different source planes are just re-scaled versions of one another, as are the \n", - "shear and the deflection angles (though not things like the magnification). So, for example: if you want to know the \n", - "mapping from image plane coordinates to (multiple) source plane coordinates over a grid of image plane positions, \n", - "you could calculate a deflection angle field for one source plane (which might be costly, involve numerical \n", - "integrals, etc.) and then the deflection angle for some other source plane can be found easily be re-scaling by the \n", - "ratio of 'Sigma_crit's between the two different source redshifts\n", - "In terms of multiple source planes, @Jack, what is the data you intend to fit to / how do you intend to do your fit? \n", - "By which I mean, people fitting clusters often treat galaxies more like point sources than people fitting to \n", - "galaxy-galaxy strong lensing (the cluster people often just want their lens model to get the different multiple \n", - "images of each multiply imaged background galaxy to map back to common positions behind the cluster, as opposed to \n", - "caring about the structure of each lensed image), but I've typically seen PyAutoLens used to fit an observed image \n", - "pixel-by-pixel (often with a pixelised source reconstruction). If you intend to do the latter (with pixelised sources) \n", - "then you would have multiple pixelised source planes, each with their own regularisation, and (I imagine) the linear \n", - "algebra to find the most likely set of pixel fluxes across all the source planes would be rather difficult.\n", - "Not sure that will have helped, but my two cents...\n", - "\n", - "WikipediaWikipedia\n", - "Friedmann equations\n", - "https://en.wikipedia.org/wiki/Friedmann_equations#Density_parameter\n", - "\n", - "\n", - "WikipediaWikipedia\n", - "Gravitational lensing formalism\n", - "https://en.wikipedia.org/wiki/Gravitational_lensing_formalism\n", - "\n", - "\n", - "WikipediaWikipedia\n", - "Navarro\u2013Frenk\u2013White profile\n", - "\n", - "Jack\n", - "\n", - "Hi \n", - "@Andrew, thanks for this! I agree with all of this: my concern is that since the critical surface density depends on zs, \n", - "if there are multiple zses, it is ambiguous which one to use.\n", - "\n", - "For now we are using an observed-position likelihood as you point out is the standard with clusters, which James has \n", - "implemented and is working fine. The system we're looking at has two sources at different redshifts, but one is \n", - "clearly visible and the other someone barely noticed an emission line in MUSE data.\n", - "\n", - "For the more obvious source, we may end up doing a pixel reconstruction. (In which case we would probably ignore the \n", - "other source). Andrew Newman was able to do so with a double sersic model in \n", - "this paper: https://ui.adsabs.harvard.edu/abs/2018ApJ...862..125N/abstract\n", - "\n", - "\n", - "Jam\n", - "\n", - "The multi-plane implementation \n", - "follows (section 2): https://arxiv.org/abs/1403.5278 [NOTE TO READER, I WAS WRONG ABOUT THIS, DIFFERENT CORRECT PAPER LINKED TO BELOW]\n", - "\n", - "I can provide links to the source code, but basically you compute scaling factors (beta in equation 5) based on \n", - "angulars of diameter distances and then apply them when doing the multi-plane tracing the image-plane to each plane one \n", - "after another. The ray-tracing is recursive in that you go from the image-plane to each source-plane one-by-one I believe.\n", - "\n", - "\n", - "Then the deflection angle for some other source plane can be found easily be re-scaling by the ratio of 'Sigma_crit's \n", - "between the two different source redshifts\n", - "\n", - "I'm going to hazard a guess that equation (5) can be rewritten as a ratio of sigma_crit values (e.g. via equation 4, \n", - "provided D_l1 = D_l2). We basically then just need the code to use the sigma_crit values computed specifically for \n", - "the NFW's (which are related to kappa_s) , when computing the beta values, instead of how the values are computed currently?\n", - "\n", - "\n", - "Jack\n", - "\n", - "(You can even use the ratios of sigma_crit as a probe of cosmology! https://arxiv.org/abs/2110.06232)\n", - "\n", - "For now, the question is: which zsource is correct to use? Is it the redshift of the first source, or the last one?\n", - "\n", - "\n", - "Jam\n", - "\n", - "@Qiuhan He When simulating lenses with many DM subhalos (e.g. for the ABC paper) did we account for how to treat the \n", - "source redshift when computing their sigma_crit but also how to set their mass parameters via the critical density \n", - "of the Universe?\n", - "\n", - "\n", - "Qiuhan He\n", - "\n", - "If we want to compute how the first source is lensed by an NFW halo, then the source redshift is the first source's \n", - "redshift. If we want to compute the lensing quantities about the second source, then the source redshift should be \n", - "the second source's.\n", - "\n", - "The critical density needed for an NFW halo is the critical density of the Universe at the redshift of the halo\n", - "\n", - "The PhysicalNFW defined should involve one more parameter called source_redshift, because the lensing quantity of an NFW halo would change with different background sources at different redshifts\n", - "\n", - "Jam\n", - "\n", - "When simulating lenses for the ABC paper, we had 100s of NFWs at different redshifts making up the line of sight.\n", - "\n", - "I am assuming we always used 'redshift_source' as the source redshift for every NFW, which was what went into \n", - "computing its sigma_crit value.\n", - "\n", - "When performing multi plane ray tracing, we used the redshift of the NFW we were computing deflections of, \n", - "as well as all other NFWs.\n", - "\n", - "Independently, I know these two calculations are valid. What I'm unclear on is whether combining them in the way \n", - "we did is valid (e.g. is the sigma_crit we compute when using collosus defined the same as the ones used for scaling \n", - "between different planes for multi plane ray tracing?l\n", - "\n", - "\n", - "Qiuhan He\n", - "\n", - "yes\n", - "\n", - "I need to find the multiplane ray tracing equation in Schneider+1992\n", - "\n", - "PDF\n", - "\n", - "Here, Qiuhan linked to section \"9.1 The multiple lens-plane theory\" of Schneider 1992: \n", - "https://ui.adsabs.harvard.edu/abs/1992grle.book.....S/abstract\n", - "\n", - "Eq. (9.7b) is the multiplane ray-tracing equation. The dimensionless deflection angle of each plane is defined as \n", - "Eq. (9.6), which is the way AutoLens implements.\n", - "\n", - "\n", - "Jack\n", - "\n", - "It sounds like what you are saying is that for a mass in plane i, AutoLens interprets it's deflection angles as \n", - "between plane i+1 and plane i. So the Sigma_crit is a function of z_i and z_{i+1}. Is that correct? Elsewhere \n", - "James said something implying that it is betwen plane i and i_{max}\n", - "\n", - "This whole approach of sampling unitless quantities is cute and simple with one lens/source plane, but with multiple \n", - "mass/source planes like this, or if you actually care about the physical quantities of the halos, I have always \n", - "thought it would make much more sense to only sample physical parameters of each halo and have your ray-tracing code \n", - "do everything correctly under the hood\n", - "\n", - "\n", - "Qiuhan He\n", - "\n", - "The sigma_crit is what James described as shown by Eq.(9.4).\n", - "\n", - "\n", - "Jack\n", - "\n", - "I understand what sigma_crit is.\n", - "\n", - "My point of confusion is the following: sigma_crit is a function of cosmology, lens redshift (z_l), and source \n", - "redshift (z_s). When you have one lens plane and one source plane, parametrizing everything with kappa is clear and \n", - "unambiguous, but when you have multiple lens and source redshifts, it becomes ambiguous.\n", - "\n", - "Choosing which z_l and z_s to use to define kappa then becomes an implementation choice: you could uses the lowest \n", - "redshift z_s, or the highest redshift z_s . I suppose you could do anything in between but that would be a pretty \n", - "pathological. My question is really what choice does PyAutoLens make here?\n", - "\n", - "AutoLens parameterizes an NFW as kappa_s and theta_scale, the angle of the scale radius. The convergence at the \n", - "scale radius kappa_s is only a meaningful quantity between a specific lens redshift, source redshift, and cosmology. \n", - "\n", - "So what sigma_crit does kappa_s correspond to? What physical mass surface density \\Sigma = \\Sigma_{crit} \\kappa is causing the lensing?\n", - "\n", - "\n", - "Qiuhan He\n", - "\n", - "The NFW profile parameterized by kappa_s and theta_scale has no physical meaning. It can be a halo of mass A \n", - "in one strong lensing system or a halo of mass 2A in another strong lensing system whose sigma_crit is twice.\n", - "\n", - "It is only a model for unitless lensing computation.\n", - "\n", - "We have a profile called NFWMCRLudlow . To use that, one needs to specify the redshift of the halo and the source by \n", - "redshift_object and redshift_source.\n", - "\n", - "In multiplane lensing, we set redshift_object to be the redshift of the i-th lensing plane and redshift_source \n", - "to be the source galaxy redshift\n", - "\n", - "\n", - "Jack\n", - "\n", - "So I have the opposite problem: We have the physical parameters (Mass, concentration) and need to get the correct \n", - "lensing from that mass.\n", - "\n", - "Ahh so the ambiguity I am struggling with is that we have multiple source galaxies at different redshifts\n", - "\n", - "James sent me a code snippet where it appears that the multi-plane lensing always interprets the deflection angles as \n", - "between the i'th plane and the highest-redshift plane. That would answer my question\n", - "\n", - "\n", - "Jam\n", - "\n", - "I don't think it is necessary to think about whether a plane has a \"source galaxy\" (in the sense that you have \n", - "observed imaging data of it and want to model or analyse it).\n", - "\n", - "The multi-plane ray tracing calculations do not care if a plane has a \"source galaxy\" or if its just another plane \n", - "with a galaxy with mass in. The deflection angles would not change if you added a \"source galaxy\" to a plane in a \n", - "multi-plane system which previously only had mass components.\n", - "\n", - "Andrew\n", - "@Jam, @Qiuhan He have we ever dealt with multiple source planes, or just multiple lens planes? \n", - "@Jack's sentiment that defining a lens by its convergence (and not it's physical mass distribution) is ambiguous when \n", - "there are multiple source planes is something I very much agree with (well, one cannot disagree, it is clearly true!)\n", - "\n", - "\n", - "Jam\n", - "\n", - "I have used autolens to model multiple source plane systems, but never with physical units throw into the mixer as well\n", - "\n", - "\n", - "Andrew\n", - "\n", - "But \"without physical units\" sounds ill defined. If we take a case of a single lens plane (with an NFW) and two \n", - "source planes, what should kappa_s and r_s be? r_s (being an angle to the lens plane) is well defined, but kappa_s \n", - "is different depending on which of the two source planes is being considered\n", - "\n", - "\n", - "Qiuhan He\n", - "\n", - "Ah. I get it. \n", - "\n", - "@Jack, for your purpose, ray-tracing for multiple source galaxies, please compute the kappa_s using the last source \n", - "galaxy's redshift\n", - "\n", - "\n", - "Jack\n", - "\n", - "great, thank you \n", - "\n", - "Qiuhan He\n", - "\n", - "Autolens will rescale the kappa_s for source galaxies infront of it. (I think autolens is actually rescaling the \n", - "deflection angles instead) (edited) \n", - "\n", - "\n", - "Jam\n", - "\n", - "Ok, I'm gonna write this somewhere permenant.. [Oh look I did lol]\n", - "\n", - "\n", - "\n", - "Qiuhan He\n", - "\n", - "Thats how I understand Autolens is doing for multiplane raytracing\n", - "\n", - "\n", - "Andrew\n", - "\n", - "Sounds like we have an answer :grinning:\n", - "\n", - "\n", - "Jack\n", - "\n", - "but yes @Andrew that is exactly my concern! in my case we are always concerned with the physical mass of the lenses. \n", - "so for us it never makes sense to sample \"lensing units\" that can then be rescaled for different redshifts: we always \n", - "want a fixed well-defined cosmology with known redshifts for each source\n", - "\n", - "(I suppose we could vary the cosmology but that doesn't change the principle here)\n", - "\n", - "\n", - "Andrew\n", - "\n", - "Perhaps best to check what \n", - "@Qiuhan He described with a few simple cases (like do some analysis with a single source plane, then re-do it adding\n", - "a second source plane at lower-z with no light and see that nothing changes; and various other variants of this)\n", - "\n", - "\n", - "Jam\n", - "\n", - "Are you assuming your source redshifts are always known exactly?\n", - "\n", - "\n", - "Jack\n", - "\n", - "In our case the source redshifts are known to high precision, we don't anticipate working with systems of unknown redshift\n", - "\n", - "\n", - "Jack\n", - "\n", - "@Andrew James sent me some example code doing something similar to what you described, I'll play with it to \n", - "double-check everything!\n", - "\n", - "and \n", - "@Jam re: source redshift precision: there may be some cluster systems where several sources have known redshift, and \n", - "some do not (or have e.g. photo-zs not spec-zs).\n", - "\n", - "in that case we would like to sample the redshifts of the sources with poorly-constrained redshift. the sources with \n", - "known redshift serve to constrain the lens mass, and if the lens mass is understood well enough it actually can \n", - "constrain the redshift of the other sources. (basically, at what z_source does this source's observed lensing line up)\n", - "\n", - "This is something that's actually done with huge clusters, where it's hard to get speczs on every object. I think \n", - "I've seen this with some of the crazy massive clusters observed with JWST\n", - "\n", - "\n", - "Andrew\n", - "\n", - "A quick comment (which doesn't actually suggest how things would be best implemented) and is written with explicit \n", - "reference to NFW profiles, but applies to defining lenses by their convergence (not their mass) more generally...\n", - "\n", - "Presumably, the reason why PyAutoLens uses (by default) kappa_s and theta_s [rs in the code, but it sounds like \n", - "it is an angle] (not physical things like M200, c and z_l; or rho_s, r_s and z_l) is that they uniquely specify the \n", - "mapping from image plane to source plane coordinates (i.e. the \"deflection angles\"). Whereas if specifying the \n", - "3 numbers describing the physical mass distribution (how much it weighs, how concentrated it is, and how far along \n", - "the line-of-sight it is) there are different combinations that give you the same deflection angles (i.e. the same \n", - "mapping from positions in the image back to positions in the source plane). This means that, particularly \n", - "when z_l and/or z_s is unknown, it makes sense to fit for kappa_s and theta_s. Of course, if we know the redshifts, \n", - "then whether we fit for (kappa_s, theta_s), (M200, c), or (rho_s, r_s) doesn't really matter. \n", - "\n", - "The point I am trying to make, is that there are (in a single lens plane, single source plane case) good reasons to \n", - "work directly with the convergence, rather than thinking first about the physical mass distribution, and then using \n", - "redshifts to turn this into kappa(theta), from which we get alpha(theta) and hence the mapping from image plane to \n", - "source plan positions that we need to \"do the lensing\".\n", - "\n", - "I think @Jack's multi-source-plane case has caused some confusion because it breaks the property we had with a s\n", - "ingle source plane, that we could have got the same lensing effect from different mass distributions at different \n", - "redshifts. As Qiuhan said earlier The NFW profile parameterized by kappa_s and theta_scale has no physical meaning. \n", - "It can be a halo of mass A in one strong lensing system or a halo of mass 2A in another strong lensing system whose \n", - "sigma_crit is twice. But with two source planes, a change in z_l that doubles sigma_crit to one source plane, \n", - "might only increase sigma_crit by a factor of 1.5 to the other source plane. So there is not a sense in which we \n", - "can just fit for some \"dimensionless lensing parameters\" and then later use the lens redshift to convert this to \n", - "physical lens parameters if we are so inclined. We need to know z_l in order to know how the convergence of the \n", - "lens for source plane 1 relates to the convergence of the lens for source plane 2. And if we know the lens and \n", - "source redshifts, we may as well parameterise the lens by physical parameters, rather than kappa_s and theta_s \n", - "(because the benefit of kappa_s and theta_s has gone).\n", - "\n", - "Coming from a simulation background, I definitely think of physical mass distributions first and then their lensing \n", - "effect second. But I think the \"kappa(theta) is primary, and we can worry about what actual mass distribution this \n", - "corresponds to later\" approach does make sense in some observational cases. And I'm hoping that hearing those two \n", - "perspectives might help someone (or maybe everyone already knew this! :melting_face:) (edited) \n", - "\n", - "\n", - "\n", - "Jack\n", - "\n", - "Thanks for the reasoned description Andrew! You are of course right that both approaches make sense in \n", - "different cases - I apologize if I was being uncharitable in describing the \"lensing units\" approach before.\n", - "\n", - "I think this is a frequent division between codes for \"galaxy-scale\" lensing and \"cluster-scale\" lensing: with \n", - "galaxy-galaxy lensing often the \"lensing units\" modeling is simple and works great, but with cluster-scale lensing, \n", - "with may different mass components and multiple source redshifts, the formalism behind the \"lensing units\" becomes \n", - "confusing just how you describe and it makes more sense to parametrize by physical parameters.\n", - "\n", - "The property that breaks in multiple-source lensing, as you describe, actually has a few interesting science corollaries:\n", - "The mass-sheet degeneracy is broken. With galaxy-scale lensing for cosmology, people fret quite a bit about the \n", - "mass sheet degeneracy, but with multiple sources probing different sigma_crit's you have a lever arm to break this \n", - "and constrain the physical mass distribution. This is a reason people don't talk much about the mass-sheet \n", - "degeneracy in cluster-scale lensing\n", - "\n", - "You can actually constrain cosmology with the lensing strength between planes! the ratio of lensing strength between \n", - "planes is expressed as beta (below) for each pair of planes\n", - "\n", - "it is a strict function of geometry (distances) - so comparing predicted to measured values of beta is a method for \n", - "doing cosmography. You can sample over a cosmology for this, and have a similar constraining mechanism to BAO \n", - "measurements. It's basically ratios of angular diameter distances\n", - "\n", - "Here is an example paper doing # 2 with a small sample of cluster-scale lenses. We are currently waiting on \n", - "spectroscopic follow-up of one golden system with many background sources, where we'd like to do this measurement.\n", - "https://arxiv.org/abs/2110.06232\n", - "\n", - "arXiv.orgarXiv.org\n", - "Galaxy cluster strong lensing cosmography: cosmological constraints from a sample of regular galaxy clusters\n", - "Cluster strong lensing cosmography is a promising probe of the background geometry of the Universe and several studies \n", - "have emerged, thanks to the increased quality of observations using space and ground-based telescopes. For the first \n", - "time, we use a sample of five cluster strong lenses to measure the values of cosmological parameters and combine them \n", - "with those from classical probes. In order to assess the degeneracies and the effectiveness of strong-lensing \n", - "cosmography in constraining the background geometry of the Universe, we adopt four cosmological scenarios. We find \n", - "good constraining power on the total matter density of the Universe ($\u03a9_{\\rm m}$) and the equation of state of the dark e\u2026 Show more" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Misc: Multi-Plane\n", + "=================\n", + "\n", + "Multi-plane ray-tracing is used when there are more planes than just an image-plane and source-plane. When tracing\n", + "from one plane to another, the redshifts of the different planes must be used to determine scaling factors that are\n", + "applied to the deflection angles.\n", + "\n", + "There are different formalisms for multi-plane ray-tracing, PyAutoLens follows the formalism described in\n", + "this paper: ?.\n", + "\n", + "Examples of multi-plane lensing systems include:\n", + "\n", + " - A standard lens galaxy and source galaxy system, but where there is also a dark matter subhalo whose redshift is\n", + " not at the redshift of the lens galaxy.\n", + "\n", + " - A strong lens system where the deflection due to many dark matter halos down the line-of-sight are included, which\n", + " may be at a large range of different redshifts.\n", + "\n", + " - A galaxy cluster, where the observed different background source galaxies are at a range of different redshifts\n", + " and their deflections due to one another must be included.\n", + "\n", + "__Contents__\n", + "\n", + "- **Example:** To illustrate multi-plane ray-tracing, we first set up a simple lens system, using a `Tracer`.\n", + "- **Ray Tracing:** Multi-plane ray tracing is implemented in the `tracer_util.py` module of the following package.\n", + "- **Profiles With Physical Units:** The above ray-tracing used dimensionless angular units (e.g.\n", + "- **SLACK:** This script was written after discussion on the PyAutoLens Slack channel, where some users modeling.\n", + "\n", + "__Example__\n", + "\n", + "To illustrate multi-plane ray-tracing, we first set up a simple lens system, using a `Tracer` object.\n", + "\n", + "We'll make things simple and assume 3 galaxies at redshifts 0.5, 1.0 and 2.0. We'll use a singular isothermal sphere\n", + "for each galaxy's mass profile." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from typing import List, Optional, Union\n", + "\n", + "import autoarray as aa\n", + "import autolens as al\n", + "\n", + "lens_0 = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0))\n", + "lens_1 = al.Galaxy(redshift=1.0, mass=al.mp.IsothermalSph(einstein_radius=1.0))\n", + "lens_2 = al.Galaxy(redshift=2.0, mass=al.mp.IsothermalSph(einstein_radius=1.0))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Multi-plane ray tracing is based on the redshifts of the planes that make up the lens system, as opposed to the\n", + "redshifts of the galaxies. \n", + "\n", + "These two things are equivalent, but it means we need to set up the above galaxies as planes in order to perform\n", + "multi-plane ray-tracing." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxies_0 = al.Galaxies(galaxies=[lens_0])\n", + "galaxies_1 = al.Galaxies(galaxies=[lens_1])\n", + "galaxies_2 = al.Galaxies(galaxies=[lens_2])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Multi-plane ray tracing is implemented in the `tracer_util.py` module of the following package:\n", + "\n", + "https://github.com/PyAutoLabs/PyAutoLens/blob/main/autolens/lens/tracer_util.py\n", + "\n", + "It uses the function `traced_grid_2d_list_from`.\n", + "\n", + "Multi-plane ray-tracing also heavily relies on the `scaling_factor_between_redshifts_from` function, which is\n", + "implemented in the `cosmology` package of autolens.\n", + "\n", + "I have copy and pasted both functions below, and put print statements in to show how they works." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def scaling_factor_between_redshifts_from(\n", + " cosmology, redshift_0: float, redshift_1: float, redshift_final: float\n", + ") -> float:\n", + " \"\"\"\n", + " For strong lens systems with more than 2 planes, the deflection angles between different planes must be scaled\n", + " by the angular diameter distances between the planes in order to properly perform multi-plane ray-tracing.\n", + "\n", + " For a system with a first lens galaxy l0 at `redshift_0`, second lens galaxy l1 at `redshift_1` and final\n", + " source galaxy at `redshift_final` this scaling factor is given by:\n", + "\n", + " (D_l0l1 * D_s) / (D_l1* D_l1s)\n", + "\n", + " The critical surface density for lensing, often written as $\\sigma_{cr}$, is given by:\n", + "\n", + " critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l)\n", + "\n", + " D_l0l1 = Angular diameter distance of first lens redshift to second lens redshift.\n", + " D_s = Angular diameter distance of source redshift to earth\n", + " D_l1 = Angular diameter distance of second lens redshift to Earth.\n", + " D_l1s = Angular diameter distance of second lens redshift to source redshift\n", + "\n", + " For systems with more planes this scaling factor is computed multiple times for the different redshift\n", + " combinations and applied recursively when scaling the deflection angles.\n", + "\n", + " Parameters\n", + " ----------\n", + " redshift_0\n", + " The redshift of the first strong lens galaxy.\n", + " redshift_1\n", + " The redshift of the second strong lens galaxy.\n", + " redshift_final\n", + " The redshift of the source galaxy.\n", + " \"\"\"\n", + " angular_diameter_distance_between_redshifts_0_and_1 = (\n", + " cosmology.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_1)\n", + " .to(\"kpc\")\n", + " .value\n", + " )\n", + "\n", + " angular_diameter_distance_to_redshift_final = (\n", + " cosmology.angular_diameter_distance(z=redshift_final).to(\"kpc\").value\n", + " )\n", + "\n", + " angular_diameter_distance_of_redshift_1_to_earth = (\n", + " cosmology.angular_diameter_distance(z=redshift_1).to(\"kpc\").value\n", + " )\n", + "\n", + " angular_diameter_distance_between_redshift_1_and_final = (\n", + " cosmology.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_final)\n", + " .to(\"kpc\")\n", + " .value\n", + " )\n", + "\n", + " return (\n", + " angular_diameter_distance_between_redshifts_0_and_1\n", + " * angular_diameter_distance_to_redshift_final\n", + " ) / (\n", + " angular_diameter_distance_of_redshift_1_to_earth\n", + " * angular_diameter_distance_between_redshift_1_and_final\n", + " )\n", + "\n", + "\n", + "def traced_grid_2d_list_from(\n", + " planes: Union[List[List[al.Galaxy]], List[al.Galaxies]],\n", + " grid: aa.type.Grid2DLike,\n", + " cosmology: al.cosmo.LensingCosmology = al.cosmo.Planck15(),\n", + " plane_index_limit: int = Optional[None],\n", + "):\n", + " \"\"\"\n", + " Returns a ray-traced grid of 2D Cartesian (y,x) coordinates which accounts for multi-plane ray-tracing.\n", + "\n", + " This uses the redshifts and mass profiles of the galaxies contained within the tracer to perform the multi-plane\n", + " ray-tracing calculation.\n", + "\n", + " This function returns a list of 2D (y,x) grids, corresponding to each redshift in the input list of planes. The\n", + " plane redshifts are determined from the redshifts of the galaxies in each plane, whereby there is a unique plane\n", + " at each redshift containing all galaxies at the same redshift.\n", + "\n", + " For example, if the `planes` list contains three lists of galaxies with `redshift`'s z0.5, z=1.0 and z=2.0, the\n", + " returned list of traced grids will contain three entries corresponding to the input grid after ray-tracing to\n", + " redshifts 0.5, 1.0 and 2.0.\n", + "\n", + " An input `AstroPy` cosmology object can change the cosmological model, which is used to compute the scaling\n", + " factors between planes (which are derived from their redshifts and angular diameter distances). It is these\n", + " scaling factors that account for multi-plane ray tracing effects.\n", + "\n", + " The calculation can be terminated early by inputting a `plane_index_limit`. All planes whose integer indexes are\n", + " above this value are omitted from the calculation and not included in the returned list of grids (the size of\n", + " this list is reduced accordingly).\n", + "\n", + " For example, if `planes` has 3 lists of galaxies, but `plane_index_limit=1`, the third plane (corresponding to\n", + " index 2) will not be calculated. The `plane_index_limit` is used to avoid uncessary ray tracing calculations\n", + " of higher redshift planes whose galaxies do not have mass profile (and only have light profiles).\n", + "\n", + " Parameters\n", + " ----------\n", + " galaxies\n", + " The galaxies whose mass profiles are used to perform multi-plane ray-tracing, where the list of galaxies\n", + " has an index for each plane, correspond to each unique redshift in the multi-plane system.\n", + " grid\n", + " The 2D (y, x) coordinates on which multi-plane ray-tracing calculations are performed.\n", + " cosmology\n", + " The cosmology used for ray-tracing from which angular diameter distances between planes are computed.\n", + " plane_index_limit\n", + " The integer index of the last plane which is used to perform ray-tracing, all planes with an index above\n", + " this value are omitted.\n", + "\n", + " Returns\n", + " -------\n", + " traced_grid_list\n", + " A list of 2D (y,x) grids each of which are the input grid ray-traced to a redshift of the input list of planes.\n", + " \"\"\"\n", + "\n", + " traced_grid_list = []\n", + " traced_deflection_list = []\n", + "\n", + " redshift_list = [galaxies[0].redshift for galaxies in planes]\n", + "\n", + " for plane_index, galaxies in enumerate(planes):\n", + " scaled_grid = grid.copy()\n", + "\n", + " if plane_index > 0:\n", + " for previous_plane_index in range(plane_index):\n", + " scaling_factor = cosmology.scaling_factor_between_redshifts_from(\n", + " redshift_0=redshift_list[previous_plane_index],\n", + " redshift_1=galaxies[0].redshift,\n", + " redshift_final=redshift_list[-1],\n", + " )\n", + "\n", + " scaled_deflections = (\n", + " scaling_factor * traced_deflection_list[previous_plane_index]\n", + " )\n", + "\n", + " scaled_grid -= scaled_deflections\n", + "\n", + " traced_grid_list.append(scaled_grid)\n", + "\n", + " if plane_index_limit is not None:\n", + " if plane_index == plane_index_limit:\n", + " return traced_grid_list\n", + "\n", + " deflections_yx_2d = sum(\n", + " map(lambda g: g.deflections_yx_2d_from(grid=scaled_grid), galaxies)\n", + " )\n", + "\n", + " traced_deflection_list.append(deflections_yx_2d)\n", + "\n", + " return traced_grid_list\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Example__\n", + "\n", + "The code below ray-traces a Cartesian coordinate y=1.0\", x=0.0\" to redshift 0.5, 1.0 and 2.0 via multi-plane\n", + "ray-tracing.\n", + "\n", + "The print statements show how the coordinates are transformed as they are ray-traced through each plane and\n", + "therefore how the multi-plane ray-tracing algorithm works." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2DIrregular(values=[(1.0, 0.0)])\n", + "\n", + "traced_grid_2d_list_from(\n", + " planes=[[galaxies_0], [galaxies_1], [galaxies_2]],\n", + " grid=grid,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Profiles With Physical Units__\n", + "\n", + "The above ray-tracing used dimensionless angular units (e.g. the grid was in arc-seconds and mass profile quantities \n", + "like the `einstein_radius` were in arc-seconds).\n", + "\n", + "For certain mass profiles, we define them in physical units (e.g. kpc, solar masses). For example, for the dark matter\n", + "NFW profile called `NFWMCRLudlow` in **PyAutoLens**, it is defined physically with a `mass_at_200` parameter,\n", + "which is the mass in solar masses at which the density profile drops to 200 times the critical density of the Universe.\n", + "\n", + "All internal **PyAutoLens** calculations use dimensionless units, irrespective of whether a mass profile is defined\n", + "in angular dimensionless units of physical units. Therefore, when a physical mass profile is set up, an internal \n", + "conversion is performed which converts its parameters to dimensionless units. This typically requires a cosmology,\n", + "the mass profile redshift and the redshift of the highest redshift plane in the multi-plane system, which are \n", + "often input parameters of physical mass profiles.\n", + "\n", + "For example, when setting up the ``NFWMCRLudlow``'s `mass_at_200`, an internal conversion of this value to the \n", + "dimensionless value used for NFWs, `kappa_s`, is performed. This uses the lens's critical surface mass density, \n", + "`sigma_crit`. This is computed using a cosmology, the NFW redshift and the redshift of the highest redshift \n", + "galaxy (`redshift_source`).\n", + "\n", + "The `scaling_factor` of multi-plane ray-tracing is based on ratios of `sigma_crit` values at different redshifts. For \n", + "NFW profiles in physical units, this can create ambiguity whether the `scaling_factor`'s being applied in multi-plane \n", + "ray-tracing systems are consistent with the `sigma_crit` values used to set up the physical mass profiles values. \n", + "\n", + "**PyAutoLens** uses a convention such that for every physical mass profile in a multi-plane system, their input\n", + "`redshift_source` parameters are the highest redshift plane in the system. When multi-plane ray-tracing algorithm \n", + "computes the `scaling_factor` between planes it correctly scales the lensing parameter (e.g. the `kappa_s` values \n", + "for NFWs), in order to produce the correct deflection angles.\n", + "\n", + "The factor which converts between the physical lens mass and it's lensing strength is sigma_crit. **PyAutoLens**, \n", + "always interprets this with `redshift_source` = `redshift_max_plane`. Therefore, for any profile, if you want the \n", + "projected mass associated with it at some point, you can multiply kappa at that point by sigma_crit(z_profile, z_max).\n", + "\n", + "__SLACK__\n", + "\n", + "This script was written after discussion on the PyAutoLens Slack channel, where some users modeling cluster-scale\n", + "lenses wanted to know how to perform multi-plane ray-tracing in physical units. The following text is the SLACK \n", + "conversion, which if you read in detail should help you fully understand the autolens implementation and details\n", + "of the issue.\n", + "\n", + "\n", + "\n", + "Jack\n", + "\n", + "Hi all - I am working with an undergrad to model cluster-scale lenses with PyAutoLens.\n", + "\n", + "We need to sample in physical units, rather than dimensionless quantities. For example, we want to sample an \n", + "NFW by (log10(M200), c200), rather than (kappa_s, theta_s).\n", + "\n", + "To do so, we will likely be making our own classes along these lines:\n", + "\n", + "class PhysicalNFW(autogalaxy.NFW):\n", + " def __init__(self, center, ell_comps, logM200m, c200m, cosmology):\n", + " kappa_s = ... computing rhos*rs ... / cosmology.critical_density()\n", + " theta_s = ... computing rs ... / cosmology.angular_diamter_distance()\n", + " super(autogalaxy.NFW, self).__init__(\n", + " center=center, ell_comps=ell_comps,\n", + " kappa_s=kappa_s, scale_radius=theta_s\n", + " )\n", + "\n", + "That seems wonderfully nice and simple!\n", + "However, the critical density is ambiguous when we have multiple lens and source planes (our clusters have multiple \n", + "sources at different redshifts). With multiple source planes, which zs should be used to compute the critical density \n", + "to give the right behavior? Is it the redshift of i.e. the next plane after the halo, or the last plane?\n", + "\n", + "\n", + "Jam\n", + "\n", + " However, the critical density is ambiguous when we have multiple lens and source planes (our clusters have multiple \n", + " sources at different redshifts). With multiple source planes, which zs should be used to compute the critical density \n", + " to give the right behavior? Is it the redshift of i.e. the next plane after the halo, or the last plane?\n", + "I have no idea, but I would guess it comes out as a lot of ratios of angular diameter distances. (edited) \n", + "\n", + "\n", + "\n", + "Jack\n", + "\n", + "I think it's an implementation detail of PyAutoLens: when a mass profile defines a bunch of deflection angles, \n", + "which planes does PyAutoLens interpret them as deflections between?\n", + "\n", + "\n", + "Jack\n", + "\n", + "If the mass is in plane i, is it between plane i and i+1?\n", + "\n", + "\n", + "Andrew\n", + "\n", + "So, without being an expert on the internals of PyAutoLens...\n", + "I'm going to use \"critical density\" to \n", + "mean 3H^2/8.pi.G (https://en.wikipedia.org/wiki/Friedmann_equations#Density_parameter), i.e. the 3D density for a \n", + "spatially flat Universe, and \"Sigma_crit\" to mean the critical surface density \n", + "for lensing (https://en.wikipedia.org/wiki/Gravitational_lensing_formalism) (from @Jack's pseudo-code, I'm guessing \n", + "cosmology.critical_density() is Sigma_crit?)\n", + "You will need the critical density at the lens redshift to convert from M200 and c to more physical NFW \n", + "parameters (i.e. rho_0 and R_s here https://en.wikipedia.org/wiki/Navarro%E2%80%93Frenk%E2%80%93White_profile, I \n", + "guess @Jack's rhos, rs).\n", + "\n", + "Projecting this, you get the surface density as a function of position in the image plane (i.e. Sigma(theta))\n", + "One would normally divide by Sigma_crit to get kappa(theta). But it doesn't really make sense to define some \n", + "convergence normalisation (i.e.) kappa_s for autogalaxy's NFW, because with multiple source planes there is no \n", + "one convergence field (because Sigma_crit depends on z_s). So @Jam, do you have a standard way to deal with multiple \n", + "source planes?\n", + "\n", + "That said, the convergences for the different source planes are just re-scaled versions of one another, as are the \n", + "shear and the deflection angles (though not things like the magnification). So, for example: if you want to know the \n", + "mapping from image plane coordinates to (multiple) source plane coordinates over a grid of image plane positions, \n", + "you could calculate a deflection angle field for one source plane (which might be costly, involve numerical \n", + "integrals, etc.) and then the deflection angle for some other source plane can be found easily be re-scaling by the \n", + "ratio of 'Sigma_crit's between the two different source redshifts\n", + "In terms of multiple source planes, @Jack, what is the data you intend to fit to / how do you intend to do your fit? \n", + "By which I mean, people fitting clusters often treat galaxies more like point sources than people fitting to \n", + "galaxy-galaxy strong lensing (the cluster people often just want their lens model to get the different multiple \n", + "images of each multiply imaged background galaxy to map back to common positions behind the cluster, as opposed to \n", + "caring about the structure of each lensed image), but I've typically seen PyAutoLens used to fit an observed image \n", + "pixel-by-pixel (often with a pixelised source reconstruction). If you intend to do the latter (with pixelised sources) \n", + "then you would have multiple pixelised source planes, each with their own regularisation, and (I imagine) the linear \n", + "algebra to find the most likely set of pixel fluxes across all the source planes would be rather difficult.\n", + "Not sure that will have helped, but my two cents...\n", + "\n", + "WikipediaWikipedia\n", + "Friedmann equations\n", + "https://en.wikipedia.org/wiki/Friedmann_equations#Density_parameter\n", + "\n", + "\n", + "WikipediaWikipedia\n", + "Gravitational lensing formalism\n", + "https://en.wikipedia.org/wiki/Gravitational_lensing_formalism\n", + "\n", + "\n", + "WikipediaWikipedia\n", + "Navarro\u2013Frenk\u2013White profile\n", + "\n", + "Jack\n", + "\n", + "Hi \n", + "@Andrew, thanks for this! I agree with all of this: my concern is that since the critical surface density depends on zs, \n", + "if there are multiple zses, it is ambiguous which one to use.\n", + "\n", + "For now we are using an observed-position likelihood as you point out is the standard with clusters, which James has \n", + "implemented and is working fine. The system we're looking at has two sources at different redshifts, but one is \n", + "clearly visible and the other someone barely noticed an emission line in MUSE data.\n", + "\n", + "For the more obvious source, we may end up doing a pixel reconstruction. (In which case we would probably ignore the \n", + "other source). Andrew Newman was able to do so with a double sersic model in \n", + "this paper: https://ui.adsabs.harvard.edu/abs/2018ApJ...862..125N/abstract\n", + "\n", + "\n", + "Jam\n", + "\n", + "The multi-plane implementation \n", + "follows (section 2): https://arxiv.org/abs/1403.5278 [NOTE TO READER, I WAS WRONG ABOUT THIS, DIFFERENT CORRECT PAPER LINKED TO BELOW]\n", + "\n", + "I can provide links to the source code, but basically you compute scaling factors (beta in equation 5) based on \n", + "angulars of diameter distances and then apply them when doing the multi-plane tracing the image-plane to each plane one \n", + "after another. The ray-tracing is recursive in that you go from the image-plane to each source-plane one-by-one I believe.\n", + "\n", + "\n", + "Then the deflection angle for some other source plane can be found easily be re-scaling by the ratio of 'Sigma_crit's \n", + "between the two different source redshifts\n", + "\n", + "I'm going to hazard a guess that equation (5) can be rewritten as a ratio of sigma_crit values (e.g. via equation 4, \n", + "provided D_l1 = D_l2). We basically then just need the code to use the sigma_crit values computed specifically for \n", + "the NFW's (which are related to kappa_s) , when computing the beta values, instead of how the values are computed currently?\n", + "\n", + "\n", + "Jack\n", + "\n", + "(You can even use the ratios of sigma_crit as a probe of cosmology! https://arxiv.org/abs/2110.06232)\n", + "\n", + "For now, the question is: which zsource is correct to use? Is it the redshift of the first source, or the last one?\n", + "\n", + "\n", + "Jam\n", + "\n", + "@Qiuhan He When simulating lenses with many DM subhalos (e.g. for the ABC paper) did we account for how to treat the \n", + "source redshift when computing their sigma_crit but also how to set their mass parameters via the critical density \n", + "of the Universe?\n", + "\n", + "\n", + "Qiuhan He\n", + "\n", + "If we want to compute how the first source is lensed by an NFW halo, then the source redshift is the first source's \n", + "redshift. If we want to compute the lensing quantities about the second source, then the source redshift should be \n", + "the second source's.\n", + "\n", + "The critical density needed for an NFW halo is the critical density of the Universe at the redshift of the halo\n", + "\n", + "The PhysicalNFW defined should involve one more parameter called source_redshift, because the lensing quantity of an NFW halo would change with different background sources at different redshifts\n", + "\n", + "Jam\n", + "\n", + "When simulating lenses for the ABC paper, we had 100s of NFWs at different redshifts making up the line of sight.\n", + "\n", + "I am assuming we always used 'redshift_source' as the source redshift for every NFW, which was what went into \n", + "computing its sigma_crit value.\n", + "\n", + "When performing multi plane ray tracing, we used the redshift of the NFW we were computing deflections of, \n", + "as well as all other NFWs.\n", + "\n", + "Independently, I know these two calculations are valid. What I'm unclear on is whether combining them in the way \n", + "we did is valid (e.g. is the sigma_crit we compute when using collosus defined the same as the ones used for scaling \n", + "between different planes for multi plane ray tracing?l\n", + "\n", + "\n", + "Qiuhan He\n", + "\n", + "yes\n", + "\n", + "I need to find the multiplane ray tracing equation in Schneider+1992\n", + "\n", + "PDF\n", + "\n", + "Here, Qiuhan linked to section \"9.1 The multiple lens-plane theory\" of Schneider 1992: \n", + "https://ui.adsabs.harvard.edu/abs/1992grle.book.....S/abstract\n", + "\n", + "Eq. (9.7b) is the multiplane ray-tracing equation. The dimensionless deflection angle of each plane is defined as \n", + "Eq. (9.6), which is the way AutoLens implements.\n", + "\n", + "\n", + "Jack\n", + "\n", + "It sounds like what you are saying is that for a mass in plane i, AutoLens interprets it's deflection angles as \n", + "between plane i+1 and plane i. So the Sigma_crit is a function of z_i and z_{i+1}. Is that correct? Elsewhere \n", + "James said something implying that it is betwen plane i and i_{max}\n", + "\n", + "This whole approach of sampling unitless quantities is cute and simple with one lens/source plane, but with multiple \n", + "mass/source planes like this, or if you actually care about the physical quantities of the halos, I have always \n", + "thought it would make much more sense to only sample physical parameters of each halo and have your ray-tracing code \n", + "do everything correctly under the hood\n", + "\n", + "\n", + "Qiuhan He\n", + "\n", + "The sigma_crit is what James described as shown by Eq.(9.4).\n", + "\n", + "\n", + "Jack\n", + "\n", + "I understand what sigma_crit is.\n", + "\n", + "My point of confusion is the following: sigma_crit is a function of cosmology, lens redshift (z_l), and source \n", + "redshift (z_s). When you have one lens plane and one source plane, parametrizing everything with kappa is clear and \n", + "unambiguous, but when you have multiple lens and source redshifts, it becomes ambiguous.\n", + "\n", + "Choosing which z_l and z_s to use to define kappa then becomes an implementation choice: you could uses the lowest \n", + "redshift z_s, or the highest redshift z_s . I suppose you could do anything in between but that would be a pretty \n", + "pathological. My question is really what choice does PyAutoLens make here?\n", + "\n", + "AutoLens parameterizes an NFW as kappa_s and theta_scale, the angle of the scale radius. The convergence at the \n", + "scale radius kappa_s is only a meaningful quantity between a specific lens redshift, source redshift, and cosmology. \n", + "\n", + "So what sigma_crit does kappa_s correspond to? What physical mass surface density \\Sigma = \\Sigma_{crit} \\kappa is causing the lensing?\n", + "\n", + "\n", + "Qiuhan He\n", + "\n", + "The NFW profile parameterized by kappa_s and theta_scale has no physical meaning. It can be a halo of mass A \n", + "in one strong lensing system or a halo of mass 2A in another strong lensing system whose sigma_crit is twice.\n", + "\n", + "It is only a model for unitless lensing computation.\n", + "\n", + "We have a profile called NFWMCRLudlow . To use that, one needs to specify the redshift of the halo and the source by \n", + "redshift_object and redshift_source.\n", + "\n", + "In multiplane lensing, we set redshift_object to be the redshift of the i-th lensing plane and redshift_source \n", + "to be the source galaxy redshift\n", + "\n", + "\n", + "Jack\n", + "\n", + "So I have the opposite problem: We have the physical parameters (Mass, concentration) and need to get the correct \n", + "lensing from that mass.\n", + "\n", + "Ahh so the ambiguity I am struggling with is that we have multiple source galaxies at different redshifts\n", + "\n", + "James sent me a code snippet where it appears that the multi-plane lensing always interprets the deflection angles as \n", + "between the i'th plane and the highest-redshift plane. That would answer my question\n", + "\n", + "\n", + "Jam\n", + "\n", + "I don't think it is necessary to think about whether a plane has a \"source galaxy\" (in the sense that you have \n", + "observed imaging data of it and want to model or analyse it).\n", + "\n", + "The multi-plane ray tracing calculations do not care if a plane has a \"source galaxy\" or if its just another plane \n", + "with a galaxy with mass in. The deflection angles would not change if you added a \"source galaxy\" to a plane in a \n", + "multi-plane system which previously only had mass components.\n", + "\n", + "Andrew\n", + "@Jam, @Qiuhan He have we ever dealt with multiple source planes, or just multiple lens planes? \n", + "@Jack's sentiment that defining a lens by its convergence (and not it's physical mass distribution) is ambiguous when \n", + "there are multiple source planes is something I very much agree with (well, one cannot disagree, it is clearly true!)\n", + "\n", + "\n", + "Jam\n", + "\n", + "I have used autolens to model multiple source plane systems, but never with physical units throw into the mixer as well\n", + "\n", + "\n", + "Andrew\n", + "\n", + "But \"without physical units\" sounds ill defined. If we take a case of a single lens plane (with an NFW) and two \n", + "source planes, what should kappa_s and r_s be? r_s (being an angle to the lens plane) is well defined, but kappa_s \n", + "is different depending on which of the two source planes is being considered\n", + "\n", + "\n", + "Qiuhan He\n", + "\n", + "Ah. I get it. \n", + "\n", + "@Jack, for your purpose, ray-tracing for multiple source galaxies, please compute the kappa_s using the last source \n", + "galaxy's redshift\n", + "\n", + "\n", + "Jack\n", + "\n", + "great, thank you \n", + "\n", + "Qiuhan He\n", + "\n", + "Autolens will rescale the kappa_s for source galaxies infront of it. (I think autolens is actually rescaling the \n", + "deflection angles instead) (edited) \n", + "\n", + "\n", + "Jam\n", + "\n", + "Ok, I'm gonna write this somewhere permenant.. [Oh look I did lol]\n", + "\n", + "\n", + "\n", + "Qiuhan He\n", + "\n", + "Thats how I understand Autolens is doing for multiplane raytracing\n", + "\n", + "\n", + "Andrew\n", + "\n", + "Sounds like we have an answer :grinning:\n", + "\n", + "\n", + "Jack\n", + "\n", + "but yes @Andrew that is exactly my concern! in my case we are always concerned with the physical mass of the lenses. \n", + "so for us it never makes sense to sample \"lensing units\" that can then be rescaled for different redshifts: we always \n", + "want a fixed well-defined cosmology with known redshifts for each source\n", + "\n", + "(I suppose we could vary the cosmology but that doesn't change the principle here)\n", + "\n", + "\n", + "Andrew\n", + "\n", + "Perhaps best to check what \n", + "@Qiuhan He described with a few simple cases (like do some analysis with a single source plane, then re-do it adding\n", + "a second source plane at lower-z with no light and see that nothing changes; and various other variants of this)\n", + "\n", + "\n", + "Jam\n", + "\n", + "Are you assuming your source redshifts are always known exactly?\n", + "\n", + "\n", + "Jack\n", + "\n", + "In our case the source redshifts are known to high precision, we don't anticipate working with systems of unknown redshift\n", + "\n", + "\n", + "Jack\n", + "\n", + "@Andrew James sent me some example code doing something similar to what you described, I'll play with it to \n", + "double-check everything!\n", + "\n", + "and \n", + "@Jam re: source redshift precision: there may be some cluster systems where several sources have known redshift, and \n", + "some do not (or have e.g. photo-zs not spec-zs).\n", + "\n", + "in that case we would like to sample the redshifts of the sources with poorly-constrained redshift. the sources with \n", + "known redshift serve to constrain the lens mass, and if the lens mass is understood well enough it actually can \n", + "constrain the redshift of the other sources. (basically, at what z_source does this source's observed lensing line up)\n", + "\n", + "This is something that's actually done with huge clusters, where it's hard to get speczs on every object. I think \n", + "I've seen this with some of the crazy massive clusters observed with JWST\n", + "\n", + "\n", + "Andrew\n", + "\n", + "A quick comment (which doesn't actually suggest how things would be best implemented) and is written with explicit \n", + "reference to NFW profiles, but applies to defining lenses by their convergence (not their mass) more generally...\n", + "\n", + "Presumably, the reason why PyAutoLens uses (by default) kappa_s and theta_s [rs in the code, but it sounds like \n", + "it is an angle] (not physical things like M200, c and z_l; or rho_s, r_s and z_l) is that they uniquely specify the \n", + "mapping from image plane to source plane coordinates (i.e. the \"deflection angles\"). Whereas if specifying the \n", + "3 numbers describing the physical mass distribution (how much it weighs, how concentrated it is, and how far along \n", + "the line-of-sight it is) there are different combinations that give you the same deflection angles (i.e. the same \n", + "mapping from positions in the image back to positions in the source plane). This means that, particularly \n", + "when z_l and/or z_s is unknown, it makes sense to fit for kappa_s and theta_s. Of course, if we know the redshifts, \n", + "then whether we fit for (kappa_s, theta_s), (M200, c), or (rho_s, r_s) doesn't really matter. \n", + "\n", + "The point I am trying to make, is that there are (in a single lens plane, single source plane case) good reasons to \n", + "work directly with the convergence, rather than thinking first about the physical mass distribution, and then using \n", + "redshifts to turn this into kappa(theta), from which we get alpha(theta) and hence the mapping from image plane to \n", + "source plan positions that we need to \"do the lensing\".\n", + "\n", + "I think @Jack's multi-source-plane case has caused some confusion because it breaks the property we had with a s\n", + "ingle source plane, that we could have got the same lensing effect from different mass distributions at different \n", + "redshifts. As Qiuhan said earlier The NFW profile parameterized by kappa_s and theta_scale has no physical meaning. \n", + "It can be a halo of mass A in one strong lensing system or a halo of mass 2A in another strong lensing system whose \n", + "sigma_crit is twice. But with two source planes, a change in z_l that doubles sigma_crit to one source plane, \n", + "might only increase sigma_crit by a factor of 1.5 to the other source plane. So there is not a sense in which we \n", + "can just fit for some \"dimensionless lensing parameters\" and then later use the lens redshift to convert this to \n", + "physical lens parameters if we are so inclined. We need to know z_l in order to know how the convergence of the \n", + "lens for source plane 1 relates to the convergence of the lens for source plane 2. And if we know the lens and \n", + "source redshifts, we may as well parameterise the lens by physical parameters, rather than kappa_s and theta_s \n", + "(because the benefit of kappa_s and theta_s has gone).\n", + "\n", + "Coming from a simulation background, I definitely think of physical mass distributions first and then their lensing \n", + "effect second. But I think the \"kappa(theta) is primary, and we can worry about what actual mass distribution this \n", + "corresponds to later\" approach does make sense in some observational cases. And I'm hoping that hearing those two \n", + "perspectives might help someone (or maybe everyone already knew this! :melting_face:) (edited) \n", + "\n", + "\n", + "\n", + "Jack\n", + "\n", + "Thanks for the reasoned description Andrew! You are of course right that both approaches make sense in \n", + "different cases - I apologize if I was being uncharitable in describing the \"lensing units\" approach before.\n", + "\n", + "I think this is a frequent division between codes for \"galaxy-scale\" lensing and \"cluster-scale\" lensing: with \n", + "galaxy-galaxy lensing often the \"lensing units\" modeling is simple and works great, but with cluster-scale lensing, \n", + "with may different mass components and multiple source redshifts, the formalism behind the \"lensing units\" becomes \n", + "confusing just how you describe and it makes more sense to parametrize by physical parameters.\n", + "\n", + "The property that breaks in multiple-source lensing, as you describe, actually has a few interesting science corollaries:\n", + "The mass-sheet degeneracy is broken. With galaxy-scale lensing for cosmology, people fret quite a bit about the \n", + "mass sheet degeneracy, but with multiple sources probing different sigma_crit's you have a lever arm to break this \n", + "and constrain the physical mass distribution. This is a reason people don't talk much about the mass-sheet \n", + "degeneracy in cluster-scale lensing\n", + "\n", + "You can actually constrain cosmology with the lensing strength between planes! the ratio of lensing strength between \n", + "planes is expressed as beta (below) for each pair of planes\n", + "\n", + "it is a strict function of geometry (distances) - so comparing predicted to measured values of beta is a method for \n", + "doing cosmography. You can sample over a cosmology for this, and have a similar constraining mechanism to BAO \n", + "measurements. It's basically ratios of angular diameter distances\n", + "\n", + "Here is an example paper doing # 2 with a small sample of cluster-scale lenses. We are currently waiting on \n", + "spectroscopic follow-up of one golden system with many background sources, where we'd like to do this measurement.\n", + "https://arxiv.org/abs/2110.06232\n", + "\n", + "arXiv.orgarXiv.org\n", + "Galaxy cluster strong lensing cosmography: cosmological constraints from a sample of regular galaxy clusters\n", + "Cluster strong lensing cosmography is a promising probe of the background geometry of the Universe and several studies \n", + "have emerged, thanks to the increased quality of observations using space and ground-based telescopes. For the first \n", + "time, we use a sample of five cluster strong lenses to measure the values of cosmological parameters and combine them \n", + "with those from classical probes. In order to assess the degeneracies and the effectiveness of strong-lensing \n", + "cosmography in constraining the background geometry of the Universe, we adopt four cosmological scenarios. We find \n", + "good constraining power on the total matter density of the Universe ($\u03a9_{\\rm m}$) and the equation of state of the dark e\u2026 Show more" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/advanced/over_sampling.ipynb b/notebooks/guides/advanced/over_sampling.ipynb index b62151eb7..880c6f431 100644 --- a/notebooks/guides/advanced/over_sampling.ipynb +++ b/notebooks/guides/advanced/over_sampling.ipynb @@ -1,734 +1,860 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Over Sampling\n", - "=============\n", - "\n", - "Throughout the workspace, we have created 2D grids of (y,x) coordinates and input them into light profiles to\n", - "compute their image.\n", - "\n", - "This calculates how much of the light profile's emission is observed with every 2D pixel defined on the grid.\n", - "\n", - "However, there is a problem. If we only input the (y,x) coordinates at the centre of every pixel, we are not\n", - "evaluating how the entire light profile is observed within that pixel. If the light profile has a very steep gradient\n", - "in intensity from one edge of the pixel to the other, only evaluating the intensity at the centre of the pixel will\n", - "not give an accurate estimate of the total amount of light that falls within that pixel.\n", - "\n", - "Over-sampling addresses this problem. Instead of evaluating the light profile at the centre of every pixel, we\n", - "evaluate it using a sub-grid of coordinates within every pixel and take the average of the intensity values.\n", - "Provided the sub-grid is high enough resolution that it \"over-samples\" the light profile within the pixel enough, this\n", - "will give an accurate estimate of the total intensity within the pixel.\n", - "\n", - "__Default Over-Sampling__\n", - "\n", - "Examples throughout the workspace use a default over-sampling set up that should ensure accurate results for any\n", - "analysis you have done. This default over-sampling is as follows:\n", - "\n", - "- When evaluating the image of a galaxy, an adaptive over sampling grid is used which uses sub grids of size 8 x 8\n", - "in the central regions of the image, 4x4 further out and 1x1 beyond that.\n", - "\n", - "- When evaluating the image of the source galaxy, no over-sampling (e.g. a 1 x 1 subgrid) is performed but instead\n", - "cored light profiles for the source are used which can be evaluated accurate without over-sampling.\n", - "\n", - "This guide will explain why these choices were made for the default over-sampling behaviour.\n", - "\n", - "__Contents__\n", - "\n", - "- **Illustration:** To illustrate over sampling, lets first create a uniform grid which does not over sample the.\n", - "- **Numerics:** Lets quickly check how the sub-grid is defined and stored numerically.\n", - "- **Images:** We now use over-sampling to compute the image of a Sersic light profile, which has a steep.\n", - "- **Adaptive Over Sampling:** We have shown that over-sampling is important for accurate image evaluation.\n", - "- **Multiple Lens Galaxies:** The analysis may contain multiple lens galaxies, each of which must be over-sampled accurately.\n", - "- **Ray Tracing:** So far, we have evaluated the image of a light profile using over-sampling on an unlensed uniform.\n", - "- **Default Ray Tracing:** By assuming the lens galaxy is near (0.0\", 0.0\"), it was simple to set up an adaptive over sampling.\n", - "- **Dataset & Modeling:** Throughout this guide, grid objects have been used to compute the image of light and mass profiles.\n", - "- **Pixelization:** Source galaxies can be reconstructed using pixelizations, which discretize the source's light onto." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Illustration__\n", - "\n", - "To illustrate over sampling, lets first create a uniform grid which does not over sample the pixels, using \n", - "the `over_sample_size` input.\n", - "\n", - "The input below uses `over_sample_size=1`, therefore each pixel is split into a sub-grid of \n", - "size `over_sample_size x over_sample_size` = `1 x 1`. This means the light profile is evaluated once at the centre of each pixel, \n", - "which is equivalent to not over-sampling the grid at all. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_sub_1 = al.Grid2D.uniform(\n", - " shape_native=(40, 40),\n", - " pixel_scales=0.1,\n", - " over_sample_size=1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now plot the grid, over laying a uniform grid of pixels to illustrate the area of each pixel within which we\n", - "want light profile intensities to be computed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_grid(grid=grid_sub_1, title=\"Grid (No Over-Sampling)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now create and plot a uniform grid which does over-sample the pixels, by inputting `over_sample_size=2`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_sub_2 = al.Grid2D.uniform(\n", - " shape_native=(40, 40),\n", - " pixel_scales=0.1,\n", - " over_sample_size=2,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we print `grid_sub_2` and its shape, we will find it is actually identical to `grid_sub_1`, despite the change\n", - "in `over_sample_size`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(grid_sub_1)\n", - "print(grid_sub_2)\n", - "print(grid_sub_1.shape)\n", - "print(grid_sub_2.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This is because the over sampled version of the grid is stored in a separate attribute, called `over_sampled`,\n", - "which we print below.\n", - "\n", - "We see that for `grid_sub_1` and `grid_sub_2` the `over_sampled` grids are different, with the over sampled grid for\n", - "`grid_sub_2` containing four times as many entries corresponding to each pixel being sub-gridded in a 2 x 2 shape." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(grid_sub_1.over_sampled)\n", - "print(grid_sub_2.over_sampled)\n", - "print(grid_sub_1.over_sampled.shape)\n", - "print(grid_sub_2.over_sampled.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now plot the over sampled grid over the image, showing that each pixel is now split into a 2x2 sub-grid of \n", - "coordinates. \n", - "\n", - "These are used to compute the intensity of the light profile and therefore more accurately estimate the total \n", - "intensity within each pixel if there is a significant gradient in intensity within the pixel.\n", - "\n", - "In the code below, it is the input `plot_over_sampled_grid=True` which ensures we plot the over sampled grid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_grid(grid=grid_sub_2, title=\"Over-Sampled Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Numerics__\n", - "\n", - "Lets quickly check how the sub-grid is defined and stored numerically.\n", - "\n", - "The first four pixels of this sub-grid correspond to the first four sub-pixels in the first pixel of the grid. \n", - "\n", - "The top-left pixel image above shows how the sub-pixels are spaced within the pixel. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"(y,x) pixel 0 of grid_sub_1:\")\n", - "print(grid_sub_1.over_sampled[0])\n", - "print(\"(y,x) pixel 0 of grid_sub_2:\")\n", - "print(grid_sub_2.over_sampled[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now confirm that the first four sub-pixels of the over-sampled grid correspond are contained within the \n", - "first pixel of the grid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"(y,x) pixel 0 (of original grid):\")\n", - "print(grid_sub_2[0])\n", - "print(\"(y,x) sub-pixel 0 (of pixel 0):\")\n", - "print(grid_sub_2.over_sampled[0])\n", - "print(\"(y,x) sub-pixel 1 (of pixel 0):\")\n", - "print(grid_sub_2.over_sampled[1])\n", - "print(\"(y,x) sub-pixel 2 (of pixel 0):\")\n", - "print(grid_sub_2.over_sampled[2])\n", - "print(\"(y,x) sub-pixel 3 (of pixel 0):\")\n", - "print(grid_sub_2.over_sampled[3])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Numerically, the over-sampled grid contains the sub-pixel coordinates of every pixel in the grid, going from the \n", - "first top-left pixel right and downwards to the bottom-right pixel. \n", - "\n", - "So the pixel to the right of the first pixel is the next 4 sub-pixels in the over-sampled grid, and so on.\n", - "\n", - "__Images__\n", - "\n", - "We now use over-sampling to compute the image of a Sersic light profile, which has a steep intensity gradient\n", - "at its centre which a lack of over-sampling does not accurately capture.\n", - "\n", - "We create the light profile, input the two grids (with `over_sample_size=1` and `over_sample_size=2`) and compute the image of the\n", - "light profile using each grid. We then plot the residuals between the two images in order to show the difference\n", - "between the two images and thus why over-sampling is important.\n", - "\n", - "Over sampling occurs automatically when a grid is input into a function like `image_2d_from`, therefore internally \n", - "the line of code, `image_sub_2 = light.image_2d_from(grid=grid_sub_2)`, is evaluating the light profile using the\n", - "2 x 2 oversampled grid and internally binning it up in to fully perform over sampling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "light = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=1.0,\n", - " effective_radius=0.2,\n", - " sersic_index=3.0,\n", - ")\n", - "\n", - "image_sub_1 = light.image_2d_from(grid=grid_sub_1)\n", - "image_sub_2 = light.image_2d_from(grid=grid_sub_2)\n", - "\n", - "aplt.plot_array(array=image_sub_1, title=\"Image of Sersic Profile\")\n", - "\n", - "residual_map = image_sub_2 - image_sub_1\n", - "\n", - "aplt.plot_array(array=residual_map, title=\"Over-Sampling Residuals\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the central 4 pixels of the image, the residuals are large due to the steep intensity gradient of the Sersic\n", - "profile at its centre. \n", - "\n", - "The gradient in these pixels is so steep that evaluating the intensity at the centre of the pixel, without over \n", - "sampling, does not accurately capture the total intensity within the pixel.\n", - "\n", - "At the edges of the image, the residuals are very small, as the intensity gradient of the Sersic profile is very \n", - "shallow and it is an accurate approximation to evaluate the intensity at the centre of the pixel.\n", - "\n", - "The absolute value of the central residuals are 0.74, however it is difficult to assess whether this is a large or\n", - "small value. We can quantify this by dividing by the evaluated value of the Sersic image in each pixel in order\n", - "to compute the fractional residuals." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fractional_residual_map = residual_map / image_sub_2\n", - "\n", - "\n", - "aplt.plot_array(\n", - " array=fractional_residual_map, title=\"Fractional Over-Sampling Residuals\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fractional residuals in the centre exceed 0.1, or 10%, which is a significant error in the image and\n", - "demonstrates why over-sampling is important.\n", - "\n", - "Lets confirm sub-griding can converge to central residuals that are very small.\n", - "\n", - "The fractional residuals with high levels of over-sampling are below 0.01, or 1%, which is sufficiently accurate\n", - "for most scientific purposes (albeit you should think carefully about the level of over-sampling you need for\n", - "your specific science case)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_sub_16 = al.Grid2D.uniform(\n", - " shape_native=(40, 40), pixel_scales=0.1, over_sample_size=16\n", - ")\n", - "grid_sub_32 = al.Grid2D.uniform(\n", - " shape_native=(40, 40), pixel_scales=0.1, over_sample_size=32\n", - ")\n", - "\n", - "image_sub_16 = light.image_2d_from(grid=grid_sub_16)\n", - "image_sub_32 = light.image_2d_from(grid=grid_sub_32)\n", - "\n", - "residual_map = image_sub_32 - image_sub_16\n", - "\n", - "aplt.plot_array(array=residual_map, title=\"Over-Sampling Reduces Residuals\")\n", - "\n", - "fractional_residual_map = residual_map / image_sub_32\n", - "\n", - "aplt.plot_array(\n", - " array=fractional_residual_map, title=\"Fractional Residuals With Over-Sampling\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Adaptive Over Sampling__\n", - "\n", - "We have shown that over-sampling is important for accurate image evaluation. However, there is a major drawback to\n", - "over-sampling, which is that it is computationally expensive. \n", - "\n", - "For example, for the 32x32 over-sampled grid above, 1024 sub-pixels are used in every pixel, which must all be \n", - "evaluated using the Sersic light profile. The calculation of the image is therefore at least 1000 times slower than if\n", - "we had not used over-sampling.\n", - "\n", - "Speeding up the calculation is crucial for model-fitting where the image is evaluated many times to fit the\n", - "model to the data.\n", - "\n", - "Fortunately, there is a solution to this problem. We saw above that the residuals rapidly decrease away\n", - "from the centre of the light profile. Therefore, we only need to over-sample the central regions of the image,\n", - "where the intensity gradient is steep. We can use lower levels of over-sampling away from the centre, which\n", - "will be fast to evaluate.\n", - "\n", - "Up to now, the `over_sample_size` input has been an integer, however it can also be an `ndarray` of values corresponding\n", - "to each pixel. We create an `ndarray` of values which are high in the centre, but reduce to 2 at the outskirts,\n", - "therefore providing high levels of over sampling where we need it whilst using lower values which are computationally\n", - "fast to evaluate at the outskirts.\n", - "\n", - "Specifically, we define a 24 x 24 sub-grid within the central 0.3\" of pixels, uses a 8 x 8 grid between\n", - "0.3\" and 0.6\" and a 2 x 2 grid beyond that. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid_sub_1, sub_size_list=[24, 8, 2], radial_list=[0.3, 0.6]\n", - ")\n", - "\n", - "grid_adaptive = al.Grid2D.no_mask(\n", - " values=grid_sub_1.native,\n", - " pixel_scales=grid_sub_1.pixel_scales,\n", - " over_sample_size=over_sample_size,\n", - ")\n", - "\n", - "\n", - "aplt.plot_grid(grid=grid_adaptive, title=\"Over-Sampled Grid\")\n", - "\n", - "print(over_sample_size)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling uses masked grids, therefore the code below shows how we would create this adaptive over sample grid via \n", - "a circular mask, which can be used for modeling.\n", - "\n", - "Throughout the modeling examples in the workspace, we use this adaptive grid to ensure that the image of the\n", - "galaxy is evaluated accurately and efficiently." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(shape_native=(40, 40), pixel_scales=0.1, radius=5.0)\n", - "\n", - "grid = al.Grid2D.from_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid, sub_size_list=[24, 8, 2], radial_list=[0.3, 0.6]\n", - ")\n", - "\n", - "grid_adaptive = al.Grid2D(values=grid, mask=mask, over_sample_size=over_sample_size)\n", - "\n", - "\n", - "aplt.plot_grid(grid=grid, title=\"Over-Sampled Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can compare this adaptive grid to the grid with over sampling of 32 x 32 to confine it produces low amounts\n", - "of residuals." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_masked_sub_32 = al.Grid2D(values=grid, mask=mask, over_sample_size=32)\n", - "\n", - "image_adaptive = light.image_2d_from(grid=grid_adaptive)\n", - "image_masked_sub_32 = light.image_2d_from(grid=grid_masked_sub_32)\n", - "\n", - "residual_map = image_adaptive - image_masked_sub_32\n", - "\n", - "fractional_residual_map = residual_map / image_masked_sub_32\n", - "\n", - "\n", - "aplt.plot_array(array=fractional_residual_map, title=\"Adaptive Fractional Residuals\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Default Over-Sampling__\n", - "\n", - "The default over-sampling scheme used by the source code is 4 x 4 uniform over sampling over the whole image. \n", - "\n", - "A uniform scheme is used, instead of the adaptive scheme above, because the adaptive scheme requires input knowledge of \n", - "where the centre of the galaxy is (e.g. above the centre is at (0.0\", 0.0\").\n", - "\n", - "Uniform over sampling is precise enough for many calculations, especially when you are simply performing quick \n", - "calculations to investigate a problem. However, for detailed calculations you must ensure that high enough\n", - "levels of over sampling are used.\n", - "\n", - "For modeling, all example scripts begin by switching to an adaptive over sampling scheme, as modeling assumes\n", - "the centre of the lens galaxy is at (0.0\", 0.0\").\n", - "\n", - "__Multiple Lens Galaxies__\n", - "\n", - "The analysis may contain multiple lens galaxies, each of which must be over-sampled accurately. \n", - "\n", - "\n", - "There are two approaches you can take to over sampling multi-galaxy systems:\n", - "\n", - "1) Use a high level of uniform over sampling over the full image.\n", - "\n", - "2) Use an adaptive over sampling scheme with multiple centres of high over sampling levels, with the API shown below\n", - " for two galaxies with centres (1.0, 0.0) and (-1.0, 0.0)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid_sub_1,\n", - " sub_size_list=[24, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(1.0, 0.0), (-1.0, 0.0)],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "So far, we have evaluated the image of a light profile using over-sampling on an unlensed uniform grid. \n", - "\n", - "For lensing calculations, the grid is ray-traced via a mass model to an irregular grid in the source plane.\n", - "\n", - "The over sampling of lensed images is therefore describe as follows: \n", - "\n", - "1) Splits each image-pixel into a sub-grid of pixels in the image-plane.\n", - "2) Ray trace this sub-grid of pixels using the mass model.\n", - "3) Evaluate the source light of each sub-pixel in the source-plane.\n", - "4) Bin up the values to evaluate the over sampled values.\n", - "\n", - "For lensing calculations, over sampling therefore requires us to ray-trace (and therefore compute the deflecitons angles\n", - "of) many more (y,x) coordinates!\n", - "\n", - "We now illustrate using over-sampling with a mass profile, noting that for lensing:\n", - "\n", - "1) The fractional residuals due to differing over-sampling levels now occur in the lensed source's brightest multiply \n", - " imaged pixels. \n", - " \n", - "2) It is the combination of a rapidly changing source light profile and the magnification pattern of the mass model\n", - " which requires over sampling. The mass model focuses many image-pixels to the source's brightest regions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = al.mp.Isothermal(centre=(0.0, 0.0), ell_comps=(0.1, 0.0), einstein_radius=1.0)\n", - "\n", - "light = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=1.0,\n", - " effective_radius=0.2,\n", - " sersic_index=3.0,\n", - ")\n", - "\n", - "lens = al.Galaxy(redshift=0.5, mass=mass)\n", - "\n", - "source = al.Galaxy(redshift=1.0, bulge=light)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "image_sub_1 = tracer.image_2d_from(grid=grid_sub_1)\n", - "image_sub_2 = tracer.image_2d_from(grid=grid_sub_2)\n", - "\n", - "aplt.plot_array(array=image_sub_1, title=\"Source Image 1x1\")\n", - "\n", - "residual_map = image_sub_2 - image_sub_1\n", - "\n", - "fractional_residual_map = residual_map / image_sub_2\n", - "\n", - "aplt.plot_array(array=fractional_residual_map, title=\"Fractional Residuals\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Default Ray Tracing__\n", - "\n", - "By assuming the lens galaxy is near (0.0\", 0.0\"), it was simple to set up an adaptive over sampling grid which is\n", - "applicable to all strong lens dataset.\n", - "\n", - "An adaptive oversampling grid cannot be defined for the lensed source because its light appears in different regions of \n", - "the image plane for each dataset. For this reason, most workspace examples utilize cored light profiles for the \n", - "source galaxy. Cored light profiles change gradually in their central regions, allowing accurate evaluation without \n", - "requiring oversampling.\n", - "\n", - "__Adaptive Over Sampling__\n", - "\n", - "There is a way to set up an adaptive over sampling grid for a lensed source, however it requries one to use and\n", - "understanding the advanced lens modeling feature search chaining.\n", - "\n", - "An example of how to use search chaining to over sample sources efficient is provided in \n", - "the `autolens_workspace/*/guides/modeling/chaining/over_sampling.ipynb` example.\n", - "\n", - "__Dataset & Modeling__\n", - "\n", - "Throughout this guide, grid objects have been used to compute the image of light and mass profiles and illustrate\n", - "over sampling.\n", - "\n", - "If you are performing calculations with imaging data or want to fit a model to the data with a specific\n", - "over-sampling level, the `apply_over_sampling` method is used to update the over sampling scheme of the dataset.\n", - "\n", - "The grid this is applied to is called `lp`, to indicate that it is the grid used to evaluate the emission of light\n", - "profiles for which this over sampling scheme is applied." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "# This can be any of the over-sampling objects we have used above.\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=4)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixelization__\n", - "\n", - "Source galaxies can be reconstructed using pixelizations, which discretize the source's light onto a mesh,\n", - "for example a Voronoi mesh.\n", - "\n", - "Over sampling is used by pixelizations in an analogous way to light profiles. By default, a 4 x 4 sub-grid is used,\n", - "whereby every image pixel is ray-traced on its 4 x 4 sub grid to the source mesh and fractional mappings are computed.\n", - "\n", - "A different grid and over sampling scheme is applied to light profiles and pixelizations, which is why\n", - "there are separate inputs called `lp` and `pix`.\n", - "\n", - "This is explained in more detail in the pixelization examples.\n", - "\n", - "Here is an example of how to change the over sampling applied to a pixelization for a lens model fit:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=4)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Over Sampling\n", + "=============\n", + "\n", + "Throughout the workspace, we have created 2D grids of (y,x) coordinates and input them into light profiles to\n", + "compute their image.\n", + "\n", + "This calculates how much of the light profile's emission is observed with every 2D pixel defined on the grid.\n", + "\n", + "However, there is a problem. If we only input the (y,x) coordinates at the centre of every pixel, we are not\n", + "evaluating how the entire light profile is observed within that pixel. If the light profile has a very steep gradient\n", + "in intensity from one edge of the pixel to the other, only evaluating the intensity at the centre of the pixel will\n", + "not give an accurate estimate of the total amount of light that falls within that pixel.\n", + "\n", + "Over-sampling addresses this problem. Instead of evaluating the light profile at the centre of every pixel, we\n", + "evaluate it using a sub-grid of coordinates within every pixel and take the average of the intensity values.\n", + "Provided the sub-grid is high enough resolution that it \"over-samples\" the light profile within the pixel enough, this\n", + "will give an accurate estimate of the total intensity within the pixel.\n", + "\n", + "__Default Over-Sampling__\n", + "\n", + "Examples throughout the workspace use a default over-sampling set up that should ensure accurate results for any\n", + "analysis you have done. This default over-sampling is as follows:\n", + "\n", + "- When evaluating the image of a galaxy, an adaptive over sampling grid is used which uses sub grids of size 8 x 8\n", + "in the central regions of the image, 4x4 further out and 1x1 beyond that.\n", + "\n", + "- When evaluating the image of the source galaxy, no over-sampling (e.g. a 1 x 1 subgrid) is performed but instead\n", + "cored light profiles for the source are used which can be evaluated accurate without over-sampling.\n", + "\n", + "This guide will explain why these choices were made for the default over-sampling behaviour.\n", + "\n", + "__Contents__\n", + "\n", + "- **Illustration:** To illustrate over sampling, lets first create a uniform grid which does not over sample the.\n", + "- **Numerics:** Lets quickly check how the sub-grid is defined and stored numerically.\n", + "- **Images:** We now use over-sampling to compute the image of a Sersic light profile, which has a steep.\n", + "- **Adaptive Over Sampling:** We have shown that over-sampling is important for accurate image evaluation.\n", + "- **Multiple Lens Galaxies:** The analysis may contain multiple lens galaxies, each of which must be over-sampled accurately.\n", + "- **Ray Tracing:** So far, we have evaluated the image of a light profile using over-sampling on an unlensed uniform.\n", + "- **Default Ray Tracing:** By assuming the lens galaxy is near (0.0\", 0.0\"), it was simple to set up an adaptive over sampling.\n", + "- **Dataset & Modeling:** Throughout this guide, grid objects have been used to compute the image of light and mass profiles.\n", + "- **Pixelization:** Source galaxies can be reconstructed using pixelizations, which discretize the source's light onto." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Illustration__\n", + "\n", + "To illustrate over sampling, lets first create a uniform grid which does not over sample the pixels, using \n", + "the `over_sample_size` input.\n", + "\n", + "The input below uses `over_sample_size=1`, therefore each pixel is split into a sub-grid of \n", + "size `over_sample_size x over_sample_size` = `1 x 1`. This means the light profile is evaluated once at the centre of each pixel, \n", + "which is equivalent to not over-sampling the grid at all. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_sub_1 = al.Grid2D.uniform(\n", + " shape_native=(40, 40),\n", + " pixel_scales=0.1,\n", + " over_sample_size=1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now plot the grid, over laying a uniform grid of pixels to illustrate the area of each pixel within which we\n", + "want light profile intensities to be computed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_grid(grid=grid_sub_1, title=\"Grid (No Over-Sampling)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now create and plot a uniform grid which does over-sample the pixels, by inputting `over_sample_size=2`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_sub_2 = al.Grid2D.uniform(\n", + " shape_native=(40, 40),\n", + " pixel_scales=0.1,\n", + " over_sample_size=2,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we print `grid_sub_2` and its shape, we will find it is actually identical to `grid_sub_1`, despite the change\n", + "in `over_sample_size`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(grid_sub_1)\n", + "print(grid_sub_2)\n", + "print(grid_sub_1.shape)\n", + "print(grid_sub_2.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This is because the over sampled version of the grid is stored in a separate attribute, called `over_sampled`,\n", + "which we print below.\n", + "\n", + "We see that for `grid_sub_1` and `grid_sub_2` the `over_sampled` grids are different, with the over sampled grid for\n", + "`grid_sub_2` containing four times as many entries corresponding to each pixel being sub-gridded in a 2 x 2 shape." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(grid_sub_1.over_sampled)\n", + "print(grid_sub_2.over_sampled)\n", + "print(grid_sub_1.over_sampled.shape)\n", + "print(grid_sub_2.over_sampled.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now plot the over sampled grid over the image, showing that each pixel is now split into a 2x2 sub-grid of \n", + "coordinates. \n", + "\n", + "These are used to compute the intensity of the light profile and therefore more accurately estimate the total \n", + "intensity within each pixel if there is a significant gradient in intensity within the pixel.\n", + "\n", + "In the code below, it is the input `plot_over_sampled_grid=True` which ensures we plot the over sampled grid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_grid(grid=grid_sub_2, title=\"Over-Sampled Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Numerics__\n", + "\n", + "Lets quickly check how the sub-grid is defined and stored numerically.\n", + "\n", + "The first four pixels of this sub-grid correspond to the first four sub-pixels in the first pixel of the grid. \n", + "\n", + "The top-left pixel image above shows how the sub-pixels are spaced within the pixel. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"(y,x) pixel 0 of grid_sub_1:\")\n", + "print(grid_sub_1.over_sampled[0])\n", + "print(\"(y,x) pixel 0 of grid_sub_2:\")\n", + "print(grid_sub_2.over_sampled[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now confirm that the first four sub-pixels of the over-sampled grid correspond are contained within the \n", + "first pixel of the grid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"(y,x) pixel 0 (of original grid):\")\n", + "print(grid_sub_2[0])\n", + "print(\"(y,x) sub-pixel 0 (of pixel 0):\")\n", + "print(grid_sub_2.over_sampled[0])\n", + "print(\"(y,x) sub-pixel 1 (of pixel 0):\")\n", + "print(grid_sub_2.over_sampled[1])\n", + "print(\"(y,x) sub-pixel 2 (of pixel 0):\")\n", + "print(grid_sub_2.over_sampled[2])\n", + "print(\"(y,x) sub-pixel 3 (of pixel 0):\")\n", + "print(grid_sub_2.over_sampled[3])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Numerically, the over-sampled grid contains the sub-pixel coordinates of every pixel in the grid, going from the \n", + "first top-left pixel right and downwards to the bottom-right pixel. \n", + "\n", + "So the pixel to the right of the first pixel is the next 4 sub-pixels in the over-sampled grid, and so on.\n", + "\n", + "__Images__\n", + "\n", + "We now use over-sampling to compute the image of a Sersic light profile, which has a steep intensity gradient\n", + "at its centre which a lack of over-sampling does not accurately capture.\n", + "\n", + "We create the light profile, input the two grids (with `over_sample_size=1` and `over_sample_size=2`) and compute the image of the\n", + "light profile using each grid. We then plot the residuals between the two images in order to show the difference\n", + "between the two images and thus why over-sampling is important.\n", + "\n", + "Over sampling occurs automatically when a grid is input into a function like `image_2d_from`, therefore internally \n", + "the line of code, `image_sub_2 = light.image_2d_from(grid=grid_sub_2)`, is evaluating the light profile using the\n", + "2 x 2 oversampled grid and internally binning it up in to fully perform over sampling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "light = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=1.0,\n", + " effective_radius=0.2,\n", + " sersic_index=3.0,\n", + ")\n", + "\n", + "image_sub_1 = light.image_2d_from(grid=grid_sub_1)\n", + "image_sub_2 = light.image_2d_from(grid=grid_sub_2)\n", + "\n", + "aplt.plot_array(array=image_sub_1, title=\"Image of Sersic Profile\")\n", + "\n", + "residual_map = image_sub_2 - image_sub_1\n", + "\n", + "aplt.plot_array(array=residual_map, title=\"Over-Sampling Residuals\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the central 4 pixels of the image, the residuals are large due to the steep intensity gradient of the Sersic\n", + "profile at its centre. \n", + "\n", + "The gradient in these pixels is so steep that evaluating the intensity at the centre of the pixel, without over \n", + "sampling, does not accurately capture the total intensity within the pixel.\n", + "\n", + "At the edges of the image, the residuals are very small, as the intensity gradient of the Sersic profile is very \n", + "shallow and it is an accurate approximation to evaluate the intensity at the centre of the pixel.\n", + "\n", + "The absolute value of the central residuals are 0.74, however it is difficult to assess whether this is a large or\n", + "small value. We can quantify this by dividing by the evaluated value of the Sersic image in each pixel in order\n", + "to compute the fractional residuals." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fractional_residual_map = residual_map / image_sub_2\n", + "\n", + "\n", + "aplt.plot_array(\n", + " array=fractional_residual_map, title=\"Fractional Over-Sampling Residuals\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fractional residuals in the centre exceed 0.1, or 10%, which is a significant error in the image and\n", + "demonstrates why over-sampling is important.\n", + "\n", + "Lets confirm sub-griding can converge to central residuals that are very small.\n", + "\n", + "The fractional residuals with high levels of over-sampling are below 0.01, or 1%, which is sufficiently accurate\n", + "for most scientific purposes (albeit you should think carefully about the level of over-sampling you need for\n", + "your specific science case)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_sub_16 = al.Grid2D.uniform(\n", + " shape_native=(40, 40), pixel_scales=0.1, over_sample_size=16\n", + ")\n", + "grid_sub_32 = al.Grid2D.uniform(\n", + " shape_native=(40, 40), pixel_scales=0.1, over_sample_size=32\n", + ")\n", + "\n", + "image_sub_16 = light.image_2d_from(grid=grid_sub_16)\n", + "image_sub_32 = light.image_2d_from(grid=grid_sub_32)\n", + "\n", + "residual_map = image_sub_32 - image_sub_16\n", + "\n", + "aplt.plot_array(array=residual_map, title=\"Over-Sampling Reduces Residuals\")\n", + "\n", + "fractional_residual_map = residual_map / image_sub_32\n", + "\n", + "aplt.plot_array(\n", + " array=fractional_residual_map, title=\"Fractional Residuals With Over-Sampling\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Adaptive Over Sampling__\n", + "\n", + "We have shown that over-sampling is important for accurate image evaluation. However, there is a major drawback to\n", + "over-sampling, which is that it is computationally expensive. \n", + "\n", + "For example, for the 32x32 over-sampled grid above, 1024 sub-pixels are used in every pixel, which must all be \n", + "evaluated using the Sersic light profile. The calculation of the image is therefore at least 1000 times slower than if\n", + "we had not used over-sampling.\n", + "\n", + "Speeding up the calculation is crucial for model-fitting where the image is evaluated many times to fit the\n", + "model to the data.\n", + "\n", + "Fortunately, there is a solution to this problem. We saw above that the residuals rapidly decrease away\n", + "from the centre of the light profile. Therefore, we only need to over-sample the central regions of the image,\n", + "where the intensity gradient is steep. We can use lower levels of over-sampling away from the centre, which\n", + "will be fast to evaluate.\n", + "\n", + "Up to now, the `over_sample_size` input has been an integer, however it can also be an `ndarray` of values corresponding\n", + "to each pixel. We create an `ndarray` of values which are high in the centre, but reduce to 2 at the outskirts,\n", + "therefore providing high levels of over sampling where we need it whilst using lower values which are computationally\n", + "fast to evaluate at the outskirts.\n", + "\n", + "Specifically, we define a 24 x 24 sub-grid within the central 0.3\" of pixels, uses a 8 x 8 grid between\n", + "0.3\" and 0.6\" and a 2 x 2 grid beyond that. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid_sub_1, sub_size_list=[24, 8, 2], radial_list=[0.3, 0.6]\n", + ")\n", + "\n", + "grid_adaptive = al.Grid2D.no_mask(\n", + " values=grid_sub_1.native,\n", + " pixel_scales=grid_sub_1.pixel_scales,\n", + " over_sample_size=over_sample_size,\n", + ")\n", + "\n", + "\n", + "aplt.plot_grid(grid=grid_adaptive, title=\"Over-Sampled Grid\")\n", + "\n", + "print(over_sample_size)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling uses masked grids, therefore the code below shows how we would create this adaptive over sample grid via \n", + "a circular mask, which can be used for modeling.\n", + "\n", + "Throughout the modeling examples in the workspace, we use this adaptive grid to ensure that the image of the\n", + "galaxy is evaluated accurately and efficiently." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(shape_native=(40, 40), pixel_scales=0.1, radius=5.0)\n", + "\n", + "grid = al.Grid2D.from_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid, sub_size_list=[24, 8, 2], radial_list=[0.3, 0.6]\n", + ")\n", + "\n", + "grid_adaptive = al.Grid2D(values=grid, mask=mask, over_sample_size=over_sample_size)\n", + "\n", + "\n", + "aplt.plot_grid(grid=grid, title=\"Over-Sampled Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can compare this adaptive grid to the grid with over sampling of 32 x 32 to confine it produces low amounts\n", + "of residuals." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_masked_sub_32 = al.Grid2D(values=grid, mask=mask, over_sample_size=32)\n", + "\n", + "image_adaptive = light.image_2d_from(grid=grid_adaptive)\n", + "image_masked_sub_32 = light.image_2d_from(grid=grid_masked_sub_32)\n", + "\n", + "residual_map = image_adaptive - image_masked_sub_32\n", + "\n", + "fractional_residual_map = residual_map / image_masked_sub_32\n", + "\n", + "\n", + "aplt.plot_array(array=fractional_residual_map, title=\"Adaptive Fractional Residuals\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Default Over-Sampling__\n", + "\n", + "The default over-sampling scheme used by the source code is 4 x 4 uniform over sampling over the whole image. \n", + "\n", + "A uniform scheme is used, instead of the adaptive scheme above, because the adaptive scheme requires input knowledge of \n", + "where the centre of the galaxy is (e.g. above the centre is at (0.0\", 0.0\").\n", + "\n", + "Uniform over sampling is precise enough for many calculations, especially when you are simply performing quick \n", + "calculations to investigate a problem. However, for detailed calculations you must ensure that high enough\n", + "levels of over sampling are used.\n", + "\n", + "For modeling, all example scripts begin by switching to an adaptive over sampling scheme, as modeling assumes\n", + "the centre of the lens galaxy is at (0.0\", 0.0\").\n", + "\n", + "__Multiple Lens Galaxies__\n", + "\n", + "The analysis may contain multiple lens galaxies, each of which must be over-sampled accurately. \n", + "\n", + "\n", + "There are two approaches you can take to over sampling multi-galaxy systems:\n", + "\n", + "1) Use a high level of uniform over sampling over the full image.\n", + "\n", + "2) Use an adaptive over sampling scheme with multiple centres of high over sampling levels, with the API shown below\n", + " for two galaxies with centres (1.0, 0.0) and (-1.0, 0.0)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid_sub_1,\n", + " sub_size_list=[24, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(1.0, 0.0), (-1.0, 0.0)],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "So far, we have evaluated the image of a light profile using over-sampling on an unlensed uniform grid. \n", + "\n", + "For lensing calculations, the grid is ray-traced via a mass model to an irregular grid in the source plane.\n", + "\n", + "The over sampling of lensed images is therefore describe as follows: \n", + "\n", + "1) Splits each image-pixel into a sub-grid of pixels in the image-plane.\n", + "2) Ray trace this sub-grid of pixels using the mass model.\n", + "3) Evaluate the source light of each sub-pixel in the source-plane.\n", + "4) Bin up the values to evaluate the over sampled values.\n", + "\n", + "For lensing calculations, over sampling therefore requires us to ray-trace (and therefore compute the deflecitons angles\n", + "of) many more (y,x) coordinates!\n", + "\n", + "We now illustrate using over-sampling with a mass profile, noting that for lensing:\n", + "\n", + "1) The fractional residuals due to differing over-sampling levels now occur in the lensed source's brightest multiply \n", + " imaged pixels. \n", + " \n", + "2) It is the combination of a rapidly changing source light profile and the magnification pattern of the mass model\n", + " which requires over sampling. The mass model focuses many image-pixels to the source's brightest regions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = al.mp.Isothermal(centre=(0.0, 0.0), ell_comps=(0.1, 0.0), einstein_radius=1.0)\n", + "\n", + "light = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=1.0,\n", + " effective_radius=0.2,\n", + " sersic_index=3.0,\n", + ")\n", + "\n", + "lens = al.Galaxy(redshift=0.5, mass=mass)\n", + "\n", + "source = al.Galaxy(redshift=1.0, bulge=light)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "image_sub_1 = tracer.image_2d_from(grid=grid_sub_1)\n", + "image_sub_2 = tracer.image_2d_from(grid=grid_sub_2)\n", + "\n", + "aplt.plot_array(array=image_sub_1, title=\"Source Image 1x1\")\n", + "\n", + "residual_map = image_sub_2 - image_sub_1\n", + "\n", + "fractional_residual_map = residual_map / image_sub_2\n", + "\n", + "aplt.plot_array(array=fractional_residual_map, title=\"Fractional Residuals\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Default Ray Tracing__\n", + "\n", + "By assuming the lens galaxy is near (0.0\", 0.0\"), it was simple to set up an adaptive over sampling grid which is\n", + "applicable to all strong lens dataset.\n", + "\n", + "An adaptive oversampling grid cannot be defined for the lensed source because its light appears in different regions of \n", + "the image plane for each dataset. For this reason, most workspace examples utilize cored light profiles for the \n", + "source galaxy. Cored light profiles change gradually in their central regions, allowing accurate evaluation without \n", + "requiring oversampling.\n", + "\n", + "__Adaptive Over Sampling__\n", + "\n", + "There is a way to set up an adaptive over sampling grid for a lensed source, however it requries one to use and\n", + "understanding the advanced lens modeling feature search chaining.\n", + "\n", + "An example of how to use search chaining to over sample sources efficient is provided in \n", + "the `autolens_workspace/*/guides/modeling/chaining/over_sampling.ipynb` example.\n", + "\n", + "__Dataset & Modeling__\n", + "\n", + "Throughout this guide, grid objects have been used to compute the image of light and mass profiles and illustrate\n", + "over sampling.\n", + "\n", + "If you are performing calculations with imaging data or want to fit a model to the data with a specific\n", + "over-sampling level, the `apply_over_sampling` method is used to update the over sampling scheme of the dataset.\n", + "\n", + "The grid this is applied to is called `lp`, to indicate that it is the grid used to evaluate the emission of light\n", + "profiles for which this over sampling scheme is applied." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "# This can be any of the over-sampling objects we have used above.\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=4)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixelization__\n", + "\n", + "Source galaxies can be reconstructed using pixelizations, which discretize the source's light onto a mesh,\n", + "for example a Voronoi mesh.\n", + "\n", + "Over sampling is used by pixelizations in an analogous way to light profiles. By default, a 4 x 4 sub-grid is used,\n", + "whereby every image pixel is ray-traced on its 4 x 4 sub grid to the source mesh and fractional mappings are computed.\n", + "\n", + "A different grid and over sampling scheme is applied to light profiles and pixelizations, which is why\n", + "there are separate inputs called `lp` and `pix`.\n", + "\n", + "This is explained in more detail in the pixelization examples.\n", + "\n", + "Here is an example of how to change the over sampling applied to a pixelization for a lens model fit:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=4)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Oversampled PSF Convolution__\n", + "\n", + "Everything above concerns over sampling the *evaluation* of light profiles \u2014 computing the image on a sub-grid\n", + "and binning it to pixel values before any instrument effect is applied. PSF convolution then blurs those\n", + "pixel values at the resolution of the image.\n", + "\n", + "Convolution itself can also be over sampled. This matters when the PSF is undersampled by the detector (its\n", + "width is comparable to the pixel scale, as for HST or Euclid VIS imaging): blurring the binned image with a\n", + "pixel-scale PSF loses the sub-pixel structure of both the image and the PSF, and the two operations (bin then\n", + "convolve, versus convolve finely then bin) do not commute.\n", + "\n", + "Supplying the PSF at a multiple of the image resolution and setting `convolve_over_sample_size` performs the\n", + "convolution on the over-sampled grid and bins the result back to image resolution:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf_fine = al.Convolver.from_gaussian(\n", + " shape_native=(21, 21), # twice the pixels of an 11x11 image-resolution kernel...\n", + " pixel_scales=0.1 / 2, # ...at half the pixel scale, so the same physical extent.\n", + " sigma=0.05,\n", + " normalize=True,\n", + " convolve_over_sample_size=2,\n", + ")\n", + "\n", + "dataset_fine = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For a dataset, the sizes are set per operation, mirroring the `lp` / `pixelization` split above. The matching\n", + "evaluation `over_sample_size` must be a uniform integer equal to the convolution size \u2014 the values convolved at\n", + "the fine resolution must first be evaluated at exactly that resolution, so adaptive over sampling cannot be\n", + "combined with oversampled convolution (the code raises a clear error rather than silently degrading):\n", + "\n", + "```python\n", + "dataset = al.Imaging(\n", + " data=data,\n", + " noise_map=noise_map,\n", + " psf=psf_fine,\n", + " over_sample_size_lp=2,\n", + " over_sample_size_pixelization=2,\n", + " convolve_over_sample_size_lp=2,\n", + " convolve_over_sample_size_pixelization=2,\n", + ")\n", + "```\n", + "\n", + "Fits then work across every model surface: standard light profiles, linear light profiles, operated light\n", + "profiles (which are added at image resolution unblurred, by definition) and pixelized source reconstructions.\n", + "Simulation is supported too \u2014 see the `__Oversampled PSF__` section of `scripts/imaging/simulator.py`.\n", + "\n", + "__Adaptive Evaluation: the k x s Coupling__\n", + "\n", + "Adaptive over sampling composes with oversampled convolution: every `over_sample_size` entry must be\n", + "*divisible* by `convolve_over_sample_size` (not equal to it). Each pixel is evaluated on its own adaptive\n", + "sub-grid of size k * s, the values are partially binned to a uniform image at the convolution resolution s,\n", + "convolved, and binned to image resolution \u2014 adaptive sampling does the integration accuracy, the uniform\n", + "intermediate does the convolution. The adaptive radial schemes above therefore work unchanged with an\n", + "oversampled PSF, provided their sizes are multiples of s (e.g. [32, 8, 2] with s=2). A non-divisible\n", + "combination raises a clear error. This holds for pixelized sources too, where the coupling is exact by\n", + "linearity of the mapping matrix.\n", + "\n", + "__Limitations__\n", + "\n", + "The following are not supported with an oversampled PSF and raise a clear error if combined with it:\n", + "\n", + " - Over sampling sizes not divisible by `convolve_over_sample_size` (e.g. 3 with s=2).\n", + " - The sparse linear-algebra formalism (`apply_sparse_operator`), whose PSF products are precomputed at image\n", + " resolution.\n", + " - The fixed-linear-function preload (`data_linear_func_matrix`) used to accelerate some fixed-MGE fits.\n", + "\n", + "The padded / unmasked visualization images (`padded_image_2d_from`, `unmasked_blurred_image_2d_from`) remain at\n", + "image resolution \u2014 visualization is unaffected by the convolution accuracy improvement.\n", + "\n", + "The numerical verification of all of the above lives in\n", + "`autolens_workspace_test/scripts/imaging/convolution_over_sampled.py`, which pins the implementation to\n", + "brute-force reference calculations.\n", + "\n", + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/advanced/over_sampling_chaining.ipynb b/notebooks/guides/advanced/over_sampling_chaining.ipynb index 33db13685..64664dc09 100644 --- a/notebooks/guides/advanced/over_sampling_chaining.ipynb +++ b/notebooks/guides/advanced/over_sampling_chaining.ipynb @@ -1,497 +1,534 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Chaining: Over Sample\n", - "=====================\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", - "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "If you are reading this example, you should be familiar with over sampling already. If this is not the case,\n", - "checkout the over sampling guide at `autolens_workspace/*/guides/over_sampling.py`.\n", - "\n", - "The guide illustrated adaptive over sampling, where the over sampling sub grid used high resolution pixels in the\n", - "centre of a light profile and lower resolution pixels further out. This reached high levels of numerical accuracy\n", - "with efficient run times.\n", - "\n", - "However, the adaptive sub grid only works for uniform grids which have not been deflected or ray-traced by the lens\n", - "mass model. This criteria is met for the lens galaxy's light, but not for the emission of the lensed source. There\n", - "is no automatic adaptive method for the lensed source, which is why the autolens workspace uses cored light profiles\n", - "throughout.\n", - "\n", - "An efficient and adaptive over sampling grid is possible. However, it requires using search chaining, where between\n", - "searches the over sampling grid is updated.\n", - "\n", - "This example shows how to combine lensed source adaptive over sampling with search chaining.\n", - "\n", - "__Contents__\n", - "\n", - "- **Start Here Notebook:** If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook.\n", - "- **Dataset + Masking:** Load, plot and mask the `Imaging` data.\n", - "- **Paths:** The path the results of all chained searches are output.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Model (Search 1):** Search 1 fits a lens model where the lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` and the source galaxy's light is an MGE.\n", - "- **Search + Analysis + Model-Fit (Search 1):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "- **Result (Search 1):** The results which are used for prior passing are summarised in the `info` attribute.\n", - "- **Over Sampling (Search 2):** We update the over sampling grid between searches using the results of search 1.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import os\n", - "import sys\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "sys.path.insert(0, os.getcwd())" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__ \n", - "\n", - "Load, plot and mask the `Imaging` data.\n", - "\n", - "The data is simulated using a Sersic without a core, unlike most datasets fitted throughout the workspace." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__\n", - "\n", - "The path the results of all chained searches are output:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"imaging\") / \"chaining\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Search 1 fits a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", - " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", - "provides a reasonably accurate lens model.\n", - "\n", - "Faint residuals around the multiple images will be present, because the simulated data used a non-cored Sersic\n", - "whereas the model fitted is a cored Sersic." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__sersic_core\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 1)__\n", - "\n", - "The results which are used for prior passing are summarised in the `info` attribute." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling (Search 2)__\n", - "\n", - "We now create an over sampling grid which applies high levels of over sampling to the brightest regions of the\n", - "lensed source in the image plane.\n", - "\n", - "This uses the result of the first lens model in the following way:\n", - "\n", - " 1) Use the lens mass model to ray-trace every deflected image pixel to the source plane, computed the traced grid.\n", - " \n", - " 2) Use the traced grid and the centre of the source light profile to compute the distance of every traced image pixel \n", - " to the source centre. \n", - " \n", - " 3) For all pixels with a distance below a threshold value of 0.1\", we set the over sampling factor to a high value of \n", - " 32, which will ensure accuracy in the evaluated of the source's light profile, even after lensing. Pixels 0.1\" to\n", - " 0.3\" from the centre use an over sampling factor of 4, and all other pixels use an over sampling factor of 2. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "tracer = result_1.max_log_likelihood_tracer\n", - "\n", - "traced_grid = tracer.traced_grid_2d_list_from(\n", - " grid=dataset.grid,\n", - ")[-1]\n", - "\n", - "source_centre = tracer.galaxies[1].bulge.centre\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=traced_grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=[source_centre],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "One thing to note is this over sampling grid, although specific to the lensed source, is also applied to the lens\n", - "galaxy's light. Other than slower computation times, this is not a problem, as the lens galaxy's light in these\n", - "pixels will still be evaluated accurately.\n", - "\n", - "The data fitted in this example omits the lens light for simplicity, however if it were present we would want\n", - "the lens galaxy's light to be over sampled in the same way as the source galaxy's light because we still require high \n", - "levels of over sampling in the lens galaxy's light to evaluate it correctly. \n", - "\n", - "We therefore create an over sampling grid which is centred on the lens galaxy's light and combine these values with \n", - "those found for the source galaxy's light, to ensure over sampling is centred on the brightest regions of both galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size_lens = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "over_sample_size = np.where(\n", - " over_sample_size > over_sample_size_lens, over_sample_size, over_sample_size_lens\n", - ")\n", - "over_sample_size = al.Array2D(values=over_sample_size, mask=mask)\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Search 1 fits a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", - " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.Sersic)\n", - "\n", - "model_2 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", - "provides a reasonably accurate lens model.\n", - "\n", - "Faint residuals around the multiple images will be present, because the simulated data used a non-cored Sersic\n", - "whereas the model fitted is a cored Sersic." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__sersic_over_sampled\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(dataset=dataset)\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 1)__\n", - "\n", - "The results which are used for prior passing are summarised in the `info` attribute." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Chaining: Over Sample\n", + "=====================\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", + "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "If you are reading this example, you should be familiar with over sampling already. If this is not the case,\n", + "checkout the over sampling guide at `autolens_workspace/*/guides/over_sampling.py`.\n", + "\n", + "The guide illustrated adaptive over sampling, where the over sampling sub grid used high resolution pixels in the\n", + "centre of a light profile and lower resolution pixels further out. This reached high levels of numerical accuracy\n", + "with efficient run times.\n", + "\n", + "However, the adaptive sub grid only works for uniform grids which have not been deflected or ray-traced by the lens\n", + "mass model. This criteria is met for the lens galaxy's light, but not for the emission of the lensed source. There\n", + "is no automatic adaptive method for the lensed source, which is why the autolens workspace uses cored light profiles\n", + "throughout.\n", + "\n", + "An efficient and adaptive over sampling grid is possible. However, it requires using search chaining, where between\n", + "searches the over sampling grid is updated.\n", + "\n", + "This example shows how to combine lensed source adaptive over sampling with search chaining.\n", + "\n", + "__Contents__\n", + "\n", + "- **Start Here Notebook:** If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook.\n", + "- **Dataset + Masking:** Load, plot and mask the `Imaging` data.\n", + "- **Paths:** The path the results of all chained searches are output.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Model (Search 1):** Search 1 fits a lens model where the lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` and the source galaxy's light is an MGE.\n", + "- **Search + Analysis + Model-Fit (Search 1):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "- **Result (Search 1):** The results which are used for prior passing are summarised in the `info` attribute.\n", + "- **Over Sampling (Search 2):** We update the over sampling grid between searches using the results of search 1.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import os\n", + "import sys\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "sys.path.insert(0, os.getcwd())" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__ \n", + "\n", + "Load, plot and mask the `Imaging` data.\n", + "\n", + "The data is simulated using a Sersic without a core, unlike most datasets fitted throughout the workspace." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__\n", + "\n", + "The path the results of all chained searches are output:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"imaging\") / \"chaining\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Search 1 fits a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", + " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", + "provides a reasonably accurate lens model.\n", + "\n", + "Faint residuals around the multiple images will be present, because the simulated data used a non-cored Sersic\n", + "whereas the model fitted is a cored Sersic." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__sersic_core\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 1)__\n", + "\n", + "The results which are used for prior passing are summarised in the `info` attribute." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling (Search 2)__\n", + "\n", + "We now create an over sampling grid which applies high levels of over sampling to the brightest regions of the\n", + "lensed source in the image plane.\n", + "\n", + "This uses the result of the first lens model in the following way:\n", + "\n", + " 1) Use the lens mass model to ray-trace every deflected image pixel to the source plane, computed the traced grid.\n", + " \n", + " 2) Use the traced grid and the centre of the source light profile to compute the distance of every traced image pixel \n", + " to the source centre. \n", + " \n", + " 3) For all pixels with a distance below a threshold value of 0.1\", we set the over sampling factor to a high value of \n", + " 32, which will ensure accuracy in the evaluated of the source's light profile, even after lensing. Pixels 0.1\" to\n", + " 0.3\" from the centre use an over sampling factor of 4, and all other pixels use an over sampling factor of 2. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "tracer = result_1.max_log_likelihood_tracer\n", + "\n", + "traced_grid = tracer.traced_grid_2d_list_from(\n", + " grid=dataset.grid,\n", + ")[-1]\n", + "\n", + "source_centre = tracer.galaxies[1].bulge.centre\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=traced_grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=[source_centre],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "One thing to note is this over sampling grid, although specific to the lensed source, is also applied to the lens\n", + "galaxy's light. Other than slower computation times, this is not a problem, as the lens galaxy's light in these\n", + "pixels will still be evaluated accurately.\n", + "\n", + "The data fitted in this example omits the lens light for simplicity, however if it were present we would want\n", + "the lens galaxy's light to be over sampled in the same way as the source galaxy's light because we still require high \n", + "levels of over sampling in the lens galaxy's light to evaluate it correctly. \n", + "\n", + "We therefore create an over sampling grid which is centred on the lens galaxy's light and combine these values with \n", + "those found for the source galaxy's light, to ensure over sampling is centred on the brightest regions of both galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size_lens = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "over_sample_size = np.where(\n", + " over_sample_size > over_sample_size_lens, over_sample_size, over_sample_size_lens\n", + ")\n", + "over_sample_size = al.Array2D(values=over_sample_size, mask=mask)\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Search 1 fits a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", + " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.Sersic)\n", + "\n", + "model_2 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", + "provides a reasonably accurate lens model.\n", + "\n", + "Faint residuals around the multiple images will be present, because the simulated data used a non-cored Sersic\n", + "whereas the model fitted is a cored Sersic." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__sersic_over_sampled\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(dataset=dataset)\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 1)__\n", + "\n", + "The results which are used for prior passing are summarised in the `info` attribute." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/data_structures.ipynb b/notebooks/guides/data_structures.ipynb index 83842ad32..7e542a038 100644 --- a/notebooks/guides/data_structures.ipynb +++ b/notebooks/guides/data_structures.ipynb @@ -1,838 +1,875 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Structures\n", - "===============\n", - "\n", - "This tutorial illustrates the data structure objects which data and results quantities are stored using, which are\n", - "extensions of NumPy arrays.\n", - "\n", - "These data structures are used because for different lensing calculations it is convenient to store the data in\n", - "different formats. For example, when ray-tracing a uniform grid of image-plane (y,x) coordinates, to an irregular\n", - "grid of source-plane (y,x) coordinates, the image-plane coordinates can be stored in 2D (because the grid is uniform)\n", - "whereas the source-plane coordinates must be stored in 1D (because after lensing it is irregular).\n", - "\n", - "These data structures use the `slim` and `native` data representations API to make it simple to map quantities from\n", - "1D dimensions to their native dimensions.\n", - "\n", - "__Contents__\n", - "\n", - "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", - "- **API:** We discuss in detail why these data structures and illustrate their functionality below.\n", - "- **Grids:** We now illustrate data structures using a `Grid2D` object, which is a set of two-dimensional.\n", - "- **Native:** This plot shows the grid in its `native` format, that is in 2D dimensions where the y and x.\n", - "- **Slim:** Every `Grid2D` object is accessible via two attributes, `native` and `slim`, which store the grid.\n", - "- **Masked Data Structures:** When a mask is applied to a grid or other data structure, this changes the `slim` and `native`.\n", - "- **Data:** Two dimensional arrays of data are stored using the `Array2D` object, which has `slim` and `native`.\n", - "- **Tracer:** The `Tracer` produces many lensing quantities all of which use the `slim` and `native` data.\n", - "- **Irregular Structures:** We may want to perform calculations at specific (y,x) coordinates which are not tied to a uniform.\n", - "- **Vector Quantities:** Many lensing quantities are vectors.\n", - "\n", - "__Units__\n", - "\n", - "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", - "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", - "\n", - "The guide `units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", - "kiloparsecs, magnitudes and solar masses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__API__\n", - "\n", - "We discuss in detail why these data structures and illustrate their functionality below.\n", - "\n", - "However, we first create the three data structures we'll use in this example, to set expectations for what they do.\n", - "\n", - "We create three data structures:\n", - "\n", - " - `Array2D`: A 2D array of data, which is used for storing an image, a noise-map, etc. \n", - "\n", - " - `Grid2D`: A 2D array of (y,x) coordinates, which is used for ray-tracing.\n", - "\n", - " -`VectorYX2D`: A 2D array of vector values, which is used for deflection angles, shear and other vector fields.\n", - "\n", - "All data structures are defined according to a uniform grid of coordinates and therefore they have a `pixel_scales`\n", - "input defining the pixel-to-arcssecond conversion factor of its grid. \n", - "\n", - "For example, for an image stored as an `Array2D`, it has a grid where each coordinate is the centre of each image pixel\n", - "and the pixel-scale is therefore the resolution of the image.\n", - "\n", - "We first create each data structure without a mask using the `no_mask` method:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "arr = al.Array2D.no_mask(\n", - " values=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], pixel_scales=1.0\n", - ")\n", - "\n", - "print(arr)\n", - "\n", - "grid = al.Grid2D.no_mask(\n", - " values=[\n", - " [[-1.0, -1.0], [-1.0, 0.0], [-1.0, 1.0]],\n", - " [[0.0, -1.0], [0.0, 0.0], [0.0, 1.0]],\n", - " [\n", - " [1.0, -1.0],\n", - " [1.0, 0.0],\n", - " [1.0, 1.0],\n", - " ],\n", - " ],\n", - " pixel_scales=1.0,\n", - ")\n", - "\n", - "print(grid)\n", - "\n", - "vector_yx = al.VectorYX2D.no_mask(\n", - " values=[\n", - " [[5.0, -5.0], [5.0, 0.0], [5.0, 5.0]],\n", - " [[0.0, -5.0], [0.0, 0.0], [0.0, 5.0]],\n", - " [\n", - " [-5.0, -5.0],\n", - " [-5.0, 0.0],\n", - " [-5.0, 5.0],\n", - " ],\n", - " ],\n", - " pixel_scales=1.0,\n", - ")\n", - "\n", - "print(vector_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grids__\n", - "\n", - "We now illustrate data structures using a `Grid2D` object, which is a set of two-dimensional $(y,x)$ coordinates\n", - "(in arc-seconds) that are deflected and traced by a strong lensing system.\n", - "\n", - "These are fundamental to all lensing calculations and drive why data structures are used.\n", - "\n", - "First, lets make a uniform 100 x 100 grid of (y,x) coordinates and plot it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "\n", - "aplt.plot_grid(grid=grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Native__\n", - "\n", - "This plot shows the grid in its `native` format, that is in 2D dimensions where the y and x coordinates are plotted\n", - "where we expect them to be on the grid.\n", - "\n", - "We can print values from the grid's `native` property to confirm this:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"(y,x) pixel 0:\")\n", - "print(grid.native[0, 0])\n", - "print(\"(y,x) pixel 1:\")\n", - "print(grid.native[0, 1])\n", - "print(\"(y,x) pixel 2:\")\n", - "print(grid.native[0, 2])\n", - "print(\"(y,x) pixel 100:\")\n", - "print(grid.native[1, 0])\n", - "print(\"etc.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Slim__\n", - "\n", - "Every `Grid2D` object is accessible via two attributes, `native` and `slim`, which store the grid as NumPy ndarrays \n", - "of two different shapes:\n", - " \n", - " - `native`: an ndarray of shape [total_y_image_pixels, total_x_image_pixels, 2] which is the native shape of the \n", - " 2D grid and corresponds to the resolution of the image datasets we pair with a grid.\n", - " \n", - " - `slim`: an ndarray of shape [total_y_image_pixels*total_x_image_pixels, 2] which is a slimmed-down representation \n", - " the grid which collapses the inner two dimensions of the native ndarray to a single dimension." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"(y,x) pixel 0 (accessed via native):\")\n", - "print(grid.native[0, 0])\n", - "print(\"(y,x) pixel 0 (accessed via slim 1D):\")\n", - "print(grid.slim[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As discussed above, the reason we need the slim representation is because when we ray-trace a grid of (y,x) coordinates\n", - "from the image-plane to the source-plane, the source-plane grid will be irregular.\n", - "\n", - "The shapes of the `Grid2D` in its `native` and `slim` formats are also available, confirming that this grid has a \n", - "`native` resolution of (100 x 100) and a `slim` resolution of 10000 coordinates." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(grid.shape_native)\n", - "print(grid.shape_slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Neither shape above include the third index of the `Grid` which has dimensions 2 (corresponding to the y and x \n", - "coordinates). \n", - "\n", - "This is accessible by using the standard numpy `shape` method on each grid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(grid.native.shape)\n", - "print(grid.slim.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can print the entire `Grid2D` in its `slim` or `native` form. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(grid.native)\n", - "print(grid.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masked Data Structures__\n", - "\n", - "When a mask is applied to a grid or other data structure, this changes the `slim` and `native` representations as \n", - "follows:\n", - "\n", - " - `slim`: only contains image-pixels that are not masked, removing all masked pixels from the 1D array.\n", - " \n", - " - `native`: retains the dimensions [total_y_image_pixels, total_x_image_pixels], but the masked pixels have values\n", - " of 0.0 or (0.0, 0.0).\n", - "\n", - "This can be seen by computing a grid via a mask and comparing the its`shape_slim` attribute to the `pixels_in_mask` of \n", - "the mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(shape_native=(100, 100), pixel_scales=0.05, radius=3.0)\n", - "\n", - "grid = al.Grid2D.from_mask(mask=mask)\n", - "\n", - "print(\"The shape_slim and number of unmasked pixels\")\n", - "print(grid.shape_slim)\n", - "print(mask.pixels_in_mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can use the `slim` attribute to print unmasked values of the grid:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"First unmasked image value:\")\n", - "print(grid.slim[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `native` representation of the `Grid2D` retains the dimensions [total_y_image_pixels, total_x_image_pixels], \n", - "however the exterior pixels have values of 0.0 indicating that they have been masked." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Example masked pixels in the grid native representation:\")\n", - "print(grid.shape_native)\n", - "print(grid.native[0, 0])\n", - "print(grid.native[2, 2])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Data__\n", - "\n", - "Two dimensional arrays of data are stored using the `Array2D` object, which has `slim` and `native` representations\n", - "analogous to the `Grid2D` object and described as follows:\n", - "\n", - " - `slim`: an ndarray of shape [total_unmasked_pixels] which is a slimmed-down representation of the data in 1D that \n", - " contains only the unmasked data points (where this mask is the one used by the model-fit above).\n", - "\n", - " - `native`: an ndarray of shape [total_y_image_pixels, total_x_image_pixels], which is the native shape of the \n", - " masked 2D grid used to fit the lens model. All masked pixels are assigned a value 0.0 in the `native` array.\n", - "\n", - "For example, the `data` and `noise_map` in an `Imaging` object are stored as `Array2D` objects.\n", - "\n", - "We load an imaging dataset and illustrate its data structures below. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "data = dataset.data" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here is what `slim` and `native` representations of the data's first pixel look like for the `data` before masking:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"First unmasked data value:\")\n", - "print(data.slim[0])\n", - "print(data.native[0, 0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By default, all arrays in **PyAutoLens** are stored as their `slim` 1D numpy array, meaning we don't need to use the\n", - "`slim` attribute to access the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(data[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By applying a mask the first value in `slim` changes and the native value becomes 0.0:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "data = dataset.data\n", - "\n", - "print(\"First unmasked data value:\")\n", - "print(data.slim[0])\n", - "print(data.native[0, 0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "The `Tracer` produces many lensing quantities all of which use the `slim` and `native` data structures.\n", - "\n", - "For example, by passing it a 2D grid of (y,x) coordinates we can return a numpy array containing its 2D image. \n", - "This includes the lens light and lensed source images.\n", - "\n", - "Below, we use the grid that is aligned with the imaging data (e.g. where each grid coordinate is at the centre of each\n", - "image pixel) to compute the galaxy image and show its data structure." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5, mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6)\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " light=al.lp.SersicCoreSph(\n", - " centre=(0.0, 0.0),\n", - " intensity=0.2,\n", - " effective_radius=0.2,\n", - " sersic_index=1.0,\n", - " radius_break=0.025,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "image = tracer.image_2d_from(grid=dataset.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we print the type of the `image` we note that it is an `Array2D`, which is a data structure that inherits \n", - "from a numpy array but is extended to include specific functionality discussed below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(type(image))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Because the image is a numpy array, we can print its shape and see that it is 1D." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(image.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Irregular Structures__\n", - "\n", - "We may want to perform calculations at specific (y,x) coordinates which are not tied to a uniform grid.\n", - "\n", - "We can use an irregular 2D (y,x) grid of coordinates for this. The grid below evaluates the image at:\n", - "\n", - "- y = 1.0, x = 1.0.\n", - "- y = 1.0, x = 2.0.\n", - "- y = 2.0, x = 2.0." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_irregular = al.Grid2DIrregular(values=[[1.0, 1.0], [1.0, 2.0], [2.0, 2.0]])\n", - "\n", - "image = tracer.image_2d_from(grid=grid_irregular)\n", - "\n", - "print(image)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The result is stored using an `ArrayIrregular` object, which is a data structure that handles irregular arrays." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(type(image))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Vector Quantities__\n", - "\n", - "Many lensing quantities are vectors. That is, they are (y,x) coordinates that have 2 values representing their\n", - "magnitudes in both the y and x directions.\n", - "\n", - "The most obvious of these is the deflection angles, which are used throughout lens modeling to ray-trace grids\n", - "from the image-plane to the source-plane via a lens galaxy mass model.\n", - "\n", - "To indicate that a quantities is a vector, **PyAutoLens** uses the label `_yx`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_yx_2d = tracer.deflections_yx_2d_from(grid=dataset.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we print the type of the `deflections_yx` we note that it is a `VectorYX2D`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(type(deflections_yx_2d))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Unlike the scalar quantities above, which were a 1D numpy array in the `slim` representation and a 2D numpy array in \n", - "the `native` representation, vectors are 2D in `slim` and 3D in `native`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(deflections_yx_2d.slim.shape)\n", - "print(deflections_yx_2d.native.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For vector quantities the has shape `2`, corresponding to the y and x vectors respectively." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(deflections_yx_2d.slim[0, :])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The role of the terms `slim` and `native` can be thought of in the same way as for scalar quantities. \n", - "\n", - "For a scalar, the `slim` property gives every scalar value as a 1D ndarray for every unmasked pixel. For a vector we \n", - "still get an ndarray of every unmasked pixel, however each entry now contains two entries: the vector of (y,x) values. \n", - "\n", - "For a `native` property these vectors are shown on an image-plane 2D grid where again each pixel\n", - "contains a (y,x) vector.\n", - "\n", - "Like we did for the convergence, we can use whatever grid we want to compute a vector and use sub-gridding to estimate\n", - "values more precisely:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(3, 3), pixel_scales=0.1)\n", - "\n", - "deflections_yx_2d = tracer.deflections_yx_2d_from(grid=grid)\n", - "\n", - "print(deflections_yx_2d.slim)\n", - "print(deflections_yx_2d.native)\n", - "\n", - "grid_irregular = al.Grid2DIrregular(values=[[1.0, 1.0], [1.0, 2.0], [2.0, 2.0]])\n", - "\n", - "deflections_yx_2d = tracer.deflections_yx_2d_from(grid=grid_irregular)\n", - "\n", - "print(deflections_yx_2d)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__JAX__\n", - "\n", - "PyAutoLens runs on either NumPy (the default) or JAX. The data structures\n", - "you've met above are *backend-polymorphic* \u2014 they're Python wrappers\n", - "around a numerical array, and that array can be a `numpy.ndarray` or a\n", - "`jax.Array` depending on how the structure was constructed and what code\n", - "path produced it.\n", - "\n", - "You can always reach the raw backing array via `.array`:\n", - "\n", - "```python\n", - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "print(type(grid.array)) # on the default path\n", - "```\n", - "\n", - "__When the backing becomes `jax.Array`__\n", - "\n", - "Three situations switch the backing to `jax.Array`:\n", - "\n", - "1. The structure comes back from a JAX-accelerated `Analysis(use_jax=True)`\n", - " fit \u2014 e.g. `fit.residual_map.array`, `fit.model_image.array` are\n", - " JAX-backed when the analysis ran with `use_jax=True` (the default).\n", - "2. The structure comes back from a `Simulator(use_jax=True)` simulation\n", - " (see `scripts/imaging/simulator.py` `__JAX Variant__`).\n", - "3. You constructed it inside a JAX-traced function and the upstream grid\n", - " was `jnp`-backed (e.g. `tracer.image_2d_from(grid=jnp_grid, xp=jnp)`).\n", - "\n", - "In all three cases the Python-level wrapper is the same `aa.Array2D` /\n", - "`aa.Grid2D` / etc. you've been using \u2014 only the underlying array type\n", - "changes. Workspace code reads identically on either backend.\n", - "\n", - "__Host transfer (the JAX \u2192 NumPy boundary)__\n", - "\n", - "Most things you do with these structures convert back to NumPy\n", - "transparently:\n", - "\n", - "- Plotting (`aplt.plot_array`, `aplt.subplot_fit_imaging`, ...) \u2014 calls\n", - " `np.asarray(...)` internally.\n", - "- `.fits` writing (`aplt.fits_imaging`).\n", - "- `.copy()`, `.tolist()`.\n", - "- Direct NumPy arithmetic \u2014 `np.sqrt(fit.residual_map.array)` transfers\n", - " the array off the GPU. Use `jnp.sqrt(...)` if you want to stay on the\n", - " GPU inside a hot loop.\n", - "\n", - "For one-off analysis code (notebooks, single-figure plotting), the\n", - "transfer is invisible. For hot loops or production fits, prefer the\n", - "JAX-native call.\n", - "\n", - "__The not-pytree rule__\n", - "\n", - "There's one place the abstraction leaks. If you write your own\n", - "`@jax.jit` function and try to return an `aa.Array2D` (or\n", - "`aa.Grid2DIrregular`) from inside it, the JIT boundary may fail with\n", - "`TypeError: ... is not a valid JAX type`. The wrapper types are not\n", - "reliably registered as JAX pytrees for return-from-JIT purposes.\n", - "\n", - "The workaround: return the raw `.array` from inside the jit and rewrap\n", - "outside on the host side.\n", - "\n", - "```python\n", - "@jax.jit\n", - "def my_image_fn(tracer, grid):\n", - " return tracer.image_2d_from(grid=grid, xp=jnp).array # raw jax.Array\n", - "\n", - "arr = my_image_fn(tracer, grid)\n", - "img_wrapped = al.Array2D(values=arr, mask=grid.mask)\n", - "```\n", - "\n", - "You only encounter this when *you* write the `@jax.jit` \u2014 the library\n", - "handles its own returns correctly (`AnalysisImaging(use_jax=True)` gives\n", - "back a proper `FitImaging`; `SimulatorImaging(use_jax=True)` gives back a\n", - "proper `Imaging`).\n", - "\n", - "For the full deep-dive on writing your own JAX-jit functions that\n", - "compose PyAutoLens library calls (decorator-on-def vs `jax.jit(bound_method)`,\n", - "cache-identity considerations, closure-captured `self` vs traced-argument\n", - "distinction, the `@jax.jit + xp=jnp` pairing rule), see\n", - "`scripts/guides/lens_calc.py` \u2014 that's the canonical advanced guide for\n", - "the \"JIT-it-yourself\" path.\n", - "\n", - "__Summary__\n", - "\n", - "| You construct / receive | Backing type |\n", - "|---|---|\n", - "| `al.Grid2D.uniform(shape_native, pixel_scales)` | `numpy.ndarray` |\n", - "| `fit = analysis.fit_from(instance)` from `AnalysisImaging(use_jax=True)` | `jax.Array` |\n", - "| `dataset = simulator.via_tracer_from(...)` from `SimulatorImaging(use_jax=True)` | `jax.Array` |\n", - "| `tracer.image_2d_from(grid=jnp_grid, xp=jnp)` inside your own `@jax.jit` | `jax.Array` |\n", - "\n", - "`.array` is the safe accessor for the raw backing in all cases.\n", - "Plotting and `.fits` writers handle the conversion transparently." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "Finish.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Structures\n", + "===============\n", + "\n", + "This tutorial illustrates the data structure objects which data and results quantities are stored using, which are\n", + "extensions of NumPy arrays.\n", + "\n", + "These data structures are used because for different lensing calculations it is convenient to store the data in\n", + "different formats. For example, when ray-tracing a uniform grid of image-plane (y,x) coordinates, to an irregular\n", + "grid of source-plane (y,x) coordinates, the image-plane coordinates can be stored in 2D (because the grid is uniform)\n", + "whereas the source-plane coordinates must be stored in 1D (because after lensing it is irregular).\n", + "\n", + "These data structures use the `slim` and `native` data representations API to make it simple to map quantities from\n", + "1D dimensions to their native dimensions.\n", + "\n", + "__Contents__\n", + "\n", + "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", + "- **API:** We discuss in detail why these data structures and illustrate their functionality below.\n", + "- **Grids:** We now illustrate data structures using a `Grid2D` object, which is a set of two-dimensional.\n", + "- **Native:** This plot shows the grid in its `native` format, that is in 2D dimensions where the y and x.\n", + "- **Slim:** Every `Grid2D` object is accessible via two attributes, `native` and `slim`, which store the grid.\n", + "- **Masked Data Structures:** When a mask is applied to a grid or other data structure, this changes the `slim` and `native`.\n", + "- **Data:** Two dimensional arrays of data are stored using the `Array2D` object, which has `slim` and `native`.\n", + "- **Tracer:** The `Tracer` produces many lensing quantities all of which use the `slim` and `native` data.\n", + "- **Irregular Structures:** We may want to perform calculations at specific (y,x) coordinates which are not tied to a uniform.\n", + "- **Vector Quantities:** Many lensing quantities are vectors.\n", + "\n", + "__Units__\n", + "\n", + "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", + "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", + "\n", + "The guide `units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", + "kiloparsecs, magnitudes and solar masses." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__API__\n", + "\n", + "We discuss in detail why these data structures and illustrate their functionality below.\n", + "\n", + "However, we first create the three data structures we'll use in this example, to set expectations for what they do.\n", + "\n", + "We create three data structures:\n", + "\n", + " - `Array2D`: A 2D array of data, which is used for storing an image, a noise-map, etc. \n", + "\n", + " - `Grid2D`: A 2D array of (y,x) coordinates, which is used for ray-tracing.\n", + "\n", + " -`VectorYX2D`: A 2D array of vector values, which is used for deflection angles, shear and other vector fields.\n", + "\n", + "All data structures are defined according to a uniform grid of coordinates and therefore they have a `pixel_scales`\n", + "input defining the pixel-to-arcssecond conversion factor of its grid. \n", + "\n", + "For example, for an image stored as an `Array2D`, it has a grid where each coordinate is the centre of each image pixel\n", + "and the pixel-scale is therefore the resolution of the image.\n", + "\n", + "We first create each data structure without a mask using the `no_mask` method:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "arr = al.Array2D.no_mask(\n", + " values=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], pixel_scales=1.0\n", + ")\n", + "\n", + "print(arr)\n", + "\n", + "grid = al.Grid2D.no_mask(\n", + " values=[\n", + " [[-1.0, -1.0], [-1.0, 0.0], [-1.0, 1.0]],\n", + " [[0.0, -1.0], [0.0, 0.0], [0.0, 1.0]],\n", + " [\n", + " [1.0, -1.0],\n", + " [1.0, 0.0],\n", + " [1.0, 1.0],\n", + " ],\n", + " ],\n", + " pixel_scales=1.0,\n", + ")\n", + "\n", + "print(grid)\n", + "\n", + "vector_yx = al.VectorYX2D.no_mask(\n", + " values=[\n", + " [[5.0, -5.0], [5.0, 0.0], [5.0, 5.0]],\n", + " [[0.0, -5.0], [0.0, 0.0], [0.0, 5.0]],\n", + " [\n", + " [-5.0, -5.0],\n", + " [-5.0, 0.0],\n", + " [-5.0, 5.0],\n", + " ],\n", + " ],\n", + " pixel_scales=1.0,\n", + ")\n", + "\n", + "print(vector_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grids__\n", + "\n", + "We now illustrate data structures using a `Grid2D` object, which is a set of two-dimensional $(y,x)$ coordinates\n", + "(in arc-seconds) that are deflected and traced by a strong lensing system.\n", + "\n", + "These are fundamental to all lensing calculations and drive why data structures are used.\n", + "\n", + "First, lets make a uniform 100 x 100 grid of (y,x) coordinates and plot it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "\n", + "aplt.plot_grid(grid=grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Native__\n", + "\n", + "This plot shows the grid in its `native` format, that is in 2D dimensions where the y and x coordinates are plotted\n", + "where we expect them to be on the grid.\n", + "\n", + "We can print values from the grid's `native` property to confirm this:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"(y,x) pixel 0:\")\n", + "print(grid.native[0, 0])\n", + "print(\"(y,x) pixel 1:\")\n", + "print(grid.native[0, 1])\n", + "print(\"(y,x) pixel 2:\")\n", + "print(grid.native[0, 2])\n", + "print(\"(y,x) pixel 100:\")\n", + "print(grid.native[1, 0])\n", + "print(\"etc.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Slim__\n", + "\n", + "Every `Grid2D` object is accessible via two attributes, `native` and `slim`, which store the grid as NumPy ndarrays \n", + "of two different shapes:\n", + " \n", + " - `native`: an ndarray of shape [total_y_image_pixels, total_x_image_pixels, 2] which is the native shape of the \n", + " 2D grid and corresponds to the resolution of the image datasets we pair with a grid.\n", + " \n", + " - `slim`: an ndarray of shape [total_y_image_pixels*total_x_image_pixels, 2] which is a slimmed-down representation \n", + " the grid which collapses the inner two dimensions of the native ndarray to a single dimension." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"(y,x) pixel 0 (accessed via native):\")\n", + "print(grid.native[0, 0])\n", + "print(\"(y,x) pixel 0 (accessed via slim 1D):\")\n", + "print(grid.slim[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As discussed above, the reason we need the slim representation is because when we ray-trace a grid of (y,x) coordinates\n", + "from the image-plane to the source-plane, the source-plane grid will be irregular.\n", + "\n", + "The shapes of the `Grid2D` in its `native` and `slim` formats are also available, confirming that this grid has a \n", + "`native` resolution of (100 x 100) and a `slim` resolution of 10000 coordinates." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(grid.shape_native)\n", + "print(grid.shape_slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Neither shape above include the third index of the `Grid` which has dimensions 2 (corresponding to the y and x \n", + "coordinates). \n", + "\n", + "This is accessible by using the standard numpy `shape` method on each grid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(grid.native.shape)\n", + "print(grid.slim.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can print the entire `Grid2D` in its `slim` or `native` form. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(grid.native)\n", + "print(grid.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masked Data Structures__\n", + "\n", + "When a mask is applied to a grid or other data structure, this changes the `slim` and `native` representations as \n", + "follows:\n", + "\n", + " - `slim`: only contains image-pixels that are not masked, removing all masked pixels from the 1D array.\n", + " \n", + " - `native`: retains the dimensions [total_y_image_pixels, total_x_image_pixels], but the masked pixels have values\n", + " of 0.0 or (0.0, 0.0).\n", + "\n", + "This can be seen by computing a grid via a mask and comparing the its`shape_slim` attribute to the `pixels_in_mask` of \n", + "the mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(shape_native=(100, 100), pixel_scales=0.05, radius=3.0)\n", + "\n", + "grid = al.Grid2D.from_mask(mask=mask)\n", + "\n", + "print(\"The shape_slim and number of unmasked pixels\")\n", + "print(grid.shape_slim)\n", + "print(mask.pixels_in_mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can use the `slim` attribute to print unmasked values of the grid:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"First unmasked image value:\")\n", + "print(grid.slim[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `native` representation of the `Grid2D` retains the dimensions [total_y_image_pixels, total_x_image_pixels], \n", + "however the exterior pixels have values of 0.0 indicating that they have been masked." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Example masked pixels in the grid native representation:\")\n", + "print(grid.shape_native)\n", + "print(grid.native[0, 0])\n", + "print(grid.native[2, 2])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Data__\n", + "\n", + "Two dimensional arrays of data are stored using the `Array2D` object, which has `slim` and `native` representations\n", + "analogous to the `Grid2D` object and described as follows:\n", + "\n", + " - `slim`: an ndarray of shape [total_unmasked_pixels] which is a slimmed-down representation of the data in 1D that \n", + " contains only the unmasked data points (where this mask is the one used by the model-fit above).\n", + "\n", + " - `native`: an ndarray of shape [total_y_image_pixels, total_x_image_pixels], which is the native shape of the \n", + " masked 2D grid used to fit the lens model. All masked pixels are assigned a value 0.0 in the `native` array.\n", + "\n", + "For example, the `data` and `noise_map` in an `Imaging` object are stored as `Array2D` objects.\n", + "\n", + "We load an imaging dataset and illustrate its data structures below. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "data = dataset.data" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here is what `slim` and `native` representations of the data's first pixel look like for the `data` before masking:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"First unmasked data value:\")\n", + "print(data.slim[0])\n", + "print(data.native[0, 0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By default, all arrays in **PyAutoLens** are stored as their `slim` 1D numpy array, meaning we don't need to use the\n", + "`slim` attribute to access the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(data[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By applying a mask the first value in `slim` changes and the native value becomes 0.0:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "data = dataset.data\n", + "\n", + "print(\"First unmasked data value:\")\n", + "print(data.slim[0])\n", + "print(data.native[0, 0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "The `Tracer` produces many lensing quantities all of which use the `slim` and `native` data structures.\n", + "\n", + "For example, by passing it a 2D grid of (y,x) coordinates we can return a numpy array containing its 2D image. \n", + "This includes the lens light and lensed source images.\n", + "\n", + "Below, we use the grid that is aligned with the imaging data (e.g. where each grid coordinate is at the centre of each\n", + "image pixel) to compute the galaxy image and show its data structure." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5, mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6)\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " light=al.lp.SersicCoreSph(\n", + " centre=(0.0, 0.0),\n", + " intensity=0.2,\n", + " effective_radius=0.2,\n", + " sersic_index=1.0,\n", + " radius_break=0.025,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "image = tracer.image_2d_from(grid=dataset.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we print the type of the `image` we note that it is an `Array2D`, which is a data structure that inherits \n", + "from a numpy array but is extended to include specific functionality discussed below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(type(image))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Because the image is a numpy array, we can print its shape and see that it is 1D." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(image.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Irregular Structures__\n", + "\n", + "We may want to perform calculations at specific (y,x) coordinates which are not tied to a uniform grid.\n", + "\n", + "We can use an irregular 2D (y,x) grid of coordinates for this. The grid below evaluates the image at:\n", + "\n", + "- y = 1.0, x = 1.0.\n", + "- y = 1.0, x = 2.0.\n", + "- y = 2.0, x = 2.0." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_irregular = al.Grid2DIrregular(values=[[1.0, 1.0], [1.0, 2.0], [2.0, 2.0]])\n", + "\n", + "image = tracer.image_2d_from(grid=grid_irregular)\n", + "\n", + "print(image)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The result is stored using an `ArrayIrregular` object, which is a data structure that handles irregular arrays." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(type(image))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Vector Quantities__\n", + "\n", + "Many lensing quantities are vectors. That is, they are (y,x) coordinates that have 2 values representing their\n", + "magnitudes in both the y and x directions.\n", + "\n", + "The most obvious of these is the deflection angles, which are used throughout lens modeling to ray-trace grids\n", + "from the image-plane to the source-plane via a lens galaxy mass model.\n", + "\n", + "To indicate that a quantities is a vector, **PyAutoLens** uses the label `_yx`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_yx_2d = tracer.deflections_yx_2d_from(grid=dataset.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we print the type of the `deflections_yx` we note that it is a `VectorYX2D`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(type(deflections_yx_2d))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Unlike the scalar quantities above, which were a 1D numpy array in the `slim` representation and a 2D numpy array in \n", + "the `native` representation, vectors are 2D in `slim` and 3D in `native`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(deflections_yx_2d.slim.shape)\n", + "print(deflections_yx_2d.native.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For vector quantities the has shape `2`, corresponding to the y and x vectors respectively." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(deflections_yx_2d.slim[0, :])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The role of the terms `slim` and `native` can be thought of in the same way as for scalar quantities. \n", + "\n", + "For a scalar, the `slim` property gives every scalar value as a 1D ndarray for every unmasked pixel. For a vector we \n", + "still get an ndarray of every unmasked pixel, however each entry now contains two entries: the vector of (y,x) values. \n", + "\n", + "For a `native` property these vectors are shown on an image-plane 2D grid where again each pixel\n", + "contains a (y,x) vector.\n", + "\n", + "Like we did for the convergence, we can use whatever grid we want to compute a vector and use sub-gridding to estimate\n", + "values more precisely:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(3, 3), pixel_scales=0.1)\n", + "\n", + "deflections_yx_2d = tracer.deflections_yx_2d_from(grid=grid)\n", + "\n", + "print(deflections_yx_2d.slim)\n", + "print(deflections_yx_2d.native)\n", + "\n", + "grid_irregular = al.Grid2DIrregular(values=[[1.0, 1.0], [1.0, 2.0], [2.0, 2.0]])\n", + "\n", + "deflections_yx_2d = tracer.deflections_yx_2d_from(grid=grid_irregular)\n", + "\n", + "print(deflections_yx_2d)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__JAX__\n", + "\n", + "PyAutoLens runs on either NumPy (the default) or JAX. The data structures\n", + "you've met above are *backend-polymorphic* \u2014 they're Python wrappers\n", + "around a numerical array, and that array can be a `numpy.ndarray` or a\n", + "`jax.Array` depending on how the structure was constructed and what code\n", + "path produced it.\n", + "\n", + "You can always reach the raw backing array via `.array`:\n", + "\n", + "```python\n", + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "print(type(grid.array)) # on the default path\n", + "```\n", + "\n", + "__When the backing becomes `jax.Array`__\n", + "\n", + "Three situations switch the backing to `jax.Array`:\n", + "\n", + "1. The structure comes back from a JAX-accelerated `Analysis(use_jax=True)`\n", + " fit \u2014 e.g. `fit.residual_map.array`, `fit.model_image.array` are\n", + " JAX-backed when the analysis ran with `use_jax=True` (the default).\n", + "2. The structure comes back from a `Simulator(use_jax=True)` simulation\n", + " (see `scripts/imaging/simulator.py` `__JAX Variant__`).\n", + "3. You constructed it inside a JAX-traced function and the upstream grid\n", + " was `jnp`-backed (e.g. `tracer.image_2d_from(grid=jnp_grid, xp=jnp)`).\n", + "\n", + "In all three cases the Python-level wrapper is the same `aa.Array2D` /\n", + "`aa.Grid2D` / etc. you've been using \u2014 only the underlying array type\n", + "changes. Workspace code reads identically on either backend.\n", + "\n", + "__Host transfer (the JAX \u2192 NumPy boundary)__\n", + "\n", + "Most things you do with these structures convert back to NumPy\n", + "transparently:\n", + "\n", + "- Plotting (`aplt.plot_array`, `aplt.subplot_fit_imaging`, ...) \u2014 calls\n", + " `np.asarray(...)` internally.\n", + "- `.fits` writing (`aplt.fits_imaging`).\n", + "- `.copy()`, `.tolist()`.\n", + "- Direct NumPy arithmetic \u2014 `np.sqrt(fit.residual_map.array)` transfers\n", + " the array off the GPU. Use `jnp.sqrt(...)` if you want to stay on the\n", + " GPU inside a hot loop.\n", + "\n", + "For one-off analysis code (notebooks, single-figure plotting), the\n", + "transfer is invisible. For hot loops or production fits, prefer the\n", + "JAX-native call.\n", + "\n", + "__The not-pytree rule__\n", + "\n", + "There's one place the abstraction leaks. If you write your own\n", + "`@jax.jit` function and try to return an `aa.Array2D` (or\n", + "`aa.Grid2DIrregular`) from inside it, the JIT boundary may fail with\n", + "`TypeError: ... is not a valid JAX type`. The wrapper types are not\n", + "reliably registered as JAX pytrees for return-from-JIT purposes.\n", + "\n", + "The workaround: return the raw `.array` from inside the jit and rewrap\n", + "outside on the host side.\n", + "\n", + "```python\n", + "@jax.jit\n", + "def my_image_fn(tracer, grid):\n", + " return tracer.image_2d_from(grid=grid, xp=jnp).array # raw jax.Array\n", + "\n", + "arr = my_image_fn(tracer, grid)\n", + "img_wrapped = al.Array2D(values=arr, mask=grid.mask)\n", + "```\n", + "\n", + "You only encounter this when *you* write the `@jax.jit` \u2014 the library\n", + "handles its own returns correctly (`AnalysisImaging(use_jax=True)` gives\n", + "back a proper `FitImaging`; `SimulatorImaging(use_jax=True)` gives back a\n", + "proper `Imaging`).\n", + "\n", + "For the full deep-dive on writing your own JAX-jit functions that\n", + "compose PyAutoLens library calls (decorator-on-def vs `jax.jit(bound_method)`,\n", + "cache-identity considerations, closure-captured `self` vs traced-argument\n", + "distinction, the `@jax.jit + xp=jnp` pairing rule), see\n", + "`scripts/guides/lens_calc.py` \u2014 that's the canonical advanced guide for\n", + "the \"JIT-it-yourself\" path.\n", + "\n", + "__Summary__\n", + "\n", + "| You construct / receive | Backing type |\n", + "|---|---|\n", + "| `al.Grid2D.uniform(shape_native, pixel_scales)` | `numpy.ndarray` |\n", + "| `fit = analysis.fit_from(instance)` from `AnalysisImaging(use_jax=True)` | `jax.Array` |\n", + "| `dataset = simulator.via_tracer_from(...)` from `SimulatorImaging(use_jax=True)` | `jax.Array` |\n", + "| `tracer.image_2d_from(grid=jnp_grid, xp=jnp)` inside your own `@jax.jit` | `jax.Array` |\n", + "\n", + "`.array` is the safe accessor for the raw backing in all cases.\n", + "Plotting and `.fits` writers handle the conversion transparently." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "Finish.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/galaxies.ipynb b/notebooks/guides/galaxies.ipynb index 4f4894374..00af5e610 100644 --- a/notebooks/guides/galaxies.ipynb +++ b/notebooks/guides/galaxies.ipynb @@ -1,422 +1,459 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Galaxies\n", - "========\n", - "\n", - "In the guide `tracer.py`, we inspected the results of a `Tracer` and computed the overall properties of the\n", - "lens model's image, convergence and other quantities.\n", - "\n", - "However, we did not compute the individual properties of each galaxy. For example, we did not compute an image of the\n", - "source galaxy on the source plane or compute individual quantities for each mass profile.\n", - "\n", - "This tutorial illustrates how to compute these more complicated results. We therefore fit a slightly more complicated\n", - "lens model, where the lens galaxy's light is composed of two components (a bulge and disk) and the source-plane\n", - "comprises two galaxies.\n", - "\n", - "__Contents__\n", - "\n", - "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", - "- **Data Structures:** Arrays inspected in this example use bespoke data structures for storing arrays, grids, vectors and.\n", - "- **Grids:** To describe the luminous emission of galaxies, **PyAutoGalaxy** uses `Grid2D` data structures.\n", - "- **Tracer:** We first set up a tracer with a lens galaxy and two source galaxies, which we will use to.\n", - "- **Individual Lens Galaxy Components:** We are able to create an image of the lens galaxy as follows, which includes the emission of both.\n", - "- **Log10:** The light distributions of galaxies are closer to a log10 distribution than a linear one.\n", - "\n", - "__Units__\n", - "\n", - "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", - "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", - "\n", - "The guide `guides/units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", - "kiloparsecs, magnitudes and solar masses.\n", - "\n", - "__Data Structures__\n", - "\n", - "Arrays inspected in this example use bespoke data structures for storing arrays, grids,\n", - "vectors and other 1D and 2D quantities. These use the `slim` and `native` API to toggle between representing the\n", - "data in 1D numpy arrays or high dimension numpy arrays.\n", - "\n", - "This tutorial will only use the `slim` properties which show results in 1D numpy arrays of\n", - "shape [total_unmasked_pixels]. This is a slimmed-down representation of the data in 1D that contains only the\n", - "unmasked data points\n", - "\n", - "These are documented fully in the `autolens_workspace/*/guides/data_structures.ipynb` guide.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `results/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grids__\n", - "\n", - "To describe the luminous emission of galaxies, **PyAutoGalaxy** uses `Grid2D` data structures, which are \n", - "two-dimensional Cartesian grids of (y,x) coordinates. \n", - "\n", - "Below, we make and plot a uniform Cartesian grid:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1, # The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "aplt.plot_grid(grid=grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "We first set up a tracer with a lens galaxy and two source galaxies, which we will use to illustrate how to extract\n", - "individual galaxy images." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " disk=al.lp.Exponential(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", - " intensity=0.1,\n", - " effective_radius=1.6,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.25, 0.15),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=120.0),\n", - " intensity=0.7,\n", - " effective_radius=0.7,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "source_galaxy_1 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.7, -0.5),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", - " intensity=0.2,\n", - " effective_radius=1.6,\n", - " sersic_index=3.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])\n", - "\n", - "aplt.subplot_tracer(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Individual Lens Galaxy Components__\n", - "\n", - "We are able to create an image of the lens galaxy as follows, which includes the emission of both the lens galaxy's\n", - "`bulge` and `disk` components." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = tracer.image_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In order to create images of the `bulge` and `disk` separately, we need to extract each individual component from the \n", - "tracer. \n", - "\n", - "To do this, we first use the tracer's `planes` attribute, which is a list of all `Planes` objects in the tracer. \n", - "\n", - "This list is in ascending order of plane redshift, such that `planes[0]` is the image-plane and `planes[1]` is the \n", - "source-plane. Had we modeled a multi-plane lens system there would be additional planes at each individual redshift \n", - "(the redshifts of the galaxies in the model determine at what redshifts planes are created)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_plane = tracer.planes[0]\n", - "source_plane = tracer.planes[1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Each plane contains a list of galaxies, which are in order of how we specify them in the `collection` above.\n", - "\n", - "In order to extract the `bulge` and `disk` we therefore need the lens galaxy, which we can extract from \n", - "the `image_plane` and print to make sure it contains the correct light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(image_plane)\n", - "\n", - "lens_galaxy = image_plane[0]\n", - "\n", - "print(lens_galaxy)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally, we can use the `lens_galaxy` to extract the `bulge` and `disk` and make the image of each." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = lens_galaxy.bulge\n", - "disk = lens_galaxy.disk\n", - "\n", - "bulge_image_2d = bulge.image_2d_from(grid=grid)\n", - "disk_image_2d = disk.image_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If you are unclear on what `slim` means, refer to the section `Data Structure` at the top of this example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(bulge_image_2d.slim[0])\n", - "print(disk_image_2d.slim[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It is more concise to extract these quantities in one line of Python.\n", - "\n", - "The way to think about index accessing of `planes`, as shown below is as follows:\n", - "\n", - "- The first index, `planes[0]` accesses the first plane (the image-plane).\n", - "- The second index, `planes[0][0]` accesses the first galaxy in the first plane (the lens galaxy)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge_image_2d = tracer.planes[0][0].bulge.image_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `aplt.plot_array` makes it straight forward to extract and plot an individual light profile component." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log10__\n", - "\n", - "The light distributions of galaxies are closer to a log10 distribution than a linear one. \n", - "\n", - "This means that when we plot an image of a light profile, its appearance is better highlighted when we take the\n", - "logarithm of its values and plot it in log10 space.\n", - "\n", - "The `plot_array`/`subplot_\\*` object has an input `use_log10`, which will do this automatically when we call the `plot_array` method.\n", - "Below, we can see that the image plotted now appears more clearly, with the outskirts of the light profile more visible.\n", - "\n", - "__JAX__\n", - "\n", - "When you write your own `@jax.jit` around a function that takes a\n", - "`Galaxy` or `Galaxies` as an argument, JAX needs to flatten and unflatten\n", - "that object across the JIT boundary \u2014 i.e. the class must be registered\n", - "as a JAX pytree. The library handles this for you automatically in two\n", - "situations:\n", - "\n", - "1. You constructed an `Analysis` with `use_jax=True` (the default for\n", - " modeling fits). `AnalysisImaging._register_fit_imaging_pytrees()`\n", - " walks the dataset on first `fit_from` call and registers every\n", - " reachable `Galaxy` / profile class.\n", - "2. You constructed a `Simulator` with `use_jax=True` and made a call \u2014\n", - " same walk happens.\n", - "\n", - "After either, every `Galaxy`, `LightProfile`, `MassProfile`, `Point`,\n", - "etc. of the same class is JIT-safe forever in the current process. You\n", - "never call `register_instance_pytree(Galaxy)` yourself.\n", - "\n", - "__The \"I have no Analysis or Simulator handy\" case__\n", - "\n", - "For a quick exploration script or a custom forward model that doesn't go\n", - "through `Simulator.via_tracer_from`, you may want to JIT a function that\n", - "takes a `Galaxy` or list of galaxies as an argument:\n", - "\n", - "```python\n", - "@jax.jit\n", - "def galaxy_image(galaxy, grid):\n", - " return galaxy.image_2d_from(grid=grid, xp=jnp).array\n", - "```\n", - "\n", - "Without prior pytree registration this fails the first time `galaxy` is\n", - "traced. The workaround: trigger registration at the top of your script.\n", - "Two equivalent approaches:\n", - "\n", - "```python\n", - "# (a) Reach for an Analysis instance \u2014 its first construction triggers\n", - "# the registration walk on the (possibly dummy) dataset.\n", - "_ = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "# (b) Use the dedicated walker on a representative tracer.\n", - "from autolens.jax import register_tracer_classes\n", - "register_tracer_classes(tracer)\n", - "```\n", - "\n", - "After either, `galaxy_image` JITs cleanly.\n", - "\n", - "__Closure-captured galaxy: registration not needed__\n", - "\n", - "There's a way to JIT a galaxy-method call that does NOT need pytree\n", - "registration: pass the galaxy as the bound method's `self`, not as an\n", - "argument.\n", - "\n", - "```python\n", - "jitted_image = jax.jit(galaxy.image_2d_from) # bound method; assign ONCE\n", - "arr = jitted_image(grid=grid, xp=jnp).array\n", - "```\n", - "\n", - "`galaxy` here is closed over inside the bound method; JAX treats it as a\n", - "closure constant, doesn't trace through it, and pytree registration\n", - "never enters the picture.\n", - "\n", - "Trade-off: you can't vary `galaxy` across calls and still hit the\n", - "compilation cache. If you want to evaluate the same function for many\n", - "different galaxies (parameter sweep), the argument form (with prior\n", - "registration) is the right choice.\n", - "\n", - "The full deep-dive on the bound-method vs traced-argument trade-off,\n", - "cache-identity footguns, and the `@jax.jit + xp=jnp` pairing rule lives\n", - "in `scripts/guides/lens_calc.py`. The `.array` host-transfer mechanics\n", - "live in `scripts/guides/data_structures.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Galaxies\n", + "========\n", + "\n", + "In the guide `tracer.py`, we inspected the results of a `Tracer` and computed the overall properties of the\n", + "lens model's image, convergence and other quantities.\n", + "\n", + "However, we did not compute the individual properties of each galaxy. For example, we did not compute an image of the\n", + "source galaxy on the source plane or compute individual quantities for each mass profile.\n", + "\n", + "This tutorial illustrates how to compute these more complicated results. We therefore fit a slightly more complicated\n", + "lens model, where the lens galaxy's light is composed of two components (a bulge and disk) and the source-plane\n", + "comprises two galaxies.\n", + "\n", + "__Contents__\n", + "\n", + "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", + "- **Data Structures:** Arrays inspected in this example use bespoke data structures for storing arrays, grids, vectors and.\n", + "- **Grids:** To describe the luminous emission of galaxies, **PyAutoGalaxy** uses `Grid2D` data structures.\n", + "- **Tracer:** We first set up a tracer with a lens galaxy and two source galaxies, which we will use to.\n", + "- **Individual Lens Galaxy Components:** We are able to create an image of the lens galaxy as follows, which includes the emission of both.\n", + "- **Log10:** The light distributions of galaxies are closer to a log10 distribution than a linear one.\n", + "\n", + "__Units__\n", + "\n", + "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", + "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", + "\n", + "The guide `guides/units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", + "kiloparsecs, magnitudes and solar masses.\n", + "\n", + "__Data Structures__\n", + "\n", + "Arrays inspected in this example use bespoke data structures for storing arrays, grids,\n", + "vectors and other 1D and 2D quantities. These use the `slim` and `native` API to toggle between representing the\n", + "data in 1D numpy arrays or high dimension numpy arrays.\n", + "\n", + "This tutorial will only use the `slim` properties which show results in 1D numpy arrays of\n", + "shape [total_unmasked_pixels]. This is a slimmed-down representation of the data in 1D that contains only the\n", + "unmasked data points\n", + "\n", + "These are documented fully in the `autolens_workspace/*/guides/data_structures.ipynb` guide.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `results/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grids__\n", + "\n", + "To describe the luminous emission of galaxies, **PyAutoGalaxy** uses `Grid2D` data structures, which are \n", + "two-dimensional Cartesian grids of (y,x) coordinates. \n", + "\n", + "Below, we make and plot a uniform Cartesian grid:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1, # The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "aplt.plot_grid(grid=grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "We first set up a tracer with a lens galaxy and two source galaxies, which we will use to illustrate how to extract\n", + "individual galaxy images." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " disk=al.lp.Exponential(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", + " intensity=0.1,\n", + " effective_radius=1.6,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.25, 0.15),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=120.0),\n", + " intensity=0.7,\n", + " effective_radius=0.7,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "source_galaxy_1 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.7, -0.5),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", + " intensity=0.2,\n", + " effective_radius=1.6,\n", + " sersic_index=3.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])\n", + "\n", + "aplt.subplot_tracer(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Individual Lens Galaxy Components__\n", + "\n", + "We are able to create an image of the lens galaxy as follows, which includes the emission of both the lens galaxy's\n", + "`bulge` and `disk` components." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = tracer.image_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In order to create images of the `bulge` and `disk` separately, we need to extract each individual component from the \n", + "tracer. \n", + "\n", + "To do this, we first use the tracer's `planes` attribute, which is a list of all `Planes` objects in the tracer. \n", + "\n", + "This list is in ascending order of plane redshift, such that `planes[0]` is the image-plane and `planes[1]` is the \n", + "source-plane. Had we modeled a multi-plane lens system there would be additional planes at each individual redshift \n", + "(the redshifts of the galaxies in the model determine at what redshifts planes are created)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_plane = tracer.planes[0]\n", + "source_plane = tracer.planes[1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each plane contains a list of galaxies, which are in order of how we specify them in the `collection` above.\n", + "\n", + "In order to extract the `bulge` and `disk` we therefore need the lens galaxy, which we can extract from \n", + "the `image_plane` and print to make sure it contains the correct light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(image_plane)\n", + "\n", + "lens_galaxy = image_plane[0]\n", + "\n", + "print(lens_galaxy)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, we can use the `lens_galaxy` to extract the `bulge` and `disk` and make the image of each." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = lens_galaxy.bulge\n", + "disk = lens_galaxy.disk\n", + "\n", + "bulge_image_2d = bulge.image_2d_from(grid=grid)\n", + "disk_image_2d = disk.image_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If you are unclear on what `slim` means, refer to the section `Data Structure` at the top of this example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(bulge_image_2d.slim[0])\n", + "print(disk_image_2d.slim[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It is more concise to extract these quantities in one line of Python.\n", + "\n", + "The way to think about index accessing of `planes`, as shown below is as follows:\n", + "\n", + "- The first index, `planes[0]` accesses the first plane (the image-plane).\n", + "- The second index, `planes[0][0]` accesses the first galaxy in the first plane (the lens galaxy)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge_image_2d = tracer.planes[0][0].bulge.image_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `aplt.plot_array` makes it straight forward to extract and plot an individual light profile component." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log10__\n", + "\n", + "The light distributions of galaxies are closer to a log10 distribution than a linear one. \n", + "\n", + "This means that when we plot an image of a light profile, its appearance is better highlighted when we take the\n", + "logarithm of its values and plot it in log10 space.\n", + "\n", + "The `plot_array`/`subplot_\\*` object has an input `use_log10`, which will do this automatically when we call the `plot_array` method.\n", + "Below, we can see that the image plotted now appears more clearly, with the outskirts of the light profile more visible.\n", + "\n", + "__JAX__\n", + "\n", + "When you write your own `@jax.jit` around a function that takes a\n", + "`Galaxy` or `Galaxies` as an argument, JAX needs to flatten and unflatten\n", + "that object across the JIT boundary \u2014 i.e. the class must be registered\n", + "as a JAX pytree. The library handles this for you automatically in two\n", + "situations:\n", + "\n", + "1. You constructed an `Analysis` with `use_jax=True` (the default for\n", + " modeling fits). `AnalysisImaging._register_fit_imaging_pytrees()`\n", + " walks the dataset on first `fit_from` call and registers every\n", + " reachable `Galaxy` / profile class.\n", + "2. You constructed a `Simulator` with `use_jax=True` and made a call \u2014\n", + " same walk happens.\n", + "\n", + "After either, every `Galaxy`, `LightProfile`, `MassProfile`, `Point`,\n", + "etc. of the same class is JIT-safe forever in the current process. You\n", + "never call `register_instance_pytree(Galaxy)` yourself.\n", + "\n", + "__The \"I have no Analysis or Simulator handy\" case__\n", + "\n", + "For a quick exploration script or a custom forward model that doesn't go\n", + "through `Simulator.via_tracer_from`, you may want to JIT a function that\n", + "takes a `Galaxy` or list of galaxies as an argument:\n", + "\n", + "```python\n", + "@jax.jit\n", + "def galaxy_image(galaxy, grid):\n", + " return galaxy.image_2d_from(grid=grid, xp=jnp).array\n", + "```\n", + "\n", + "Without prior pytree registration this fails the first time `galaxy` is\n", + "traced. The workaround: trigger registration at the top of your script.\n", + "Two equivalent approaches:\n", + "\n", + "```python\n", + "# (a) Reach for an Analysis instance \u2014 its first construction triggers\n", + "# the registration walk on the (possibly dummy) dataset.\n", + "_ = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "# (b) Use the dedicated walker on a representative tracer.\n", + "from autolens.jax import register_tracer_classes\n", + "register_tracer_classes(tracer)\n", + "```\n", + "\n", + "After either, `galaxy_image` JITs cleanly.\n", + "\n", + "__Closure-captured galaxy: registration not needed__\n", + "\n", + "There's a way to JIT a galaxy-method call that does NOT need pytree\n", + "registration: pass the galaxy as the bound method's `self`, not as an\n", + "argument.\n", + "\n", + "```python\n", + "jitted_image = jax.jit(galaxy.image_2d_from) # bound method; assign ONCE\n", + "arr = jitted_image(grid=grid, xp=jnp).array\n", + "```\n", + "\n", + "`galaxy` here is closed over inside the bound method; JAX treats it as a\n", + "closure constant, doesn't trace through it, and pytree registration\n", + "never enters the picture.\n", + "\n", + "Trade-off: you can't vary `galaxy` across calls and still hit the\n", + "compilation cache. If you want to evaluate the same function for many\n", + "different galaxies (parameter sweep), the argument form (with prior\n", + "registration) is the right choice.\n", + "\n", + "The full deep-dive on the bound-method vs traced-argument trade-off,\n", + "cache-identity footguns, and the `@jax.jit + xp=jnp` pairing rule lives\n", + "in `scripts/guides/lens_calc.py`. The `.array` host-transfer mechanics\n", + "live in `scripts/guides/data_structures.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/hpc/example_cpu.ipynb b/notebooks/guides/hpc/example_cpu.ipynb index 3ff4737ca..254c0a33a 100644 --- a/notebooks/guides/hpc/example_cpu.ipynb +++ b/notebooks/guides/hpc/example_cpu.ipynb @@ -1,610 +1,647 @@ { - "cells": [ - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# %%\n", - "'''\n", - "HPC: Example CPU\n", - "================\n", - "\n", - "This example illustrates how to set up lens modeling on a High Performance Computing (HPC) system using multiple CPUs.\n", - "\n", - "It illustrates two different forms of parallelization:\n", - "\n", - "1) Set off many single CPU jobs in a single HPC submission script, where each job fits a different dataset using the\n", - "same lens model analysis. This form of parallelization is efficient when we have many datasets we wish to fit\n", - "simultaneously, but each individual fit only uses one CPU so overall run times are slower.\n", - "\n", - "2) Fit a single dataset using a parallelized Nautilus model-fit, where the non-linear search distributes the model-fit\n", - "over multiple CPUs. This form of parallelization is efficient when we have a single dataset to fit, but we wish to\n", - "speed up the overall run time of the model-fit by using multiple CPUs. However, parallelizing over multiple CPUs\n", - "have communication overheads, and so this form of parallelization is less efficient than fitting many single CPU jobs.\n", - "\n", - "The example assumes the HPC environment uses slurm for job management, which is standard for many academic HPCs but\n", - "may not necessarily be the case for your HPC. If your HPC does not use slurm, you should still be able to adapt this\n", - "example to your HPC`s job management system.\n", - "\n", - "This example will likely require adaptation for you to run it on your HPC enviroment, its goal is to simply\n", - "illustrate the general principles of how to set up lens modeling on an HPC.\n", - "\n", - "__Contents__\n", - "\n", - "- **HPC Output Path:** We first set the `hpc_output_path`, where the results of lens modeling are output on your HPC.\n", - "- **HPC Dataset Path:** We next set the `hpc_dataset_path`, which is the path where datasets are stored on the hpc.\n", - "- **HPC Home Path:** The `home_path` is in your the hpc home directory, which again may be different from your output.\n", - "- **Batch Script Many Lenses:** HPC submissions require a batch script, which tells the HPC the CPU hardware you want the job to.\n", - "- **Batch Script One Lenses:** Lets now look at the second example batch script, `example_cpu_one_dataset_parallel`, which fits a.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Nautilus CPUs:** The final change we need to make is to set the number of CPUs Nautilus uses in the model-fit.\n", - "- **Lens Modeling:** Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "# %%\n", - "'''\n", - "__HPC Output Path__\n", - "\n", - "We first set the `hpc_output_path`, where the results of lens modeling are output on your HPC. \n", - " \n", - "On certain HPCs this may be different from your home directory or where you store data, because lens modeling has more \n", - "IO and outputs many individual files. \n", - "\n", - "This example assumes results are output to the directory, where `hpc_username` is your hpc username:\n", - "\n", - " `/hpc/data/hpc_username/output`.\n", - " \n", - "You will need to update `hpc_username` to your hpc username below.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from pathlib import Path\n", - "\n", - "hpc_output_path = Path(\"/\") / \"hpc\" / \"data\" / \"hpc_username\" / \"output\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__HPC Dataset Path__\n", - "\n", - "We next set the `hpc_dataset_path`, which is the path where datasets are stored on the hpc.\n", - "\n", - "This may be the same as your output path, or you may have been advised to store datasets in a different location,\n", - "especially if they are large in file size.\n", - "\n", - "We therefore define it separately from the `hpc_output_path`.\n", - "\n", - "Below, we set `hpc_dataset_path=/hpc/data/hpc_username/dataset/example/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_folder = \"example\"\n", - "dataset_name = \"simple__no_lens_light\"\n", - "\n", - "hpc_dataset_path = (\n", - " Path(\"/\")\n", - " / \"hpc\"\n", - " / \"data\"\n", - " / \"hpc_username\"\n", - " / \"dataset\"\n", - " / dataset_folder\n", - " / dataset_name\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__HPC Home Path__\n", - "\n", - "The `home_path` is in your the hpc home directory, which again may be different from your output and dataset paths.\n", - "\n", - "The home path often has signficant storage restrictions, so is not a good location to store datasets or output results.\n", - "But may be where you store the python lens modeling scripts you run on the HPC, the config files, batch scripts\n", - "and other files you use to set up lens modeling on the hpc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "home_path = Path(\"/\", \"hpc\", \"home\", \"hpc_username\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "On the HPC, most likely in your home directory, you should have a config folder which contains the config files used by \n", - "modeling.\n", - "\n", - "This `config_path` sets the path to the config files that are used in this analysis, which are contained within the `hpc` \n", - "directory of the example project in your the hpc home directory." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "config_path = Path(home_path, \"hpc\", \"config\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Set the config and output paths using autoconf, as you would for a laptop run." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoconf import conf\n", - "\n", - "conf.instance.push(new_path=config_path, output_path=hpc_output_path)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Above, we set up many different paths required to run modeling on the hpc. You should basically determine where\n", - "all the different paths are on your HPC are which correspond to the paths above, and update the code accordingly.\n", - "\n", - "__Batch Script Many Lenses__\n", - "\n", - "HPC submissions require a batch script, which tells the HPC the CPU hardware you want the job to run on and the \n", - "PyAutoLens Python script you want it to execute. This script then distributes the job to nodes and CPUs. \n", - "\n", - "Lets look at the batch script \n", - "\n", - " `autolens_workspace/*/guides/hpc/batch/example_cpu_one_dataset_parallel\n", - " \n", - "The following options are worth noting:\n", - "\n", - " `#SBATCH -N 1` - The number of nodes we require, where 1 node contains 28 CPUs on the hpc.\n", - " `#SBATCH --ntasks=16` - The total number of task we are submitting.\n", - " `#SBATCH --cpus-per-task=1` - The number of tasks per CPU.\n", - " `#SBATCH -J example` - The name of the job, which is how it appears on hpc when you inspect it.\n", - " `#SBATCH -o output/output.%A.out` - Python interpreter output is placed in a file in the `output` folder.\n", - " `#SBATCH -o error/error.%A.out` - Python interpreter errors are placed in a file in the `error` folder.\n", - " `#SBATCH -p hpc` - Signifies we are running the job on the hpc.\n", - " `#SBATCH -A dp004` - The project code of the submission.\n", - " `#SBATCH -t 48:00:00` - The job will terminate after this length of time (if it does not end naturally).\n", - " `#SBATCH --mail-type=END` - If you input your email, when you`ll get an email about the job (END means once finished).\n", - " `#SBATCH --mail-user=fill@me.co.uk` - The email address the hpc sends the email too.\n", - "\n", - "The following line activates the PyAutoLens virtual environment we set up on hpc for this run:\n", - "\n", - " `source /hpc/home/hpc_username/activate.sh`\n", - "\n", - "These lines prevent the NumPy linear algebra libraries from overloading the CPUs during calculations.\n", - " \n", - "export CPUS_PER_TASK=1\n", - "\n", - "export OPENBLAS_NUM_THREADS=$CPUS_PER_TASK\n", - "export MKL_NUM_THREADS=$CPUS_PER_TASK\n", - "export OMP_NUM_THREADS=$CPUS_PER_TASK\n", - "export VECLIB_MAXIMUM_THREADS=$CPUS_PER_TASK\n", - "export NUMEXPR_NUM_THREADS=$CPUS_PER_TASK\n", - "\n", - "This line sets off the job:\n", - "\n", - " srun -n 16 --multi-prog conf/example.conf\n", - "\n", - "Lets checkout the file `example_cpu_many_datasets.conf`:\n", - "\n", - " 0 python3 /hpc/home/hpc_username/runners/example.py 0\n", - " 1 python3 /hpc/home/hpc_username/runners/example.py 1\n", - " 2 python3 /hpc/home/hpc_username/runners/example.py 2\n", - " 3 python3 /hpc/home/hpc_username/runners/example.py 3\n", - " 4 python3 /hpc/home/hpc_username/runners/example.py 4\n", - " 5 python3 /hpc/home/hpc_username/runners/example.py 5\n", - " 6 python3 /hpc/home/hpc_username/runners/example.py 6\n", - " 7 python3 /hpc/home/hpc_username/runners/example.py 7\n", - " 8 python3 /hpc/home/hpc_username/runners/example.py 8\n", - " 9 python3 /hpc/home/hpc_username/runners/example.py 9\n", - " 10 python3 /hpc/home/hpc_username/runners/example.py 10\n", - " 11 python3 /hpc/home/hpc_username/runners/example.py 11\n", - " 12 python3 /hpc/home/hpc_username/runners/example.py 12\n", - " 13 python3 /hpc/home/hpc_username/runners/example.py 13\n", - " 14 python3 /hpc/home/hpc_username/runners/example.py 14\n", - " 15 python3 /hpc/home/hpc_username/runners/example.py 15\n", - " \n", - "This file contains lines of python3 commands which set off our modeling script script(s)! It is now clear how to set \n", - "off many hpc jobs; just add each modeling script you want to run to this script. \n", - "\n", - "The numbers on the left running from 0-15 specify the CPU number and should always run from 0. \n", - "\n", - "The numbers on the right are inputting an integer, which is then used to load a specific dataset. Below, using \n", - "the `sys.argv[1]` command, we load each integer into the Python script. For example, the first job loads the integer\n", - "0, the second job the integer 1 and so forth. Each job will therefore have a unique integer value in the `hpc_id` \n", - "variable." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import sys\n", - "\n", - "hpc_id = int(sys.argv[1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now use this variable to load a specific piece of data for this run!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "dataset_type = \"imaging\"\n", - "pixel_scales = 0.1\n", - "\n", - "dataset_name = []\n", - "dataset_name.append(\"example_image_1\") # Index 0\n", - "dataset_name.append(\"example_image_2\") # Index 1\n", - "dataset_name.append(\"example_image_3\") # Index 2\n", - "dataset_name.append(\"example_image_4\") # Index 3\n", - "dataset_name.append(\"example_image_5\") # Index 4\n", - "dataset_name.append(\"example_image_6\") # Index 5\n", - "dataset_name.append(\"example_image_7\") # Index 6\n", - "dataset_name.append(\"example_image_8\") # Index 7\n", - "# ...and so on." - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now extract the dataset name specific to this hpc id, meaning every CPU run will load and fit a different dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = dataset_name[hpc_id]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now create the overall path to the dataset this specific call of the script fits, which for the first line in the \n", - "`.conf` file above (which has integer input 0) is: \n", - "\n", - " `/hpc/data/hpc_username/dataset/imaging/example_image_1`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(hpc_dataset_path, dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You now have all the code you need to set up many single-CPU jobs on the hpc!\n", - "\n", - "You would simply append the batch scripts and Python code aboves to the lens modeling script script you are using,\n", - "which is given below for completeness.\n", - "\n", - "However, first we describe how to set up a single multi-CPU Nautilus job on the hpc.\n", - "\n", - "__Batch Script One Lenses__\n", - "\n", - "Lets now look at the second example batch script, `example_cpu_one_dataset_parallel`, which fits a single dataset\n", - "using multiple CPUs.\n", - "\n", - "#!/bin/bash -l\n", - "\n", - " `#SBATCH -N 1` - The number of nodes we require, where 1 node contains 28 CPUs on the hpc.\n", - " `#SBATCH --ntasks=1` - The total number of task we are submitting.\n", - " `#SBATCH --cpus-per-task=16` - The number of tasks per CPU.\n", - " `#SBATCH -J example` - The name of the job, which is how it appears on hpc when you inspect it.\n", - " `#SBATCH -o output/output.%A.out` - Python interpreter output is placed in a file in the `output` folder.\n", - " `#SBATCH -o error/error.%A.out` - Python interpreter errors are placed in a file in the `error` folder.\n", - " `#SBATCH -p hpc` - Signifies we are running the job on the hpc.\n", - " `#SBATCH -A dp004` - The project code of the submission.\n", - " `#SBATCH -t 48:00:00` - The job will terminate after this length of time (if it does not end naturally).\n", - " `#SBATCH --mail-type=END` - If you input your email, when you`ll get an email about the job (END means once finished).\n", - " `#SBATCH --mail-user=fill@me.co.uk` - The email address the hpc sends the email too.\n", - "\n", - "source /hpc/home/hpc_username/activate.sh\n", - "\n", - "export CPUS_PER_TASK=1\n", - "\n", - "export OPENBLAS_NUM_THREADS=$CPUS_PER_TASK\n", - "export MKL_NUM_THREADS=$CPUS_PER_TASK\n", - "export OMP_NUM_THREADS=$CPUS_PER_TASK\n", - "export VECLIB_MAXIMUM_THREADS=$CPUS_PER_TASK\n", - "export NUMEXPR_NUM_THREADS=$CPUS_PER_TASK\n", - "\n", - "python3 /hpc/home/hpc_username/runners/example.py 16\n", - "\n", - "The key difference in this batch script is the line which sets off the job:\n", - "\n", - "- ntasks=1 - We only require one task because we are only fitting one dataset.\n", - "\n", - "- cpus-per-task=16 - We require 16 CPUs for this single task, which the Nautilus non-linear search will use\n", - "to parallelize the model-fit over.\n", - "\n", - "- python3 ... example.py 16 - We input the integer 16, which is used below to set the number of CPUs Nautilus\n", - "\n", - "Lets now look at the beginning of the modeling script script again, which now does not use a list of datasets \n", - "to load, but instead has the dataset name hard coded." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `example_image_1` via .fits files, which we will fit with the lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_folder = \"example\"\n", - "dataset_name = \"simple__no_lens_light\"\n", - "\n", - "dataset_path = Path(hpc_dataset_path, dataset_folder, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Nautilus CPUs__\n", - "\n", - "The final change we need to make is to set the number of CPUs Nautilus uses in the model-fit.\n", - "\n", - "We do this by loading the integer input form the batch script, which we set to be the number of CPUs Nautilus uses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "number_of_cores = int(sys.argv[1])\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=\"hpc\",\n", - " name=\"example\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " number_of_cores=number_of_cores,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now have everything we need to fit a single dataset using multiple CPUs on the hpc!\n", - "\n", - "The code below performs standard lens modeling, which is unchanged from normal modeling on a laptop. It can be\n", - "used for either many single-CPU jobs or a single multi-CPU Nautilus job.\n", - "\n", - "__Lens Modeling__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - " \n", - " - The source galaxy's light is an MGE [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.SersicCore)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# %%\n", + "'''\n", + "HPC: Example CPU\n", + "================\n", + "\n", + "This example illustrates how to set up lens modeling on a High Performance Computing (HPC) system using multiple CPUs.\n", + "\n", + "It illustrates two different forms of parallelization:\n", + "\n", + "1) Set off many single CPU jobs in a single HPC submission script, where each job fits a different dataset using the\n", + "same lens model analysis. This form of parallelization is efficient when we have many datasets we wish to fit\n", + "simultaneously, but each individual fit only uses one CPU so overall run times are slower.\n", + "\n", + "2) Fit a single dataset using a parallelized Nautilus model-fit, where the non-linear search distributes the model-fit\n", + "over multiple CPUs. This form of parallelization is efficient when we have a single dataset to fit, but we wish to\n", + "speed up the overall run time of the model-fit by using multiple CPUs. However, parallelizing over multiple CPUs\n", + "have communication overheads, and so this form of parallelization is less efficient than fitting many single CPU jobs.\n", + "\n", + "The example assumes the HPC environment uses slurm for job management, which is standard for many academic HPCs but\n", + "may not necessarily be the case for your HPC. If your HPC does not use slurm, you should still be able to adapt this\n", + "example to your HPC`s job management system.\n", + "\n", + "This example will likely require adaptation for you to run it on your HPC enviroment, its goal is to simply\n", + "illustrate the general principles of how to set up lens modeling on an HPC.\n", + "\n", + "__Contents__\n", + "\n", + "- **HPC Output Path:** We first set the `hpc_output_path`, where the results of lens modeling are output on your HPC.\n", + "- **HPC Dataset Path:** We next set the `hpc_dataset_path`, which is the path where datasets are stored on the hpc.\n", + "- **HPC Home Path:** The `home_path` is in your the hpc home directory, which again may be different from your output.\n", + "- **Batch Script Many Lenses:** HPC submissions require a batch script, which tells the HPC the CPU hardware you want the job to.\n", + "- **Batch Script One Lenses:** Lets now look at the second example batch script, `example_cpu_one_dataset_parallel`, which fits a.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Nautilus CPUs:** The final change we need to make is to set the number of CPUs Nautilus uses in the model-fit.\n", + "- **Lens Modeling:** Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "# %%\n", + "'''\n", + "__HPC Output Path__\n", + "\n", + "We first set the `hpc_output_path`, where the results of lens modeling are output on your HPC. \n", + " \n", + "On certain HPCs this may be different from your home directory or where you store data, because lens modeling has more \n", + "IO and outputs many individual files. \n", + "\n", + "This example assumes results are output to the directory, where `hpc_username` is your hpc username:\n", + "\n", + " `/hpc/data/hpc_username/output`.\n", + " \n", + "You will need to update `hpc_username` to your hpc username below.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from pathlib import Path\n", + "\n", + "hpc_output_path = Path(\"/\") / \"hpc\" / \"data\" / \"hpc_username\" / \"output\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__HPC Dataset Path__\n", + "\n", + "We next set the `hpc_dataset_path`, which is the path where datasets are stored on the hpc.\n", + "\n", + "This may be the same as your output path, or you may have been advised to store datasets in a different location,\n", + "especially if they are large in file size.\n", + "\n", + "We therefore define it separately from the `hpc_output_path`.\n", + "\n", + "Below, we set `hpc_dataset_path=/hpc/data/hpc_username/dataset/example/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_folder = \"example\"\n", + "dataset_name = \"simple__no_lens_light\"\n", + "\n", + "hpc_dataset_path = (\n", + " Path(\"/\")\n", + " / \"hpc\"\n", + " / \"data\"\n", + " / \"hpc_username\"\n", + " / \"dataset\"\n", + " / dataset_folder\n", + " / dataset_name\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__HPC Home Path__\n", + "\n", + "The `home_path` is in your the hpc home directory, which again may be different from your output and dataset paths.\n", + "\n", + "The home path often has signficant storage restrictions, so is not a good location to store datasets or output results.\n", + "But may be where you store the python lens modeling scripts you run on the HPC, the config files, batch scripts\n", + "and other files you use to set up lens modeling on the hpc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "home_path = Path(\"/\", \"hpc\", \"home\", \"hpc_username\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "On the HPC, most likely in your home directory, you should have a config folder which contains the config files used by \n", + "modeling.\n", + "\n", + "This `config_path` sets the path to the config files that are used in this analysis, which are contained within the `hpc` \n", + "directory of the example project in your the hpc home directory." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "config_path = Path(home_path, \"hpc\", \"config\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set the config and output paths using autoconf, as you would for a laptop run." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoconf import conf\n", + "\n", + "conf.instance.push(new_path=config_path, output_path=hpc_output_path)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Above, we set up many different paths required to run modeling on the hpc. You should basically determine where\n", + "all the different paths are on your HPC are which correspond to the paths above, and update the code accordingly.\n", + "\n", + "__Batch Script Many Lenses__\n", + "\n", + "HPC submissions require a batch script, which tells the HPC the CPU hardware you want the job to run on and the \n", + "PyAutoLens Python script you want it to execute. This script then distributes the job to nodes and CPUs. \n", + "\n", + "Lets look at the batch script \n", + "\n", + " `autolens_workspace/*/guides/hpc/batch/example_cpu_one_dataset_parallel\n", + " \n", + "The following options are worth noting:\n", + "\n", + " `#SBATCH -N 1` - The number of nodes we require, where 1 node contains 28 CPUs on the hpc.\n", + " `#SBATCH --ntasks=16` - The total number of task we are submitting.\n", + " `#SBATCH --cpus-per-task=1` - The number of tasks per CPU.\n", + " `#SBATCH -J example` - The name of the job, which is how it appears on hpc when you inspect it.\n", + " `#SBATCH -o output/output.%A.out` - Python interpreter output is placed in a file in the `output` folder.\n", + " `#SBATCH -o error/error.%A.out` - Python interpreter errors are placed in a file in the `error` folder.\n", + " `#SBATCH -p hpc` - Signifies we are running the job on the hpc.\n", + " `#SBATCH -A dp004` - The project code of the submission.\n", + " `#SBATCH -t 48:00:00` - The job will terminate after this length of time (if it does not end naturally).\n", + " `#SBATCH --mail-type=END` - If you input your email, when you`ll get an email about the job (END means once finished).\n", + " `#SBATCH --mail-user=fill@me.co.uk` - The email address the hpc sends the email too.\n", + "\n", + "The following line activates the PyAutoLens virtual environment we set up on hpc for this run:\n", + "\n", + " `source /hpc/home/hpc_username/activate.sh`\n", + "\n", + "These lines prevent the NumPy linear algebra libraries from overloading the CPUs during calculations.\n", + " \n", + "export CPUS_PER_TASK=1\n", + "\n", + "export OPENBLAS_NUM_THREADS=$CPUS_PER_TASK\n", + "export MKL_NUM_THREADS=$CPUS_PER_TASK\n", + "export OMP_NUM_THREADS=$CPUS_PER_TASK\n", + "export VECLIB_MAXIMUM_THREADS=$CPUS_PER_TASK\n", + "export NUMEXPR_NUM_THREADS=$CPUS_PER_TASK\n", + "\n", + "This line sets off the job:\n", + "\n", + " srun -n 16 --multi-prog conf/example.conf\n", + "\n", + "Lets checkout the file `example_cpu_many_datasets.conf`:\n", + "\n", + " 0 python3 /hpc/home/hpc_username/runners/example.py 0\n", + " 1 python3 /hpc/home/hpc_username/runners/example.py 1\n", + " 2 python3 /hpc/home/hpc_username/runners/example.py 2\n", + " 3 python3 /hpc/home/hpc_username/runners/example.py 3\n", + " 4 python3 /hpc/home/hpc_username/runners/example.py 4\n", + " 5 python3 /hpc/home/hpc_username/runners/example.py 5\n", + " 6 python3 /hpc/home/hpc_username/runners/example.py 6\n", + " 7 python3 /hpc/home/hpc_username/runners/example.py 7\n", + " 8 python3 /hpc/home/hpc_username/runners/example.py 8\n", + " 9 python3 /hpc/home/hpc_username/runners/example.py 9\n", + " 10 python3 /hpc/home/hpc_username/runners/example.py 10\n", + " 11 python3 /hpc/home/hpc_username/runners/example.py 11\n", + " 12 python3 /hpc/home/hpc_username/runners/example.py 12\n", + " 13 python3 /hpc/home/hpc_username/runners/example.py 13\n", + " 14 python3 /hpc/home/hpc_username/runners/example.py 14\n", + " 15 python3 /hpc/home/hpc_username/runners/example.py 15\n", + " \n", + "This file contains lines of python3 commands which set off our modeling script script(s)! It is now clear how to set \n", + "off many hpc jobs; just add each modeling script you want to run to this script. \n", + "\n", + "The numbers on the left running from 0-15 specify the CPU number and should always run from 0. \n", + "\n", + "The numbers on the right are inputting an integer, which is then used to load a specific dataset. Below, using \n", + "the `sys.argv[1]` command, we load each integer into the Python script. For example, the first job loads the integer\n", + "0, the second job the integer 1 and so forth. Each job will therefore have a unique integer value in the `hpc_id` \n", + "variable." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import sys\n", + "\n", + "hpc_id = int(sys.argv[1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now use this variable to load a specific piece of data for this run!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "dataset_type = \"imaging\"\n", + "pixel_scales = 0.1\n", + "\n", + "dataset_name = []\n", + "dataset_name.append(\"example_image_1\") # Index 0\n", + "dataset_name.append(\"example_image_2\") # Index 1\n", + "dataset_name.append(\"example_image_3\") # Index 2\n", + "dataset_name.append(\"example_image_4\") # Index 3\n", + "dataset_name.append(\"example_image_5\") # Index 4\n", + "dataset_name.append(\"example_image_6\") # Index 5\n", + "dataset_name.append(\"example_image_7\") # Index 6\n", + "dataset_name.append(\"example_image_8\") # Index 7\n", + "# ...and so on." + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now extract the dataset name specific to this hpc id, meaning every CPU run will load and fit a different dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = dataset_name[hpc_id]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now create the overall path to the dataset this specific call of the script fits, which for the first line in the \n", + "`.conf` file above (which has integer input 0) is: \n", + "\n", + " `/hpc/data/hpc_username/dataset/imaging/example_image_1`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(hpc_dataset_path, dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You now have all the code you need to set up many single-CPU jobs on the hpc!\n", + "\n", + "You would simply append the batch scripts and Python code aboves to the lens modeling script script you are using,\n", + "which is given below for completeness.\n", + "\n", + "However, first we describe how to set up a single multi-CPU Nautilus job on the hpc.\n", + "\n", + "__Batch Script One Lenses__\n", + "\n", + "Lets now look at the second example batch script, `example_cpu_one_dataset_parallel`, which fits a single dataset\n", + "using multiple CPUs.\n", + "\n", + "#!/bin/bash -l\n", + "\n", + " `#SBATCH -N 1` - The number of nodes we require, where 1 node contains 28 CPUs on the hpc.\n", + " `#SBATCH --ntasks=1` - The total number of task we are submitting.\n", + " `#SBATCH --cpus-per-task=16` - The number of tasks per CPU.\n", + " `#SBATCH -J example` - The name of the job, which is how it appears on hpc when you inspect it.\n", + " `#SBATCH -o output/output.%A.out` - Python interpreter output is placed in a file in the `output` folder.\n", + " `#SBATCH -o error/error.%A.out` - Python interpreter errors are placed in a file in the `error` folder.\n", + " `#SBATCH -p hpc` - Signifies we are running the job on the hpc.\n", + " `#SBATCH -A dp004` - The project code of the submission.\n", + " `#SBATCH -t 48:00:00` - The job will terminate after this length of time (if it does not end naturally).\n", + " `#SBATCH --mail-type=END` - If you input your email, when you`ll get an email about the job (END means once finished).\n", + " `#SBATCH --mail-user=fill@me.co.uk` - The email address the hpc sends the email too.\n", + "\n", + "source /hpc/home/hpc_username/activate.sh\n", + "\n", + "export CPUS_PER_TASK=1\n", + "\n", + "export OPENBLAS_NUM_THREADS=$CPUS_PER_TASK\n", + "export MKL_NUM_THREADS=$CPUS_PER_TASK\n", + "export OMP_NUM_THREADS=$CPUS_PER_TASK\n", + "export VECLIB_MAXIMUM_THREADS=$CPUS_PER_TASK\n", + "export NUMEXPR_NUM_THREADS=$CPUS_PER_TASK\n", + "\n", + "python3 /hpc/home/hpc_username/runners/example.py 16\n", + "\n", + "The key difference in this batch script is the line which sets off the job:\n", + "\n", + "- ntasks=1 - We only require one task because we are only fitting one dataset.\n", + "\n", + "- cpus-per-task=16 - We require 16 CPUs for this single task, which the Nautilus non-linear search will use\n", + "to parallelize the model-fit over.\n", + "\n", + "- python3 ... example.py 16 - We input the integer 16, which is used below to set the number of CPUs Nautilus\n", + "\n", + "Lets now look at the beginning of the modeling script script again, which now does not use a list of datasets \n", + "to load, but instead has the dataset name hard coded." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `example_image_1` via .fits files, which we will fit with the lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_folder = \"example\"\n", + "dataset_name = \"simple__no_lens_light\"\n", + "\n", + "dataset_path = Path(hpc_dataset_path, dataset_folder, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Nautilus CPUs__\n", + "\n", + "The final change we need to make is to set the number of CPUs Nautilus uses in the model-fit.\n", + "\n", + "We do this by loading the integer input form the batch script, which we set to be the number of CPUs Nautilus uses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "number_of_cores = int(sys.argv[1])\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=\"hpc\",\n", + " name=\"example\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " number_of_cores=number_of_cores,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now have everything we need to fit a single dataset using multiple CPUs on the hpc!\n", + "\n", + "The code below performs standard lens modeling, which is unchanged from normal modeling on a laptop. It can be\n", + "used for either many single-CPU jobs or a single multi-CPU Nautilus job.\n", + "\n", + "__Lens Modeling__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + " \n", + " - The source galaxy's light is an MGE [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.SersicCore)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/lens_calc.ipynb b/notebooks/guides/lens_calc.ipynb index c2294fb9a..59b661c63 100644 --- a/notebooks/guides/lens_calc.ipynb +++ b/notebooks/guides/lens_calc.ipynb @@ -1,893 +1,930 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lens Calc\n", - "=========\n", - "\n", - "This guide explains the ``LensCalc`` class, which computes a comprehensive set of lensing quantities from the\n", - "deflection angles of a mass distribution.\n", - "\n", - "Given any mass model (a ``MassProfile``, ``Galaxy``, or ``Tracer``), ``LensCalc`` derives:\n", - "\n", - "- **Convergence** \u2014 the projected surface mass density of the lens, normalised by the critical density.\n", - "- **Shear** \u2014 the tidal stretching and squeezing of lensed images.\n", - "- **Magnification** \u2014 how much brighter (or fainter) a lensed image appears compared to the unlensed source.\n", - "- **Critical curves** \u2014 special curves in the image plane where magnification diverges to infinity.\n", - "- **Caustics** \u2014 the source-plane counterparts of critical curves, which delimit regions of multiple imaging.\n", - "- **Einstein radius** \u2014 the characteristic angular scale of the lens, derived from the critical curves.\n", - "- **Fermat potential** \u2014 the time-delay surface, whose stationary points correspond to the observed image positions.\n", - "\n", - "All of these are derived from the **deflection angles** of the lens. If you are new to gravitational lensing, this\n", - "guide walks through each quantity from first principles, with equations and code examples.\n", - "\n", - "__Contents__\n", - "\n", - "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", - "- **Data Structures:** Arrays inspected in this example use bespoke data structures for storing arrays, grids, vectors and.\n", - "- **Grids:** To describe the deflection of light, **PyAutoLens** uses `Grid2D` data structures.\n", - "- **Mass Profile and Galaxy:** We create a simple elliptical isothermal mass profile and wrap it in a Galaxy.\n", - "- **Tracer:** We create a two-plane Tracer from a lens and source galaxy.\n", - "- **LensCalc:** We introduce the ``LensCalc`` object and how to construct one from a Tracer.\n", - "- **The Lens Equation:** The fundamental equation of gravitational lensing.\n", - "- **Deflection Angles:** The deflection angles are the input to every other lensing quantity.\n", - "- **Hessian:** The matrix of second derivatives of the lensing potential.\n", - "- **Convergence:** The projected surface mass density normalised by the critical density.\n", - "- **Shear:** The tidal distortion field that stretches lensed images.\n", - "- **Magnification:** How much a lensed image is brightened (or dimmed) relative to the unlensed source.\n", - "- **Critical Curves and Caustics:** Where magnification formally diverges and how this maps to the source plane.\n", - "- **Einstein Radius:** The characteristic angular size of a lens.\n", - "- **Fermat Potential:** The time-delay surface whose extrema locate lensed images.\n", - "- **Wrap Up:** Summary and pointers to further reading.\n", - "\n", - "__Units__\n", - "\n", - "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", - "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", - "\n", - "The guide ``guides/units/cosmology.ipynb`` illustrates how to convert these quantities to physical units like\n", - "kiloparsecs, magnitudes and solar masses.\n", - "\n", - "__Data Structures__\n", - "\n", - "Arrays inspected in this example use bespoke data structures for storing arrays, grids,\n", - "vectors and other 1D and 2D quantities. These use the ``slim`` and ``native`` API to toggle between representing the\n", - "data in 1D numpy arrays or high dimension numpy arrays.\n", - "\n", - "This tutorial will only use the ``slim`` properties which show results in 1D numpy arrays of\n", - "shape [total_unmasked_pixels]. This is a slimmed-down representation of the data in 1D that contains only the\n", - "unmasked data points.\n", - "\n", - "These are documented fully in the ``autolens_workspace/*/guides/data_structures.ipynb`` guide." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import autolens as al\n", - "import autoarray as aa\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grids__\n", - "\n", - "To describe the deflection of light, **PyAutoLens** uses `Grid2D` data structures, which are two-dimensional\n", - "Cartesian grids of (y,x) coordinates.\n", - "\n", - "Below, we make a uniform Cartesian grid in units of arcseconds. This grid will be used throughout this guide\n", - "to evaluate every lensing quantity." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(200, 200),\n", - " pixel_scales=0.05,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mass Profile and Galaxy__\n", - "\n", - "We create a simple elliptical isothermal mass profile (`Isothermal`). This is one of the most commonly used\n", - "mass models in strong lensing \u2014 it describes a singular isothermal ellipsoid (SIE), a good first approximation\n", - "for the mass distribution of an early-type galaxy.\n", - "\n", - "We then wrap it in a `Galaxy` at redshift 0.5, which represents the foreground lens galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass_profile = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " einstein_radius=1.6,\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass_profile)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We also create a simple source galaxy at redshift 1.0. The source does not need a mass profile \u2014 it is the\n", - "background object whose light is being lensed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.ExponentialCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=(0.1, 0.0),\n", - " intensity=0.1,\n", - " effective_radius=0.5,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "A `Tracer` combines the lens and source galaxies with a cosmological model to perform ray-tracing.\n", - "\n", - "Ray-tracing means computing how light rays from the source are deflected by the lens galaxy's gravity, producing\n", - "the distorted, magnified images we observe. The `Tracer` handles all of the cosmological distance calculations\n", - "behind the scenes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, source_galaxy],\n", - " cosmology=al.cosmo.Planck15(),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here is the lensed image of the source galaxy, showing the characteristic arcs and multiple images produced by\n", - "strong gravitational lensing." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = tracer.image_2d_from(grid=grid)\n", - "aplt.plot_array(array=image, title=\"Lensed Image of the Source Galaxy\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LensCalc__\n", - "\n", - "The `LensCalc` class is a calculator that derives all secondary lensing quantities from a deflection-angle\n", - "callable. You can construct one from a `Tracer`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_calc = al.LensCalc.from_tracer(tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can also construct one directly from a mass profile or galaxy:\n", - "\n", - " lens_calc = al.LensCalc.from_mass_obj(mass_obj=lens_galaxy)\n", - "\n", - "Both approaches give you the same interface. Using `from_tracer` is recommended when your lens system has\n", - "multiple planes or when you want to include all galaxies in the deflection calculation.\n", - "\n", - "Now let's walk through each lensing quantity that `LensCalc` can compute.\n", - "\n", - "__The Lens Equation__\n", - "\n", - "The fundamental equation of gravitational lensing relates the **observed image position** to the **true source\n", - "position**. In the simplest case of a thin lens, this is:\n", - "\n", - " beta = theta - alpha(theta)\n", - "\n", - "where:\n", - "\n", - "- theta is the observed (image-plane) position of a light ray, in arcseconds.\n", - "- alpha(theta) is the deflection angle \u2014 how much the light ray is bent by the lens's gravity at position theta.\n", - "- beta is the true (source-plane) position, i.e. where the source would appear if there were no lens.\n", - "\n", - "If the lens is strong enough, multiple image-plane positions theta can map to the same source position beta.\n", - "This is why we see multiple images of the same source in strong lensing systems.\n", - "\n", - "Everything that `LensCalc` computes starts from the deflection angles alpha(theta).\n", - "\n", - "__Deflection Angles__\n", - "\n", - "The deflection angles describe how much light is bent at each point in the image plane. They are a 2D vector\n", - "field \u2014 at every (y, x) coordinate, there is a deflection in both the y and x directions.\n", - "\n", - "We can compute them directly from the tracer:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections = tracer.deflections_yx_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The deflection angles are a `VectorYX2D` data structure with shape [N, 2], where the first column is the\n", - "y-deflection and the second column is the x-deflection.\n", - "\n", - "Let's print the deflection at the central pixel:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Deflection at centre (y, x):\", deflections[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And plot the y-component and x-component of the deflection field:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=grid.mask)\n", - "aplt.plot_array(array=deflections_y, title=\"Deflection Angles (y-component)\")\n", - "\n", - "deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=grid.mask)\n", - "aplt.plot_array(array=deflections_x, title=\"Deflection Angles (x-component)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "These deflection angles are the only input that `LensCalc` needs. Every other quantity \u2014 convergence, shear,\n", - "magnification, critical curves, etc. \u2014 is derived from them by taking derivatives or solving equations.\n", - "\n", - "__Hessian__\n", - "\n", - "The **Hessian** is the 2x2 matrix of second partial derivatives of the lensing potential psi(theta).\n", - "Equivalently, it is the matrix of first partial derivatives of the deflection angles:\n", - "\n", - " H_yy = d(alpha_y) / d(theta_y)\n", - " H_xy = d(alpha_x) / d(theta_y)\n", - " H_yx = d(alpha_y) / d(theta_x)\n", - " H_xx = d(alpha_x) / d(theta_x)\n", - "\n", - "Written as a matrix:\n", - "\n", - " H = | H_yy H_xy |\n", - " | H_yx H_xx |\n", - "\n", - "The Hessian captures how the deflection angles *change* across the image plane. This is the key to understanding\n", - "how images are distorted: convergence, shear, and magnification are all simple combinations of Hessian components.\n", - "\n", - "`LensCalc` computes the Hessian by **finite differences** (nudging the grid positions slightly and measuring\n", - "how the deflections change). If JAX is available, it can alternatively use automatic differentiation for\n", - "exact derivatives." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "hessian_yy, hessian_xy, hessian_yx, hessian_xx = lens_calc.hessian_from(grid=grid)\n", - "\n", - "print(\"Hessian_yy at centre:\", hessian_yy[0])\n", - "print(\"Hessian_xx at centre:\", hessian_xx[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Convergence__\n", - "\n", - "The **convergence** (kappa) is the projected surface mass density of the lens, normalised by the **critical\n", - "surface density**. It tells you how much mass is concentrated along the line of sight at each point.\n", - "\n", - "Physically:\n", - "\n", - "- kappa = 1 means the projected mass density equals the critical density \u2014 this is roughly the threshold\n", - " for strong lensing.\n", - "- kappa > 1 means the lens is \"super-critical\" at that point.\n", - "- kappa < 1 means the lens is \"sub-critical\".\n", - "\n", - "The convergence is computed from the Hessian as:\n", - "\n", - " kappa = 0.5 * (H_yy + H_xx)\n", - "\n", - "This is the trace of the Hessian divided by 2. It measures the isotropic part of the image distortion \u2014\n", - "convergence magnifies images uniformly (making them bigger and brighter) without changing their shape." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "convergence = lens_calc.convergence_2d_via_hessian_from(grid=grid)\n", - "\n", - "print(\"Convergence at centre:\", convergence[0])\n", - "\n", - "convergence_array = aa.Array2D(values=convergence, mask=grid.mask)\n", - "aplt.plot_array(array=convergence_array, title=\"Convergence (kappa)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The convergence computed this way (via the Hessian) is independent of any analytic formula \u2014 it works for\n", - "any mass distribution, as long as you can compute deflection angles.\n", - "\n", - "__Shear__\n", - "\n", - "The **shear** (gamma) describes the tidal stretching of lensed images. Unlike convergence (which magnifies\n", - "images isotropically), shear distorts images *anisotropically* \u2014 it stretches them along one axis and\n", - "compresses them along the perpendicular axis.\n", - "\n", - "Shear has two components:\n", - "\n", - " gamma_1 = 0.5 * (H_xx - H_yy)\n", - " gamma_2 = H_xy\n", - "\n", - "The total shear magnitude is:\n", - "\n", - " |gamma| = sqrt(gamma_1^2 + gamma_2^2)\n", - "\n", - "Physically, the shear direction tells you the orientation of the tidal stretching: images near a lens are\n", - "elongated tangentially (forming the characteristic arcs of strong lensing), while images far from the lens\n", - "experience weaker, more radially oriented distortion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "shear = lens_calc.shear_yx_2d_via_hessian_from(grid=grid)\n", - "\n", - "print(\"Shear gamma_2 at centre:\", shear[0, 0])\n", - "print(\"Shear gamma_1 at centre:\", shear[0, 1])\n", - "print(\"Shear magnitude at centre:\", shear.magnitudes[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Magnification__\n", - "\n", - "The **magnification** (mu) tells you how much brighter (or fainter) a lensed image appears compared to the\n", - "unlensed source. It is defined as the inverse of the determinant of the **lensing Jacobian matrix**:\n", - "\n", - " A = I - H = | 1 - H_yy -H_xy |\n", - " | -H_yx 1 - H_xx |\n", - "\n", - " mu = 1 / det(A) = 1 / [(1 - H_yy)(1 - H_xx) - H_xy * H_yx]\n", - "\n", - "Equivalently, using convergence and shear:\n", - "\n", - " mu = 1 / [(1 - kappa)^2 - |gamma|^2]\n", - "\n", - "Key points:\n", - "\n", - "- |mu| > 1 means the image is magnified (brighter and larger than the unlensed source).\n", - "- |mu| < 1 means the image is demagnified.\n", - "- mu < 0 means the image has flipped parity (it is a mirror image of the source).\n", - "- Where det(A) = 0, the magnification diverges to infinity \u2014 these special locations are the **critical curves**." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnification = lens_calc.magnification_2d_from(grid=grid)\n", - "\n", - "print(\"Magnification at centre:\", magnification[0])\n", - "\n", - "magnification_array = aa.Array2D(values=magnification, mask=grid.mask)\n", - "aplt.plot_array(array=magnification_array, title=\"Magnification (mu)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The magnification map shows extremely high values near the critical curves, where images are stretched into\n", - "the bright arcs that make strong lensing systems so visually striking.\n", - "\n", - "__Critical Curves and Caustics__\n", - "\n", - "**Critical curves** are closed curves in the image plane where the magnification formally diverges (det(A) = 0).\n", - "\n", - "There are two types:\n", - "\n", - "- **Tangential critical curves** \u2014 found where the tangential eigenvalue (1 - kappa - |gamma|) = 0. These\n", - " are the ones that produce the bright, highly magnified arcs seen in strong lensing systems.\n", - "\n", - "- **Radial critical curves** \u2014 found where the radial eigenvalue (1 - kappa + |gamma|) = 0. These produce\n", - " fainter, radially oriented counter-images.\n", - "\n", - "`LensCalc` finds critical curves by evaluating the eigenvalues on a fine grid and tracing the zero-contours\n", - "using a marching-squares algorithm." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tangential_critical_curves = lens_calc.tangential_critical_curve_list_from(grid=grid)\n", - "radial_critical_curves = lens_calc.radial_critical_curve_list_from(grid=grid)\n", - "\n", - "print(\"Number of tangential critical curves:\", len(tangential_critical_curves))\n", - "print(\"Number of radial critical curves:\", len(radial_critical_curves))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**Caustics** are the source-plane images of the critical curves. They are found by ray-tracing each critical\n", - "curve through the lens equation (subtracting the deflection angles):\n", - "\n", - " caustic = critical_curve - alpha(critical_curve)\n", - "\n", - "Caustics divide the source plane into regions with different numbers of images. A source inside the tangential\n", - "caustic produces multiple (typically 4) images, while a source outside produces fewer (typically 2)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tangential_caustics = lens_calc.tangential_caustic_list_from(grid=grid)\n", - "radial_caustics = lens_calc.radial_caustic_list_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Let's plot the convergence map. The critical curves trace the boundary between highly magnified and weakly\n", - "magnified regions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "convergence_for_plot = tracer.convergence_2d_from(grid=grid)\n", - "aplt.plot_array(array=convergence_for_plot, title=\"Convergence with Critical Curves\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Einstein Radius__\n", - "\n", - "The **Einstein radius** is the characteristic angular scale of a strong lens. It is defined as the radius of\n", - "the circle that encloses the same area as the tangential critical curve:\n", - "\n", - " theta_E = sqrt(A_crit / pi)\n", - "\n", - "where A_crit is the area enclosed by the tangential critical curve.\n", - "\n", - "This is sometimes called the \"effective Einstein radius\" in the literature. For a circular lens, the tangential\n", - "critical curve is a perfect circle and the Einstein radius equals its geometric radius. For an elliptical lens,\n", - "the critical curve is not circular, so the Einstein radius is an effective average.\n", - "\n", - "The Einstein radius sets the scale of the lensing system \u2014 the separation between multiple images, the size\n", - "of the arcs, and the enclosed mass are all closely related to it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "einstein_radius = lens_calc.einstein_radius_from(grid=grid)\n", - "\n", - "print(\"Einstein radius:\", einstein_radius, \"arcsec\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The angular Einstein mass (in arcseconds squared) is:\n", - "\n", - " M_E = pi * theta_E^2\n", - "\n", - "To convert this to physical mass (e.g. solar masses), you need the critical surface density, which depends\n", - "on the cosmological distances to the lens and source. See the ``guides/units/cosmology.ipynb`` guide." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "einstein_mass = lens_calc.einstein_mass_angular_from(grid=grid)\n", - "\n", - "print(\"Angular Einstein mass:\", einstein_mass, \"arcsec^2\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fermat Potential__\n", - "\n", - "The **Fermat potential** (also called the time-delay surface or arrival-time surface) is a scalar field in the\n", - "image plane that encodes the light travel time from source to observer via each image-plane position.\n", - "\n", - "It is given by:\n", - "\n", - " phi(theta) = 0.5 * |theta - beta|^2 - psi(theta)\n", - "\n", - "where:\n", - "\n", - "- theta is the image-plane position.\n", - "- beta is the source-plane position (= theta - alpha(theta)).\n", - "- psi(theta) is the lensing potential (the scalar potential whose gradient gives the deflection angles).\n", - "\n", - "The first term, 0.5 * |theta - beta|^2, is the **geometric delay** \u2014 the extra path length due to the\n", - "bending of light. The second term, psi(theta), is the **gravitational (Shapiro) delay** \u2014 the slowing of\n", - "light in the gravitational potential of the lens.\n", - "\n", - "Fermat's principle tells us that observed images form at the stationary points (extrema and saddle points) of\n", - "this surface. This is a powerful result: it means the positions of lensed images are determined by the topology\n", - "of the Fermat potential.\n", - "\n", - "The *differences* in the Fermat potential between image positions are proportional to the time delays between\n", - "images. Measuring these time delays (e.g. from a variable quasar) can constrain the Hubble constant." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fermat_potential = lens_calc.fermat_potential_from(grid=grid)\n", - "\n", - "fermat_array = aa.Array2D(values=fermat_potential, mask=grid.mask)\n", - "aplt.plot_array(array=fermat_array, title=\"Fermat Potential\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The geometric delay term alone can also be inspected:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "geometric_delay = lens_calc.time_delay_geometry_term_from(grid=grid)\n", - "\n", - "geometric_array = aa.Array2D(values=geometric_delay, mask=grid.mask)\n", - "aplt.plot_array(array=geometric_array, title=\"Geometric Time Delay Term\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__JAX (JIT-it-yourself)__\n", - "\n", - "This is the canonical home for the \"JIT-it-yourself\" advanced path.\n", - "Audience: users building custom forward models or scientific tools\n", - "using PyAutoLens primitives directly \u2014 not running standard fits\n", - "(which go through `Analysis(use_jax=True)` and need zero JAX code from\n", - "you) or standard simulations (which use `Simulator(use_jax=True)`\n", - "similarly).\n", - "\n", - "If you're writing `@jax.jit` yourself around library calls like\n", - "`tracer.image_2d_from`, `LensCalc.magnification_2d_via_hessian_from`,\n", - "or your own `log_likelihood(instance)` function, this section is for you.\n", - "\n", - "__The pairing rule: `@jax.jit` + `xp=jnp`__\n", - "\n", - "The single rule to remember: when you decorate a function with\n", - "`@jax.jit` that calls a PyAutoLens library method internally, **pass\n", - "`xp=jnp` to that method inside the function body**.\n", - "\n", - "```python\n", - "import jax\n", - "import jax.numpy as jnp\n", - "\n", - "@jax.jit\n", - "def magnification_fn(lens_calc, grid):\n", - " return lens_calc.magnification_2d_via_hessian_from(grid=grid, xp=jnp)\n", - "```\n", - "\n", - "The `xp=jnp` is what tells the library \"you're inside a JAX trace \u2014\n", - "route calls through `jax.numpy` and don't wrap the return in an\n", - "autoarray type (it would fail to cross the JIT boundary)\".\n", - "\n", - "__The footgun: forgetting `xp=jnp`__\n", - "\n", - "The default for `xp` in library method signatures is `np`. If you\n", - "forget to pass `xp=jnp` inside `@jax.jit`, one of two things happens:\n", - "\n", - "- The function body does `np.sqrt(jax_array)` \u2014 NumPy routes through\n", - " `__array__()` on the JAX tracer, host-transferring off the GPU.\n", - " Your fit runs, invisibly slower than the NumPy path.\n", - "- The `if xp is np:` guard inside library functions fires and wraps\n", - " the result in `aa.Array2D`, which fails at the JIT boundary with\n", - " `TypeError: ... is not a valid JAX type`.\n", - "\n", - "The library raises a clear `ValueError` on the easy mismatch \u2014 when\n", - "you pass `xp=np` with a `grid.use_jax=True` input:\n", - "\n", - "```\n", - "ValueError: Called magnification_2d_via_hessian_from with xp=np but\n", - "the input grid is JAX-backed (grid.use_jax=True). Inside @jax.jit,\n", - "pass xp=jnp explicitly.\n", - "```\n", - "\n", - "If you see this error: add `xp=jnp` to the call site. Done.\n", - "\n", - "__Decorator-on-def vs `jax.jit(bound_method)`__\n", - "\n", - "JAX accepts any callable \u2014 `@jax.jit` is sugar for `fn = jax.jit(fn)`.\n", - "You can JIT a standalone function (canonical) or a bound method\n", - "(shortcut). Both work:\n", - "\n", - "```python\n", - "# Form 1 (canonical): decorator on def\n", - "@jax.jit\n", - "def image_fn(tracer, grid):\n", - " return tracer.image_2d_from(grid=grid, xp=jnp).array\n", - "\n", - "# Form 2: jit on bound method, assign-to-variable\n", - "jitted = jax.jit(tracer.image_2d_from)\n", - "arr = jitted(grid=grid, xp=jnp).array\n", - "```\n", - "\n", - "Form 2 is shorter for interactive use. **Footgun:** bound methods are\n", - "fresh objects on every attribute access, so this silently misses the\n", - "JIT cache every iteration:\n", - "\n", - "```python\n", - "# DON'T DO THIS \u2014 fresh jax.jit closure every iteration\n", - "for grid in many_grids:\n", - " arr = jax.jit(tracer.image_2d_from)(grid=grid, xp=jnp).array\n", - "```\n", - "\n", - "If you're calling a JITted method in a loop: assign once outside the\n", - "loop, or use the decorator-on-def form.\n", - "\n", - "__Closure-captured `self` vs traced argument__\n", - "\n", - "Form 1 and Form 2 differ in *semantics*, not just syntax:\n", - "\n", - "- **Form 2 (`jax.jit(tracer.image_2d_from)`):** `tracer` is the bound\n", - " method's `self`; JAX captures it as a closure constant and doesn't\n", - " trace through it. **`Tracer` does NOT need to be pytree-registered.**\n", - " Trade-off: a different tracer means a fresh bound-method object and a\n", - " fresh JIT cache key.\n", - "- **Form 1 (`@jax.jit def image_fn(tracer, grid)`):** `tracer` is a\n", - " traced argument. **`Tracer` DOES need pytree registration** (which\n", - " `Analysis(use_jax=True)` does for you, or you can trigger via\n", - " `autolens.jax.register_tracer_classes(tracer)`). Trade-off: cache\n", - " reuse across different tracers \u2014 parameter sweeps and `jax.vmap`\n", - " work naturally.\n", - "\n", - "Pick based on whether you want to vary the tracer across calls:\n", - "\n", - "- Parameter sweep / `jax.vmap` over models? Form 1.\n", - "- Quick one-off / interactive exploration? Form 2 (assign once).\n", - "\n", - "__`LensCalc` and the wrapped-vs-raw return type__\n", - "\n", - "`LensCalc.magnification_2d_via_hessian_from`,\n", - "`.shear_yx_2d_via_hessian_from`, `.convergence_2d_via_hessian_from`,\n", - "and the eigen-value methods all implement the `if xp is np:` guard\n", - "inside:\n", - "\n", - "```python\n", - "def magnification_2d_via_hessian_from(self, grid, xp=np):\n", - " ...\n", - " if xp is np:\n", - " return aa.Array2D(values=mag, mask=grid.mask) # numpy: wrapped\n", - " return mag # jax: raw jax.Array\n", - "```\n", - "\n", - "This is intentional. Inside `@jax.jit` (where you pass `xp=jnp`), you\n", - "get back a raw `jax.Array`. On the NumPy path (no `xp` or `xp=np`),\n", - "you get an `aa.Array2D` wrapper. The function adapts to which path\n", - "you're on.\n", - "\n", - "Implication: when you JIT-wrap a `LensCalc` method, expect a raw\n", - "`jax.Array` back. Rewrap with `aa.Array2D(values=..., mask=...)` on\n", - "the host if you want the wrapper for downstream plotting:\n", - "\n", - "```python\n", - "@jax.jit\n", - "def magnification_fn(lens_calc, grid):\n", - " return lens_calc.magnification_2d_via_hessian_from(grid=grid, xp=jnp)\n", - "\n", - "mag_raw = magnification_fn(lens_calc, grid)\n", - "mag_wrapped = aa.Array2D(values=mag_raw, mask=grid.mask)\n", - "aplt.plot_array(array=mag_wrapped)\n", - "```\n", - "\n", - "For `Tracer` and `Galaxy` methods that don't have the guard internally\n", - "(e.g. `tracer.image_2d_from`), the `.array` unwrap inside the jit +\n", - "rewrap outside discipline applies \u2014 see `scripts/guides/data_structures.py`\n", - "`__JAX__` section.\n", - "\n", - "__Summary \u2014 the three rules__\n", - "\n", - "The \"JIT-it-yourself\" path is bounded by three rules:\n", - "\n", - "1. **`@jax.jit` and `xp=jnp` are paired.** Forgetting `xp=jnp` either\n", - " silently host-transfers or fails at the boundary.\n", - "2. **`.array` unwrap inside the jit; rewrap on the host** if you want\n", - " the autoarray wrapper.\n", - "3. **Tracer-as-argument needs pytree registration** (via\n", - " `register_tracer_classes(tracer)` or any `Analysis(use_jax=True)`\n", - " construction); tracer-as-closure (bound-method form) doesn't.\n", - "\n", - "For the standard `Analysis` / `Simulator` paths \u2014 where you do none\n", - "of this and JAX runs implicitly \u2014 see the top-level\n", - "`autolens_workspace/start_here.py` `__JAX__` section." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Wrap Up__\n", - "\n", - "This guide introduced the `LensCalc` class and the key lensing quantities it computes:\n", - "\n", - "1. **Deflection angles** \u2014 the bending of light by the lens's gravity.\n", - "2. **Hessian** \u2014 the matrix of second derivatives, from which all other quantities are derived.\n", - "3. **Convergence** \u2014 the projected mass density (isotropic magnification).\n", - "4. **Shear** \u2014 the tidal distortion (anisotropic stretching).\n", - "5. **Magnification** \u2014 the brightness and size change of lensed images.\n", - "6. **Critical curves and caustics** \u2014 where magnification diverges, and the corresponding source-plane boundaries.\n", - "7. **Einstein radius** \u2014 the characteristic angular scale of the lens.\n", - "8. **Fermat potential** \u2014 the time-delay surface whose stationary points give image positions.\n", - "\n", - "For further reading:\n", - "\n", - "- ``guides/tracer.py`` \u2014 ray-tracing and image formation.\n", - "- ``guides/galaxies.py`` \u2014 working with individual galaxy components.\n", - "- ``guides/units/cosmology.ipynb`` \u2014 converting to physical units.\n", - "- ``guides/data_structures.py`` \u2014 the ``Array2D``, ``Grid2D`` and ``VectorYX2D`` data structures.\n", - "\n", - "The `LensCalc` class is also available via `al.LensCalc` \u2014 it lives in ``PyAutoGalaxy/autogalaxy/operate/lens_calc.py``.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lens Calc\n", + "=========\n", + "\n", + "This guide explains the ``LensCalc`` class, which computes a comprehensive set of lensing quantities from the\n", + "deflection angles of a mass distribution.\n", + "\n", + "Given any mass model (a ``MassProfile``, ``Galaxy``, or ``Tracer``), ``LensCalc`` derives:\n", + "\n", + "- **Convergence** \u2014 the projected surface mass density of the lens, normalised by the critical density.\n", + "- **Shear** \u2014 the tidal stretching and squeezing of lensed images.\n", + "- **Magnification** \u2014 how much brighter (or fainter) a lensed image appears compared to the unlensed source.\n", + "- **Critical curves** \u2014 special curves in the image plane where magnification diverges to infinity.\n", + "- **Caustics** \u2014 the source-plane counterparts of critical curves, which delimit regions of multiple imaging.\n", + "- **Einstein radius** \u2014 the characteristic angular scale of the lens, derived from the critical curves.\n", + "- **Fermat potential** \u2014 the time-delay surface, whose stationary points correspond to the observed image positions.\n", + "\n", + "All of these are derived from the **deflection angles** of the lens. If you are new to gravitational lensing, this\n", + "guide walks through each quantity from first principles, with equations and code examples.\n", + "\n", + "__Contents__\n", + "\n", + "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", + "- **Data Structures:** Arrays inspected in this example use bespoke data structures for storing arrays, grids, vectors and.\n", + "- **Grids:** To describe the deflection of light, **PyAutoLens** uses `Grid2D` data structures.\n", + "- **Mass Profile and Galaxy:** We create a simple elliptical isothermal mass profile and wrap it in a Galaxy.\n", + "- **Tracer:** We create a two-plane Tracer from a lens and source galaxy.\n", + "- **LensCalc:** We introduce the ``LensCalc`` object and how to construct one from a Tracer.\n", + "- **The Lens Equation:** The fundamental equation of gravitational lensing.\n", + "- **Deflection Angles:** The deflection angles are the input to every other lensing quantity.\n", + "- **Hessian:** The matrix of second derivatives of the lensing potential.\n", + "- **Convergence:** The projected surface mass density normalised by the critical density.\n", + "- **Shear:** The tidal distortion field that stretches lensed images.\n", + "- **Magnification:** How much a lensed image is brightened (or dimmed) relative to the unlensed source.\n", + "- **Critical Curves and Caustics:** Where magnification formally diverges and how this maps to the source plane.\n", + "- **Einstein Radius:** The characteristic angular size of a lens.\n", + "- **Fermat Potential:** The time-delay surface whose extrema locate lensed images.\n", + "- **Wrap Up:** Summary and pointers to further reading.\n", + "\n", + "__Units__\n", + "\n", + "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", + "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", + "\n", + "The guide ``guides/units/cosmology.ipynb`` illustrates how to convert these quantities to physical units like\n", + "kiloparsecs, magnitudes and solar masses.\n", + "\n", + "__Data Structures__\n", + "\n", + "Arrays inspected in this example use bespoke data structures for storing arrays, grids,\n", + "vectors and other 1D and 2D quantities. These use the ``slim`` and ``native`` API to toggle between representing the\n", + "data in 1D numpy arrays or high dimension numpy arrays.\n", + "\n", + "This tutorial will only use the ``slim`` properties which show results in 1D numpy arrays of\n", + "shape [total_unmasked_pixels]. This is a slimmed-down representation of the data in 1D that contains only the\n", + "unmasked data points.\n", + "\n", + "These are documented fully in the ``autolens_workspace/*/guides/data_structures.ipynb`` guide." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import autolens as al\n", + "import autoarray as aa\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grids__\n", + "\n", + "To describe the deflection of light, **PyAutoLens** uses `Grid2D` data structures, which are two-dimensional\n", + "Cartesian grids of (y,x) coordinates.\n", + "\n", + "Below, we make a uniform Cartesian grid in units of arcseconds. This grid will be used throughout this guide\n", + "to evaluate every lensing quantity." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(200, 200),\n", + " pixel_scales=0.05,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mass Profile and Galaxy__\n", + "\n", + "We create a simple elliptical isothermal mass profile (`Isothermal`). This is one of the most commonly used\n", + "mass models in strong lensing \u2014 it describes a singular isothermal ellipsoid (SIE), a good first approximation\n", + "for the mass distribution of an early-type galaxy.\n", + "\n", + "We then wrap it in a `Galaxy` at redshift 0.5, which represents the foreground lens galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass_profile = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " einstein_radius=1.6,\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass_profile)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We also create a simple source galaxy at redshift 1.0. The source does not need a mass profile \u2014 it is the\n", + "background object whose light is being lensed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.ExponentialCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=(0.1, 0.0),\n", + " intensity=0.1,\n", + " effective_radius=0.5,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "A `Tracer` combines the lens and source galaxies with a cosmological model to perform ray-tracing.\n", + "\n", + "Ray-tracing means computing how light rays from the source are deflected by the lens galaxy's gravity, producing\n", + "the distorted, magnified images we observe. The `Tracer` handles all of the cosmological distance calculations\n", + "behind the scenes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, source_galaxy],\n", + " cosmology=al.cosmo.Planck15(),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here is the lensed image of the source galaxy, showing the characteristic arcs and multiple images produced by\n", + "strong gravitational lensing." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = tracer.image_2d_from(grid=grid)\n", + "aplt.plot_array(array=image, title=\"Lensed Image of the Source Galaxy\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LensCalc__\n", + "\n", + "The `LensCalc` class is a calculator that derives all secondary lensing quantities from a deflection-angle\n", + "callable. You can construct one from a `Tracer`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_calc = al.LensCalc.from_tracer(tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can also construct one directly from a mass profile or galaxy:\n", + "\n", + " lens_calc = al.LensCalc.from_mass_obj(mass_obj=lens_galaxy)\n", + "\n", + "Both approaches give you the same interface. Using `from_tracer` is recommended when your lens system has\n", + "multiple planes or when you want to include all galaxies in the deflection calculation.\n", + "\n", + "Now let's walk through each lensing quantity that `LensCalc` can compute.\n", + "\n", + "__The Lens Equation__\n", + "\n", + "The fundamental equation of gravitational lensing relates the **observed image position** to the **true source\n", + "position**. In the simplest case of a thin lens, this is:\n", + "\n", + " beta = theta - alpha(theta)\n", + "\n", + "where:\n", + "\n", + "- theta is the observed (image-plane) position of a light ray, in arcseconds.\n", + "- alpha(theta) is the deflection angle \u2014 how much the light ray is bent by the lens's gravity at position theta.\n", + "- beta is the true (source-plane) position, i.e. where the source would appear if there were no lens.\n", + "\n", + "If the lens is strong enough, multiple image-plane positions theta can map to the same source position beta.\n", + "This is why we see multiple images of the same source in strong lensing systems.\n", + "\n", + "Everything that `LensCalc` computes starts from the deflection angles alpha(theta).\n", + "\n", + "__Deflection Angles__\n", + "\n", + "The deflection angles describe how much light is bent at each point in the image plane. They are a 2D vector\n", + "field \u2014 at every (y, x) coordinate, there is a deflection in both the y and x directions.\n", + "\n", + "We can compute them directly from the tracer:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections = tracer.deflections_yx_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The deflection angles are a `VectorYX2D` data structure with shape [N, 2], where the first column is the\n", + "y-deflection and the second column is the x-deflection.\n", + "\n", + "Let's print the deflection at the central pixel:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Deflection at centre (y, x):\", deflections[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "And plot the y-component and x-component of the deflection field:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=grid.mask)\n", + "aplt.plot_array(array=deflections_y, title=\"Deflection Angles (y-component)\")\n", + "\n", + "deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=grid.mask)\n", + "aplt.plot_array(array=deflections_x, title=\"Deflection Angles (x-component)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "These deflection angles are the only input that `LensCalc` needs. Every other quantity \u2014 convergence, shear,\n", + "magnification, critical curves, etc. \u2014 is derived from them by taking derivatives or solving equations.\n", + "\n", + "__Hessian__\n", + "\n", + "The **Hessian** is the 2x2 matrix of second partial derivatives of the lensing potential psi(theta).\n", + "Equivalently, it is the matrix of first partial derivatives of the deflection angles:\n", + "\n", + " H_yy = d(alpha_y) / d(theta_y)\n", + " H_xy = d(alpha_x) / d(theta_y)\n", + " H_yx = d(alpha_y) / d(theta_x)\n", + " H_xx = d(alpha_x) / d(theta_x)\n", + "\n", + "Written as a matrix:\n", + "\n", + " H = | H_yy H_xy |\n", + " | H_yx H_xx |\n", + "\n", + "The Hessian captures how the deflection angles *change* across the image plane. This is the key to understanding\n", + "how images are distorted: convergence, shear, and magnification are all simple combinations of Hessian components.\n", + "\n", + "`LensCalc` computes the Hessian by **finite differences** (nudging the grid positions slightly and measuring\n", + "how the deflections change). If JAX is available, it can alternatively use automatic differentiation for\n", + "exact derivatives." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "hessian_yy, hessian_xy, hessian_yx, hessian_xx = lens_calc.hessian_from(grid=grid)\n", + "\n", + "print(\"Hessian_yy at centre:\", hessian_yy[0])\n", + "print(\"Hessian_xx at centre:\", hessian_xx[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Convergence__\n", + "\n", + "The **convergence** (kappa) is the projected surface mass density of the lens, normalised by the **critical\n", + "surface density**. It tells you how much mass is concentrated along the line of sight at each point.\n", + "\n", + "Physically:\n", + "\n", + "- kappa = 1 means the projected mass density equals the critical density \u2014 this is roughly the threshold\n", + " for strong lensing.\n", + "- kappa > 1 means the lens is \"super-critical\" at that point.\n", + "- kappa < 1 means the lens is \"sub-critical\".\n", + "\n", + "The convergence is computed from the Hessian as:\n", + "\n", + " kappa = 0.5 * (H_yy + H_xx)\n", + "\n", + "This is the trace of the Hessian divided by 2. It measures the isotropic part of the image distortion \u2014\n", + "convergence magnifies images uniformly (making them bigger and brighter) without changing their shape." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "convergence = lens_calc.convergence_2d_via_hessian_from(grid=grid)\n", + "\n", + "print(\"Convergence at centre:\", convergence[0])\n", + "\n", + "convergence_array = aa.Array2D(values=convergence, mask=grid.mask)\n", + "aplt.plot_array(array=convergence_array, title=\"Convergence (kappa)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The convergence computed this way (via the Hessian) is independent of any analytic formula \u2014 it works for\n", + "any mass distribution, as long as you can compute deflection angles.\n", + "\n", + "__Shear__\n", + "\n", + "The **shear** (gamma) describes the tidal stretching of lensed images. Unlike convergence (which magnifies\n", + "images isotropically), shear distorts images *anisotropically* \u2014 it stretches them along one axis and\n", + "compresses them along the perpendicular axis.\n", + "\n", + "Shear has two components:\n", + "\n", + " gamma_1 = 0.5 * (H_xx - H_yy)\n", + " gamma_2 = H_xy\n", + "\n", + "The total shear magnitude is:\n", + "\n", + " |gamma| = sqrt(gamma_1^2 + gamma_2^2)\n", + "\n", + "Physically, the shear direction tells you the orientation of the tidal stretching: images near a lens are\n", + "elongated tangentially (forming the characteristic arcs of strong lensing), while images far from the lens\n", + "experience weaker, more radially oriented distortion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "shear = lens_calc.shear_yx_2d_via_hessian_from(grid=grid)\n", + "\n", + "print(\"Shear gamma_2 at centre:\", shear[0, 0])\n", + "print(\"Shear gamma_1 at centre:\", shear[0, 1])\n", + "print(\"Shear magnitude at centre:\", shear.magnitudes[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Magnification__\n", + "\n", + "The **magnification** (mu) tells you how much brighter (or fainter) a lensed image appears compared to the\n", + "unlensed source. It is defined as the inverse of the determinant of the **lensing Jacobian matrix**:\n", + "\n", + " A = I - H = | 1 - H_yy -H_xy |\n", + " | -H_yx 1 - H_xx |\n", + "\n", + " mu = 1 / det(A) = 1 / [(1 - H_yy)(1 - H_xx) - H_xy * H_yx]\n", + "\n", + "Equivalently, using convergence and shear:\n", + "\n", + " mu = 1 / [(1 - kappa)^2 - |gamma|^2]\n", + "\n", + "Key points:\n", + "\n", + "- |mu| > 1 means the image is magnified (brighter and larger than the unlensed source).\n", + "- |mu| < 1 means the image is demagnified.\n", + "- mu < 0 means the image has flipped parity (it is a mirror image of the source).\n", + "- Where det(A) = 0, the magnification diverges to infinity \u2014 these special locations are the **critical curves**." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnification = lens_calc.magnification_2d_from(grid=grid)\n", + "\n", + "print(\"Magnification at centre:\", magnification[0])\n", + "\n", + "magnification_array = aa.Array2D(values=magnification, mask=grid.mask)\n", + "aplt.plot_array(array=magnification_array, title=\"Magnification (mu)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The magnification map shows extremely high values near the critical curves, where images are stretched into\n", + "the bright arcs that make strong lensing systems so visually striking.\n", + "\n", + "__Critical Curves and Caustics__\n", + "\n", + "**Critical curves** are closed curves in the image plane where the magnification formally diverges (det(A) = 0).\n", + "\n", + "There are two types:\n", + "\n", + "- **Tangential critical curves** \u2014 found where the tangential eigenvalue (1 - kappa - |gamma|) = 0. These\n", + " are the ones that produce the bright, highly magnified arcs seen in strong lensing systems.\n", + "\n", + "- **Radial critical curves** \u2014 found where the radial eigenvalue (1 - kappa + |gamma|) = 0. These produce\n", + " fainter, radially oriented counter-images.\n", + "\n", + "`LensCalc` finds critical curves by evaluating the eigenvalues on a fine grid and tracing the zero-contours\n", + "using a marching-squares algorithm." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tangential_critical_curves = lens_calc.tangential_critical_curve_list_from(grid=grid)\n", + "radial_critical_curves = lens_calc.radial_critical_curve_list_from(grid=grid)\n", + "\n", + "print(\"Number of tangential critical curves:\", len(tangential_critical_curves))\n", + "print(\"Number of radial critical curves:\", len(radial_critical_curves))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Caustics** are the source-plane images of the critical curves. They are found by ray-tracing each critical\n", + "curve through the lens equation (subtracting the deflection angles):\n", + "\n", + " caustic = critical_curve - alpha(critical_curve)\n", + "\n", + "Caustics divide the source plane into regions with different numbers of images. A source inside the tangential\n", + "caustic produces multiple (typically 4) images, while a source outside produces fewer (typically 2)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tangential_caustics = lens_calc.tangential_caustic_list_from(grid=grid)\n", + "radial_caustics = lens_calc.radial_caustic_list_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let's plot the convergence map. The critical curves trace the boundary between highly magnified and weakly\n", + "magnified regions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "convergence_for_plot = tracer.convergence_2d_from(grid=grid)\n", + "aplt.plot_array(array=convergence_for_plot, title=\"Convergence with Critical Curves\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Einstein Radius__\n", + "\n", + "The **Einstein radius** is the characteristic angular scale of a strong lens. It is defined as the radius of\n", + "the circle that encloses the same area as the tangential critical curve:\n", + "\n", + " theta_E = sqrt(A_crit / pi)\n", + "\n", + "where A_crit is the area enclosed by the tangential critical curve.\n", + "\n", + "This is sometimes called the \"effective Einstein radius\" in the literature. For a circular lens, the tangential\n", + "critical curve is a perfect circle and the Einstein radius equals its geometric radius. For an elliptical lens,\n", + "the critical curve is not circular, so the Einstein radius is an effective average.\n", + "\n", + "The Einstein radius sets the scale of the lensing system \u2014 the separation between multiple images, the size\n", + "of the arcs, and the enclosed mass are all closely related to it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "einstein_radius = lens_calc.einstein_radius_from(grid=grid)\n", + "\n", + "print(\"Einstein radius:\", einstein_radius, \"arcsec\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The angular Einstein mass (in arcseconds squared) is:\n", + "\n", + " M_E = pi * theta_E^2\n", + "\n", + "To convert this to physical mass (e.g. solar masses), you need the critical surface density, which depends\n", + "on the cosmological distances to the lens and source. See the ``guides/units/cosmology.ipynb`` guide." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "einstein_mass = lens_calc.einstein_mass_angular_from(grid=grid)\n", + "\n", + "print(\"Angular Einstein mass:\", einstein_mass, \"arcsec^2\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fermat Potential__\n", + "\n", + "The **Fermat potential** (also called the time-delay surface or arrival-time surface) is a scalar field in the\n", + "image plane that encodes the light travel time from source to observer via each image-plane position.\n", + "\n", + "It is given by:\n", + "\n", + " phi(theta) = 0.5 * |theta - beta|^2 - psi(theta)\n", + "\n", + "where:\n", + "\n", + "- theta is the image-plane position.\n", + "- beta is the source-plane position (= theta - alpha(theta)).\n", + "- psi(theta) is the lensing potential (the scalar potential whose gradient gives the deflection angles).\n", + "\n", + "The first term, 0.5 * |theta - beta|^2, is the **geometric delay** \u2014 the extra path length due to the\n", + "bending of light. The second term, psi(theta), is the **gravitational (Shapiro) delay** \u2014 the slowing of\n", + "light in the gravitational potential of the lens.\n", + "\n", + "Fermat's principle tells us that observed images form at the stationary points (extrema and saddle points) of\n", + "this surface. This is a powerful result: it means the positions of lensed images are determined by the topology\n", + "of the Fermat potential.\n", + "\n", + "The *differences* in the Fermat potential between image positions are proportional to the time delays between\n", + "images. Measuring these time delays (e.g. from a variable quasar) can constrain the Hubble constant." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fermat_potential = lens_calc.fermat_potential_from(grid=grid)\n", + "\n", + "fermat_array = aa.Array2D(values=fermat_potential, mask=grid.mask)\n", + "aplt.plot_array(array=fermat_array, title=\"Fermat Potential\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The geometric delay term alone can also be inspected:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "geometric_delay = lens_calc.time_delay_geometry_term_from(grid=grid)\n", + "\n", + "geometric_array = aa.Array2D(values=geometric_delay, mask=grid.mask)\n", + "aplt.plot_array(array=geometric_array, title=\"Geometric Time Delay Term\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__JAX (JIT-it-yourself)__\n", + "\n", + "This is the canonical home for the \"JIT-it-yourself\" advanced path.\n", + "Audience: users building custom forward models or scientific tools\n", + "using PyAutoLens primitives directly \u2014 not running standard fits\n", + "(which go through `Analysis(use_jax=True)` and need zero JAX code from\n", + "you) or standard simulations (which use `Simulator(use_jax=True)`\n", + "similarly).\n", + "\n", + "If you're writing `@jax.jit` yourself around library calls like\n", + "`tracer.image_2d_from`, `LensCalc.magnification_2d_via_hessian_from`,\n", + "or your own `log_likelihood(instance)` function, this section is for you.\n", + "\n", + "__The pairing rule: `@jax.jit` + `xp=jnp`__\n", + "\n", + "The single rule to remember: when you decorate a function with\n", + "`@jax.jit` that calls a PyAutoLens library method internally, **pass\n", + "`xp=jnp` to that method inside the function body**.\n", + "\n", + "```python\n", + "import jax\n", + "import jax.numpy as jnp\n", + "\n", + "@jax.jit\n", + "def magnification_fn(lens_calc, grid):\n", + " return lens_calc.magnification_2d_via_hessian_from(grid=grid, xp=jnp)\n", + "```\n", + "\n", + "The `xp=jnp` is what tells the library \"you're inside a JAX trace \u2014\n", + "route calls through `jax.numpy` and don't wrap the return in an\n", + "autoarray type (it would fail to cross the JIT boundary)\".\n", + "\n", + "__The footgun: forgetting `xp=jnp`__\n", + "\n", + "The default for `xp` in library method signatures is `np`. If you\n", + "forget to pass `xp=jnp` inside `@jax.jit`, one of two things happens:\n", + "\n", + "- The function body does `np.sqrt(jax_array)` \u2014 NumPy routes through\n", + " `__array__()` on the JAX tracer, host-transferring off the GPU.\n", + " Your fit runs, invisibly slower than the NumPy path.\n", + "- The `if xp is np:` guard inside library functions fires and wraps\n", + " the result in `aa.Array2D`, which fails at the JIT boundary with\n", + " `TypeError: ... is not a valid JAX type`.\n", + "\n", + "The library raises a clear `ValueError` on the easy mismatch \u2014 when\n", + "you pass `xp=np` with a `grid.use_jax=True` input:\n", + "\n", + "```\n", + "ValueError: Called magnification_2d_via_hessian_from with xp=np but\n", + "the input grid is JAX-backed (grid.use_jax=True). Inside @jax.jit,\n", + "pass xp=jnp explicitly.\n", + "```\n", + "\n", + "If you see this error: add `xp=jnp` to the call site. Done.\n", + "\n", + "__Decorator-on-def vs `jax.jit(bound_method)`__\n", + "\n", + "JAX accepts any callable \u2014 `@jax.jit` is sugar for `fn = jax.jit(fn)`.\n", + "You can JIT a standalone function (canonical) or a bound method\n", + "(shortcut). Both work:\n", + "\n", + "```python\n", + "# Form 1 (canonical): decorator on def\n", + "@jax.jit\n", + "def image_fn(tracer, grid):\n", + " return tracer.image_2d_from(grid=grid, xp=jnp).array\n", + "\n", + "# Form 2: jit on bound method, assign-to-variable\n", + "jitted = jax.jit(tracer.image_2d_from)\n", + "arr = jitted(grid=grid, xp=jnp).array\n", + "```\n", + "\n", + "Form 2 is shorter for interactive use. **Footgun:** bound methods are\n", + "fresh objects on every attribute access, so this silently misses the\n", + "JIT cache every iteration:\n", + "\n", + "```python\n", + "# DON'T DO THIS \u2014 fresh jax.jit closure every iteration\n", + "for grid in many_grids:\n", + " arr = jax.jit(tracer.image_2d_from)(grid=grid, xp=jnp).array\n", + "```\n", + "\n", + "If you're calling a JITted method in a loop: assign once outside the\n", + "loop, or use the decorator-on-def form.\n", + "\n", + "__Closure-captured `self` vs traced argument__\n", + "\n", + "Form 1 and Form 2 differ in *semantics*, not just syntax:\n", + "\n", + "- **Form 2 (`jax.jit(tracer.image_2d_from)`):** `tracer` is the bound\n", + " method's `self`; JAX captures it as a closure constant and doesn't\n", + " trace through it. **`Tracer` does NOT need to be pytree-registered.**\n", + " Trade-off: a different tracer means a fresh bound-method object and a\n", + " fresh JIT cache key.\n", + "- **Form 1 (`@jax.jit def image_fn(tracer, grid)`):** `tracer` is a\n", + " traced argument. **`Tracer` DOES need pytree registration** (which\n", + " `Analysis(use_jax=True)` does for you, or you can trigger via\n", + " `autolens.jax.register_tracer_classes(tracer)`). Trade-off: cache\n", + " reuse across different tracers \u2014 parameter sweeps and `jax.vmap`\n", + " work naturally.\n", + "\n", + "Pick based on whether you want to vary the tracer across calls:\n", + "\n", + "- Parameter sweep / `jax.vmap` over models? Form 1.\n", + "- Quick one-off / interactive exploration? Form 2 (assign once).\n", + "\n", + "__`LensCalc` and the wrapped-vs-raw return type__\n", + "\n", + "`LensCalc.magnification_2d_via_hessian_from`,\n", + "`.shear_yx_2d_via_hessian_from`, `.convergence_2d_via_hessian_from`,\n", + "and the eigen-value methods all implement the `if xp is np:` guard\n", + "inside:\n", + "\n", + "```python\n", + "def magnification_2d_via_hessian_from(self, grid, xp=np):\n", + " ...\n", + " if xp is np:\n", + " return aa.Array2D(values=mag, mask=grid.mask) # numpy: wrapped\n", + " return mag # jax: raw jax.Array\n", + "```\n", + "\n", + "This is intentional. Inside `@jax.jit` (where you pass `xp=jnp`), you\n", + "get back a raw `jax.Array`. On the NumPy path (no `xp` or `xp=np`),\n", + "you get an `aa.Array2D` wrapper. The function adapts to which path\n", + "you're on.\n", + "\n", + "Implication: when you JIT-wrap a `LensCalc` method, expect a raw\n", + "`jax.Array` back. Rewrap with `aa.Array2D(values=..., mask=...)` on\n", + "the host if you want the wrapper for downstream plotting:\n", + "\n", + "```python\n", + "@jax.jit\n", + "def magnification_fn(lens_calc, grid):\n", + " return lens_calc.magnification_2d_via_hessian_from(grid=grid, xp=jnp)\n", + "\n", + "mag_raw = magnification_fn(lens_calc, grid)\n", + "mag_wrapped = aa.Array2D(values=mag_raw, mask=grid.mask)\n", + "aplt.plot_array(array=mag_wrapped)\n", + "```\n", + "\n", + "For `Tracer` and `Galaxy` methods that don't have the guard internally\n", + "(e.g. `tracer.image_2d_from`), the `.array` unwrap inside the jit +\n", + "rewrap outside discipline applies \u2014 see `scripts/guides/data_structures.py`\n", + "`__JAX__` section.\n", + "\n", + "__Summary \u2014 the three rules__\n", + "\n", + "The \"JIT-it-yourself\" path is bounded by three rules:\n", + "\n", + "1. **`@jax.jit` and `xp=jnp` are paired.** Forgetting `xp=jnp` either\n", + " silently host-transfers or fails at the boundary.\n", + "2. **`.array` unwrap inside the jit; rewrap on the host** if you want\n", + " the autoarray wrapper.\n", + "3. **Tracer-as-argument needs pytree registration** (via\n", + " `register_tracer_classes(tracer)` or any `Analysis(use_jax=True)`\n", + " construction); tracer-as-closure (bound-method form) doesn't.\n", + "\n", + "For the standard `Analysis` / `Simulator` paths \u2014 where you do none\n", + "of this and JAX runs implicitly \u2014 see the top-level\n", + "`autolens_workspace/start_here.py` `__JAX__` section." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Wrap Up__\n", + "\n", + "This guide introduced the `LensCalc` class and the key lensing quantities it computes:\n", + "\n", + "1. **Deflection angles** \u2014 the bending of light by the lens's gravity.\n", + "2. **Hessian** \u2014 the matrix of second derivatives, from which all other quantities are derived.\n", + "3. **Convergence** \u2014 the projected mass density (isotropic magnification).\n", + "4. **Shear** \u2014 the tidal distortion (anisotropic stretching).\n", + "5. **Magnification** \u2014 the brightness and size change of lensed images.\n", + "6. **Critical curves and caustics** \u2014 where magnification diverges, and the corresponding source-plane boundaries.\n", + "7. **Einstein radius** \u2014 the characteristic angular scale of the lens.\n", + "8. **Fermat potential** \u2014 the time-delay surface whose stationary points give image positions.\n", + "\n", + "For further reading:\n", + "\n", + "- ``guides/tracer.py`` \u2014 ray-tracing and image formation.\n", + "- ``guides/galaxies.py`` \u2014 working with individual galaxy components.\n", + "- ``guides/units/cosmology.ipynb`` \u2014 converting to physical units.\n", + "- ``guides/data_structures.py`` \u2014 the ``Array2D``, ``Grid2D`` and ``VectorYX2D`` data structures.\n", + "\n", + "The `LensCalc` class is also available via `al.LensCalc` \u2014 it lives in ``PyAutoGalaxy/autogalaxy/operate/lens_calc.py``.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/advanced/expectation_propagation.ipynb b/notebooks/guides/modeling/advanced/expectation_propagation.ipynb index 577c39a07..4abbe5bdb 100644 --- a/notebooks/guides/modeling/advanced/expectation_propagation.ipynb +++ b/notebooks/guides/modeling/advanced/expectation_propagation.ipynb @@ -1,575 +1,612 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Expectation Propagation\n", - "=================================\n", - "\n", - "In the `hierarchical` example, we fitted graphical models to a dataset comprising 3 images of strong lenses, which had a\n", - "hierarchical parameter, the power-law `slope`. This provides the basis of composing and fitting complex graphical\n", - "models to large datasets.\n", - "\n", - "The challenge is that we will soon hit a ceiling scaling these graphical models up to extremely large datasets.\n", - "One would soon find that the parameter space is too complex to sample, and computational limits would ultimately\n", - "cap how many datasets one could feasible fit.\n", - "\n", - "This example introduces expectation propagation (EP), the solution to this problem, which inspects a factor graph\n", - "and partitions the model-fit into many simpler fits of sub-components of the graph to individual datasets. This\n", - "overcomes the challenge of model complexity, and mitigates computational restrictions that may occur if one tries to\n", - "fit every dataset simultaneously.\n", - "\n", - "__Contents__\n", - "\n", - "- **Sample Simulation:** The dataset fitted in this example script is simulated imaging data of a sample of 3 galaxies.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Model Individual Factors:** We first set up a model for each lens, with an `PowerLawSph` mass and `ExponentialSph` bulge, which.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Paths:** Overview of paths for this example.\n", - "- **Analysis Factors:** Now we have our `Analysis` classes and graphical model, we can compose our `AnalysisFactor`'s.\n", - "- **Factor Graph:** We combine our `AnalysisFactors` into one, to compose the factor graph.\n", - "- **Expectation Propagation:** In the previous tutorials, we used the `global_prior_model` of the `factor_graph` to fit the global.\n", - "- **Cyclic Fitting:** After every `AnalysisFactor` has been fitted (e.g.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Output:** The results of the factor graph, using the EP framework and message passing, are contained in the.\n", - "- **Results:** The `MeanField` object represent the posterior of the entire factor graph and is used to infer.\n", - "\n", - "__Sample Simulation__\n", - "\n", - "The dataset fitted in this example script is simulated imaging data of a sample of 3 galaxies.\n", - "\n", - "This data is not automatically provided with the autogalaxy workspace, and must be first simulated by running the\n", - "script `autolens_workspace/scripts/advanced/graphical/simulator/samples/simple__no_lens_light.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import autolens as al\n", - "import autofit as af\n", - "\n", - "import numpy as np\n", - "from pathlib import Path" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "The following steps repeat all the initial steps performed in tutorial 2 and 3:\n", - "\n", - "This data is not automatically provided with the autogalaxy workspace, and must be first simulated by running the \n", - "script `autolens_workspace/scripts/advanced/graphical/simulator/samples/simple__no_lens_light.py`. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"samples\"\n", - "dataset_type = \"imaging\"\n", - "dataset_sample_name = \"simple\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_label, dataset_sample_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator_sample.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "total_datasets = 3\n", - "\n", - "dataset_list = []\n", - "\n", - "for dataset_index in range(total_datasets):\n", - " dataset_sample_path = Path(dataset_path, f\"dataset_{dataset_index}\")\n", - "\n", - " dataset_list.append(\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_sample_path, \"data.fits\"),\n", - " psf_path=Path(dataset_sample_path, \"psf.fits\"),\n", - " noise_map_path=Path(dataset_sample_path, \"noise_map.fits\"),\n", - " pixel_scales=0.1,\n", - " )\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_dataset_list = []\n", - "\n", - "for dataset in dataset_list:\n", - " mask_radius = 3.0\n", - "\n", - " mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=mask)\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " masked_dataset_list.append(dataset)\n", - "\n", - "dataset_list = masked_dataset_list" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Individual Factors__\n", - "\n", - "We first set up a model for each lens, with an `PowerLawSph` mass and `ExponentialSph` bulge, which we will use to \n", - "fit the hierarchical model.\n", - "\n", - "Note that the `PowerLawSph` mass model has a `slope` parameter, which we will assume is drawn from a shared parent\n", - "Gaussian distribution, albeit building this into the model is done later in this script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_list = []\n", - "\n", - "for dataset_index in range(total_datasets):\n", - "\n", - " lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.PowerLawSph)\n", - " lens.mass.centre = (0.0, 0.0)\n", - "\n", - " source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.ExponentialCoreSph)\n", - "\n", - " model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - " model_list.append(model)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "For each dataset we now create a corresponding `AnalysisImaging` class, as we are used to doing for `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = []\n", - "\n", - "for dataset in dataset_list:\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - "\n", - " analysis_list.append(analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We now compose the hierarchical model that we fit, using the individual model components created above.\n", - "\n", - "This uses the same API as the `hierarchical` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "hierarchical_factor = af.HierarchicalFactor(\n", - " af.GaussianPrior,\n", - " mean=af.TruncatedGaussianPrior(\n", - " mean=2.0, sigma=1.0, lower_limit=0.0, upper_limit=100.0\n", - " ),\n", - " sigma=af.TruncatedGaussianPrior(\n", - " mean=0.5, sigma=0.5, lower_limit=0.0, upper_limit=100.0\n", - " ),\n", - " use_jax=False,\n", - ")\n", - "\n", - "for model in model_list:\n", - " hierarchical_factor.add_drawn_variable(model.galaxies.lens.mass.slope)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"imaging\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis Factors__\n", - "\n", - "Now we have our `Analysis` classes and graphical model, we can compose our `AnalysisFactor`'s.\n", - "\n", - "However, unlike the previous tutorials, each `AnalysisFactor` is now assigned its own `search`. This is because the EP \n", - "framework performs a model-fit to each node on the factor graph (e.g. each `AnalysisFactor`). Therefore, each node \n", - "requires its own non-linear search, and in this tutorial we use `dynesty`. For complex graphs consisting of many \n", - "nodes, one could easily use different searches for different nodes on the factor graph.\n", - "\n", - "Each `AnalysisFactor` is also given a `name`, corresponding to the name of the dataset it fits. These names are used\n", - "to name the folders containing the results in the output directory." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "paths = af.DirectoryPaths(\n", - " path_prefix=path_prefix,\n", - " name=\"expectation_propagation\",\n", - ")\n", - "\n", - "search = af.Nautilus(paths=paths, n_live=100)\n", - "\n", - "analysis_factor_list = []\n", - "\n", - "dataset_index = 0\n", - "\n", - "for model, analysis in zip(model_list, analysis_list):\n", - "\n", - " dataset_name = f\"dataset_{dataset_index}\"\n", - " dataset_index += 1\n", - "\n", - " analysis_factor = af.AnalysisFactor(\n", - " prior_model=model, analysis=analysis, optimiser=search, name=dataset_name\n", - " )\n", - "\n", - " analysis_factor_list.append(analysis_factor)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Factor Graph__\n", - "\n", - "We combine our `AnalysisFactors` into one, to compose the factor graph." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "factor_graph = af.FactorGraphModel(\n", - " *analysis_factor_list, hierarchical_factor, use_jax=False\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The factor graph model `info` attribute shows the model which we fit via expectaton propagation (note that we do\n", - "not use `global_prior_model` below when performing the fit)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Expectation Propagation__\n", - "\n", - "In the previous tutorials, we used the `global_prior_model` of the `factor_graph` to fit the global model. In this \n", - "tutorial, we instead fit the `factor_graph` using the EP framework, which fits the graphical model composed in this \n", - "tutorial as follows:\n", - "\n", - "1) Go to the first node on the factor graph (e.g. `analysis_factor_list[0]`) and fit its model to its dataset. This is \n", - "simply a fit of the `Gaussian` model to the first 1D Gaussian dataset, the model-fit we are used to performing by now.\n", - "\n", - "2) Once the model-fit is complete, inspect the model for parameters that are shared with other nodes on the factor\n", - "graph. In this example, the `centre` of the `Gaussian` fitted to the first dataset is global, and therefore connects\n", - "to the other nodes on the factor graph (the `AnalysisFactor`'s) of the second and first `Gaussian` datasets.\n", - "\n", - "3) The EP framework now creates a 'message' that is to be passed to the connecting nodes on the factor graph. This\n", - "message informs them of the results of the model-fit, so they can update their priors on the `Gaussian`'s centre \n", - "accordingly and, more importantly, update their posterior inference and therefore estimate of the global centre.\n", - "\n", - "For example, the model fitted to the first Gaussian dataset includes the global centre. Therefore, after the model is \n", - "fitted, the EP framework creates a 'message' informing the factor graph about its inference on that Gaussians's centre,\n", - "thereby updating our overall inference on this shared parameter. This is termed 'message passing'.\n", - "\n", - "__Cyclic Fitting__\n", - "\n", - "After every `AnalysisFactor` has been fitted (e.g. after each fit to each of the 5 datasets in this example), we have a \n", - "new estimate of the shared parameter `centre`. This updates our priors on the shared parameter `centre`, which needs \n", - "to be reflected in each model-fit we perform on each `AnalysisFactor`. \n", - "\n", - "The EP framework therefore performs a second iteration of model-fits. It again cycles through each `AnalysisFactor` \n", - "and refits the model, using updated priors on shared parameters like the `centre`. At the end of each fit, we again \n", - "create messages that update our knowledge about other parameters on the graph.\n", - "\n", - "This process is repeated multiple times, until a convergence criteria is met whereby continued cycles are expected to\n", - "produce the same estimate of the shared parameter `centre`. \n", - "\n", - "When we fit the factor graph a `name` is passed, which determines the folder all results of the factor graph are\n", - "stored in." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "laplace = af.LaplaceOptimiser()\n", - "\n", - "factor_graph_result = factor_graph.optimise(\n", - " optimiser=laplace, paths=paths, ep_history=af.EPHistory(kl_tol=0.05), max_steps=5\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "An `info` attribute for the result of a factor graph fitted via EP does not exist yet, its on the to do list!\n", - "\n", - "The result can be seen in the `graph.result` file output to hard-disk." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "### print(factor_graph_result.info)##" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "The results of the factor graph, using the EP framework and message passing, are contained in the folder \n", - "`output/howtofit/chapter_graphical_models/tutorial_5_expectation_propagation`. \n", - "\n", - "The following folders and files are worth of note:\n", - "\n", - " - `graph.info`: this provides an overall summary of the graphical model that is fitted, including every parameter, \n", - " how parameters are shared across `AnalysisFactor`'s and the priors associated to each individual parameter.\n", - "\n", - " - The 3 folders titled `gaussian_x1_#__low_snr` correspond to the three `AnalysisFactor`'s and therefore signify \n", - " repeated non-linear searches that are performed to fit each dataset.\n", - "\n", - " - Inside each of these folders are `optimization_#` folders, corresponding to each model-fit performed over cycles of\n", - " the EP fit. A careful inspection of the `model.info` files inside each folder reveals how the priors are updated\n", - " over each cycle, whereas the `model.results` file should indicate the improved estimate of model parameters over each\n", - " cycle.\n", - "\n", - "__Results__\n", - "\n", - "The `MeanField` object represent the posterior of the entire factor graph and is used to infer estimates of the \n", - "values and error of each parameter in the graph." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mean_field = factor_graph_result.updated_ep_mean_field.mean_field\n", - "print(mean_field)\n", - "print()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The object has a `variables` property which lists every variable in the factor graph, which is essentially all of the \n", - "free parameters on the graph.\n", - "\n", - "This includes the parameters specific to each data (E.g. each node on the graph) as well as the shared centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(mean_field.variables)\n", - "print()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The variables above use the priors on each parameter as their key. \n", - "\n", - "Therefore to estimate mean-field quantities of the shared centre, we can simply use the `centre_shared_prior` defined\n", - "above.\n", - "\n", - "Each parameter estimate is given by the mean of its value in the `MeanField`. Below, we use the `centred_shared_prior` \n", - "as a key to the `MeanField.mean` dictionary to print the estimated value of the shared centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# prior = hierarchical_factor.drawn_variables[0]\n", - "#\n", - "# print(f\"Centre Mean Parameter Estimate = {mean_field.mean[prior]}\")\n", - "# print()\n", - "#\n", - "# \"\"\"\n", - "# If we want the parameter estimate of another parameter in the model, we can use the `model_list` that we composed\n", - "# above to pass a parameter prior to the mean field dictionary.\n", - "# \"\"\"\n", - "# print(\n", - "# f\"Einstein Radius Dataset 0 Mean = {mean_field.mean[model_list[0].galaxies.lens.mass.einstein_radius]}\"\n", - "# )\n", - "#\n", - "# \"\"\"\n", - "# The mean-field mean dictionary contains the estimate value of every parameter.\n", - "# \"\"\"\n", - "# print(f\"All Parameter Estimates = {mean_field.mean}\")\n", - "# print()\n", - "#\n", - "# \"\"\"\n", - "# The mean-field also contains a `variance` dictionary, which has the same keys as the `mean` dictionary above.\n", - "#\n", - "# This is the easier way to estimate the error on every parameter, for example that of the shared centre.\n", - "# \"\"\"\n", - "# print(f\"Centre Variance = {mean_field.variance[prior]}\")\n", - "# print()\n", - "#\n", - "# \"\"\"\n", - "# The standard deviation (or error at one sigma confidence interval) is given by the square root of the variance.\n", - "# \"\"\"\n", - "# print(f\"Centre 1 Sigma = {np.sqrt(mean_field.variance[prior])}\")\n", - "# print()\n", - "#\n", - "# \"\"\"\n", - "# The mean field object also contains a dictionary of the s.d./variance**0.5.\n", - "# \"\"\"\n", - "# print(f\"Centre SD/sqrt(variance) = {mean_field.scale[prior]}\")\n", - "# print()\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Expectation Propagation\n", + "=================================\n", + "\n", + "In the `hierarchical` example, we fitted graphical models to a dataset comprising 3 images of strong lenses, which had a\n", + "hierarchical parameter, the power-law `slope`. This provides the basis of composing and fitting complex graphical\n", + "models to large datasets.\n", + "\n", + "The challenge is that we will soon hit a ceiling scaling these graphical models up to extremely large datasets.\n", + "One would soon find that the parameter space is too complex to sample, and computational limits would ultimately\n", + "cap how many datasets one could feasible fit.\n", + "\n", + "This example introduces expectation propagation (EP), the solution to this problem, which inspects a factor graph\n", + "and partitions the model-fit into many simpler fits of sub-components of the graph to individual datasets. This\n", + "overcomes the challenge of model complexity, and mitigates computational restrictions that may occur if one tries to\n", + "fit every dataset simultaneously.\n", + "\n", + "__Contents__\n", + "\n", + "- **Sample Simulation:** The dataset fitted in this example script is simulated imaging data of a sample of 3 galaxies.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Model Individual Factors:** We first set up a model for each lens, with an `PowerLawSph` mass and `ExponentialSph` bulge, which.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Paths:** Overview of paths for this example.\n", + "- **Analysis Factors:** Now we have our `Analysis` classes and graphical model, we can compose our `AnalysisFactor`'s.\n", + "- **Factor Graph:** We combine our `AnalysisFactors` into one, to compose the factor graph.\n", + "- **Expectation Propagation:** In the previous tutorials, we used the `global_prior_model` of the `factor_graph` to fit the global.\n", + "- **Cyclic Fitting:** After every `AnalysisFactor` has been fitted (e.g.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Output:** The results of the factor graph, using the EP framework and message passing, are contained in the.\n", + "- **Results:** The `MeanField` object represent the posterior of the entire factor graph and is used to infer.\n", + "\n", + "__Sample Simulation__\n", + "\n", + "The dataset fitted in this example script is simulated imaging data of a sample of 3 galaxies.\n", + "\n", + "This data is not automatically provided with the autogalaxy workspace, and must be first simulated by running the\n", + "script `autolens_workspace/scripts/advanced/graphical/simulator/samples/simple__no_lens_light.py`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import autolens as al\n", + "import autofit as af\n", + "\n", + "import numpy as np\n", + "from pathlib import Path" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "The following steps repeat all the initial steps performed in tutorial 2 and 3:\n", + "\n", + "This data is not automatically provided with the autogalaxy workspace, and must be first simulated by running the \n", + "script `autolens_workspace/scripts/advanced/graphical/simulator/samples/simple__no_lens_light.py`. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"samples\"\n", + "dataset_type = \"imaging\"\n", + "dataset_sample_name = \"simple\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_label, dataset_sample_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator_sample.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "total_datasets = 3\n", + "\n", + "dataset_list = []\n", + "\n", + "for dataset_index in range(total_datasets):\n", + " dataset_sample_path = Path(dataset_path, f\"dataset_{dataset_index}\")\n", + "\n", + " dataset_list.append(\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_sample_path, \"data.fits\"),\n", + " psf_path=Path(dataset_sample_path, \"psf.fits\"),\n", + " noise_map_path=Path(dataset_sample_path, \"noise_map.fits\"),\n", + " pixel_scales=0.1,\n", + " )\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_dataset_list = []\n", + "\n", + "for dataset in dataset_list:\n", + " mask_radius = 3.0\n", + "\n", + " mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=mask)\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " masked_dataset_list.append(dataset)\n", + "\n", + "dataset_list = masked_dataset_list" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Individual Factors__\n", + "\n", + "We first set up a model for each lens, with an `PowerLawSph` mass and `ExponentialSph` bulge, which we will use to \n", + "fit the hierarchical model.\n", + "\n", + "Note that the `PowerLawSph` mass model has a `slope` parameter, which we will assume is drawn from a shared parent\n", + "Gaussian distribution, albeit building this into the model is done later in this script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_list = []\n", + "\n", + "for dataset_index in range(total_datasets):\n", + "\n", + " lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.PowerLawSph)\n", + " lens.mass.centre = (0.0, 0.0)\n", + "\n", + " source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.ExponentialCoreSph)\n", + "\n", + " model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + " model_list.append(model)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "For each dataset we now create a corresponding `AnalysisImaging` class, as we are used to doing for `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = []\n", + "\n", + "for dataset in dataset_list:\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + "\n", + " analysis_list.append(analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We now compose the hierarchical model that we fit, using the individual model components created above.\n", + "\n", + "This uses the same API as the `hierarchical` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "hierarchical_factor = af.HierarchicalFactor(\n", + " af.GaussianPrior,\n", + " mean=af.TruncatedGaussianPrior(\n", + " mean=2.0, sigma=1.0, lower_limit=0.0, upper_limit=100.0\n", + " ),\n", + " sigma=af.TruncatedGaussianPrior(\n", + " mean=0.5, sigma=0.5, lower_limit=0.0, upper_limit=100.0\n", + " ),\n", + " use_jax=False,\n", + ")\n", + "\n", + "for model in model_list:\n", + " hierarchical_factor.add_drawn_variable(model.galaxies.lens.mass.slope)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"imaging\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis Factors__\n", + "\n", + "Now we have our `Analysis` classes and graphical model, we can compose our `AnalysisFactor`'s.\n", + "\n", + "However, unlike the previous tutorials, each `AnalysisFactor` is now assigned its own `search`. This is because the EP \n", + "framework performs a model-fit to each node on the factor graph (e.g. each `AnalysisFactor`). Therefore, each node \n", + "requires its own non-linear search, and in this tutorial we use `dynesty`. For complex graphs consisting of many \n", + "nodes, one could easily use different searches for different nodes on the factor graph.\n", + "\n", + "Each `AnalysisFactor` is also given a `name`, corresponding to the name of the dataset it fits. These names are used\n", + "to name the folders containing the results in the output directory." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "paths = af.DirectoryPaths(\n", + " path_prefix=path_prefix,\n", + " name=\"expectation_propagation\",\n", + ")\n", + "\n", + "search = af.Nautilus(paths=paths, n_live=100)\n", + "\n", + "analysis_factor_list = []\n", + "\n", + "dataset_index = 0\n", + "\n", + "for model, analysis in zip(model_list, analysis_list):\n", + "\n", + " dataset_name = f\"dataset_{dataset_index}\"\n", + " dataset_index += 1\n", + "\n", + " analysis_factor = af.AnalysisFactor(\n", + " prior_model=model, analysis=analysis, optimiser=search, name=dataset_name\n", + " )\n", + "\n", + " analysis_factor_list.append(analysis_factor)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Factor Graph__\n", + "\n", + "We combine our `AnalysisFactors` into one, to compose the factor graph." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph = af.FactorGraphModel(\n", + " *analysis_factor_list, hierarchical_factor, use_jax=False\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The factor graph model `info` attribute shows the model which we fit via expectaton propagation (note that we do\n", + "not use `global_prior_model` below when performing the fit)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Expectation Propagation__\n", + "\n", + "In the previous tutorials, we used the `global_prior_model` of the `factor_graph` to fit the global model. In this \n", + "tutorial, we instead fit the `factor_graph` using the EP framework, which fits the graphical model composed in this \n", + "tutorial as follows:\n", + "\n", + "1) Go to the first node on the factor graph (e.g. `analysis_factor_list[0]`) and fit its model to its dataset. This is \n", + "simply a fit of the `Gaussian` model to the first 1D Gaussian dataset, the model-fit we are used to performing by now.\n", + "\n", + "2) Once the model-fit is complete, inspect the model for parameters that are shared with other nodes on the factor\n", + "graph. In this example, the `centre` of the `Gaussian` fitted to the first dataset is global, and therefore connects\n", + "to the other nodes on the factor graph (the `AnalysisFactor`'s) of the second and first `Gaussian` datasets.\n", + "\n", + "3) The EP framework now creates a 'message' that is to be passed to the connecting nodes on the factor graph. This\n", + "message informs them of the results of the model-fit, so they can update their priors on the `Gaussian`'s centre \n", + "accordingly and, more importantly, update their posterior inference and therefore estimate of the global centre.\n", + "\n", + "For example, the model fitted to the first Gaussian dataset includes the global centre. Therefore, after the model is \n", + "fitted, the EP framework creates a 'message' informing the factor graph about its inference on that Gaussians's centre,\n", + "thereby updating our overall inference on this shared parameter. This is termed 'message passing'.\n", + "\n", + "__Cyclic Fitting__\n", + "\n", + "After every `AnalysisFactor` has been fitted (e.g. after each fit to each of the 5 datasets in this example), we have a \n", + "new estimate of the shared parameter `centre`. This updates our priors on the shared parameter `centre`, which needs \n", + "to be reflected in each model-fit we perform on each `AnalysisFactor`. \n", + "\n", + "The EP framework therefore performs a second iteration of model-fits. It again cycles through each `AnalysisFactor` \n", + "and refits the model, using updated priors on shared parameters like the `centre`. At the end of each fit, we again \n", + "create messages that update our knowledge about other parameters on the graph.\n", + "\n", + "This process is repeated multiple times, until a convergence criteria is met whereby continued cycles are expected to\n", + "produce the same estimate of the shared parameter `centre`. \n", + "\n", + "When we fit the factor graph a `name` is passed, which determines the folder all results of the factor graph are\n", + "stored in." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "laplace = af.LaplaceOptimiser()\n", + "\n", + "factor_graph_result = factor_graph.optimise(\n", + " optimiser=laplace, paths=paths, ep_history=af.EPHistory(kl_tol=0.05), max_steps=5\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "An `info` attribute for the result of a factor graph fitted via EP does not exist yet, its on the to do list!\n", + "\n", + "The result can be seen in the `graph.result` file output to hard-disk." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "### print(factor_graph_result.info)##" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "The results of the factor graph, using the EP framework and message passing, are contained in the folder \n", + "`output/howtofit/chapter_graphical_models/tutorial_5_expectation_propagation`. \n", + "\n", + "The following folders and files are worth of note:\n", + "\n", + " - `graph.info`: this provides an overall summary of the graphical model that is fitted, including every parameter, \n", + " how parameters are shared across `AnalysisFactor`'s and the priors associated to each individual parameter.\n", + "\n", + " - The 3 folders titled `gaussian_x1_#__low_snr` correspond to the three `AnalysisFactor`'s and therefore signify \n", + " repeated non-linear searches that are performed to fit each dataset.\n", + "\n", + " - Inside each of these folders are `optimization_#` folders, corresponding to each model-fit performed over cycles of\n", + " the EP fit. A careful inspection of the `model.info` files inside each folder reveals how the priors are updated\n", + " over each cycle, whereas the `model.results` file should indicate the improved estimate of model parameters over each\n", + " cycle.\n", + "\n", + "__Results__\n", + "\n", + "The `MeanField` object represent the posterior of the entire factor graph and is used to infer estimates of the \n", + "values and error of each parameter in the graph." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mean_field = factor_graph_result.updated_ep_mean_field.mean_field\n", + "print(mean_field)\n", + "print()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The object has a `variables` property which lists every variable in the factor graph, which is essentially all of the \n", + "free parameters on the graph.\n", + "\n", + "This includes the parameters specific to each data (E.g. each node on the graph) as well as the shared centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(mean_field.variables)\n", + "print()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The variables above use the priors on each parameter as their key. \n", + "\n", + "Therefore to estimate mean-field quantities of the shared centre, we can simply use the `centre_shared_prior` defined\n", + "above.\n", + "\n", + "Each parameter estimate is given by the mean of its value in the `MeanField`. Below, we use the `centred_shared_prior` \n", + "as a key to the `MeanField.mean` dictionary to print the estimated value of the shared centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# prior = hierarchical_factor.drawn_variables[0]\n", + "#\n", + "# print(f\"Centre Mean Parameter Estimate = {mean_field.mean[prior]}\")\n", + "# print()\n", + "#\n", + "# \"\"\"\n", + "# If we want the parameter estimate of another parameter in the model, we can use the `model_list` that we composed\n", + "# above to pass a parameter prior to the mean field dictionary.\n", + "# \"\"\"\n", + "# print(\n", + "# f\"Einstein Radius Dataset 0 Mean = {mean_field.mean[model_list[0].galaxies.lens.mass.einstein_radius]}\"\n", + "# )\n", + "#\n", + "# \"\"\"\n", + "# The mean-field mean dictionary contains the estimate value of every parameter.\n", + "# \"\"\"\n", + "# print(f\"All Parameter Estimates = {mean_field.mean}\")\n", + "# print()\n", + "#\n", + "# \"\"\"\n", + "# The mean-field also contains a `variance` dictionary, which has the same keys as the `mean` dictionary above.\n", + "#\n", + "# This is the easier way to estimate the error on every parameter, for example that of the shared centre.\n", + "# \"\"\"\n", + "# print(f\"Centre Variance = {mean_field.variance[prior]}\")\n", + "# print()\n", + "#\n", + "# \"\"\"\n", + "# The standard deviation (or error at one sigma confidence interval) is given by the square root of the variance.\n", + "# \"\"\"\n", + "# print(f\"Centre 1 Sigma = {np.sqrt(mean_field.variance[prior])}\")\n", + "# print()\n", + "#\n", + "# \"\"\"\n", + "# The mean field object also contains a dictionary of the s.d./variance**0.5.\n", + "# \"\"\"\n", + "# print(f\"Centre SD/sqrt(variance) = {mean_field.scale[prior]}\")\n", + "# print()\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/advanced/graphical.ipynb b/notebooks/guides/modeling/advanced/graphical.ipynb index eb77fcc4d..d48ff5109 100644 --- a/notebooks/guides/modeling/advanced/graphical.ipynb +++ b/notebooks/guides/modeling/advanced/graphical.ipynb @@ -1,431 +1,468 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Graphical\n", - "===================\n", - "\n", - "The example scripts throughout the workspace have focused on modeling **individual strong lens datasets**.\n", - "From each fit, you may have inspected lens properties (e.g., Einstein radius) and source properties\n", - "(e.g., magnification). You may even have analyzed many lenses one-by-one and combined their results to study\n", - "global trends in galaxy formation or cosmology.\n", - "\n", - "However, fitting each lens independently does **not** make full use of the information contained in a large\n", - "sample. Many properties are expected to be **shared across lenses** (e.g., population-level mass slopes,\n", - "cosmological parameters), and treating them independently ignores this shared structure.\n", - "\n", - "In this example, we demonstrate how to fit **multiple lenses simultaneously** using a **graphical model**.\n", - "A graphical model links parameters across separate lens fits, explicitly defining which parameters are unique\n", - "to each dataset and which are shared. These links can be arbitrarily complex, enabling joint analysis across\n", - "diverse datasets with structured relationships between model components.\n", - "\n", - "Here, we illustrate a cosmological application: inferring the **Hubble constant (H0)** from time-delay lenses.\n", - "A graphical model links the mass models across multiple lenses and includes a **shared H0 parameter**, allowing\n", - "a joint inference that improves cosmological constraints compared to individual fits.\n", - "\n", - "Graphical models form the foundation of **hierarchical modeling**, where the parameters of individual lenses are\n", - "assumed to be drawn from a parent distribution (see `guides/modeling/hierarchical`). Hierarchical approaches can\n", - "extract significantly more information from large samples than fitting each dataset independently. The example\n", - "shows how the power-law `slope` of each lens's mass distribution is modeled as being drawn from a shared parent\n", - "Gaussian distribution, whose hyper-parameters (mean and variance) are inferred from the data.\n", - "\n", - "This example illustrates graphical models using point-source datasets and the Hubble Constant, but it is a clear\n", - "and intuitive model whereby a single shared parameter (H0) links multiple lenses. The API and concepts demonstrated\n", - "here can be directly applied to imaging and interferometer datasets, and more complex models with many shared can\n", - "be composed and fitted using the same framework.\n", - "\n", - "__Contents__\n", - "\n", - "- **Initialization:** Load 3 simulated time-delay lens datasets which are all simulated with different mass models but.\n", - "- **Point Solver:** We set up the `PointSolver`, which is used to compute the multiple images of the point source in.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Analysis Factors:** Above, we created a `model_list` containing three lens models, each sharing the same prior on `H0`.\n", - "- **Factor Graph:** We now combine our `AnalysisFactor` objects to form a **factor graph**.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autofit as af" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Initialization__\n", - "\n", - "Load 3 simulated time-delay lens datasets which are all simulated with different mass models but \n", - "the same Hubble constant." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"samples\"\n", - "dataset_type = \"point_source\"\n", - "dataset_sample_name = \"hubble_constant_time_delays\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_sample_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/point_source/simulator_sample.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "total_datasets = 3\n", - "\n", - "dataset_list = []\n", - "\n", - "for dataset_index in range(total_datasets):\n", - " dataset_sample_path = dataset_path / f\"dataset_{dataset_index}\"\n", - "\n", - " dataset = al.from_json(\n", - " file_path=dataset_sample_path / \"point_dataset_with_time_delays.json\",\n", - " )\n", - "\n", - " dataset_list.append(dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "We set up the `PointSolver`, which is used to compute the multiple images of the point source in the image-plane.\n", - "\n", - "There are no special settings or inputs for the fitting of time_delays, therefore the `PointSolver` is set up in the same way\n", - "as in the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose our model using `Model` objects, which represent the lenses we fit to our data.\n", - "\n", - "This graphical model creates a non-linear parameter space that has parameters for every lens mass and source galaxy point\n", - "source in our sample. In this example, there are 3 lenses each with their own model, therefore:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with fixec centre [3 parameters].\n", - "\n", - " - The source galaxy's light is a point `Point` [2 parameters].\n", - "\n", - " - There is a single cosmological shared free parameter, `H0` [1 parameter]\n", - "\n", - " - There are 3 strong lenses in our graphical model [(3 x 5) + 1 = 16 parameters]. \n", - "\n", - "The overall dimensionality of parameter space is therefore N=16." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "cosmology = af.Model(al.cosmo.FlatLambdaCDM)\n", - "\n", - "cosmology.H0 = af.UniformPrior(lower_limit=0.0, upper_limit=150.0)\n", - "\n", - "model_list = []\n", - "\n", - "for model_index in range(total_datasets):\n", - " # Lens:\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre.centre_0 = 0.0\n", - " mass.centre.centre_1 = 0.0\n", - "\n", - " lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - "\n", - " # Source:\n", - "\n", - " point_0 = af.Model(al.ps.Point)\n", - "\n", - " source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)\n", - "\n", - " # Overall Lens Model:\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source=source),\n", - " cosmology=cosmology,\n", - " )\n", - "\n", - " model_list.append(model)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "For each dataset we now create a corresponding `AnalysisPoint` class." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = []\n", - "\n", - "for dataset in dataset_list:\n", - " analysis = al.AnalysisPoint(dataset=dataset, solver=solver)\n", - "\n", - " analysis_list.append(analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis Factors__\n", - "\n", - "Above, we created a `model_list` containing three lens models, each sharing the same prior on `H0`. \n", - "We also loaded three datasets and assigned each one an `Analysis` class, which defines how a model\n", - "is evaluated against that dataset.\n", - "\n", - "We now pair each model with its corresponding `Analysis` object, telling **PyAutoLens** that:\n", - "\n", - "- `model_list[0]` is fit to `dataset_list[0]` using `analysis_list[0]`\n", - "- `model_list[1]` is fit to `dataset_list[1]` using `analysis_list[1]`\n", - "- `model_list[2]` is fit to `dataset_list[2]` using `analysis_list[2]`\n", - "\n", - "The point where a `Model` and an `Analysis` meet is called an **`AnalysisFactor`**.\n", - "\n", - "This terminology reflects that we are building a **factor graph**: \n", - "each *factor* corresponds to a node that contains (i) a dataset, (ii) a model for that dataset, and (iii) the\n", - "process that fits them together. The links between these nodes define the global structure of the graphical model\n", - "we are fitting, including shared parameters such as `H0`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = []\n", - "\n", - "for model, analysis in zip(model_list, analysis_list):\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - "\n", - " analysis_factor_list.append(analysis_factor)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Factor Graph__\n", - "\n", - "We now combine our `AnalysisFactor` objects to form a **factor graph**.\n", - "\n", - "What is a factor graph? \n", - "A factor graph is the explicit representation of our graphical model. It defines:\n", - "\n", - "- the individual model components used to fit each dataset (e.g., the three `Collection` lens + source models), and\n", - "- how their parameters are linked or shared (e.g., each lens has its own mass distribution, but all share the same\n", - " cosmological parameter `H0`).\n", - "\n", - "Although PyAutoFit does not yet visualize factor graphs, the conceptual structure is straightforward. A factor graph\n", - "consists of:\n", - "\n", - "- **Nodes** \u2014 each node corresponds to an `AnalysisFactor`, meaning a specific dataset paired with a model used to fit it.\n", - "\n", - "- **Links** \u2014 these represent shared model components or parameters across nodes (e.g., a single `H0` value shared\n", - " across all lenses), ensuring they retain the same value when fitting multiple datasets.\n", - "\n", - "Together, the nodes and links define the full, coupled model that is fit across all datasets simultaneously." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit will use the factor graph's `global_prior_model`, which uses the models contained in every analysis factor \n", - "to contrast the overall global model that is fitted.\n", - "\n", - "Printing the `info` attribute of this model reveals the overall structure of the model, which is grouped in terms\n", - "of the analysis factors and therefore datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "We can now create a non-linear search and used it to the fit the factor graph, using its `global_prior_model` property." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"modeling\"),\n", - " name=\"graphical\",\n", - " n_live=150,\n", - " n_batch=10, # GPU batching and VRAM use explained in `modeling` examples.\n", - ")\n", - "\n", - "result = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result's `info` attribute shows that the result is expressed following the same structure of analysis factors\n", - "that the `global_prior_model.info` attribute revealed above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The graphical model estimated the Hubble constant by fitting all three lenses **simultaneously**, making full use of\n", - "the information shared across the sample.\n", - "\n", - "If you instead fit each lens independently and compute `H0` manually from each result, the combined uncertainty on\n", - "`H0` will be larger than the uncertainty from the graphical model. This demonstrates the power of **joint inference**\n", - "using graphical models.\n", - "\n", - "Even if you tried to combine the independent fits using importance sampling, you would still not recover the same\n", - "precision or accuracy. In addition, the prior on `H0` would be applied three times (once per lens), biasing the\n", - "inference.\n", - "\n", - "This example used a simple shared parameter (the Hubble constant) across multiple lenses. The same framework can be\n", - "extended to far more complex models, with many shared or linked parameters, enabling powerful hierarchical and\n", - "population-level inference for large lens samples." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Graphical\n", + "===================\n", + "\n", + "The example scripts throughout the workspace have focused on modeling **individual strong lens datasets**.\n", + "From each fit, you may have inspected lens properties (e.g., Einstein radius) and source properties\n", + "(e.g., magnification). You may even have analyzed many lenses one-by-one and combined their results to study\n", + "global trends in galaxy formation or cosmology.\n", + "\n", + "However, fitting each lens independently does **not** make full use of the information contained in a large\n", + "sample. Many properties are expected to be **shared across lenses** (e.g., population-level mass slopes,\n", + "cosmological parameters), and treating them independently ignores this shared structure.\n", + "\n", + "In this example, we demonstrate how to fit **multiple lenses simultaneously** using a **graphical model**.\n", + "A graphical model links parameters across separate lens fits, explicitly defining which parameters are unique\n", + "to each dataset and which are shared. These links can be arbitrarily complex, enabling joint analysis across\n", + "diverse datasets with structured relationships between model components.\n", + "\n", + "Here, we illustrate a cosmological application: inferring the **Hubble constant (H0)** from time-delay lenses.\n", + "A graphical model links the mass models across multiple lenses and includes a **shared H0 parameter**, allowing\n", + "a joint inference that improves cosmological constraints compared to individual fits.\n", + "\n", + "Graphical models form the foundation of **hierarchical modeling**, where the parameters of individual lenses are\n", + "assumed to be drawn from a parent distribution (see `guides/modeling/hierarchical`). Hierarchical approaches can\n", + "extract significantly more information from large samples than fitting each dataset independently. The example\n", + "shows how the power-law `slope` of each lens's mass distribution is modeled as being drawn from a shared parent\n", + "Gaussian distribution, whose hyper-parameters (mean and variance) are inferred from the data.\n", + "\n", + "This example illustrates graphical models using point-source datasets and the Hubble Constant, but it is a clear\n", + "and intuitive model whereby a single shared parameter (H0) links multiple lenses. The API and concepts demonstrated\n", + "here can be directly applied to imaging and interferometer datasets, and more complex models with many shared can\n", + "be composed and fitted using the same framework.\n", + "\n", + "__Contents__\n", + "\n", + "- **Initialization:** Load 3 simulated time-delay lens datasets which are all simulated with different mass models but.\n", + "- **Point Solver:** We set up the `PointSolver`, which is used to compute the multiple images of the point source in.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Analysis Factors:** Above, we created a `model_list` containing three lens models, each sharing the same prior on `H0`.\n", + "- **Factor Graph:** We now combine our `AnalysisFactor` objects to form a **factor graph**.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autofit as af" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Initialization__\n", + "\n", + "Load 3 simulated time-delay lens datasets which are all simulated with different mass models but \n", + "the same Hubble constant." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"samples\"\n", + "dataset_type = \"point_source\"\n", + "dataset_sample_name = \"hubble_constant_time_delays\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_sample_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/point_source/simulator_sample.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "total_datasets = 3\n", + "\n", + "dataset_list = []\n", + "\n", + "for dataset_index in range(total_datasets):\n", + " dataset_sample_path = dataset_path / f\"dataset_{dataset_index}\"\n", + "\n", + " dataset = al.from_json(\n", + " file_path=dataset_sample_path / \"point_dataset_with_time_delays.json\",\n", + " )\n", + "\n", + " dataset_list.append(dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "We set up the `PointSolver`, which is used to compute the multiple images of the point source in the image-plane.\n", + "\n", + "There are no special settings or inputs for the fitting of time_delays, therefore the `PointSolver` is set up in the same way\n", + "as in the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose our model using `Model` objects, which represent the lenses we fit to our data.\n", + "\n", + "This graphical model creates a non-linear parameter space that has parameters for every lens mass and source galaxy point\n", + "source in our sample. In this example, there are 3 lenses each with their own model, therefore:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with fixec centre [3 parameters].\n", + "\n", + " - The source galaxy's light is a point `Point` [2 parameters].\n", + "\n", + " - There is a single cosmological shared free parameter, `H0` [1 parameter]\n", + "\n", + " - There are 3 strong lenses in our graphical model [(3 x 5) + 1 = 16 parameters]. \n", + "\n", + "The overall dimensionality of parameter space is therefore N=16." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cosmology = af.Model(al.cosmo.FlatLambdaCDM)\n", + "\n", + "cosmology.H0 = af.UniformPrior(lower_limit=0.0, upper_limit=150.0)\n", + "\n", + "model_list = []\n", + "\n", + "for model_index in range(total_datasets):\n", + " # Lens:\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre.centre_0 = 0.0\n", + " mass.centre.centre_1 = 0.0\n", + "\n", + " lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + " # Source:\n", + "\n", + " point_0 = af.Model(al.ps.Point)\n", + "\n", + " source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)\n", + "\n", + " # Overall Lens Model:\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source=source),\n", + " cosmology=cosmology,\n", + " )\n", + "\n", + " model_list.append(model)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "For each dataset we now create a corresponding `AnalysisPoint` class." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = []\n", + "\n", + "for dataset in dataset_list:\n", + " analysis = al.AnalysisPoint(dataset=dataset, solver=solver)\n", + "\n", + " analysis_list.append(analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis Factors__\n", + "\n", + "Above, we created a `model_list` containing three lens models, each sharing the same prior on `H0`. \n", + "We also loaded three datasets and assigned each one an `Analysis` class, which defines how a model\n", + "is evaluated against that dataset.\n", + "\n", + "We now pair each model with its corresponding `Analysis` object, telling **PyAutoLens** that:\n", + "\n", + "- `model_list[0]` is fit to `dataset_list[0]` using `analysis_list[0]`\n", + "- `model_list[1]` is fit to `dataset_list[1]` using `analysis_list[1]`\n", + "- `model_list[2]` is fit to `dataset_list[2]` using `analysis_list[2]`\n", + "\n", + "The point where a `Model` and an `Analysis` meet is called an **`AnalysisFactor`**.\n", + "\n", + "This terminology reflects that we are building a **factor graph**: \n", + "each *factor* corresponds to a node that contains (i) a dataset, (ii) a model for that dataset, and (iii) the\n", + "process that fits them together. The links between these nodes define the global structure of the graphical model\n", + "we are fitting, including shared parameters such as `H0`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = []\n", + "\n", + "for model, analysis in zip(model_list, analysis_list):\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + "\n", + " analysis_factor_list.append(analysis_factor)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Factor Graph__\n", + "\n", + "We now combine our `AnalysisFactor` objects to form a **factor graph**.\n", + "\n", + "What is a factor graph? \n", + "A factor graph is the explicit representation of our graphical model. It defines:\n", + "\n", + "- the individual model components used to fit each dataset (e.g., the three `Collection` lens + source models), and\n", + "- how their parameters are linked or shared (e.g., each lens has its own mass distribution, but all share the same\n", + " cosmological parameter `H0`).\n", + "\n", + "Although PyAutoFit does not yet visualize factor graphs, the conceptual structure is straightforward. A factor graph\n", + "consists of:\n", + "\n", + "- **Nodes** \u2014 each node corresponds to an `AnalysisFactor`, meaning a specific dataset paired with a model used to fit it.\n", + "\n", + "- **Links** \u2014 these represent shared model components or parameters across nodes (e.g., a single `H0` value shared\n", + " across all lenses), ensuring they retain the same value when fitting multiple datasets.\n", + "\n", + "Together, the nodes and links define the full, coupled model that is fit across all datasets simultaneously." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit will use the factor graph's `global_prior_model`, which uses the models contained in every analysis factor \n", + "to contrast the overall global model that is fitted.\n", + "\n", + "Printing the `info` attribute of this model reveals the overall structure of the model, which is grouped in terms\n", + "of the analysis factors and therefore datasets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "We can now create a non-linear search and used it to the fit the factor graph, using its `global_prior_model` property." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"modeling\"),\n", + " name=\"graphical\",\n", + " n_live=150,\n", + " n_batch=10, # GPU batching and VRAM use explained in `modeling` examples.\n", + ")\n", + "\n", + "result = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result's `info` attribute shows that the result is expressed following the same structure of analysis factors\n", + "that the `global_prior_model.info` attribute revealed above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The graphical model estimated the Hubble constant by fitting all three lenses **simultaneously**, making full use of\n", + "the information shared across the sample.\n", + "\n", + "If you instead fit each lens independently and compute `H0` manually from each result, the combined uncertainty on\n", + "`H0` will be larger than the uncertainty from the graphical model. This demonstrates the power of **joint inference**\n", + "using graphical models.\n", + "\n", + "Even if you tried to combine the independent fits using importance sampling, you would still not recover the same\n", + "precision or accuracy. In addition, the prior on `H0` would be applied three times (once per lens), biasing the\n", + "inference.\n", + "\n", + "This example used a simple shared parameter (the Hubble constant) across multiple lenses. The same framework can be\n", + "extended to far more complex models, with many shared or linked parameters, enabling powerful hierarchical and\n", + "population-level inference for large lens samples." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/advanced/hierarchical.ipynb b/notebooks/guides/modeling/advanced/hierarchical.ipynb index 7cffd2a07..fac01a3e4 100644 --- a/notebooks/guides/modeling/advanced/hierarchical.ipynb +++ b/notebooks/guides/modeling/advanced/hierarchical.ipynb @@ -1,507 +1,544 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Hierarchical\n", - "======================\n", - "\n", - "A hierarchical model assumes that certain model parameters are drawn from a **shared parent distribution**\n", - "(e.g. a Gaussian). When we fit such a model, the parameters of this parent distribution (such as its `mean` and\n", - "`sigma`) become explicit free parameters in the inference. Scientifically, these parent-distribution parameters\n", - "are often of greatest interest because they describe **population-level trends**, rather than the properties of\n", - "individual lenses.\n", - "\n", - "In this example, we fit a hierarchical model to a sample of three strong gravitational lenses. We assume that\n", - "the **power-law slope** of each lens\u2019s mass distribution is drawn from a shared Gaussian distribution. This is\n", - "well motivated: observational studies find that the slopes of early-type lens galaxies are well approximated by\n", - "a Gaussian with mean \u2248 2.06 and sigma \u2248 0.20.\n", - "\n", - "To perform this fit, we use a graphical model (see `guides/modeling/advanced/graphical`). The model-composition\n", - "API makes it straightforward to fit multiple datasets simultaneously while linking parameters via a shared\n", - "parent distribution.\n", - "\n", - "Note that hierarchical models **do not have to be fit through graphical models**\u2014the same API can be applied to\n", - "single-object problems where multiple components of a lens share a parent distribution. For example, a single\n", - "lens could contain multiple mass components whose parameters are drawn from a common parent distribution. While\n", - "no such example is included in the current workspace, the structure shown here could be adapted easily for that case.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Model Individual Factors:** We first set up a model for each lens, with an `PowerLawSph` mass and `ExponentialSph` bulge, which.\n", - "- **Analysis Factors:** Now we have our `Analysis` classes and model components, we can compose our `AnalysisFactor`'s.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Factor Graph:** We now create the factor graph for this model, using the list of `AnalysisFactor`'s and the.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Concept:** A hierarchical model yields more precise and accurate estimates of the parent distribution\u2019s.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import autolens as al\n", - "import autofit as af\n", - "from pathlib import Path" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "For each lens dataset in our sample we set up the correct path and load it by iterating over a for loop. \n", - "\n", - "We are loading a different dataset to the previous tutorials, where the lenses only have a single bulge component\n", - "which each have different Sersic indexes which are drawn from a parent Gaussian distribution with a mean value \n", - "of 2.0 and sigma of 0.5.\n", - "\n", - "This data is not automatically provided with the autolens workspace, and must be first simulated by running the \n", - "script `autolens_workspace/scripts/advanced/graphical/simulator/samples/advanced/mass_power_law.py`. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"samples\"\n", - "dataset_type = \"imaging\"\n", - "dataset_sample_name = \"simple\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_label, dataset_sample_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator_sample.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "total_datasets = 3\n", - "\n", - "dataset_list = []\n", - "\n", - "for dataset_index in range(total_datasets):\n", - " dataset_sample_path = Path(dataset_path, f\"dataset_{dataset_index}\")\n", - "\n", - " dataset_list.append(\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_sample_path, \"data.fits\"),\n", - " psf_path=Path(dataset_sample_path, \"psf.fits\"),\n", - " noise_map_path=Path(dataset_sample_path, \"noise_map.fits\"),\n", - " pixel_scales=0.1,\n", - " )\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_dataset_list = []\n", - "\n", - "for dataset in dataset_list:\n", - " mask_radius = 3.0\n", - "\n", - " mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=mask)\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " masked_dataset_list.append(dataset)\n", - "\n", - "dataset_list = masked_dataset_list" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "For each dataset we now create a corresponding `AnalysisImaging` class, as we are used to doing for `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = []\n", - "\n", - "for dataset in dataset_list:\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - "\n", - " analysis_list.append(analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Individual Factors__\n", - "\n", - "We first set up a model for each lens, with an `PowerLawSph` mass and `ExponentialSph` bulge, which we will use to \n", - "fit the hierarchical model.\n", - "\n", - "Note that the `PowerLawSph` mass model has a `slope` parameter, which we will assume is drawn from a shared parent\n", - "Gaussian distribution, albeit building this into the model is done later in this script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_list = []\n", - "\n", - "for dataset_index in range(total_datasets):\n", - " lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.PowerLawSph)\n", - " lens.mass.centre = (0.0, 0.0)\n", - "\n", - " source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.ExponentialCoreSph)\n", - "\n", - " model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - " model_list.append(model)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis Factors__\n", - "\n", - "Now we have our `Analysis` classes and model components, we can compose our `AnalysisFactor`'s.\n", - "\n", - "These are composed in the same way as for the graphical model and are described in detail in the\n", - "`guides/modeling/advanced/graphical` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = []\n", - "\n", - "for model, analysis in zip(model_list, analysis_list):\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - "\n", - " analysis_factor_list.append(analysis_factor)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We now compose the hierarchical model that we fit, using the individual model components created above.\n", - "\n", - "We first create a `HierarchicalFactor`, which represents the parent Gaussian distribution from which we will assume \n", - "that the `slope` of each individual lens mass model is drawn. \n", - "\n", - "For this parent `Gaussian`, we have to place priors on its `mean` and `sigma`, given that they are parameters in our\n", - "model we are ultimately fitting for." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "hierarchical_factor = af.HierarchicalFactor(\n", - " af.GaussianPrior,\n", - " mean=af.TruncatedGaussianPrior(\n", - " mean=2.0, sigma=1.0, lower_limit=0.0, upper_limit=100.0\n", - " ),\n", - " sigma=af.TruncatedGaussianPrior(\n", - " mean=0.5, sigma=0.5, lower_limit=0.0, upper_limit=100.0\n", - " ),\n", - " use_jax=False,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now add each of the individual mass `slope` parameters to the `hierarchical_factor`.\n", - "\n", - "This composes the hierarchical model whereby the individual `slope` of every light model in our dataset is now \n", - "assumed to be drawn from a shared parent distribution. It is the `mean` and `sigma` of this distribution we are hoping \n", - "to estimate.\n", - "\n", - "The code below is not specific to graphical models and could be applied to any model where certain parameters are\n", - "assumed to be drawn from a shared parent distribution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for model in model_list:\n", - " hierarchical_factor.add_drawn_variable(model.galaxies.lens.mass.slope)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Factor Graph__\n", - "\n", - "We now create the factor graph for this model, using the list of `AnalysisFactor`'s and the hierarchical factor.\n", - "\n", - "Again, this code is described in detail in the `guides/modeling/advanced/graphical` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "factor_graph = af.FactorGraphModel(\n", - " *analysis_factor_list, hierarchical_factor, use_jax=False\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The factor graph model `info` attribute shows that the hierarchical factor's parameters are included in the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "We can now create a non-linear search and used it to the fit the factor graph, using its `global_prior_model` property." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"modeling\"),\n", - " name=\"hierarchical\",\n", - " n_live=150,\n", - " n_batch=10, # GPU batching and VRAM use explained in `modeling` examples.\n", - ")\n", - "\n", - "result = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result's `info` attribute shows the result, including the hierarchical factor's parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now inspect the inferred value of hierarchical factor's mean and sigma.\n", - "\n", - "We see that they are consistent with the input values of `mean=2.0` and `sigma=0.2`, which are\n", - "the values used to simulate the lens dataset sample." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "mean = samples.median_pdf(as_instance=False)[-2]\n", - "\n", - "u1_error = samples.values_at_upper_sigma(sigma=1.0)[-2]\n", - "l1_error = samples.values_at_lower_sigma(sigma=1.0)[-2]\n", - "\n", - "u3_error = samples.values_at_upper_sigma(sigma=3.0)[-2]\n", - "l3_error = samples.values_at_lower_sigma(sigma=3.0)[-2]\n", - "\n", - "print(\n", - " \"Inferred value of the mean of the parent hierarchical distribution for the mass model slopes: \\n\"\n", - ")\n", - "print(f\"{mean} ({l1_error} {u1_error}) [1.0 sigma confidence intervals]\")\n", - "print(f\"{mean} ({l3_error} {u3_error}) [3.0 sigma confidence intervals]\")\n", - "\n", - "scatter = samples.median_pdf(as_instance=False)[-2]\n", - "\n", - "u1_error = samples.values_at_upper_sigma(sigma=1.0)[-1]\n", - "l1_error = samples.values_at_lower_sigma(sigma=1.0)[-1]\n", - "\n", - "u3_error = samples.values_at_upper_sigma(sigma=3.0)[-1]\n", - "l3_error = samples.values_at_lower_sigma(sigma=3.0)[-1]\n", - "\n", - "print(\n", - " \"Inferred value of the scatter (the sigma value of the Gassuain) of the parent hierarchical distribution for the mass model slopes: \\n\"\n", - ")\n", - "print(f\"{scatter} ({l1_error} {u1_error}) [1.0 sigma confidence intervals]\")\n", - "print(f\"{scatter} ({l3_error} {u3_error}) [3.0 sigma confidence intervals]\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Concept__\n", - "\n", - "A hierarchical model yields more precise and accurate estimates of the parent distribution\u2019s parameters, but also\n", - "the individual parameters fit to each lens. \n", - "\n", - "This happens because **each dataset informs the shared distribution**, and the distribution in turn constrains \n", - "each individual dataset. This can be described as **\u201cthe datasets talking to one another.\u201d**\n", - "\n", - "For example, suppose that when fit alone, `dataset_0` yields a weak constraint on the mass\u2013slope parameter, spanning\n", - "1.3 \u2192 2.7 (1\u03c3). Now imagine that, when we include the other datasets, the hierarchical distribution is well constrained\n", - "to `mean = 2.0 \u00b1 0.1` and `sigma = 0.10 \u00b1 0.05`. This shared information tells us that values far from ~2.0 are unlikely,\n", - "so `dataset_0` will be **forced toward physically plausible solutions**, even though it could not infer this on its own.\n", - "\n", - "In large hierarchical fits with many lenses, this \u201ccommunication\u201d between datasets can break degeneracies and extract\n", - "substantially more information from the sample than independent fits ever could. For inference on parameters like\n", - "cosmology, this shrinkage of uncertainties on lens mass model parameters can lead to significantly tighter constraints\n", - "on the cosmological parameters themselves.\n", - "\n", - "__Wrap Up__\n", - "\n", - "Hierarchical models enable us to infer **population-level trends** from large lens samples. Using graphical modeling,\n", - "we can easily compose complex models with shared parameters and hierarchical structure across many datasets.\n", - "\n", - "However, scaling to large graphs introduces challenges. As models grow in size, poor sampling can lead to local maxima,\n", - "and fitting many datasets simultaneously can become computationally expensive (in both CPU time and memory).\n", - "\n", - "The next tutorial introduces **Expectation Propagation (EP)**, a framework that partitions the graphical model into\n", - "many small sub-fits\u2014one for each node in the factor graph. Each node fit passes information to its neighbors, allowing\n", - "us to fit graphs with hundreds of components and tens of thousands of parameters as a series of manageable, low-dimensional\n", - "optimizations.\n", - "\n", - "This makes hierarchical graphical modeling **scalable to the largest datasets and most complex models.**" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Hierarchical\n", + "======================\n", + "\n", + "A hierarchical model assumes that certain model parameters are drawn from a **shared parent distribution**\n", + "(e.g. a Gaussian). When we fit such a model, the parameters of this parent distribution (such as its `mean` and\n", + "`sigma`) become explicit free parameters in the inference. Scientifically, these parent-distribution parameters\n", + "are often of greatest interest because they describe **population-level trends**, rather than the properties of\n", + "individual lenses.\n", + "\n", + "In this example, we fit a hierarchical model to a sample of three strong gravitational lenses. We assume that\n", + "the **power-law slope** of each lens\u2019s mass distribution is drawn from a shared Gaussian distribution. This is\n", + "well motivated: observational studies find that the slopes of early-type lens galaxies are well approximated by\n", + "a Gaussian with mean \u2248 2.06 and sigma \u2248 0.20.\n", + "\n", + "To perform this fit, we use a graphical model (see `guides/modeling/advanced/graphical`). The model-composition\n", + "API makes it straightforward to fit multiple datasets simultaneously while linking parameters via a shared\n", + "parent distribution.\n", + "\n", + "Note that hierarchical models **do not have to be fit through graphical models**\u2014the same API can be applied to\n", + "single-object problems where multiple components of a lens share a parent distribution. For example, a single\n", + "lens could contain multiple mass components whose parameters are drawn from a common parent distribution. While\n", + "no such example is included in the current workspace, the structure shown here could be adapted easily for that case.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Model Individual Factors:** We first set up a model for each lens, with an `PowerLawSph` mass and `ExponentialSph` bulge, which.\n", + "- **Analysis Factors:** Now we have our `Analysis` classes and model components, we can compose our `AnalysisFactor`'s.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Factor Graph:** We now create the factor graph for this model, using the list of `AnalysisFactor`'s and the.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Concept:** A hierarchical model yields more precise and accurate estimates of the parent distribution\u2019s.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import autolens as al\n", + "import autofit as af\n", + "from pathlib import Path" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "For each lens dataset in our sample we set up the correct path and load it by iterating over a for loop. \n", + "\n", + "We are loading a different dataset to the previous tutorials, where the lenses only have a single bulge component\n", + "which each have different Sersic indexes which are drawn from a parent Gaussian distribution with a mean value \n", + "of 2.0 and sigma of 0.5.\n", + "\n", + "This data is not automatically provided with the autolens workspace, and must be first simulated by running the \n", + "script `autolens_workspace/scripts/advanced/graphical/simulator/samples/advanced/mass_power_law.py`. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"samples\"\n", + "dataset_type = \"imaging\"\n", + "dataset_sample_name = \"simple\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_label, dataset_sample_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator_sample.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "total_datasets = 3\n", + "\n", + "dataset_list = []\n", + "\n", + "for dataset_index in range(total_datasets):\n", + " dataset_sample_path = Path(dataset_path, f\"dataset_{dataset_index}\")\n", + "\n", + " dataset_list.append(\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_sample_path, \"data.fits\"),\n", + " psf_path=Path(dataset_sample_path, \"psf.fits\"),\n", + " noise_map_path=Path(dataset_sample_path, \"noise_map.fits\"),\n", + " pixel_scales=0.1,\n", + " )\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_dataset_list = []\n", + "\n", + "for dataset in dataset_list:\n", + " mask_radius = 3.0\n", + "\n", + " mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=mask)\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " masked_dataset_list.append(dataset)\n", + "\n", + "dataset_list = masked_dataset_list" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "For each dataset we now create a corresponding `AnalysisImaging` class, as we are used to doing for `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = []\n", + "\n", + "for dataset in dataset_list:\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + "\n", + " analysis_list.append(analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Individual Factors__\n", + "\n", + "We first set up a model for each lens, with an `PowerLawSph` mass and `ExponentialSph` bulge, which we will use to \n", + "fit the hierarchical model.\n", + "\n", + "Note that the `PowerLawSph` mass model has a `slope` parameter, which we will assume is drawn from a shared parent\n", + "Gaussian distribution, albeit building this into the model is done later in this script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_list = []\n", + "\n", + "for dataset_index in range(total_datasets):\n", + " lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.PowerLawSph)\n", + " lens.mass.centre = (0.0, 0.0)\n", + "\n", + " source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.ExponentialCoreSph)\n", + "\n", + " model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + " model_list.append(model)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis Factors__\n", + "\n", + "Now we have our `Analysis` classes and model components, we can compose our `AnalysisFactor`'s.\n", + "\n", + "These are composed in the same way as for the graphical model and are described in detail in the\n", + "`guides/modeling/advanced/graphical` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = []\n", + "\n", + "for model, analysis in zip(model_list, analysis_list):\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + "\n", + " analysis_factor_list.append(analysis_factor)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We now compose the hierarchical model that we fit, using the individual model components created above.\n", + "\n", + "We first create a `HierarchicalFactor`, which represents the parent Gaussian distribution from which we will assume \n", + "that the `slope` of each individual lens mass model is drawn. \n", + "\n", + "For this parent `Gaussian`, we have to place priors on its `mean` and `sigma`, given that they are parameters in our\n", + "model we are ultimately fitting for." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "hierarchical_factor = af.HierarchicalFactor(\n", + " af.GaussianPrior,\n", + " mean=af.TruncatedGaussianPrior(\n", + " mean=2.0, sigma=1.0, lower_limit=0.0, upper_limit=100.0\n", + " ),\n", + " sigma=af.TruncatedGaussianPrior(\n", + " mean=0.5, sigma=0.5, lower_limit=0.0, upper_limit=100.0\n", + " ),\n", + " use_jax=False,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now add each of the individual mass `slope` parameters to the `hierarchical_factor`.\n", + "\n", + "This composes the hierarchical model whereby the individual `slope` of every light model in our dataset is now \n", + "assumed to be drawn from a shared parent distribution. It is the `mean` and `sigma` of this distribution we are hoping \n", + "to estimate.\n", + "\n", + "The code below is not specific to graphical models and could be applied to any model where certain parameters are\n", + "assumed to be drawn from a shared parent distribution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for model in model_list:\n", + " hierarchical_factor.add_drawn_variable(model.galaxies.lens.mass.slope)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Factor Graph__\n", + "\n", + "We now create the factor graph for this model, using the list of `AnalysisFactor`'s and the hierarchical factor.\n", + "\n", + "Again, this code is described in detail in the `guides/modeling/advanced/graphical` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph = af.FactorGraphModel(\n", + " *analysis_factor_list, hierarchical_factor, use_jax=False\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The factor graph model `info` attribute shows that the hierarchical factor's parameters are included in the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "We can now create a non-linear search and used it to the fit the factor graph, using its `global_prior_model` property." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"modeling\"),\n", + " name=\"hierarchical\",\n", + " n_live=150,\n", + " n_batch=10, # GPU batching and VRAM use explained in `modeling` examples.\n", + ")\n", + "\n", + "result = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result's `info` attribute shows the result, including the hierarchical factor's parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now inspect the inferred value of hierarchical factor's mean and sigma.\n", + "\n", + "We see that they are consistent with the input values of `mean=2.0` and `sigma=0.2`, which are\n", + "the values used to simulate the lens dataset sample." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "mean = samples.median_pdf(as_instance=False)[-2]\n", + "\n", + "u1_error = samples.values_at_upper_sigma(sigma=1.0)[-2]\n", + "l1_error = samples.values_at_lower_sigma(sigma=1.0)[-2]\n", + "\n", + "u3_error = samples.values_at_upper_sigma(sigma=3.0)[-2]\n", + "l3_error = samples.values_at_lower_sigma(sigma=3.0)[-2]\n", + "\n", + "print(\n", + " \"Inferred value of the mean of the parent hierarchical distribution for the mass model slopes: \\n\"\n", + ")\n", + "print(f\"{mean} ({l1_error} {u1_error}) [1.0 sigma confidence intervals]\")\n", + "print(f\"{mean} ({l3_error} {u3_error}) [3.0 sigma confidence intervals]\")\n", + "\n", + "scatter = samples.median_pdf(as_instance=False)[-2]\n", + "\n", + "u1_error = samples.values_at_upper_sigma(sigma=1.0)[-1]\n", + "l1_error = samples.values_at_lower_sigma(sigma=1.0)[-1]\n", + "\n", + "u3_error = samples.values_at_upper_sigma(sigma=3.0)[-1]\n", + "l3_error = samples.values_at_lower_sigma(sigma=3.0)[-1]\n", + "\n", + "print(\n", + " \"Inferred value of the scatter (the sigma value of the Gassuain) of the parent hierarchical distribution for the mass model slopes: \\n\"\n", + ")\n", + "print(f\"{scatter} ({l1_error} {u1_error}) [1.0 sigma confidence intervals]\")\n", + "print(f\"{scatter} ({l3_error} {u3_error}) [3.0 sigma confidence intervals]\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Concept__\n", + "\n", + "A hierarchical model yields more precise and accurate estimates of the parent distribution\u2019s parameters, but also\n", + "the individual parameters fit to each lens. \n", + "\n", + "This happens because **each dataset informs the shared distribution**, and the distribution in turn constrains \n", + "each individual dataset. This can be described as **\u201cthe datasets talking to one another.\u201d**\n", + "\n", + "For example, suppose that when fit alone, `dataset_0` yields a weak constraint on the mass\u2013slope parameter, spanning\n", + "1.3 \u2192 2.7 (1\u03c3). Now imagine that, when we include the other datasets, the hierarchical distribution is well constrained\n", + "to `mean = 2.0 \u00b1 0.1` and `sigma = 0.10 \u00b1 0.05`. This shared information tells us that values far from ~2.0 are unlikely,\n", + "so `dataset_0` will be **forced toward physically plausible solutions**, even though it could not infer this on its own.\n", + "\n", + "In large hierarchical fits with many lenses, this \u201ccommunication\u201d between datasets can break degeneracies and extract\n", + "substantially more information from the sample than independent fits ever could. For inference on parameters like\n", + "cosmology, this shrinkage of uncertainties on lens mass model parameters can lead to significantly tighter constraints\n", + "on the cosmological parameters themselves.\n", + "\n", + "__Wrap Up__\n", + "\n", + "Hierarchical models enable us to infer **population-level trends** from large lens samples. Using graphical modeling,\n", + "we can easily compose complex models with shared parameters and hierarchical structure across many datasets.\n", + "\n", + "However, scaling to large graphs introduces challenges. As models grow in size, poor sampling can lead to local maxima,\n", + "and fitting many datasets simultaneously can become computationally expensive (in both CPU time and memory).\n", + "\n", + "The next tutorial introduces **Expectation Propagation (EP)**, a framework that partitions the graphical model into\n", + "many small sub-fits\u2014one for each node in the factor graph. Each node fit passes information to its neighbors, allowing\n", + "us to fit graphs with hundreds of components and tens of thousands of parameters as a series of manageable, low-dimensional\n", + "optimizations.\n", + "\n", + "This makes hierarchical graphical modeling **scalable to the largest datasets and most complex models.**" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/bug_fix.ipynb b/notebooks/guides/modeling/bug_fix.ipynb index 2f5ac4393..717c29abd 100644 --- a/notebooks/guides/modeling/bug_fix.ipynb +++ b/notebooks/guides/modeling/bug_fix.ipynb @@ -1,233 +1,270 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Parallel Bug Fix\n", - "==========================\n", - "\n", - "Depending on the operating system (e.g. Linux, Mac, Windows) and Python version, running a Python script or Jupyter\n", - "notebook may lead to a error being raised when the search begins.\n", - "\n", - "The root cause of this error is that Python parallelization and JAX may only work when the script is run in a\n", - "particular format, which this script illustrates.\n", - "\n", - "The code in this script is identical to the `autolens_workspace/scripts/imaging/modeling.py` script.\n", - "Comments have therefore been removed to avoid repetition and make the script more concise.\n", - "\n", - "__Contents__\n", - "\n", - "- **The Fix:** The fix which makes parallelization work is at the end of the script, where we use the following.\n", - "- **Trouble Shooting:** If you still cannot get parallelization to work, please ask to be added to the SLACK channel (by.\n", - "\n", - "__The Fix__\n", - "\n", - "The fix which makes parallelization work is at the end of the script, where we use the following code:\n", - "\n", - "`if __name__ == \"__main__\":`\n", - "\n", - " `fit()`\n", - "\n", - "The reason this fixes parallelization is beyond the scope of this tutorial. However, if you are curious, a quick\n", - "Google search will provide you with a detailed explanation! For example, the stack overflow page below has\n", - "some good answers:\n", - "\n", - " https://stackoverflow.com/questions/20360686/compulsory-usage-of-if-name-main-in-windows-while-using-multiprocessi\n", - "\n", - "This fix will work for all dataset formats (e.g. `imaging`, `interferometer`) and should therefore be adopted\n", - "for any modeling script you write that has the error described above.\n", - "\n", - "__Trouble Shooting__\n", - "\n", - "If you still cannot get parallelization to work, please ask to be added to the SLACK\n", - "channel (by emailing me https://github.com/Jammy2211), where we will be able to provide support." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def fit():\n", - " from autoconf import (\n", - " jax_wrapper,\n", - " ) # Ensures JAX environment variables are set before other imports\n", - "\n", - " from autoconf import setup_notebook; setup_notebook()\n", - "\n", - " from pathlib import Path\n", - " import autofit as af\n", - " import autolens as al\n", - " import autolens.plot as aplt\n", - "\n", - " \"\"\"\n", - " __Dataset__\n", - " \"\"\"\n", - " dataset_name = \"simple\"\n", - " dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - " \"\"\"\n", - " __Dataset Auto-Simulation__\n", - "\n", - " If the dataset does not already exist on your system, it will be created by running the corresponding\n", - " simulator script. This ensures that all example scripts can be run without manually simulating data first.\n", - " \"\"\"\n", - " if not Path(dataset_path).exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - " dataset = al.Imaging.from_fits(\n", - " data_path=Path(dataset_path) / \"data.fits\",\n", - " psf_path=Path(dataset_path) / \"psf.fits\",\n", - " noise_map_path=Path(dataset_path) / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - " )\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " \"\"\"\n", - " __Mask__\n", - " \"\"\"\n", - " mask_radius = 3.0\n", - "\n", - " mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=mask)\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " \"\"\"\n", - " __Over Sampling__\n", - " \"\"\"\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " \"\"\"\n", - " __Model__\n", - " \"\"\"\n", - " # Lens:\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - "\n", - " shear = af.Model(al.mp.ExternalShear)\n", - "\n", - " lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - " # Source:\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - " # Overall Lens Model:\n", - "\n", - " model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - " \"\"\"\n", - " __Search__ \n", - " \"\"\"\n", - " search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\"),\n", - " name=\"modeling\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU batching and VRAM use explained in `modeling` examples.\n", - " iterations_per_quick_update=100000,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Analysis__\n", - " \"\"\"\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", - " )\n", - "\n", - " \"\"\"\n", - " __Model-Fit__\n", - " \"\"\"\n", - " result = search.fit(model=model, analysis=analysis)\n", - "\n", - " \"\"\"\n", - " __Output Folder__\n", - " \"\"\"\n", - " print(result.info)\n", - "\n", - " print(result.max_log_likelihood_instance)\n", - "\n", - " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - " aplt.corner_anesthetic(samples=result.samples)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This small change in how the code is run fixes parallelization issues." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if __name__ == \"__main__\":\n", - " fit()\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Parallel Bug Fix\n", + "==========================\n", + "\n", + "Depending on the operating system (e.g. Linux, Mac, Windows) and Python version, running a Python script or Jupyter\n", + "notebook may lead to a error being raised when the search begins.\n", + "\n", + "The root cause of this error is that Python parallelization and JAX may only work when the script is run in a\n", + "particular format, which this script illustrates.\n", + "\n", + "The code in this script is identical to the `autolens_workspace/scripts/imaging/modeling.py` script.\n", + "Comments have therefore been removed to avoid repetition and make the script more concise.\n", + "\n", + "__Contents__\n", + "\n", + "- **The Fix:** The fix which makes parallelization work is at the end of the script, where we use the following.\n", + "- **Trouble Shooting:** If you still cannot get parallelization to work, please ask to be added to the SLACK channel (by.\n", + "\n", + "__The Fix__\n", + "\n", + "The fix which makes parallelization work is at the end of the script, where we use the following code:\n", + "\n", + "`if __name__ == \"__main__\":`\n", + "\n", + " `fit()`\n", + "\n", + "The reason this fixes parallelization is beyond the scope of this tutorial. However, if you are curious, a quick\n", + "Google search will provide you with a detailed explanation! For example, the stack overflow page below has\n", + "some good answers:\n", + "\n", + " https://stackoverflow.com/questions/20360686/compulsory-usage-of-if-name-main-in-windows-while-using-multiprocessi\n", + "\n", + "This fix will work for all dataset formats (e.g. `imaging`, `interferometer`) and should therefore be adopted\n", + "for any modeling script you write that has the error described above.\n", + "\n", + "__Trouble Shooting__\n", + "\n", + "If you still cannot get parallelization to work, please ask to be added to the SLACK\n", + "channel (by emailing me https://github.com/Jammy2211), where we will be able to provide support." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def fit():\n", + " from autoconf import (\n", + " jax_wrapper,\n", + " ) # Ensures JAX environment variables are set before other imports\n", + "\n", + " from autoconf import setup_notebook; setup_notebook()\n", + "\n", + " from pathlib import Path\n", + " import autofit as af\n", + " import autolens as al\n", + " import autolens.plot as aplt\n", + "\n", + " \"\"\"\n", + " __Dataset__\n", + " \"\"\"\n", + " dataset_name = \"simple\"\n", + " dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + " \"\"\"\n", + " __Dataset Auto-Simulation__\n", + "\n", + " If the dataset does not already exist on your system, it will be created by running the corresponding\n", + " simulator script. This ensures that all example scripts can be run without manually simulating data first.\n", + " \"\"\"\n", + " if not Path(dataset_path).exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + " dataset = al.Imaging.from_fits(\n", + " data_path=Path(dataset_path) / \"data.fits\",\n", + " psf_path=Path(dataset_path) / \"psf.fits\",\n", + " noise_map_path=Path(dataset_path) / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + " )\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " \"\"\"\n", + " __Mask__\n", + " \"\"\"\n", + " mask_radius = 3.0\n", + "\n", + " mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=mask)\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " \"\"\"\n", + " __Over Sampling__\n", + " \"\"\"\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " \"\"\"\n", + " __Model__\n", + " \"\"\"\n", + " # Lens:\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + "\n", + " shear = af.Model(al.mp.ExternalShear)\n", + "\n", + " lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + " # Source:\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + " # Overall Lens Model:\n", + "\n", + " model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + " \"\"\"\n", + " __Search__ \n", + " \"\"\"\n", + " search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\"),\n", + " name=\"modeling\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU batching and VRAM use explained in `modeling` examples.\n", + " iterations_per_quick_update=100000,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Analysis__\n", + " \"\"\"\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", + " )\n", + "\n", + " \"\"\"\n", + " __Model-Fit__\n", + " \"\"\"\n", + " result = search.fit(model=model, analysis=analysis)\n", + "\n", + " \"\"\"\n", + " __Output Folder__\n", + " \"\"\"\n", + " print(result.info)\n", + "\n", + " print(result.max_log_likelihood_instance)\n", + "\n", + " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + " aplt.corner_anesthetic(samples=result.samples)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This small change in how the code is run fixes parallelization issues." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if __name__ == \"__main__\":\n", + " fit()\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/chaining.ipynb b/notebooks/guides/modeling/chaining.ipynb index b0eb29a64..47a28992f 100644 --- a/notebooks/guides/modeling/chaining.ipynb +++ b/notebooks/guides/modeling/chaining.ipynb @@ -1,577 +1,614 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Chaining\n", - "==================\n", - "\n", - "Non-linear search chaining is an advanced model-fitting approach which breaks the model-fitting procedure down into\n", - "multiple non-linear searches, using the results of the initial searches to initialization parameter\n", - "sampling in subsequent searches. This contrasts the `modeling` examples which each compose and fit a single lens\n", - "model-fit using one non-linear search.\n", - "\n", - "The benefits of non-linear search chaining are:\n", - "\n", - " - Earlier searches fit simpler lens models than the later searches, which have a less complex non-linear parameter\n", - " space that can be sampled more efficiently, with a reduced chance of inferring an incorrect local maxima solution.\n", - "\n", - " - Earlier searches can use faster non-linear search settings which infer the highest log likelihood models but not\n", - " precisely quantify the parameter errors, with only the final searches using slow settings to robustly estimate errors.\n", - "\n", - " - Earlier searches can augment the data or alter the fitting-procedure in ways that speed up the computational run\n", - " time. These may impact the quality of the model-fit overall, but they can be reverted to the more accurate but more\n", - " computationally expense setting in the final searches.\n", - "\n", - "__Contents__\n", - "\n", - "- **Concise Model Composition API:** Chaining uses the concise `Model` API to compose lens models, which is nearly identical to the standard API but avoids the need to use `Model` objects to compose the lens model when a light or mass profile is passed to a `Collection` object.\n", - "- **This Example:** This script gives an overview of the API for search chaining, a description of how the priors on parameters are used to pass information between searches as well as tools for customizing prior passing.\n", - "- **Dataset + Masking:** Load, plot and mask the `Imaging` data.\n", - "- **Paths:** The path the results of all chained searches are output.\n", - "- **Model (Search 1):** We compose our lens model using `Model` objects, which represent the galaxies we fit to our data.\n", - "- **Search + Analysis + Model-Fit (Search 1):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "- **Result (Search 1):** The results which are used for prior passing are summarised in the `info` attribute.\n", - "- **Model Chaining:** We use the results of search 1 to create the `Model` components that we fit in search 2.\n", - "- **Model Centred Chaining:** We use the results of search 1 to create the `Model` components that we fit in search 2, using `model_centred` to pass priors centered on the maximum likelihood parameter values.\n", - "- **Search + Analysis + Model-Fit (Search 2):** We now create the non-linear search, analysis and perform the model-fit using the chained model.\n", - "- **Result (Search 2):** The final results of the chained model-fit.\n", - "- **How is Search Chaining Used?:** Overview of how search chaining is used in practice for automated lens modeling and SLaM pipelines.\n", - "- **Detailed Explanation Of Prior Passing:** Detailed overview of how prior passing works and tools to customize its behaviour.\n", - "- **EXAMPLE:** Lets go through an example using a real parameter.\n", - "\n", - "__Concise Model Composition API__\n", - "\n", - "Chaining uses the concise `Model` API to compose lens models, which is nearly identical to\n", - "the standard API but avoids the need to use `Model` objects to compose the lens model when a light or mass\n", - "profile is passed to a `Collection` object.\n", - "\n", - "__This Example__\n", - "\n", - "This script gives an overview of the API for search chaining, a description of how the priors on parameters are used\n", - "to pass information between searches as well as tools for customizing prior passing.\n", - "\n", - "There are examples throughout the workspace where search chaining improves and helps automate lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__ \n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__\n", - "\n", - "The path the results of all chained searches are output:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"imaging\") / \"chaining\" / \"start_here\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In the first\n", - "search our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", - " \n", - " - an MGE with 1 x 20 Gaussians for the source galaxy's light [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal)\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", - "provides a reasonably accurate lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__start_here\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 1)__\n", - "\n", - "The results which are used for prior passing are summarised in the `info` attribute." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Chaining__\n", - "\n", - "We use the results of search 1 to create the `Model` components that we fit in search 2.\n", - "\n", - "The term `model` below passes the lens and source models as model-components that are to be fitted\n", - "for by the non-linear search. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_2 = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=result_1.model.galaxies.lens, source=result_1.model.galaxies.source\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The priors on the model components are the same as their original priors in `model_1`, that is the priors loaded\n", - "from the default configuration files. \n", - "\n", - "This makes model and search chaining pointless, as the second fit has the same initialization to sample parameter \n", - "space as the first.\n", - "\n", - "We instead want the model to be passed in a way where the priors are updated to reflect the inferred high likelihood \n", - "regions of the first in the first model. To do this we use the `model_centred` attribute of the result, which passes \n", - "the model components with priors that are centered on the maximum likelihood parameter values of the first search.\n", - "\n", - "__Model Centred Chaining__\n", - "\n", - "We use the results of search 1 to create the `Model` components that we fit in search 2.\n", - "\n", - "The term `model_centred` below passes the lens and source models as model-components that are to be fitted\n", - "for by the non-linear search. In other chaining examples, we'll see other ways to pass prior results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_2 = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=result_1.model_centred.galaxies.lens,\n", - " source=result_1.model_centred.galaxies.source,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`.\n", - "\n", - "This now contains `GaussianPrior`'s that are centered on the maximum likelihood parameter values of the first search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 2)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the `model.info` file of the search 2 model-fit to ensure the priors were passed correctly, as \n", - "well as the checkout the results to ensure an accurate power-law mass model is inferred." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__start_here\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 2)__\n", - "\n", - "The final results can be summarised via printing `info`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We will expand on this API in the following tutorials. The main thing to note is that we can pass entire profiles or\n", - "galaxies using prior passing, if their model does not change. \n", - "\n", - "The API to pass a whole profile or galaxy is as follows:\n", - " \n", - " bulge = result_1.model_centred.galaxies.lens.bulge\n", - " lens = result_1.model_centred.galaxies.lens\n", - " source = result_1.model_centred.galaxies.source\n", - " \n", - "We can also pass priors using an `instance` instead of a `model_centred`. When an `instance` is used, the maximum likelihood\n", - "parameter values are passed as fixed values that are therefore not fitted for by the non-linear search (reducing its\n", - "dimensionality). \n", - "\n", - "We will use this in other examples to \"split up\" the components of the model we fit, for example fit the lens light, \n", - "and then fix it to the best-fit model in a second search which fits the mass and source.\n", - " \n", - "__How is Search Chaining Used?__\n", - " \n", - "For a lot of the history of **PyAutoLens***, the passing of priors illustrated in this tutorial has been a key\n", - "aspect of how search chaining is used to model strong lenses.\n", - "\n", - "However, since ~2024, the direct passing of priors between searches has become less important. This is because\n", - "the non-linear search (Nautilus) is a lot more robust and run times are a lot faster since the inclusion of JAX. Thus,\n", - "the need to change priors via `model_centred` has become less important, but can still offer more reliable results\n", - "and efficient lens modeling.\n", - "\n", - "In fact, most example chaining examples and the Source, Light and Mass pipeline which are built are this feature\n", - "**no longer use prior passing**, always using the `model` attribute to pass models between searches which resets\n", - "the prior on parameters to their default priors.\n", - "\n", - "However, search chaining is still integral to automated lens modeling and the SLaM pipelines, but it is now used\n", - "to:\n", - "\n", - "- Split up the components of the lens model to be fitted in different searches (e.g. fit lens light and the \n", - " mass + source in different searches using `instance`).\n", - "\n", - "- Fit simpler models (e.g. MGE source) in earlier searches to get a fast estimate of the lens light and mass models,\n", - " which are then used as fixed components in later searches.\n", - " \n", - "- To pass `adapt_images` through the pipelines, which use the results of earlier searches to adapt aspects of the\n", - " lens model to the data in later searches. This is especially important for pixelized sources and described in the \n", - " `features/pixelization/adaptive` example.\n", - "\n", - "The most robust automated lens modeling is therefore still built using search chaining, but the direct passing of priors\n", - "between searches is no longer the main tool by which this is achieved.\n", - " \n", - "__Detailed Explanation Of Prior Passing__\n", - " \n", - "Lets now think about how priors are passed. Checkout the `model.info` file of the second search of this tutorial. The \n", - "parameters do not use the default priors we saw in search 1 (which are typically broad UniformPriors). Instead, \n", - "they use TruncatedGaussianPrior`s where:\n", - "\n", - " - The mean values are the median PDF results of every parameter in search 1.\n", - " - The sigma values are specified in the `width_modifier` field of the profile's entry in the `priors.yaml' config \n", - " file (we will discuss why this is used in a moment).\n", - "\n", - "Like the manual `TruncatedGaussianPrior`'s that were used in tutorial 1, the prior passing API sets up the prior on each \n", - "parameter with a `TruncatedGaussianPrior` centred on the high likelihood regions of parameter space!\n", - "\n", - "To end, I provide a detailed overview of how prior passing works and illustrate tools that can be used to customize\n", - "its behaviour. It is up to you whether you want read this, or go ahead to the next tutorial!\n", - "\n", - "Lets say I chain two parameters as follows:\n", - "\n", - " `mass.einstein_radius = result_1.model_centred.galaxies.lens.mass.einstein_radius`\n", - "\n", - "By invoking the `model_centred` attribute, the prior is passed following 3 rules:\n", - "\n", - " 1) The new parameter, in this case the einstein radius, uses a `TruncatedGaussianPrior`.This is ideal, as the 1D pdf results \n", - " we compute at the end of a search are easily summarised as a Gaussian.\n", - "\n", - " 2) The mean of the `TruncatedGaussianPrior` is the median PDF value of the parameter estimated in search 1.\n", - "\n", - " This ensures that the initial sampling of the new search's non-linear starts by searching the region of non-linear \n", - " parameter space that correspond to highest log likelihood solutions in the previous search. Our priors therefore \n", - " correspond to the `correct` regions of parameter space.\n", - "\n", - " 3) The sigma of the Gaussian uses the value specified for the profile in the `config/priors/*.yaml` config file's \n", - " `width_modifer` field (check these files out now).\n", - "\n", - "The idea here is simple. We want a value of sigma that gives a `TruncatedGaussianPrior` wide enough to search a broad \n", - "region of parameter space, so that the lens model can change if a better solution is nearby. However, we want it \n", - "to be narrow enough that we don't search too much of parameter space, as this will be slow or risk leading us \n", - "into an incorrect solution! \n", - "\n", - "The `width_modifier` values in the priors config file have been chosen based on our experience as being a good\n", - "balance broadly sampling parameter space but not being so narrow important solutions are missed.\n", - "\n", - "There are two ways a value is specified using the priors/width file:\n", - "\n", - " 1) Absolute: In this case, the error assumed on the parameter is the value given in the config file. \n", - " For example, if for the width on centre_0 of a light profile, the width modifier reads \"Absolute\" with a value \n", - " 0.05. This means if the error on the parameter centre_0 was less than 0.05 in the previous search, the sigma of \n", - " its `GaussianPrior` in this search will be 0.05.\n", - "\n", - " 2) Relative: In this case, the error assumed on the parameter is the % of the value of the estimated value given in \n", - " the config file. For example, if the intensity estimated in the previous search was 2.0, and the relative error in \n", - " the config file reads \"Relative\" with a value 0.5, then the sigma of the `GaussianPrior` will be 50% of this \n", - " value, i.e. sigma = 0.5 * 2.0 = 1.0.\n", - "\n", - "We use absolute and relative values for different parameters, depending on their properties. For example, using the \n", - "relative value of a parameter like the `Profile` centre makes no sense. If our lens galaxy is centred at (0.0, 0.0), \n", - "the relative error will always be tiny and thus poorly defined. Therefore, the default configs in **PyAutoLens** use \n", - "absolute errors on the centre.\n", - "\n", - "However, there are parameters where using an absolute value does not make sense. Intensity is a good example of this. \n", - "The intensity of an image depends on its units, S/N, galaxy brightness, etc. There is no single absolute value that \n", - "one can use to generically chain the intensity of any two profiles. Thus, it makes more sense to chain them using \n", - "the relative value from a previous search.\n", - "\n", - "We can customize how priors are passed from the results of a search and non-linear search by editing the\n", - " `prior_passer` settings in the `general.yaml` config file.\n", - "\n", - "__EXAMPLE__\n", - "\n", - "Lets go through an example using a real parameter. Lets say in search 1 we fit the lens galaxy's light with an \n", - "elliptical Sersic profile, and we estimate that its sersic index is equal to 4.0.\n", - "\n", - "To pass this as a prior to search 2 we write:\n", - "\n", - " lens.bulge.sersic_index = result_1.model_centred.lens.bulge.sersic_index\n", - "\n", - "The prior on the lens galaxy's bulge sersic index in search 2 would thus be a `GaussianPrior` with mean=4.0. \n", - "\n", - "The value of the Sersic index `width_modifier` in the priors config file sets sigma. The prior config file specifies \n", - "that we use an \"Absolute\" value of 0.8 to chain this prior. Thus, the `GaussianPrior` in search 2 would have a \n", - "mean=4.0 and sigma=0.8.\n", - "\n", - "If the prior config file had specified that we use an relative value of 0.8, the GaussianPrior in search 2 would have a \n", - "mean=4.0 and sigma = 4.0 * 0.8 = 3.2.\n", - "\n", - "And with that, we're done. Chaining priors is a bit of an art form, but one that works really well. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Chaining\n", + "==================\n", + "\n", + "Non-linear search chaining is an advanced model-fitting approach which breaks the model-fitting procedure down into\n", + "multiple non-linear searches, using the results of the initial searches to initialization parameter\n", + "sampling in subsequent searches. This contrasts the `modeling` examples which each compose and fit a single lens\n", + "model-fit using one non-linear search.\n", + "\n", + "The benefits of non-linear search chaining are:\n", + "\n", + " - Earlier searches fit simpler lens models than the later searches, which have a less complex non-linear parameter\n", + " space that can be sampled more efficiently, with a reduced chance of inferring an incorrect local maxima solution.\n", + "\n", + " - Earlier searches can use faster non-linear search settings which infer the highest log likelihood models but not\n", + " precisely quantify the parameter errors, with only the final searches using slow settings to robustly estimate errors.\n", + "\n", + " - Earlier searches can augment the data or alter the fitting-procedure in ways that speed up the computational run\n", + " time. These may impact the quality of the model-fit overall, but they can be reverted to the more accurate but more\n", + " computationally expense setting in the final searches.\n", + "\n", + "__Contents__\n", + "\n", + "- **Concise Model Composition API:** Chaining uses the concise `Model` API to compose lens models, which is nearly identical to the standard API but avoids the need to use `Model` objects to compose the lens model when a light or mass profile is passed to a `Collection` object.\n", + "- **This Example:** This script gives an overview of the API for search chaining, a description of how the priors on parameters are used to pass information between searches as well as tools for customizing prior passing.\n", + "- **Dataset + Masking:** Load, plot and mask the `Imaging` data.\n", + "- **Paths:** The path the results of all chained searches are output.\n", + "- **Model (Search 1):** We compose our lens model using `Model` objects, which represent the galaxies we fit to our data.\n", + "- **Search + Analysis + Model-Fit (Search 1):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "- **Result (Search 1):** The results which are used for prior passing are summarised in the `info` attribute.\n", + "- **Model Chaining:** We use the results of search 1 to create the `Model` components that we fit in search 2.\n", + "- **Model Centred Chaining:** We use the results of search 1 to create the `Model` components that we fit in search 2, using `model_centred` to pass priors centered on the maximum likelihood parameter values.\n", + "- **Search + Analysis + Model-Fit (Search 2):** We now create the non-linear search, analysis and perform the model-fit using the chained model.\n", + "- **Result (Search 2):** The final results of the chained model-fit.\n", + "- **How is Search Chaining Used?:** Overview of how search chaining is used in practice for automated lens modeling and SLaM pipelines.\n", + "- **Detailed Explanation Of Prior Passing:** Detailed overview of how prior passing works and tools to customize its behaviour.\n", + "- **EXAMPLE:** Lets go through an example using a real parameter.\n", + "\n", + "__Concise Model Composition API__\n", + "\n", + "Chaining uses the concise `Model` API to compose lens models, which is nearly identical to\n", + "the standard API but avoids the need to use `Model` objects to compose the lens model when a light or mass\n", + "profile is passed to a `Collection` object.\n", + "\n", + "__This Example__\n", + "\n", + "This script gives an overview of the API for search chaining, a description of how the priors on parameters are used\n", + "to pass information between searches as well as tools for customizing prior passing.\n", + "\n", + "There are examples throughout the workspace where search chaining improves and helps automate lens modeling." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__ \n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__\n", + "\n", + "The path the results of all chained searches are output:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"imaging\") / \"chaining\" / \"start_here\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In the first\n", + "search our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", + " \n", + " - an MGE with 1 x 20 Gaussians for the source galaxy's light [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal)\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", + "provides a reasonably accurate lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__start_here\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 1)__\n", + "\n", + "The results which are used for prior passing are summarised in the `info` attribute." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Chaining__\n", + "\n", + "We use the results of search 1 to create the `Model` components that we fit in search 2.\n", + "\n", + "The term `model` below passes the lens and source models as model-components that are to be fitted\n", + "for by the non-linear search. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_2 = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=result_1.model.galaxies.lens, source=result_1.model.galaxies.source\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The priors on the model components are the same as their original priors in `model_1`, that is the priors loaded\n", + "from the default configuration files. \n", + "\n", + "This makes model and search chaining pointless, as the second fit has the same initialization to sample parameter \n", + "space as the first.\n", + "\n", + "We instead want the model to be passed in a way where the priors are updated to reflect the inferred high likelihood \n", + "regions of the first in the first model. To do this we use the `model_centred` attribute of the result, which passes \n", + "the model components with priors that are centered on the maximum likelihood parameter values of the first search.\n", + "\n", + "__Model Centred Chaining__\n", + "\n", + "We use the results of search 1 to create the `Model` components that we fit in search 2.\n", + "\n", + "The term `model_centred` below passes the lens and source models as model-components that are to be fitted\n", + "for by the non-linear search. In other chaining examples, we'll see other ways to pass prior results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_2 = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=result_1.model_centred.galaxies.lens,\n", + " source=result_1.model_centred.galaxies.source,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`.\n", + "\n", + "This now contains `GaussianPrior`'s that are centered on the maximum likelihood parameter values of the first search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 2)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the `model.info` file of the search 2 model-fit to ensure the priors were passed correctly, as \n", + "well as the checkout the results to ensure an accurate power-law mass model is inferred." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__start_here\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 2)__\n", + "\n", + "The final results can be summarised via printing `info`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We will expand on this API in the following tutorials. The main thing to note is that we can pass entire profiles or\n", + "galaxies using prior passing, if their model does not change. \n", + "\n", + "The API to pass a whole profile or galaxy is as follows:\n", + " \n", + " bulge = result_1.model_centred.galaxies.lens.bulge\n", + " lens = result_1.model_centred.galaxies.lens\n", + " source = result_1.model_centred.galaxies.source\n", + " \n", + "We can also pass priors using an `instance` instead of a `model_centred`. When an `instance` is used, the maximum likelihood\n", + "parameter values are passed as fixed values that are therefore not fitted for by the non-linear search (reducing its\n", + "dimensionality). \n", + "\n", + "We will use this in other examples to \"split up\" the components of the model we fit, for example fit the lens light, \n", + "and then fix it to the best-fit model in a second search which fits the mass and source.\n", + " \n", + "__How is Search Chaining Used?__\n", + " \n", + "For a lot of the history of **PyAutoLens***, the passing of priors illustrated in this tutorial has been a key\n", + "aspect of how search chaining is used to model strong lenses.\n", + "\n", + "However, since ~2024, the direct passing of priors between searches has become less important. This is because\n", + "the non-linear search (Nautilus) is a lot more robust and run times are a lot faster since the inclusion of JAX. Thus,\n", + "the need to change priors via `model_centred` has become less important, but can still offer more reliable results\n", + "and efficient lens modeling.\n", + "\n", + "In fact, most example chaining examples and the Source, Light and Mass pipeline which are built are this feature\n", + "**no longer use prior passing**, always using the `model` attribute to pass models between searches which resets\n", + "the prior on parameters to their default priors.\n", + "\n", + "However, search chaining is still integral to automated lens modeling and the SLaM pipelines, but it is now used\n", + "to:\n", + "\n", + "- Split up the components of the lens model to be fitted in different searches (e.g. fit lens light and the \n", + " mass + source in different searches using `instance`).\n", + "\n", + "- Fit simpler models (e.g. MGE source) in earlier searches to get a fast estimate of the lens light and mass models,\n", + " which are then used as fixed components in later searches.\n", + " \n", + "- To pass `adapt_images` through the pipelines, which use the results of earlier searches to adapt aspects of the\n", + " lens model to the data in later searches. This is especially important for pixelized sources and described in the \n", + " `features/pixelization/adaptive` example.\n", + "\n", + "The most robust automated lens modeling is therefore still built using search chaining, but the direct passing of priors\n", + "between searches is no longer the main tool by which this is achieved.\n", + " \n", + "__Detailed Explanation Of Prior Passing__\n", + " \n", + "Lets now think about how priors are passed. Checkout the `model.info` file of the second search of this tutorial. The \n", + "parameters do not use the default priors we saw in search 1 (which are typically broad UniformPriors). Instead, \n", + "they use TruncatedGaussianPrior`s where:\n", + "\n", + " - The mean values are the median PDF results of every parameter in search 1.\n", + " - The sigma values are specified in the `width_modifier` field of the profile's entry in the `priors.yaml' config \n", + " file (we will discuss why this is used in a moment).\n", + "\n", + "Like the manual `TruncatedGaussianPrior`'s that were used in tutorial 1, the prior passing API sets up the prior on each \n", + "parameter with a `TruncatedGaussianPrior` centred on the high likelihood regions of parameter space!\n", + "\n", + "To end, I provide a detailed overview of how prior passing works and illustrate tools that can be used to customize\n", + "its behaviour. It is up to you whether you want read this, or go ahead to the next tutorial!\n", + "\n", + "Lets say I chain two parameters as follows:\n", + "\n", + " `mass.einstein_radius = result_1.model_centred.galaxies.lens.mass.einstein_radius`\n", + "\n", + "By invoking the `model_centred` attribute, the prior is passed following 3 rules:\n", + "\n", + " 1) The new parameter, in this case the einstein radius, uses a `TruncatedGaussianPrior`.This is ideal, as the 1D pdf results \n", + " we compute at the end of a search are easily summarised as a Gaussian.\n", + "\n", + " 2) The mean of the `TruncatedGaussianPrior` is the median PDF value of the parameter estimated in search 1.\n", + "\n", + " This ensures that the initial sampling of the new search's non-linear starts by searching the region of non-linear \n", + " parameter space that correspond to highest log likelihood solutions in the previous search. Our priors therefore \n", + " correspond to the `correct` regions of parameter space.\n", + "\n", + " 3) The sigma of the Gaussian uses the value specified for the profile in the `config/priors/*.yaml` config file's \n", + " `width_modifer` field (check these files out now).\n", + "\n", + "The idea here is simple. We want a value of sigma that gives a `TruncatedGaussianPrior` wide enough to search a broad \n", + "region of parameter space, so that the lens model can change if a better solution is nearby. However, we want it \n", + "to be narrow enough that we don't search too much of parameter space, as this will be slow or risk leading us \n", + "into an incorrect solution! \n", + "\n", + "The `width_modifier` values in the priors config file have been chosen based on our experience as being a good\n", + "balance broadly sampling parameter space but not being so narrow important solutions are missed.\n", + "\n", + "There are two ways a value is specified using the priors/width file:\n", + "\n", + " 1) Absolute: In this case, the error assumed on the parameter is the value given in the config file. \n", + " For example, if for the width on centre_0 of a light profile, the width modifier reads \"Absolute\" with a value \n", + " 0.05. This means if the error on the parameter centre_0 was less than 0.05 in the previous search, the sigma of \n", + " its `GaussianPrior` in this search will be 0.05.\n", + "\n", + " 2) Relative: In this case, the error assumed on the parameter is the % of the value of the estimated value given in \n", + " the config file. For example, if the intensity estimated in the previous search was 2.0, and the relative error in \n", + " the config file reads \"Relative\" with a value 0.5, then the sigma of the `GaussianPrior` will be 50% of this \n", + " value, i.e. sigma = 0.5 * 2.0 = 1.0.\n", + "\n", + "We use absolute and relative values for different parameters, depending on their properties. For example, using the \n", + "relative value of a parameter like the `Profile` centre makes no sense. If our lens galaxy is centred at (0.0, 0.0), \n", + "the relative error will always be tiny and thus poorly defined. Therefore, the default configs in **PyAutoLens** use \n", + "absolute errors on the centre.\n", + "\n", + "However, there are parameters where using an absolute value does not make sense. Intensity is a good example of this. \n", + "The intensity of an image depends on its units, S/N, galaxy brightness, etc. There is no single absolute value that \n", + "one can use to generically chain the intensity of any two profiles. Thus, it makes more sense to chain them using \n", + "the relative value from a previous search.\n", + "\n", + "We can customize how priors are passed from the results of a search and non-linear search by editing the\n", + " `prior_passer` settings in the `general.yaml` config file.\n", + "\n", + "__EXAMPLE__\n", + "\n", + "Lets go through an example using a real parameter. Lets say in search 1 we fit the lens galaxy's light with an \n", + "elliptical Sersic profile, and we estimate that its sersic index is equal to 4.0.\n", + "\n", + "To pass this as a prior to search 2 we write:\n", + "\n", + " lens.bulge.sersic_index = result_1.model_centred.lens.bulge.sersic_index\n", + "\n", + "The prior on the lens galaxy's bulge sersic index in search 2 would thus be a `GaussianPrior` with mean=4.0. \n", + "\n", + "The value of the Sersic index `width_modifier` in the priors config file sets sigma. The prior config file specifies \n", + "that we use an \"Absolute\" value of 0.8 to chain this prior. Thus, the `GaussianPrior` in search 2 would have a \n", + "mean=4.0 and sigma=0.8.\n", + "\n", + "If the prior config file had specified that we use an relative value of 0.8, the GaussianPrior in search 2 would have a \n", + "mean=4.0 and sigma = 4.0 * 0.8 = 3.2.\n", + "\n", + "And with that, we're done. Chaining priors is a bit of an art form, but one that works really well. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/cookbook.ipynb b/notebooks/guides/modeling/cookbook.ipynb index b00a672b3..d72c04f44 100644 --- a/notebooks/guides/modeling/cookbook.ipynb +++ b/notebooks/guides/modeling/cookbook.ipynb @@ -1,621 +1,658 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Model Cookbook\n", - "==============\n", - "\n", - "The model cookbook provides a concise reference to lens model composition tools, specifically the `Model` and\n", - "`Collection` objects.\n", - "\n", - "Examples using different PyAutoLens API\u2019s for model composition are provided, which produce more concise and\n", - "readable code for different use-cases.\n", - "\n", - "__Contents__\n", - "\n", - "- **Simple Lens Model:** Compose a simple lens model with a lens galaxy and source galaxy.\n", - "- **More Complex Lens Models:** Extend the simple model to have multiple light or mass profiles and multiple galaxies.\n", - "- **Concise API:** Compose a lens model using the concise API, which is more readable and concise.\n", - "- **Prior Customization:** Customize the priors of individual lens model parameters using uniform, log-uniform and Gaussian priors.\n", - "- **Model Customization:** Customize the lens model parameters, including parameter pairing, fixing and offsets.\n", - "- **Redshift Free:** Make the redshift of a galaxy a free parameter in the model-fit.\n", - "- **Available Model Components:** List the available light profiles, mass profiles and other components that can be used for lens modeling.\n", - "\n", - "Advanced Features:\n", - "\n", - "**JSon Outputs:** Output a model to a .json file on hard-disk, which can be loaded and modified.\n", - "**Many Profile Models:** Compose and fit models with many light profiles, such as the Multi Gaussian Expansion (MGE) and shapelets.\n", - "**Model Linking:** Link the inferred model of one phase to the model in a non-linear search chain.\n", - "**Across Datasets:** Compose models where the same model component is used across multiple datasets, with certain parameters free to vary.\n", - "**Relations:** Compose models where the free parameter(s) vary according to a user-specified function.\n", - "**PyAutoFit API:** Use the PyAutoFit API to compose lens models in more advanced ways.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simple Lens Model__\n", - "\n", - "A simple lens model has a lens galaxy with a Sersic light profile, Isothermal mass profile and source galaxy with \n", - "a Sersic light profile:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - ")\n", - "\n", - "# Source:\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The redshifts in the above model are used to determine which galaxy is the lens and which is the source.\n", - "\n", - "The model `total_free_parameters` tells us the total number of free parameters (which are fitted for via a \n", - "non-linear search), which in this case is 19 (7 from the lens `Sersic`, 5 from the lens `Isothermal` and 7 from the \n", - "source `Sersic`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Model Total Free Parameters = {model.total_free_parameters}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we print the `info` attribute of the model we get information on all of the parameters and their priors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__More Complex Lens Models__\n", - "\n", - "The API above can be easily extended to compose lens models where each galaxy has multiple light or mass profiles:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "disk = af.Model(al.lp_linear.Exponential)\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " disk=disk,\n", - " mass=mass,\n", - " shear=shear,\n", - ")\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "disk = af.Model(al.lp_linear.ExponentialCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=disk)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The use of the words `bulge`, `disk`, `mass` and `shear` above are arbitrary. They can be replaced with any name you\n", - "like, e.g. `bulge_0`, `bulge_1`, `mass_0`, `mass_1`, and the model will still behave in the same way.\n", - "\n", - "The API can also be extended to compose lens models where there are multiple galaxies:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - ")\n", - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens_1 = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - ")\n", - "\n", - "# Source 0:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source_0 = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Source 1 :\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source_1 = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens_0=lens_0, lens_1=lens_1, source_0=source_0, source_1=source_1\n", - " ),\n", - ")\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The above lens model consists of only two planes (an image-plane and source-plane), but has four galaxies in total.\n", - "This is because the lens galaxies have the same redshift and the source galaxies have the same redshift.\n", - "\n", - "If we gave one of the lens galaxies a different redshift, it would be included in a third plane, and the model would\n", - "perform multi-plane ray tracing when the model-fit is performed.\n", - "\n", - "__Concise API__\n", - "\n", - "If a light or mass profile is passed directly to the `af.Model` of a galaxy, it is automatically assigned to be a\n", - "`af.Model` component of the galaxy.\n", - "\n", - "This means we can write the model above comprising multiple light and mass profiles more concisely as follows (also\n", - "removing the comments reading Lens / Source / Overall Lens Model to make the code more readable):" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.Sersic,\n", - " disk=al.lp_linear.Sersic,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - ")\n", - "\n", - "source = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore,\n", - " disk=al.lp_linear.ExponentialCore,\n", - ")\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Prior Customization__\n", - "\n", - "We can customize the priors of the lens model component individual parameters, using the following three types\n", - "of priors:\n", - "\n", - "`UniformPrior`: The values of a parameter are randomly drawn between a `lower_limit` and `upper_limit`. For example,\n", - " the effective radius of ellipitical Sersic profiles typically assumes a uniform prior between 0.0\" and 30.0\".\n", - "\n", - "`LogUniformPrior`: Like a `UniformPrior` this randomly draws values between a `limit_limit` and `upper_limit`, but the\n", - " values are drawn from a distribution with base 10. This is used for the `intensity` of a light profile, as the\n", - " luminosity of galaxies follows a log10 distribution.\n", - "\n", - "`GaussianPrior`: The values of a parameter are randomly drawn from a Gaussian distribution with a `mean` and width\n", - " `sigma`. For example, the $y$ and $x$ centre values in a light profile typically assume a mean of 0.0\" and a\n", - " sigma of 0.3\", indicating that we expect the profile centre to be located near the centre of the image.\n", - " " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "bulge.sersic_index = af.TruncatedGaussianPrior(\n", - " mean=4.0, sigma=1.0, lower_limit=1.0, upper_limit=8.0\n", - ")\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "mass.centre.centre_0 = af.TruncatedGaussianPrior(\n", - " mean=0.0, sigma=0.1, lower_limit=-0.5, upper_limit=0.5\n", - ")\n", - "mass.centre.centre_1 = af.TruncatedGaussianPrior(\n", - " mean=0.0, sigma=0.1, lower_limit=-0.5, upper_limit=0.5\n", - ")\n", - "mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=8.0)\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - ")\n", - "\n", - "# Source\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "source.effective_radius = af.TruncatedGaussianPrior(\n", - " mean=0.1, sigma=0.05, lower_limit=0.0, upper_limit=1.0\n", - ")\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Customization__\n", - "\n", - "We can customize the lens model parameters in a number of different ways, as shown below:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "disk = af.Model(al.lp_linear.Exponential)\n", - "\n", - "# Parameter Pairing: Pair the centre of the bulge and disk together, reducing\n", - "# the complexity of non-linear parameter space by N = 2\n", - "\n", - "bulge.centre = disk.centre\n", - "\n", - "# Parameter Fixing: Fix the sersic_index of the bulge to a value of 4, reducing\n", - "# the complexity of non-linear parameter space by N = 1\n", - "\n", - "bulge.sersic_index = 4.0\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "# Parameter Offsets: Make the mass model centre parameters the same value as\n", - "# the bulge / disk but with an offset.\n", - "\n", - "mass.centre.centre_0 = bulge.centre.centre_0 + 0.1\n", - "mass.centre.centre_1 = bulge.centre.centre_1 + 0.1\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " disk=disk,\n", - " mass=mass,\n", - " shear=shear,\n", - ")\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "disk = af.Model(al.lp_linear.ExponentialCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=disk)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "# Assert that the effective radius of the bulge is larger than that of the disk.\n", - "# (Assertions can only be added at the end of model composition, after all components\n", - "# have been brought together in a `Collection`.\n", - "model.add_assertion(\n", - " model.galaxies.lens.bulge.effective_radius\n", - " > model.galaxies.lens.disk.effective_radius\n", - ")\n", - "\n", - "# Assert that the Einstein Radius is below 3.0\":\n", - "model.add_assertion(model.galaxies.lens.mass.einstein_radius < 3.0)\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshift Free__\n", - "\n", - "The redshift of a galaxy can be treated as a free parameter in the model-fit by using the following API:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift = af.Model(al.Redshift)\n", - "redshift.redshift = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=redshift, mass=al.mp.Isothermal)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The model-fit will automatically enable multi-plane ray tracing and alter the ordering of the planes depending on the\n", - "redshifts of the galaxies.\n", - "\n", - "NOTE: For strong lenses with just two planes (an image-plane and source-plane) the redshifts of the galaxies do not\n", - "impact the model-fit. You should therefore never make the redshifts free if you are only modeling a two-plane lens\n", - "system. This is because lensing calculations can be defined in arc-second coordinates, which do not change as a\n", - "function of redshift.\n", - "\n", - "Redshifts should be made free when modeling three or more planes, as the mulit-plane ray-tracing calculations have an\n", - "obvious dependence on the redshifts of the galaxies which could be inferred by the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Available Model Components__\n", - "\n", - "The light profiles, mass profiles and other components that can be used for lens modeling are given at the following\n", - "API documentation pages:\n", - "\n", - " - https://pyautolens.readthedocs.io/en/latest/api/light.html\n", - " - https://pyautolens.readthedocs.io/en/latest/api/mass.html\n", - " - https://pyautolens.readthedocs.io/en/latest/api/pixelization.html\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__JSon Outputs__\n", - "\n", - "After a model is composed, it can easily be output to a .json file on hard-disk in a readable structure:\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import os\n", - "import json\n", - "\n", - "model_path = Path(\"path\", \"to\", \"model\", \"json\")\n", - "\n", - "os.makedirs(model_path, exist_ok=True)\n", - "\n", - "model_file = Path(model_path, \"model.json\")\n", - "\n", - "with open(model_file, \"w+\") as f:\n", - " json.dump(model.dict(), f, indent=4)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can load the model from its `.json` file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model = af.Model.from_json(file=model_file)\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This means in **PyAutoLens** one can write a model in a script, save it to hard disk and load it elsewhere, as well\n", - "as manually customize it in the .json file directory.\n", - "\n", - "This is used for composing complex models of group scale lenses.\n", - "\n", - "__Many Profile Models (Advanced)__\n", - "\n", - "Features such as the Multi Gaussian Expansion (MGE) and shapelets compose models consisting of 50 - 500+ light\n", - "profiles.\n", - "\n", - "The following example notebooks show how to compose and fit these models:\n", - "\n", - "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/modeling/features/multi_gaussian_expansion.ipynb\n", - "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/modeling/features/shapelets.ipynb\n", - "\n", - "__Model Linking (Advanced)__\n", - "\n", - "When performing non-linear search chaining, the inferred model of one phase can be linked to the model.\n", - "\n", - "The following example notebooks show how to compose and fit these models:\n", - "\n", - "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/imaging/advanced/guides/modeling/chaining.ipynb\n", - "\n", - "__Across Datasets (Advanced)__\n", - "\n", - "When fitting multiple datasets, model can be composed where the same model component are used across the datasets\n", - "but certain parameters are free to vary across the datasets.\n", - "\n", - "The following example notebooks show how to compose and fit these models:\n", - "\n", - "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/multi/start_here.ipynb\n", - "\n", - "__Relations (Advanced)__\n", - "\n", - "We can compose models where the free parameter(s) vary according to a user-specified function \n", - "(e.g. y = mx +c -> effective_radius = (m * wavelength) + c across the datasets.\n", - "\n", - "The following example notebooks show how to compose and fit these models:\n", - "\n", - "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/multi/features/wavelength_dependence/modeling.ipynb\n", - "\n", - "__PyAutoFit API__\n", - "\n", - "**PyAutoFit** is a general model composition library which offers even more ways to compose lens models not\n", - "detailed in this cookbook.\n", - "\n", - "The **PyAutoFit** model composition cookbooks detail this API in more detail:\n", - "\n", - "https://pyautofit.readthedocs.io/en/latest/cookbooks/model.html\n", - "https://pyautofit.readthedocs.io/en/latest/cookbooks/multi_level_model.html\n", - "\n", - "__Wrap Up__\n", - "\n", - "This cookbook shows how to compose simple lens models using the `af.Model()` and `af.Collection()` objects." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Model Cookbook\n", + "==============\n", + "\n", + "The model cookbook provides a concise reference to lens model composition tools, specifically the `Model` and\n", + "`Collection` objects.\n", + "\n", + "Examples using different PyAutoLens API\u2019s for model composition are provided, which produce more concise and\n", + "readable code for different use-cases.\n", + "\n", + "__Contents__\n", + "\n", + "- **Simple Lens Model:** Compose a simple lens model with a lens galaxy and source galaxy.\n", + "- **More Complex Lens Models:** Extend the simple model to have multiple light or mass profiles and multiple galaxies.\n", + "- **Concise API:** Compose a lens model using the concise API, which is more readable and concise.\n", + "- **Prior Customization:** Customize the priors of individual lens model parameters using uniform, log-uniform and Gaussian priors.\n", + "- **Model Customization:** Customize the lens model parameters, including parameter pairing, fixing and offsets.\n", + "- **Redshift Free:** Make the redshift of a galaxy a free parameter in the model-fit.\n", + "- **Available Model Components:** List the available light profiles, mass profiles and other components that can be used for lens modeling.\n", + "\n", + "Advanced Features:\n", + "\n", + "**JSon Outputs:** Output a model to a .json file on hard-disk, which can be loaded and modified.\n", + "**Many Profile Models:** Compose and fit models with many light profiles, such as the Multi Gaussian Expansion (MGE) and shapelets.\n", + "**Model Linking:** Link the inferred model of one phase to the model in a non-linear search chain.\n", + "**Across Datasets:** Compose models where the same model component is used across multiple datasets, with certain parameters free to vary.\n", + "**Relations:** Compose models where the free parameter(s) vary according to a user-specified function.\n", + "**PyAutoFit API:** Use the PyAutoFit API to compose lens models in more advanced ways.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simple Lens Model__\n", + "\n", + "A simple lens model has a lens galaxy with a Sersic light profile, Isothermal mass profile and source galaxy with \n", + "a Sersic light profile:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + ")\n", + "\n", + "# Source:\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The redshifts in the above model are used to determine which galaxy is the lens and which is the source.\n", + "\n", + "The model `total_free_parameters` tells us the total number of free parameters (which are fitted for via a \n", + "non-linear search), which in this case is 19 (7 from the lens `Sersic`, 5 from the lens `Isothermal` and 7 from the \n", + "source `Sersic`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Model Total Free Parameters = {model.total_free_parameters}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we print the `info` attribute of the model we get information on all of the parameters and their priors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__More Complex Lens Models__\n", + "\n", + "The API above can be easily extended to compose lens models where each galaxy has multiple light or mass profiles:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "disk = af.Model(al.lp_linear.Exponential)\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " disk=disk,\n", + " mass=mass,\n", + " shear=shear,\n", + ")\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "disk = af.Model(al.lp_linear.ExponentialCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=disk)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The use of the words `bulge`, `disk`, `mass` and `shear` above are arbitrary. They can be replaced with any name you\n", + "like, e.g. `bulge_0`, `bulge_1`, `mass_0`, `mass_1`, and the model will still behave in the same way.\n", + "\n", + "The API can also be extended to compose lens models where there are multiple galaxies:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + ")\n", + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens_1 = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + ")\n", + "\n", + "# Source 0:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source_0 = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Source 1 :\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source_1 = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens_0=lens_0, lens_1=lens_1, source_0=source_0, source_1=source_1\n", + " ),\n", + ")\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The above lens model consists of only two planes (an image-plane and source-plane), but has four galaxies in total.\n", + "This is because the lens galaxies have the same redshift and the source galaxies have the same redshift.\n", + "\n", + "If we gave one of the lens galaxies a different redshift, it would be included in a third plane, and the model would\n", + "perform multi-plane ray tracing when the model-fit is performed.\n", + "\n", + "__Concise API__\n", + "\n", + "If a light or mass profile is passed directly to the `af.Model` of a galaxy, it is automatically assigned to be a\n", + "`af.Model` component of the galaxy.\n", + "\n", + "This means we can write the model above comprising multiple light and mass profiles more concisely as follows (also\n", + "removing the comments reading Lens / Source / Overall Lens Model to make the code more readable):" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.Sersic,\n", + " disk=al.lp_linear.Sersic,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + ")\n", + "\n", + "source = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore,\n", + " disk=al.lp_linear.ExponentialCore,\n", + ")\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Prior Customization__\n", + "\n", + "We can customize the priors of the lens model component individual parameters, using the following three types\n", + "of priors:\n", + "\n", + "`UniformPrior`: The values of a parameter are randomly drawn between a `lower_limit` and `upper_limit`. For example,\n", + " the effective radius of ellipitical Sersic profiles typically assumes a uniform prior between 0.0\" and 30.0\".\n", + "\n", + "`LogUniformPrior`: Like a `UniformPrior` this randomly draws values between a `limit_limit` and `upper_limit`, but the\n", + " values are drawn from a distribution with base 10. This is used for the `intensity` of a light profile, as the\n", + " luminosity of galaxies follows a log10 distribution.\n", + "\n", + "`GaussianPrior`: The values of a parameter are randomly drawn from a Gaussian distribution with a `mean` and width\n", + " `sigma`. For example, the $y$ and $x$ centre values in a light profile typically assume a mean of 0.0\" and a\n", + " sigma of 0.3\", indicating that we expect the profile centre to be located near the centre of the image.\n", + " " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "bulge.sersic_index = af.TruncatedGaussianPrior(\n", + " mean=4.0, sigma=1.0, lower_limit=1.0, upper_limit=8.0\n", + ")\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "mass.centre.centre_0 = af.TruncatedGaussianPrior(\n", + " mean=0.0, sigma=0.1, lower_limit=-0.5, upper_limit=0.5\n", + ")\n", + "mass.centre.centre_1 = af.TruncatedGaussianPrior(\n", + " mean=0.0, sigma=0.1, lower_limit=-0.5, upper_limit=0.5\n", + ")\n", + "mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=8.0)\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + ")\n", + "\n", + "# Source\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "source.effective_radius = af.TruncatedGaussianPrior(\n", + " mean=0.1, sigma=0.05, lower_limit=0.0, upper_limit=1.0\n", + ")\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Customization__\n", + "\n", + "We can customize the lens model parameters in a number of different ways, as shown below:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "disk = af.Model(al.lp_linear.Exponential)\n", + "\n", + "# Parameter Pairing: Pair the centre of the bulge and disk together, reducing\n", + "# the complexity of non-linear parameter space by N = 2\n", + "\n", + "bulge.centre = disk.centre\n", + "\n", + "# Parameter Fixing: Fix the sersic_index of the bulge to a value of 4, reducing\n", + "# the complexity of non-linear parameter space by N = 1\n", + "\n", + "bulge.sersic_index = 4.0\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "# Parameter Offsets: Make the mass model centre parameters the same value as\n", + "# the bulge / disk but with an offset.\n", + "\n", + "mass.centre.centre_0 = bulge.centre.centre_0 + 0.1\n", + "mass.centre.centre_1 = bulge.centre.centre_1 + 0.1\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " disk=disk,\n", + " mass=mass,\n", + " shear=shear,\n", + ")\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "disk = af.Model(al.lp_linear.ExponentialCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=disk)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "# Assert that the effective radius of the bulge is larger than that of the disk.\n", + "# (Assertions can only be added at the end of model composition, after all components\n", + "# have been brought together in a `Collection`.\n", + "model.add_assertion(\n", + " model.galaxies.lens.bulge.effective_radius\n", + " > model.galaxies.lens.disk.effective_radius\n", + ")\n", + "\n", + "# Assert that the Einstein Radius is below 3.0\":\n", + "model.add_assertion(model.galaxies.lens.mass.einstein_radius < 3.0)\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshift Free__\n", + "\n", + "The redshift of a galaxy can be treated as a free parameter in the model-fit by using the following API:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift = af.Model(al.Redshift)\n", + "redshift.redshift = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=redshift, mass=al.mp.Isothermal)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The model-fit will automatically enable multi-plane ray tracing and alter the ordering of the planes depending on the\n", + "redshifts of the galaxies.\n", + "\n", + "NOTE: For strong lenses with just two planes (an image-plane and source-plane) the redshifts of the galaxies do not\n", + "impact the model-fit. You should therefore never make the redshifts free if you are only modeling a two-plane lens\n", + "system. This is because lensing calculations can be defined in arc-second coordinates, which do not change as a\n", + "function of redshift.\n", + "\n", + "Redshifts should be made free when modeling three or more planes, as the mulit-plane ray-tracing calculations have an\n", + "obvious dependence on the redshifts of the galaxies which could be inferred by the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Available Model Components__\n", + "\n", + "The light profiles, mass profiles and other components that can be used for lens modeling are given at the following\n", + "API documentation pages:\n", + "\n", + " - https://pyautolens.readthedocs.io/en/latest/api/light.html\n", + " - https://pyautolens.readthedocs.io/en/latest/api/mass.html\n", + " - https://pyautolens.readthedocs.io/en/latest/api/pixelization.html\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__JSon Outputs__\n", + "\n", + "After a model is composed, it can easily be output to a .json file on hard-disk in a readable structure:\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "import json\n", + "\n", + "model_path = Path(\"path\", \"to\", \"model\", \"json\")\n", + "\n", + "os.makedirs(model_path, exist_ok=True)\n", + "\n", + "model_file = Path(model_path, \"model.json\")\n", + "\n", + "with open(model_file, \"w+\") as f:\n", + " json.dump(model.dict(), f, indent=4)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can load the model from its `.json` file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model = af.Model.from_json(file=model_file)\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This means in **PyAutoLens** one can write a model in a script, save it to hard disk and load it elsewhere, as well\n", + "as manually customize it in the .json file directory.\n", + "\n", + "This is used for composing complex models of group scale lenses.\n", + "\n", + "__Many Profile Models (Advanced)__\n", + "\n", + "Features such as the Multi Gaussian Expansion (MGE) and shapelets compose models consisting of 50 - 500+ light\n", + "profiles.\n", + "\n", + "The following example notebooks show how to compose and fit these models:\n", + "\n", + "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/modeling/features/multi_gaussian_expansion.ipynb\n", + "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/modeling/features/shapelets.ipynb\n", + "\n", + "__Model Linking (Advanced)__\n", + "\n", + "When performing non-linear search chaining, the inferred model of one phase can be linked to the model.\n", + "\n", + "The following example notebooks show how to compose and fit these models:\n", + "\n", + "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/imaging/advanced/guides/modeling/chaining.ipynb\n", + "\n", + "__Across Datasets (Advanced)__\n", + "\n", + "When fitting multiple datasets, model can be composed where the same model component are used across the datasets\n", + "but certain parameters are free to vary across the datasets.\n", + "\n", + "The following example notebooks show how to compose and fit these models:\n", + "\n", + "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/multi/start_here.ipynb\n", + "\n", + "__Relations (Advanced)__\n", + "\n", + "We can compose models where the free parameter(s) vary according to a user-specified function \n", + "(e.g. y = mx +c -> effective_radius = (m * wavelength) + c across the datasets.\n", + "\n", + "The following example notebooks show how to compose and fit these models:\n", + "\n", + "https://github.com/PyAutoLabs/autolens_workspace/blob/main/notebooks/multi/features/wavelength_dependence/modeling.ipynb\n", + "\n", + "__PyAutoFit API__\n", + "\n", + "**PyAutoFit** is a general model composition library which offers even more ways to compose lens models not\n", + "detailed in this cookbook.\n", + "\n", + "The **PyAutoFit** model composition cookbooks detail this API in more detail:\n", + "\n", + "https://pyautofit.readthedocs.io/en/latest/cookbooks/model.html\n", + "https://pyautofit.readthedocs.io/en/latest/cookbooks/multi_level_model.html\n", + "\n", + "__Wrap Up__\n", + "\n", + "This cookbook shows how to compose simple lens models using the `af.Model()` and `af.Collection()` objects." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/customize.ipynb b/notebooks/guides/modeling/customize.ipynb index f364a14cb..c07ea4a10 100644 --- a/notebooks/guides/modeling/customize.ipynb +++ b/notebooks/guides/modeling/customize.ipynb @@ -1,440 +1,477 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Customize\n", - "===================\n", - "\n", - "This script gives a run through of all the different ways the analysis can be customized for lens modeling, with\n", - "reasons explaining why each customization is useful.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Positions:** Before fitting a strong lens, we can manually specify a grid of image-plane coordinates.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "All customizations in this script are applied to the strong lens dataset `simple__no_lens_light`, which is a\n", - "simple strong lens with no lens light emission.\n", - "\n", - "We therefore load and plot the strong lens dataset `simple__no_lens_light` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "if not (dataset_path / \"positions.json\").exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "All example default scritps use a circular mask to model a lens, which contains the lens and source emission.\n", - "However, the mask can be customized to better suit the lens and source emission, for example by using an annular\n", - "mask to only contain the emission of the Einstein ring itelf.\n", - "\n", - "Advantages: Strong lenses with complex and difficult-to-subtract foreground lens galaxies can leave residuals that\n", - "bias the mass and source models, which this custom mask can remove from the model-fit. The custom mask can also provide\n", - "faster run times, as the removal of large large regions of the image (which contain no signal) no longer need to be\n", - "processed and fitted.\n", - "\n", - "Disadvantages: Pixels containing no source emission may still constrain the lens model, if a mass model incorrectly\n", - "predicts that flux will appear in these image pixels. By using a custom mask, the model-fit will not be penalized for\n", - "incorrectly predicting flux in these image-pixels (As the mask has removed them from the fit).\n", - "\n", - "We first show an example using an annular masks, which because the data does not contain lens light can be\n", - "used to remove the central pixels containing no emission and thus only fit the source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular_annular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " inner_radius=0.5,\n", - " outer_radius=2.5,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask) # <----- The custom mask is used here!\n", - "\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also load the mask from a .fits file, which could have been produced in a way which is even more customized\n", - "to the source emission than the annular masks above.. \n", - "\n", - "To create the .fits file of a mask, we use a GUI tool which is described in the following script:\n", - "\n", - " `autolens_workspace/*/imaging/data_preparation/gui/mask.py`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " mask = al.Mask2D.from_fits(\n", - " file_path=Path(dataset_path, \"mask_gui.fits\"),\n", - " hdu=0,\n", - " pixel_scales=dataset.pixel_scales,\n", - " )\n", - "except FileNotFoundError:\n", - " mask = al.Mask2D.circular_annular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " inner_radius=0.5,\n", - " outer_radius=2.5,\n", - " )\n", - "\n", - "dataset = dataset.apply_mask(mask=mask) # <----- The custom mask is used here!\n", - "\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", - "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of\n", - "over sampling to ensure the lensed images are evaluated accurately.\n", - "\n", - "This is why throughout the workspace the cored Sersic profile is used, instead of the regular Sersic profile which\n", - "you may be more familiar with from the literature. In this example we will increase the over sampling level and\n", - "therefore fit a regular Sersic profile to the data, instead of a cored Sersic profile.\n", - "\n", - "This example demonstrates how to change the over sampling used to compute the surface brightness of every image-pixel,\n", - "whereby a higher sub-grid resolution better oversamples the image of the light profile so as to provide a more accurate\n", - "model of its image.\n", - "\n", - "**Benefit**: Higher level of over sampling provide a more accurate estimate of the surface brightness in every image-pixel.\n", - "**Downside**: Higher levels of over sampling require longer calculations and higher memory usage.\n", - "\n", - "Over sampling is applied separately to the light profiles which compute the surface brightness of the lens galaxy,\n", - "which are on a `uniform` grid, and the light profiles which compute the surface brightness of the source galaxy,\n", - "which are on a `non-uniform` grid.\n", - "\n", - "Prequisites: You should read `autolens_workspace/*/guides/advanced/over_sampling.ipynb` before running this script, which\n", - "introduces the concept of over sampling in PyAutoLens and explains why the lens and source galaxy are evaluated\n", - "on different grids.\n", - "\n", - "\n", - "\n", - "The over sampling used to fit the data is customized using the `apply_over_sampling` method, which you may have\n", - "seen in example `modeling` scripts.\n", - "\n", - "To apply uniform over sampling of degree 4x4, we simply input the integer 4.\n", - "\n", - "The grid this is applied to is called `lp`, to indicate that it is the grid used to evaluate the emission of light\n", - "profiles for which this over sampling scheme is applied." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_over_sampling(over_sample_size_lp=4)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Above, the `over_sample_size` input has been an integer, however it can also be an `ndarray` of values corresponding\n", - "to each pixel. \n", - "\n", - "We create an `ndarray` of values which are high in the centre, but reduce to 2 at the outskirts, therefore \n", - "providing high levels of over sampling where we need it whilst using lower values which are computationally fast to \n", - "evaluate at the outskirts.\n", - "\n", - "Specifically, we define a 24 x 24 sub-grid within the central 0.3\" of pixels, uses a 8 x 8 grid between\n", - "0.3\" and 0.6\" and a 2 x 2 grid beyond that. \n", - "\n", - "This will provide high levels of over sampling for the lens galaxy, whose emission peaks at the centre of the\n", - "image near (0.0\", 0.0\"), but will not produce high levels of over sampling for the lensed source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[24, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By assuming the lens galaxy is near (0.0\", 0.0\"), it was simple to set up an adaptive over sampling grid which is\n", - "applicable to all strong lens dataset.\n", - "\n", - "There is no analogous define an adaptive over sampling grid for the lensed source, because for every dataset the\n", - "source's light will appear in different regions of the image plane.\n", - "\n", - "This is why the majority of workspace examples use cored light profiles for the source galaxy. A cored light profile\n", - "does not rapidly change in its central regions, and therefore can be evaluated accurately without over-sampling.\n", - "\n", - "There is a way to set up an adaptive over sampling grid for a lensed source, however it requries one to use and\n", - "understanding the advanced lens modeling feature search chaining.\n", - "\n", - "An example of how to use search chaining to over sample sources efficient is provided in \n", - "the `autolens_workspace/*/guides/modeling/chaining/over_sampling.ipynb` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Positions__\n", - "\n", - "Before fitting a strong lens, we can manually specify a grid of image-plane coordinates corresponding to the multiple\n", - "images of the lensed source-galaxy(s). During the model-fit, **PyAutoLens** will check that these coordinates trace\n", - "within a specified arc-second threshold of one another in the source-plane. If they do not meet this threshold, the\n", - "mass model is discarded and a new sample is generated by the non-linear search.\n", - "\n", - "Advantages: The model-fit is faster, as the non-linear search avoids regions of parameter space where the mass-model\n", - "is clearly not accurate. Removing these unphysical solutions may also mean that the global-maximum solution is inferred\n", - "instead of a local-maxima, given that removing unphysical mass models makes non-linear parameter space less complex.\n", - "\n", - "Disadvantages: If the positions are inaccurate or threshold is set too low, one may inadvertantly remove the correct\n", - "mass model!\n", - "\n", - "The positions are associated with the and they are loaded from a `positions.json` file which is in the same folder as \n", - "the dataset itself. \n", - "\n", - "To create this file, we used a GUI to `draw on` the positions with our mouse. This GUI can be found in the script:\n", - "\n", - " `autolens_workspace/*/imaging/data_preparation/gui/positions.py`\n", - "\n", - "If you wish to use positions for modeling your own lens data, you should use this script to draw on the positions of\n", - "every lens in you dataset.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Alternatively, the positions can be specified manually in the modeling script script using the `Grid2DIrregular`object.\n", - "\n", - "Below, we specify a list of (y,x) coordinates (that are not on a uniform or regular grid) which correspond to the \n", - "arc-second (y,x) coordinates ot he lensed source's brightest pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(\n", - " [(0.4, 1.6), (1.58, -0.35), (-0.43, -1.59), (-1.45, 0.2)]\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To use positions in lens modeling, we pass the `AnalysisImaging` object a `PositionsLH` object, which includes the \n", - "positions we loaded above, alongside a `threshold`.\n", - "\n", - "The threshold of 0.3\" is large. For an accurate lens model we would anticipate the positions trace within < 0.01\" of\n", - "one another. The high threshold ensures only the initial mass models at the start of the fit are penalized." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The analysis receives a list of `PositionsLH` objects, which allows us to use multiple sets of positions to apply this \n", - "penalty, for example if there source has multiple sets of multiple images which we know how they map to one another." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset, positions_likelihood_list=[positions_likelihood], use_jax=True\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Customize\n", + "===================\n", + "\n", + "This script gives a run through of all the different ways the analysis can be customized for lens modeling, with\n", + "reasons explaining why each customization is useful.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Positions:** Before fitting a strong lens, we can manually specify a grid of image-plane coordinates.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "All customizations in this script are applied to the strong lens dataset `simple__no_lens_light`, which is a\n", + "simple strong lens with no lens light emission.\n", + "\n", + "We therefore load and plot the strong lens dataset `simple__no_lens_light` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "if not (dataset_path / \"positions.json\").exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "All example default scritps use a circular mask to model a lens, which contains the lens and source emission.\n", + "However, the mask can be customized to better suit the lens and source emission, for example by using an annular\n", + "mask to only contain the emission of the Einstein ring itelf.\n", + "\n", + "Advantages: Strong lenses with complex and difficult-to-subtract foreground lens galaxies can leave residuals that\n", + "bias the mass and source models, which this custom mask can remove from the model-fit. The custom mask can also provide\n", + "faster run times, as the removal of large large regions of the image (which contain no signal) no longer need to be\n", + "processed and fitted.\n", + "\n", + "Disadvantages: Pixels containing no source emission may still constrain the lens model, if a mass model incorrectly\n", + "predicts that flux will appear in these image pixels. By using a custom mask, the model-fit will not be penalized for\n", + "incorrectly predicting flux in these image-pixels (As the mask has removed them from the fit).\n", + "\n", + "We first show an example using an annular masks, which because the data does not contain lens light can be\n", + "used to remove the central pixels containing no emission and thus only fit the source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular_annular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " inner_radius=0.5,\n", + " outer_radius=2.5,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask) # <----- The custom mask is used here!\n", + "\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also load the mask from a .fits file, which could have been produced in a way which is even more customized\n", + "to the source emission than the annular masks above.. \n", + "\n", + "To create the .fits file of a mask, we use a GUI tool which is described in the following script:\n", + "\n", + " `autolens_workspace/*/imaging/data_preparation/gui/mask.py`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " mask = al.Mask2D.from_fits(\n", + " file_path=Path(dataset_path, \"mask_gui.fits\"),\n", + " hdu=0,\n", + " pixel_scales=dataset.pixel_scales,\n", + " )\n", + "except FileNotFoundError:\n", + " mask = al.Mask2D.circular_annular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " inner_radius=0.5,\n", + " outer_radius=2.5,\n", + " )\n", + "\n", + "dataset = dataset.apply_mask(mask=mask) # <----- The custom mask is used here!\n", + "\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated\n", + "on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of\n", + "over sampling to ensure the lensed images are evaluated accurately.\n", + "\n", + "This is why throughout the workspace the cored Sersic profile is used, instead of the regular Sersic profile which\n", + "you may be more familiar with from the literature. In this example we will increase the over sampling level and\n", + "therefore fit a regular Sersic profile to the data, instead of a cored Sersic profile.\n", + "\n", + "This example demonstrates how to change the over sampling used to compute the surface brightness of every image-pixel,\n", + "whereby a higher sub-grid resolution better oversamples the image of the light profile so as to provide a more accurate\n", + "model of its image.\n", + "\n", + "**Benefit**: Higher level of over sampling provide a more accurate estimate of the surface brightness in every image-pixel.\n", + "**Downside**: Higher levels of over sampling require longer calculations and higher memory usage.\n", + "\n", + "Over sampling is applied separately to the light profiles which compute the surface brightness of the lens galaxy,\n", + "which are on a `uniform` grid, and the light profiles which compute the surface brightness of the source galaxy,\n", + "which are on a `non-uniform` grid.\n", + "\n", + "Prequisites: You should read `autolens_workspace/*/guides/advanced/over_sampling.ipynb` before running this script, which\n", + "introduces the concept of over sampling in PyAutoLens and explains why the lens and source galaxy are evaluated\n", + "on different grids.\n", + "\n", + "\n", + "\n", + "The over sampling used to fit the data is customized using the `apply_over_sampling` method, which you may have\n", + "seen in example `modeling` scripts.\n", + "\n", + "To apply uniform over sampling of degree 4x4, we simply input the integer 4.\n", + "\n", + "The grid this is applied to is called `lp`, to indicate that it is the grid used to evaluate the emission of light\n", + "profiles for which this over sampling scheme is applied." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_over_sampling(over_sample_size_lp=4)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Above, the `over_sample_size` input has been an integer, however it can also be an `ndarray` of values corresponding\n", + "to each pixel. \n", + "\n", + "We create an `ndarray` of values which are high in the centre, but reduce to 2 at the outskirts, therefore \n", + "providing high levels of over sampling where we need it whilst using lower values which are computationally fast to \n", + "evaluate at the outskirts.\n", + "\n", + "Specifically, we define a 24 x 24 sub-grid within the central 0.3\" of pixels, uses a 8 x 8 grid between\n", + "0.3\" and 0.6\" and a 2 x 2 grid beyond that. \n", + "\n", + "This will provide high levels of over sampling for the lens galaxy, whose emission peaks at the centre of the\n", + "image near (0.0\", 0.0\"), but will not produce high levels of over sampling for the lensed source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[24, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By assuming the lens galaxy is near (0.0\", 0.0\"), it was simple to set up an adaptive over sampling grid which is\n", + "applicable to all strong lens dataset.\n", + "\n", + "There is no analogous define an adaptive over sampling grid for the lensed source, because for every dataset the\n", + "source's light will appear in different regions of the image plane.\n", + "\n", + "This is why the majority of workspace examples use cored light profiles for the source galaxy. A cored light profile\n", + "does not rapidly change in its central regions, and therefore can be evaluated accurately without over-sampling.\n", + "\n", + "There is a way to set up an adaptive over sampling grid for a lensed source, however it requries one to use and\n", + "understanding the advanced lens modeling feature search chaining.\n", + "\n", + "An example of how to use search chaining to over sample sources efficient is provided in \n", + "the `autolens_workspace/*/guides/modeling/chaining/over_sampling.ipynb` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Positions__\n", + "\n", + "Before fitting a strong lens, we can manually specify a grid of image-plane coordinates corresponding to the multiple\n", + "images of the lensed source-galaxy(s). During the model-fit, **PyAutoLens** will check that these coordinates trace\n", + "within a specified arc-second threshold of one another in the source-plane. If they do not meet this threshold, the\n", + "mass model is discarded and a new sample is generated by the non-linear search.\n", + "\n", + "Advantages: The model-fit is faster, as the non-linear search avoids regions of parameter space where the mass-model\n", + "is clearly not accurate. Removing these unphysical solutions may also mean that the global-maximum solution is inferred\n", + "instead of a local-maxima, given that removing unphysical mass models makes non-linear parameter space less complex.\n", + "\n", + "Disadvantages: If the positions are inaccurate or threshold is set too low, one may inadvertantly remove the correct\n", + "mass model!\n", + "\n", + "The positions are associated with the and they are loaded from a `positions.json` file which is in the same folder as \n", + "the dataset itself. \n", + "\n", + "To create this file, we used a GUI to `draw on` the positions with our mouse. This GUI can be found in the script:\n", + "\n", + " `autolens_workspace/*/imaging/data_preparation/gui/positions.py`\n", + "\n", + "If you wish to use positions for modeling your own lens data, you should use this script to draw on the positions of\n", + "every lens in you dataset.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Alternatively, the positions can be specified manually in the modeling script script using the `Grid2DIrregular`object.\n", + "\n", + "Below, we specify a list of (y,x) coordinates (that are not on a uniform or regular grid) which correspond to the \n", + "arc-second (y,x) coordinates ot he lensed source's brightest pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(\n", + " [(0.4, 1.6), (1.58, -0.35), (-0.43, -1.59), (-1.45, 0.2)]\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To use positions in lens modeling, we pass the `AnalysisImaging` object a `PositionsLH` object, which includes the \n", + "positions we loaded above, alongside a `threshold`.\n", + "\n", + "The threshold of 0.3\" is large. For an accurate lens model we would anticipate the positions trace within < 0.01\" of\n", + "one another. The high threshold ensures only the initial mass models at the start of the fit are penalized." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The analysis receives a list of `PositionsLH` objects, which allows us to use multiple sets of positions to apply this \n", + "penalty, for example if there source has multiple sets of multiple images which we know how they map to one another." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset, positions_likelihood_list=[positions_likelihood], use_jax=True\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/searches.ipynb b/notebooks/guides/modeling/searches.ipynb index 34b4f648b..77edcb882 100644 --- a/notebooks/guides/modeling/searches.ipynb +++ b/notebooks/guides/modeling/searches.ipynb @@ -1,411 +1,448 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Searches\n", - "==================\n", - "\n", - "This script gives a run through of all non-linear searches that are available for modeling.\n", - "\n", - "Extensive testing of lens modeling has shown that the default search used throughout all modeling examples,\n", - "`Nautilus`, is the most accurate and fastest search available. For users familiar with statistical inference, this may\n", - "be surprising, as nested samplers are traditionally slower than MCMC methods such as Emcee and maximum likelihood\n", - "methods such as LBFGS. A description of why Nautilus performs better than these other searches is beyond the scope\n", - "of this script, but if you add me on SLACk I'd be happy to have a discussion about it!\n", - "\n", - "Therefore, unless you really know what you are doing or want to use an alternative search, it is strongly recommended\n", - "you stick to Nautilus.\n", - "\n", - "Three different categories of searches are available, nested samplers (E.g. Nautilus, Dynesty), MCMC (E.g. Emcee) and\n", - "maximum likelihood (e.g. LBFGS). MCMC and MLE methods can often optionally use a \"starting point\" to initialize the\n", - "model-fit with the parameters where it should begin. Nested samplers do not use a starting point, but a similar\n", - "approach can be applied by putting tight priors on certain parameters.\n", - "\n", - "To perform a model-fit, a fully modeling script will include steps which compose a model, create an `Analysis`\n", - "object and pass these to the search to perform the fit. We skip these steps for brevity.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dynesty:** Dynesty (https://github.com/joshspeagle/dynesty) is a nested sampling algorithm.\n", - "- **Emcee:** Emcee (https://github.com/dfm/emcee) is an ensemble MCMC sampler that is commonly used in Astronomy.\n", - "- **Zeus:** Zeus (https://zeus-mcmc.readthedocs.io/en/latest/) is an ensemble MCMC slice sampler.\n", - "- **LBFGS:** LBFGS is a quasi-Newton optimization algorithm from scipy.\n", - "- **Start Point:** For maximum likelihood estimator (MLE) and Markov Chain Monte Carlo (MCMC) non-linear searches.\n", - "- **Search Cookbook:** There are a number of other searches supported by **PyAutoFit** and therefore which can be used.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dynesty__\n", - "\n", - "Dynesty (https://github.com/joshspeagle/dynesty) is a nested sampling algorithm.\n", - "\n", - "Dynesty used to be the default model-fitting algorithm, before Nautilus was found to be better. However, Dynesty with\n", - "random walk nested sampling is still an effective method for modeling and worth using if you want to check your\n", - "results with an alternative to Nautilus.\n", - "\n", - "Dynesty itself supports a wide variety of different nested sampling methods, including static \n", - "sampling (`DynestyStatic` where the number of live point is fixed), dynamic sampling (`DynestyDynamic` where the number \n", - "of live points varies with the fit) and different approaches to point sampling (e.g. slice sampling, uniform sampling). \n", - "\n", - "If you are familiar with nested sampling you can use all dynesty's different options by customizing the code below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.DynestyStatic(\n", - " path_prefix=Path(\"searches\"),\n", - " name=\"DynestyStatic\",\n", - " unique_tag=\"example\",\n", - " iterations_per_quick_update=2500,\n", - " # search specific settings\n", - " nlive=50,\n", - " sample=\"rwalk\",\n", - " walks=10,\n", - " bound=\"multi\",\n", - " bootstrap=None,\n", - " enlarge=None,\n", - " update_interval=None,\n", - " facc=0.5,\n", - " slices=5,\n", - " fmove=0.9,\n", - " max_move=100,\n", - ")\n", - "\n", - "search = af.DynestyDynamic(\n", - " path_prefix=Path(\"searches\"),\n", - " name=\"DynestyDynamic\",\n", - " unique_tag=\"example\",\n", - " iterations_per_quick_update=2500,\n", - " # search specific settings\n", - " nlive=50,\n", - " sample=\"rwalk\",\n", - " walks=10,\n", - " bound=\"multi\",\n", - " bootstrap=None,\n", - " enlarge=None,\n", - " update_interval=None,\n", - " facc=0.5,\n", - " slices=5,\n", - " fmove=0.9,\n", - " max_move=100,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Emcee__\n", - "\n", - "Emcee (https://github.com/dfm/emcee) is an ensemble MCMC sampler that is commonly used in Astronomy and Astrophysics.\n", - "\n", - "The wrapper with **PyAutoFit** supports different initialization methods, including a ball around the center of the\n", - "priors on the model parameters, which is the recommend initialization method for Emcee.\n", - "\n", - "It also includes functionality which checks the auto correlations of the chains, and terminates the search early\n", - "if they meet certain convergence criteria. This is useful for ensuring that the chains have converged.\n", - "\n", - "Whilst Emcee is a popular choice of MCMC method in astrophsyics, note that the MCMC method `Zeus`, described next, has\n", - "proven better as lens modeling for our tests." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Emcee(\n", - " path_prefix=Path(\"imaging\", \"searches\"),\n", - " name=\"Emcee\",\n", - " unique_tag=\"example\",\n", - " iterations_per_quick_update=5000,\n", - " # search specific settings\n", - " nwalkers=30,\n", - " nsteps=500,\n", - " initializer=af.InitializerBall(lower_limit=0.49, upper_limit=0.51),\n", - " auto_correlations_settings=af.AutoCorrelationsSettings(\n", - " check_for_convergence=True,\n", - " check_size=100,\n", - " required_length=50,\n", - " change_threshold=0.01,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Zeus__\n", - "\n", - "Zeus (https://zeus-mcmc.readthedocs.io/en/latest/) is an ensemble MCMC slice sampler.\n", - "\n", - "The wrapper with **PyAutoFit** supports different initialization methods, including a ball around the center of the\n", - "priors on the model parameters, which is the recommend initialization method for Emcee.\n", - "\n", - "It also includes functionality which checks the auto correlations of the chains, and terminates the search early\n", - "if they meet certain convergence criteria. This is useful for ensuring that the chains have converged.\n", - "\n", - "Zeus is the most effective MCMC method for lens modeling that we have tested, and is the recommended MCMC method,\n", - "however its performance is not as good as Nautilus." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Zeus(\n", - " path_prefix=Path(\"imaging\", \"searches\"),\n", - " name=\"Zeus\",\n", - " unique_tag=\"example\",\n", - " iterations_per_quick_update=5000,\n", - " # search specific settings\n", - " nwalkers=30,\n", - " nsteps=20,\n", - " initializer=af.InitializerBall(lower_limit=0.49, upper_limit=0.51),\n", - " auto_correlations_settings=af.AutoCorrelationsSettings(\n", - " check_for_convergence=True,\n", - " check_size=100,\n", - " required_length=50,\n", - " change_threshold=0.01,\n", - " ),\n", - " tune=False,\n", - " tolerance=0.05,\n", - " patience=5,\n", - " maxsteps=10000,\n", - " mu=1.0,\n", - " maxiter=10000,\n", - " vectorize=False,\n", - " check_walkers=True,\n", - " shuffle_ensemble=True,\n", - " light_mode=False,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LBFGS__\n", - "\n", - "LBFGS is a quasi-Newton optimization algorithm from scipy.\n", - "\n", - "An optimizer only seeks to find the maximum likelihood lens model, unlike MCMC or nested sampling algorithms\n", - "like Zeus and Nautilus, which aim to map out parameter space and infer errors on the parameters. Therefore, in\n", - "principle, an optimizer like LBFGS should fit a lens model very fast.\n", - "\n", - "In our experience, the parameter spaces fitted by lens models are often too complex for optimizers to be used without\n", - "careful initialization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.LBFGS(\n", - " path_prefix=Path(\"imaging\", \"searches\"),\n", - " name=\"LBFGS\",\n", - " unique_tag=\"example\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Start Point__\n", - "\n", - "For maximum likelihood estimator (MLE) and Markov Chain Monte Carlo (MCMC) non-linear searches, parameter space\n", - "sampling is built around having a \"location\" in parameter space.\n", - "\n", - "This could simply be the parameters of the current maximum likelihood model in an MLE fit, or the locations of many\n", - "walkers in parameter space (e.g. MCMC).\n", - "\n", - "For many model-fitting problems, we may have an expectation of where correct solutions lie in parameter space and\n", - "therefore want our non-linear search to start near that location of parameter space. Alternatively, we may want to\n", - "sample a specific region of parameter space, to determine what solutions look like there.\n", - "\n", - "The start-point API allows us to do this, by manually specifying the start-point of an MLE fit or the start-point of\n", - "the walkers in an MCMC fit. Because nested sampling draws from priors, it cannot use the start-point API.\n", - "\n", - "Similar behaviour can be achieved by customizing the priors of a model-fit. We could place `GaussianPrior`'s\n", - "centred on the regions of parameter space we want to sample, or we could place tight `UniformPrior`'s on regions\n", - "of parameter space we believe the correct answer lies.\n", - "\n", - "The downside of using priors is that our priors have a direct influence on the parameters we infer and the size\n", - "of the inferred parameter errors. By using priors to control the location of our model-fit, we therefore risk\n", - "inferring a non-representative model.\n", - "\n", - "For users more familiar with statistical inference, adjusting ones priors in the way described above leads to\n", - "changes in the posterior, which therefore impacts the model inferred." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "mass.centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "mass.centre_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "mass.ell_comps.ell_comps_0 = af.UniformPrior(lower_limit=-0.5, upper_limit=0.5)\n", - "mass.ell_comps.ell_comps_1 = af.UniformPrior(lower_limit=-0.5, upper_limit=0.5)\n", - "mass.einstein_radius = af.UniformPrior(lower_limit=0.2, upper_limit=3.0)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "shear.gamma_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "shear.gamma_2 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now define the start point of certain parameters in the model:\n", - "\n", - " - The galaxy is centred near (0.0, 0.0), so we set a start point for its mass distribution there.\n", - "\n", - " - The size of the lensed source galaxy is around 1.6\" thus we set the `einstein_radius` to start here.\n", - "\n", - " - We know the source galaxy is a disk galaxy, thus we set its `sersic_index` to start around 1.0.\n", - "\n", - "For all parameters where the start-point is not specified (in this case the `ell_comps`, their \n", - "parameter values are drawn randomly from the prior when determining the initial locations of the parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "initializer = af.InitializerParamBounds(\n", - " {\n", - " model.galaxies.lens.mass.centre_0: (-0.01, 0.01),\n", - " model.galaxies.lens.mass.centre_1: (-0.01, 0.01),\n", - " model.galaxies.lens.mass.einstein_radius: (1.58, 1.62),\n", - " model.galaxies.source.bulge.sersic_index: (0.95, 1.05),\n", - " }\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `initializer` is passed to the search (e.g. the MCMC method Emcee below), which uses it to set the start-point of \n", - "the walkers in parameter space. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Emcee(\n", - " path_prefix=Path(\"imaging\", \"customize\"),\n", - " name=\"start_point\",\n", - " nwalkers=50,\n", - " nsteps=500,\n", - " initializer=initializer,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Cookbook__\n", - "\n", - "There are a number of other searches supported by **PyAutoFit** and therefore which can be used, which are not \n", - "explictly documented here. These include LBFGS.\n", - "\n", - "The **PyAutoFit** search cookbook documents all searches that are available, including those not documented here,\n", - "and provides the code you can easily copy and paste to use these methods.\n", - "\n", - "https://pyautofit.readthedocs.io/en/latest/cookbooks/search.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Searches\n", + "==================\n", + "\n", + "This script gives a run through of all non-linear searches that are available for modeling.\n", + "\n", + "Extensive testing of lens modeling has shown that the default search used throughout all modeling examples,\n", + "`Nautilus`, is the most accurate and fastest search available. For users familiar with statistical inference, this may\n", + "be surprising, as nested samplers are traditionally slower than MCMC methods such as Emcee and maximum likelihood\n", + "methods such as LBFGS. A description of why Nautilus performs better than these other searches is beyond the scope\n", + "of this script, but if you add me on SLACk I'd be happy to have a discussion about it!\n", + "\n", + "Therefore, unless you really know what you are doing or want to use an alternative search, it is strongly recommended\n", + "you stick to Nautilus.\n", + "\n", + "Three different categories of searches are available, nested samplers (E.g. Nautilus, Dynesty), MCMC (E.g. Emcee) and\n", + "maximum likelihood (e.g. LBFGS). MCMC and MLE methods can often optionally use a \"starting point\" to initialize the\n", + "model-fit with the parameters where it should begin. Nested samplers do not use a starting point, but a similar\n", + "approach can be applied by putting tight priors on certain parameters.\n", + "\n", + "To perform a model-fit, a fully modeling script will include steps which compose a model, create an `Analysis`\n", + "object and pass these to the search to perform the fit. We skip these steps for brevity.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dynesty:** Dynesty (https://github.com/joshspeagle/dynesty) is a nested sampling algorithm.\n", + "- **Emcee:** Emcee (https://github.com/dfm/emcee) is an ensemble MCMC sampler that is commonly used in Astronomy.\n", + "- **Zeus:** Zeus (https://zeus-mcmc.readthedocs.io/en/latest/) is an ensemble MCMC slice sampler.\n", + "- **LBFGS:** LBFGS is a quasi-Newton optimization algorithm from scipy.\n", + "- **Start Point:** For maximum likelihood estimator (MLE) and Markov Chain Monte Carlo (MCMC) non-linear searches.\n", + "- **Search Cookbook:** There are a number of other searches supported by **PyAutoFit** and therefore which can be used.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dynesty__\n", + "\n", + "Dynesty (https://github.com/joshspeagle/dynesty) is a nested sampling algorithm.\n", + "\n", + "Dynesty used to be the default model-fitting algorithm, before Nautilus was found to be better. However, Dynesty with\n", + "random walk nested sampling is still an effective method for modeling and worth using if you want to check your\n", + "results with an alternative to Nautilus.\n", + "\n", + "Dynesty itself supports a wide variety of different nested sampling methods, including static \n", + "sampling (`DynestyStatic` where the number of live point is fixed), dynamic sampling (`DynestyDynamic` where the number \n", + "of live points varies with the fit) and different approaches to point sampling (e.g. slice sampling, uniform sampling). \n", + "\n", + "If you are familiar with nested sampling you can use all dynesty's different options by customizing the code below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.DynestyStatic(\n", + " path_prefix=Path(\"searches\"),\n", + " name=\"DynestyStatic\",\n", + " unique_tag=\"example\",\n", + " iterations_per_quick_update=2500,\n", + " # search specific settings\n", + " nlive=50,\n", + " sample=\"rwalk\",\n", + " walks=10,\n", + " bound=\"multi\",\n", + " bootstrap=None,\n", + " enlarge=None,\n", + " update_interval=None,\n", + " facc=0.5,\n", + " slices=5,\n", + " fmove=0.9,\n", + " max_move=100,\n", + ")\n", + "\n", + "search = af.DynestyDynamic(\n", + " path_prefix=Path(\"searches\"),\n", + " name=\"DynestyDynamic\",\n", + " unique_tag=\"example\",\n", + " iterations_per_quick_update=2500,\n", + " # search specific settings\n", + " nlive=50,\n", + " sample=\"rwalk\",\n", + " walks=10,\n", + " bound=\"multi\",\n", + " bootstrap=None,\n", + " enlarge=None,\n", + " update_interval=None,\n", + " facc=0.5,\n", + " slices=5,\n", + " fmove=0.9,\n", + " max_move=100,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Emcee__\n", + "\n", + "Emcee (https://github.com/dfm/emcee) is an ensemble MCMC sampler that is commonly used in Astronomy and Astrophysics.\n", + "\n", + "The wrapper with **PyAutoFit** supports different initialization methods, including a ball around the center of the\n", + "priors on the model parameters, which is the recommend initialization method for Emcee.\n", + "\n", + "It also includes functionality which checks the auto correlations of the chains, and terminates the search early\n", + "if they meet certain convergence criteria. This is useful for ensuring that the chains have converged.\n", + "\n", + "Whilst Emcee is a popular choice of MCMC method in astrophsyics, note that the MCMC method `Zeus`, described next, has\n", + "proven better as lens modeling for our tests." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Emcee(\n", + " path_prefix=Path(\"imaging\", \"searches\"),\n", + " name=\"Emcee\",\n", + " unique_tag=\"example\",\n", + " iterations_per_quick_update=5000,\n", + " # search specific settings\n", + " nwalkers=30,\n", + " nsteps=500,\n", + " initializer=af.InitializerBall(lower_limit=0.49, upper_limit=0.51),\n", + " auto_correlations_settings=af.AutoCorrelationsSettings(\n", + " check_for_convergence=True,\n", + " check_size=100,\n", + " required_length=50,\n", + " change_threshold=0.01,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Zeus__\n", + "\n", + "Zeus (https://zeus-mcmc.readthedocs.io/en/latest/) is an ensemble MCMC slice sampler.\n", + "\n", + "The wrapper with **PyAutoFit** supports different initialization methods, including a ball around the center of the\n", + "priors on the model parameters, which is the recommend initialization method for Emcee.\n", + "\n", + "It also includes functionality which checks the auto correlations of the chains, and terminates the search early\n", + "if they meet certain convergence criteria. This is useful for ensuring that the chains have converged.\n", + "\n", + "Zeus is the most effective MCMC method for lens modeling that we have tested, and is the recommended MCMC method,\n", + "however its performance is not as good as Nautilus." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Zeus(\n", + " path_prefix=Path(\"imaging\", \"searches\"),\n", + " name=\"Zeus\",\n", + " unique_tag=\"example\",\n", + " iterations_per_quick_update=5000,\n", + " # search specific settings\n", + " nwalkers=30,\n", + " nsteps=20,\n", + " initializer=af.InitializerBall(lower_limit=0.49, upper_limit=0.51),\n", + " auto_correlations_settings=af.AutoCorrelationsSettings(\n", + " check_for_convergence=True,\n", + " check_size=100,\n", + " required_length=50,\n", + " change_threshold=0.01,\n", + " ),\n", + " tune=False,\n", + " tolerance=0.05,\n", + " patience=5,\n", + " maxsteps=10000,\n", + " mu=1.0,\n", + " maxiter=10000,\n", + " vectorize=False,\n", + " check_walkers=True,\n", + " shuffle_ensemble=True,\n", + " light_mode=False,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LBFGS__\n", + "\n", + "LBFGS is a quasi-Newton optimization algorithm from scipy.\n", + "\n", + "An optimizer only seeks to find the maximum likelihood lens model, unlike MCMC or nested sampling algorithms\n", + "like Zeus and Nautilus, which aim to map out parameter space and infer errors on the parameters. Therefore, in\n", + "principle, an optimizer like LBFGS should fit a lens model very fast.\n", + "\n", + "In our experience, the parameter spaces fitted by lens models are often too complex for optimizers to be used without\n", + "careful initialization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.LBFGS(\n", + " path_prefix=Path(\"imaging\", \"searches\"),\n", + " name=\"LBFGS\",\n", + " unique_tag=\"example\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Start Point__\n", + "\n", + "For maximum likelihood estimator (MLE) and Markov Chain Monte Carlo (MCMC) non-linear searches, parameter space\n", + "sampling is built around having a \"location\" in parameter space.\n", + "\n", + "This could simply be the parameters of the current maximum likelihood model in an MLE fit, or the locations of many\n", + "walkers in parameter space (e.g. MCMC).\n", + "\n", + "For many model-fitting problems, we may have an expectation of where correct solutions lie in parameter space and\n", + "therefore want our non-linear search to start near that location of parameter space. Alternatively, we may want to\n", + "sample a specific region of parameter space, to determine what solutions look like there.\n", + "\n", + "The start-point API allows us to do this, by manually specifying the start-point of an MLE fit or the start-point of\n", + "the walkers in an MCMC fit. Because nested sampling draws from priors, it cannot use the start-point API.\n", + "\n", + "Similar behaviour can be achieved by customizing the priors of a model-fit. We could place `GaussianPrior`'s\n", + "centred on the regions of parameter space we want to sample, or we could place tight `UniformPrior`'s on regions\n", + "of parameter space we believe the correct answer lies.\n", + "\n", + "The downside of using priors is that our priors have a direct influence on the parameters we infer and the size\n", + "of the inferred parameter errors. By using priors to control the location of our model-fit, we therefore risk\n", + "inferring a non-representative model.\n", + "\n", + "For users more familiar with statistical inference, adjusting ones priors in the way described above leads to\n", + "changes in the posterior, which therefore impacts the model inferred." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "mass.centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "mass.centre_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "mass.ell_comps.ell_comps_0 = af.UniformPrior(lower_limit=-0.5, upper_limit=0.5)\n", + "mass.ell_comps.ell_comps_1 = af.UniformPrior(lower_limit=-0.5, upper_limit=0.5)\n", + "mass.einstein_radius = af.UniformPrior(lower_limit=0.2, upper_limit=3.0)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "shear.gamma_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "shear.gamma_2 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now define the start point of certain parameters in the model:\n", + "\n", + " - The galaxy is centred near (0.0, 0.0), so we set a start point for its mass distribution there.\n", + "\n", + " - The size of the lensed source galaxy is around 1.6\" thus we set the `einstein_radius` to start here.\n", + "\n", + " - We know the source galaxy is a disk galaxy, thus we set its `sersic_index` to start around 1.0.\n", + "\n", + "For all parameters where the start-point is not specified (in this case the `ell_comps`, their \n", + "parameter values are drawn randomly from the prior when determining the initial locations of the parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "initializer = af.InitializerParamBounds(\n", + " {\n", + " model.galaxies.lens.mass.centre_0: (-0.01, 0.01),\n", + " model.galaxies.lens.mass.centre_1: (-0.01, 0.01),\n", + " model.galaxies.lens.mass.einstein_radius: (1.58, 1.62),\n", + " model.galaxies.source.bulge.sersic_index: (0.95, 1.05),\n", + " }\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `initializer` is passed to the search (e.g. the MCMC method Emcee below), which uses it to set the start-point of \n", + "the walkers in parameter space. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Emcee(\n", + " path_prefix=Path(\"imaging\", \"customize\"),\n", + " name=\"start_point\",\n", + " nwalkers=50,\n", + " nsteps=500,\n", + " initializer=initializer,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Cookbook__\n", + "\n", + "There are a number of other searches supported by **PyAutoFit** and therefore which can be used, which are not \n", + "explictly documented here. These include LBFGS.\n", + "\n", + "The **PyAutoFit** search cookbook documents all searches that are available, including those not documented here,\n", + "and provides the code you can easily copy and paste to use these methods.\n", + "\n", + "https://pyautofit.readthedocs.io/en/latest/cookbooks/search.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/modeling/slam_start_here.ipynb b/notebooks/guides/modeling/slam_start_here.ipynb index d0a70b445..8d93e2007 100644 --- a/notebooks/guides/modeling/slam_start_here.ipynb +++ b/notebooks/guides/modeling/slam_start_here.ipynb @@ -1,803 +1,840 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM (Source, Light and Mass): Start Here\n", - "=========================================\n", - "\n", - "This scripts gives an introduce to the Source, (lens) Light and Mass (SLaM) pipelines. These are advanced modeling\n", - "pipelines which use many aspects of core PyAutoLens functionality to automate the modeling of strong lenses.\n", - "\n", - "__Contents__\n", - "\n", - "- **Preqrequisites:** Before using SLaM, you should understand.\n", - "- **Overview:** SLaM chains together four or more sequential modeling searches, with each stage passing its results.\n", - "- **Design Choices:** The structure of the SLaM pipelines is driven by the requirements of **adaptive pixelized source.\n", - "- **This Script:** Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this.\n", - "- **SOURCE LP PIPELINE:** The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light.\n", - "- **SOURCE PIX PIPELINE 1:** The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source.\n", - "- **Positions:** Image positions are used to prevent unphysical source reconstructions (see `features/pixelization`.\n", - "- **Adapt Images:** An adapt image is computed from the SOURCE LP result and passed to the analysis.\n", - "- **SOURCE PIX PIPELINE 2:** The second search of the SOURCE PIX PIPELINE fits the final pixelized source model using the.\n", - "- **LIGHT LP PIPELINE:** The LIGHT LP PIPELINE uses one search to fit a complex lens light model to a high level of.\n", - "- **MASS TOTAL PIPELINE:** The MASS TOTAL PIPELINE uses one search to fit a complex lens mass model to a high level of.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Preqrequisites__\n", - "\n", - "Before using SLaM, you should understand:\n", - "\n", - "- **pixelizations** (`*/features/pixelization`)\n", - " Methods that reconstruct the source galaxy using a pixel grid.\n", - "\n", - "- **Search Chaining** (`guides/modeling/chaining`)\n", - " Fitting lens models in stages of increasing complexity, e.g. first a light-profile source,\n", - " then a pixelized source.\n", - "\n", - "- **Adaptive pixelizations** (`features/pixelization/adaptive`)\n", - " pixelizations that adapt their mesh and regularization to the unlensed source morphology.\n", - "\n", - "- **Multi-Gaussian Expansions (MGEs)** (`features/multi_gaussian_expansion`)\n", - " Galaxy light modeled as a sum of Gaussians, enabling accurate lens-light subtraction.\n", - "\n", - "You can still run the script without fully understanding these concepts; however, reviewing the\n", - "referenced examples later will clarify why SLaM pipelines are structured as they are.\n", - "\n", - "__Overview__\n", - "\n", - "SLaM chains together four or more sequential modeling searches, with each stage passing its results\n", - "forward to the next. This strategy enables fully automated modeling suitable for large samples of\n", - "strong lenses.\n", - "\n", - "Each pipeline targets a specific part of the lens model:\n", - "\n", - "1. **Source Pipeline**:\n", - " Builds a reliable model of the source galaxy. For pixelized sources, this includes determining\n", - " stable mesh and regularization parameters.\n", - "\n", - "2. **Light Pipeline**:\n", - " Models the lens galaxy's light using the fixed source model from step 1. Accurate subtraction\n", - " of lens light is essential for robust mass modeling.\n", - "\n", - "3. **Mass Pipeline**:\n", - " Fits the lens mass distribution (often complex), using the refined source and lens-light models\n", - " from previous pipelines.\n", - "\n", - "\n", - "The SLaM workflow is flexible\u2014you can swap MGE light profiles for other light models if desired. Models set up in\n", - "earlier pipelines guide those used in later ones. For example, if the Source Pipeline uses a `RectangularAdaptDensity`\n", - "mesh, the same mesh type is carried into later pipelines for consistency.\n", - "\n", - "__Design Choices__\n", - "\n", - "The structure of the SLaM pipelines is driven by the requirements of **adaptive pixelized source modeling**,\n", - "which is essential for fitting complex light and mass distributions for the lens.\n", - "\n", - "Although SLaM also supports light-profile sources, pixelized sources are at the core of the pipeline design and\n", - "enable fully automated modeling of realistic, high-complexity mass models.\n", - "\n", - "Below are the key design considerations that determine the ordering of SLaM pipelines:\n", - "\n", - "- **Source First**\n", - " Complex mass models (e.g., `PowerLaw`, or composite stellar + dark matter models) require pixelized\n", - " source reconstruction, not simple light profiles. Therefore, SLaM begins with a source model using a\n", - " simpler mass profile (e.g., `Isothermal` + `ExternalShear`) to provide a stable basis for later stages.\n", - "\n", - "- **Image Positions**\n", - " Pixelized modeling needs robust multiple image-position estimates to prevent unphysical source reconstructions.\n", - " SLaM automatically determines these positions from the results of the Source Light Profile Pipeline.\n", - "\n", - "- **Adapt Images**\n", - " Advanced pixelizations use lens-light-subtracted \"adapt images\" to adapt the source pixelization mesh and\n", - " regularization to the unlensed source morphology. These are only set once a sufficiently accurate source model\n", - " is available from earlier stages in the Source Pipeline.\n", - "\n", - "- **Lens Light Before Mass**\n", - " Accurate lens-light subtraction is required before fitting complex mass models, especially for mass models\n", - " fitting stellar and dark matter components simultanoeusly. Pixelized source modeling enables reliable\n", - " deblending of the lens and source, so the lens light model is refined after the adaptive pixelized source is\n", - " accurate but before fitting more complex mass models.\n", - "\n", - "- **Mass Model Last**\n", - " The most flexible and complex mass models are fit only after high-quality source and lens-light models\n", - " are established, ensuring stable priors and accurate mass inference.\n", - "\n", - "Together, these design choices allow SLaM to perform precise, automated strong-lens modeling while maintaining\n", - "robustness and efficiency at each stage.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this SLaM modeling\n", - "script fits `Imaging` dataset of a strong lens system where in the final model:\n", - "\n", - " - The lens galaxy's light is a bulge with Multiple Gaussian Expansion (MGE) light profile.\n", - " - The lens galaxy's total mass distribution is an `PowerLaw` plus an `ExternalShear`.\n", - " - The source galaxy's light is a `Pixelization`.\n", - "\n", - "Each SLaM pipeline is implemented as a Python function below (e.g. `source_lp`, `source_pix_1`), with a\n", - "documentation string above each function describing the pipeline in detail. The full pipeline is run at the\n", - "bottom of the script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light, which in\n", - "this example:\n", - "\n", - " - Models the lens galaxy's light as an MGE with 2 x 20 Gaussians.\n", - " - Uses an `Isothermal` model for the lens's total mass distribution with an `ExternalShear`.\n", - " - Models the source galaxy's light as an MGE with 1 x 20 Gaussians.\n", - "\n", - "The mass and source models from this search initialize the SOURCE PIX PIPELINE searches that follow." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source galaxy.\n", - "This pixelization adapts its resolution to the source morphology, assigning more pixels to brighter,\n", - "more detailed regions.\n", - "\n", - "To build an adaptive pixelization, we require an **adapt image**: a lens-light-subtracted image in\n", - "which only the lensed source emission remains. This image determines how both the mesh density and\n", - "regularization weights adapt to the source structure.\n", - "\n", - "The SOURCE LP Pipeline does not provide a sufficiently accurate source model for computing this adapt\n", - "image (e.g., the true source may be more complex than a simple light profile). Therefore, the first\n", - "search of the SOURCE PIX PIPELINE fits a model using a pixelization whose purpose is to generate a\n", - "high-quality adapt image used in search 2.\n", - "\n", - "__Positions__\n", - "\n", - "Image positions are used to prevent unphysical source reconstructions (see `features/pixelization` for\n", - "details). Rather than being input manually, they are computed automatically from the SOURCE LP result.\n", - "This is a key automation feature of SLaM.\n", - "\n", - "__Adapt Images__\n", - "\n", - "An adapt image is computed from the SOURCE LP result and passed to the analysis. This provides an initial\n", - "estimate of the source morphology for the `Adapt` regularization scheme, even though the MGE source model\n", - "may not fully capture the source structure. Search 2 improves upon this using a pixelized adapt image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "The second search of the SOURCE PIX PIPELINE fits the final pixelized source model using the improved\n", - "adapt images computed from search 1's pixelized source reconstruction.\n", - "\n", - "The `RectangularAdaptImage` mesh and `Adapt` regularization adapt the source pixels and regularization\n", - "weights to the source's morphology using the high-quality adapt images from search 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "The LIGHT LP PIPELINE uses one search to fit a complex lens light model to a high level of accuracy, with\n", - "the lens mass model and source light model fixed to the maximum log likelihood result of the SOURCE PIX\n", - "PIPELINE.\n", - "\n", - "In this example:\n", - "\n", - " - The lens galaxy's light is a MGE with 2 x 20 Gaussians.\n", - " - Uses an `Isothermal` mass model with `ExternalShear` for the lens's total mass distribution [fixed from\n", - " SOURCE PIX PIPELINE].\n", - " - Uses a `Pixelization` for the source's light [fixed from SOURCE PIX PIPELINE].\n", - "\n", - "This search aims to produce an accurate model of the lens galaxy's light, which may not have been possible\n", - "in the SOURCE PIPELINE as the mass and source models were not yet precisely estimated. The adapt images\n", - "from SOURCE PIX PIPELINE search 1 are reused, providing a stable basis for the lens-light subtraction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " )\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "The MASS TOTAL PIPELINE uses one search to fit a complex lens mass model to a high level of accuracy,\n", - "using the lens mass model and source model of the SOURCE PIX PIPELINE to initialize model priors, and the\n", - "lens light model of the LIGHT LP PIPELINE.\n", - "\n", - "In this example:\n", - "\n", - " - Uses a linear MGE bulge [fixed from LIGHT LP PIPELINE].\n", - " - Uses a `PowerLaw` model for the lens's total mass distribution [priors initialized from SOURCE PIX\n", - " PIPELINE].\n", - " - Uses a `Pixelization` for the source's light [fixed from SOURCE PIX PIPELINE].\n", - "\n", - "__Positions__\n", - "\n", - "Positions are computed from the SOURCE PIX PIPELINE search 2 result, which provides more precise multiple\n", - "image positions than the SOURCE LP PIPELINE (as the pixelized source gives a better source-plane\n", - "reconstruction)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " bulge = light_result.instance.galaxies.lens.bulge\n", - " disk = light_result.instance.galaxies.lens.disk\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=bulge,\n", - " disk=disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM (Source, Light and Mass): Start Here\n", + "=========================================\n", + "\n", + "This scripts gives an introduce to the Source, (lens) Light and Mass (SLaM) pipelines. These are advanced modeling\n", + "pipelines which use many aspects of core PyAutoLens functionality to automate the modeling of strong lenses.\n", + "\n", + "__Contents__\n", + "\n", + "- **Preqrequisites:** Before using SLaM, you should understand.\n", + "- **Overview:** SLaM chains together four or more sequential modeling searches, with each stage passing its results.\n", + "- **Design Choices:** The structure of the SLaM pipelines is driven by the requirements of **adaptive pixelized source.\n", + "- **This Script:** Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this.\n", + "- **SOURCE LP PIPELINE:** The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light.\n", + "- **SOURCE PIX PIPELINE 1:** The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source.\n", + "- **Positions:** Image positions are used to prevent unphysical source reconstructions (see `features/pixelization`.\n", + "- **Adapt Images:** An adapt image is computed from the SOURCE LP result and passed to the analysis.\n", + "- **SOURCE PIX PIPELINE 2:** The second search of the SOURCE PIX PIPELINE fits the final pixelized source model using the.\n", + "- **LIGHT LP PIPELINE:** The LIGHT LP PIPELINE uses one search to fit a complex lens light model to a high level of.\n", + "- **MASS TOTAL PIPELINE:** The MASS TOTAL PIPELINE uses one search to fit a complex lens mass model to a high level of.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Preqrequisites__\n", + "\n", + "Before using SLaM, you should understand:\n", + "\n", + "- **pixelizations** (`*/features/pixelization`)\n", + " Methods that reconstruct the source galaxy using a pixel grid.\n", + "\n", + "- **Search Chaining** (`guides/modeling/chaining`)\n", + " Fitting lens models in stages of increasing complexity, e.g. first a light-profile source,\n", + " then a pixelized source.\n", + "\n", + "- **Adaptive pixelizations** (`features/pixelization/adaptive`)\n", + " pixelizations that adapt their mesh and regularization to the unlensed source morphology.\n", + "\n", + "- **Multi-Gaussian Expansions (MGEs)** (`features/multi_gaussian_expansion`)\n", + " Galaxy light modeled as a sum of Gaussians, enabling accurate lens-light subtraction.\n", + "\n", + "You can still run the script without fully understanding these concepts; however, reviewing the\n", + "referenced examples later will clarify why SLaM pipelines are structured as they are.\n", + "\n", + "__Overview__\n", + "\n", + "SLaM chains together four or more sequential modeling searches, with each stage passing its results\n", + "forward to the next. This strategy enables fully automated modeling suitable for large samples of\n", + "strong lenses.\n", + "\n", + "Each pipeline targets a specific part of the lens model:\n", + "\n", + "1. **Source Pipeline**:\n", + " Builds a reliable model of the source galaxy. For pixelized sources, this includes determining\n", + " stable mesh and regularization parameters.\n", + "\n", + "2. **Light Pipeline**:\n", + " Models the lens galaxy's light using the fixed source model from step 1. Accurate subtraction\n", + " of lens light is essential for robust mass modeling.\n", + "\n", + "3. **Mass Pipeline**:\n", + " Fits the lens mass distribution (often complex), using the refined source and lens-light models\n", + " from previous pipelines.\n", + "\n", + "\n", + "The SLaM workflow is flexible\u2014you can swap MGE light profiles for other light models if desired. Models set up in\n", + "earlier pipelines guide those used in later ones. For example, if the Source Pipeline uses a `RectangularAdaptDensity`\n", + "mesh, the same mesh type is carried into later pipelines for consistency.\n", + "\n", + "__Design Choices__\n", + "\n", + "The structure of the SLaM pipelines is driven by the requirements of **adaptive pixelized source modeling**,\n", + "which is essential for fitting complex light and mass distributions for the lens.\n", + "\n", + "Although SLaM also supports light-profile sources, pixelized sources are at the core of the pipeline design and\n", + "enable fully automated modeling of realistic, high-complexity mass models.\n", + "\n", + "Below are the key design considerations that determine the ordering of SLaM pipelines:\n", + "\n", + "- **Source First**\n", + " Complex mass models (e.g., `PowerLaw`, or composite stellar + dark matter models) require pixelized\n", + " source reconstruction, not simple light profiles. Therefore, SLaM begins with a source model using a\n", + " simpler mass profile (e.g., `Isothermal` + `ExternalShear`) to provide a stable basis for later stages.\n", + "\n", + "- **Image Positions**\n", + " Pixelized modeling needs robust multiple image-position estimates to prevent unphysical source reconstructions.\n", + " SLaM automatically determines these positions from the results of the Source Light Profile Pipeline.\n", + "\n", + "- **Adapt Images**\n", + " Advanced pixelizations use lens-light-subtracted \"adapt images\" to adapt the source pixelization mesh and\n", + " regularization to the unlensed source morphology. These are only set once a sufficiently accurate source model\n", + " is available from earlier stages in the Source Pipeline.\n", + "\n", + "- **Lens Light Before Mass**\n", + " Accurate lens-light subtraction is required before fitting complex mass models, especially for mass models\n", + " fitting stellar and dark matter components simultanoeusly. Pixelized source modeling enables reliable\n", + " deblending of the lens and source, so the lens light model is refined after the adaptive pixelized source is\n", + " accurate but before fitting more complex mass models.\n", + "\n", + "- **Mass Model Last**\n", + " The most flexible and complex mass models are fit only after high-quality source and lens-light models\n", + " are established, ensuring stable priors and accurate mass inference.\n", + "\n", + "Together, these design choices allow SLaM to perform precise, automated strong-lens modeling while maintaining\n", + "robustness and efficiency at each stage.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this SLaM modeling\n", + "script fits `Imaging` dataset of a strong lens system where in the final model:\n", + "\n", + " - The lens galaxy's light is a bulge with Multiple Gaussian Expansion (MGE) light profile.\n", + " - The lens galaxy's total mass distribution is an `PowerLaw` plus an `ExternalShear`.\n", + " - The source galaxy's light is a `Pixelization`.\n", + "\n", + "Each SLaM pipeline is implemented as a Python function below (e.g. `source_lp`, `source_pix_1`), with a\n", + "documentation string above each function describing the pipeline in detail. The full pipeline is run at the\n", + "bottom of the script." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light, which in\n", + "this example:\n", + "\n", + " - Models the lens galaxy's light as an MGE with 2 x 20 Gaussians.\n", + " - Uses an `Isothermal` model for the lens's total mass distribution with an `ExternalShear`.\n", + " - Models the source galaxy's light as an MGE with 1 x 20 Gaussians.\n", + "\n", + "The mass and source models from this search initialize the SOURCE PIX PIPELINE searches that follow." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source galaxy.\n", + "This pixelization adapts its resolution to the source morphology, assigning more pixels to brighter,\n", + "more detailed regions.\n", + "\n", + "To build an adaptive pixelization, we require an **adapt image**: a lens-light-subtracted image in\n", + "which only the lensed source emission remains. This image determines how both the mesh density and\n", + "regularization weights adapt to the source structure.\n", + "\n", + "The SOURCE LP Pipeline does not provide a sufficiently accurate source model for computing this adapt\n", + "image (e.g., the true source may be more complex than a simple light profile). Therefore, the first\n", + "search of the SOURCE PIX PIPELINE fits a model using a pixelization whose purpose is to generate a\n", + "high-quality adapt image used in search 2.\n", + "\n", + "__Positions__\n", + "\n", + "Image positions are used to prevent unphysical source reconstructions (see `features/pixelization` for\n", + "details). Rather than being input manually, they are computed automatically from the SOURCE LP result.\n", + "This is a key automation feature of SLaM.\n", + "\n", + "__Adapt Images__\n", + "\n", + "An adapt image is computed from the SOURCE LP result and passed to the analysis. This provides an initial\n", + "estimate of the source morphology for the `Adapt` regularization scheme, even though the MGE source model\n", + "may not fully capture the source structure. Search 2 improves upon this using a pixelized adapt image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "The second search of the SOURCE PIX PIPELINE fits the final pixelized source model using the improved\n", + "adapt images computed from search 1's pixelized source reconstruction.\n", + "\n", + "The `RectangularAdaptImage` mesh and `Adapt` regularization adapt the source pixels and regularization\n", + "weights to the source's morphology using the high-quality adapt images from search 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "The LIGHT LP PIPELINE uses one search to fit a complex lens light model to a high level of accuracy, with\n", + "the lens mass model and source light model fixed to the maximum log likelihood result of the SOURCE PIX\n", + "PIPELINE.\n", + "\n", + "In this example:\n", + "\n", + " - The lens galaxy's light is a MGE with 2 x 20 Gaussians.\n", + " - Uses an `Isothermal` mass model with `ExternalShear` for the lens's total mass distribution [fixed from\n", + " SOURCE PIX PIPELINE].\n", + " - Uses a `Pixelization` for the source's light [fixed from SOURCE PIX PIPELINE].\n", + "\n", + "This search aims to produce an accurate model of the lens galaxy's light, which may not have been possible\n", + "in the SOURCE PIPELINE as the mass and source models were not yet precisely estimated. The adapt images\n", + "from SOURCE PIX PIPELINE search 1 are reused, providing a stable basis for the lens-light subtraction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " )\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "The MASS TOTAL PIPELINE uses one search to fit a complex lens mass model to a high level of accuracy,\n", + "using the lens mass model and source model of the SOURCE PIX PIPELINE to initialize model priors, and the\n", + "lens light model of the LIGHT LP PIPELINE.\n", + "\n", + "In this example:\n", + "\n", + " - Uses a linear MGE bulge [fixed from LIGHT LP PIPELINE].\n", + " - Uses a `PowerLaw` model for the lens's total mass distribution [priors initialized from SOURCE PIX\n", + " PIPELINE].\n", + " - Uses a `Pixelization` for the source's light [fixed from SOURCE PIX PIPELINE].\n", + "\n", + "__Positions__\n", + "\n", + "Positions are computed from the SOURCE PIX PIPELINE search 2 result, which provides more precise multiple\n", + "image positions than the SOURCE LP PIPELINE (as the pixelized source gives a better source-plane\n", + "reconstruction)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " bulge = light_result.instance.galaxies.lens.bulge\n", + " disk = light_result.instance.galaxies.lens.disk\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=bulge,\n", + " disk=disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/plot/advanced/plotters_double_einstein_ring.ipynb b/notebooks/guides/plot/advanced/plotters_double_einstein_ring.ipynb index 00d17aaf3..7f2c00fce 100644 --- a/notebooks/guides/plot/advanced/plotters_double_einstein_ring.ipynb +++ b/notebooks/guides/plot/advanced/plotters_double_einstein_ring.ipynb @@ -1,373 +1,410 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plots: Double Einstein Ring\n", - "===========================\n", - "\n", - "This example illustrates the plotting API for double Einstein ring systems, which have\n", - "more than two planes at different redshifts.\n", - "\n", - "The new API uses:\n", - "\n", - " - `aplt.plot_array()` \u2014 plot any 2D array.\n", - " - `aplt.subplot_fit_imaging()` \u2014 multi-panel fit overview.\n", - "\n", - "For pixelized source reconstructions, inversion quantities are accessed via `fit.inversion`\n", - "and plotted with `aplt.plot_array()`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "Refer to `plots/start_here.ipynb` for an introduction to the new plotting API.\n", - "\n", - "__Contents__\n", - "\n", - "- **Setup:** General setup for the analysis.\n", - "- **Fit Imaging:** Plot individual fit attributes with `aplt.plot_array()`.\n", - "- **Full Subplot:** A multi-panel subplot overview is produced with `aplt.subplot_fit_imaging()`.\n", - "- **Pixelized Source Reconstruction:** Now set up a double Einstein ring fit using pixelized source reconstructions.\n", - "- **Inversion:** The inversion is computed directly from a `Tracer` using `TracerToInversion`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Setup__\n", - "\n", - "Set up the double Einstein ring dataset and fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.5\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.5,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.ExponentialCoreSph(\n", - " centre=(-0.15, -0.15), intensity=1.2, effective_radius=0.1\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-0.15, -0.15), einstein_radius=0.3),\n", - ")\n", - "\n", - "source_galaxy_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " bulge=al.lp.ExponentialCoreSph(\n", - " centre=(-0.45, 0.45), intensity=0.6, effective_radius=0.07\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit Imaging__\n", - "\n", - "Plot individual fit attributes with `aplt.plot_array()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=fit.data, title=\"Data\")\n", - "aplt.plot_array(array=fit.noise_map, title=\"Noise Map\")\n", - "aplt.plot_array(array=fit.signal_to_noise_map, title=\"Signal-to-Noise Map\")\n", - "aplt.plot_array(array=fit.model_data, title=\"Model Image\")\n", - "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", - "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", - "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi-Squared Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Plane Images__\n", - "\n", - "For a double Einstein ring (3-plane system), per-plane images are accessed via\n", - "`model_images_of_planes_list`, which has one entry per plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=fit.model_images_of_planes_list[0], title=\"Plane 0 Model Image\")\n", - "aplt.plot_array(array=fit.model_images_of_planes_list[1], title=\"Plane 1 Model Image\")\n", - "aplt.plot_array(array=fit.model_images_of_planes_list[2], title=\"Plane 2 Model Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Full Subplot__\n", - "\n", - "A multi-panel subplot overview is produced with `aplt.subplot_fit_imaging()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixelized Source Reconstruction__\n", - "\n", - "Now set up a double Einstein ring fit using pixelized source reconstructions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " pixelization=al.Pixelization(\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=(24, 24)),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " pixelization=al.Pixelization(\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=(24, 24)),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the fit overview." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixelized source reconstructions are accessed via the `inversion` property of the fit.\n", - "\n", - "For a double Einstein ring there are two reconstructions (one per source plane), indexed\n", - "by their position in `fit.inversion.reconstruction_dict`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=fit.model_images_of_planes_list[1],\n", - " title=\"Plane 1 Model Image (Pixelized)\",\n", - ")\n", - "aplt.plot_array(\n", - " array=fit.model_images_of_planes_list[2],\n", - " title=\"Plane 2 Model Image (Pixelized)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Inversion__\n", - "\n", - "The inversion is computed directly from a `Tracer` using `TracerToInversion`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(\n", - " tracer=tracer,\n", - " dataset=dataset,\n", - ")\n", - "\n", - "inversion = tracer_to_inversion.inversion" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the reconstructed source for each pixelization index." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=fit.model_images_of_planes_list[1],\n", - " title=\"Inversion Reconstruction (Plane 1)\",\n", - ")\n", - "aplt.plot_array(\n", - " array=fit.model_images_of_planes_list[2],\n", - " title=\"Inversion Reconstruction (Plane 2)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plots: Double Einstein Ring\n", + "===========================\n", + "\n", + "This example illustrates the plotting API for double Einstein ring systems, which have\n", + "more than two planes at different redshifts.\n", + "\n", + "The new API uses:\n", + "\n", + " - `aplt.plot_array()` \u2014 plot any 2D array.\n", + " - `aplt.subplot_fit_imaging()` \u2014 multi-panel fit overview.\n", + "\n", + "For pixelized source reconstructions, inversion quantities are accessed via `fit.inversion`\n", + "and plotted with `aplt.plot_array()`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "Refer to `plots/start_here.ipynb` for an introduction to the new plotting API.\n", + "\n", + "__Contents__\n", + "\n", + "- **Setup:** General setup for the analysis.\n", + "- **Fit Imaging:** Plot individual fit attributes with `aplt.plot_array()`.\n", + "- **Full Subplot:** A multi-panel subplot overview is produced with `aplt.subplot_fit_imaging()`.\n", + "- **Pixelized Source Reconstruction:** Now set up a double Einstein ring fit using pixelized source reconstructions.\n", + "- **Inversion:** The inversion is computed directly from a `Tracer` using `TracerToInversion`." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Setup__\n", + "\n", + "Set up the double Einstein ring dataset and fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.5\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.5,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.ExponentialCoreSph(\n", + " centre=(-0.15, -0.15), intensity=1.2, effective_radius=0.1\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-0.15, -0.15), einstein_radius=0.3),\n", + ")\n", + "\n", + "source_galaxy_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " bulge=al.lp.ExponentialCoreSph(\n", + " centre=(-0.45, 0.45), intensity=0.6, effective_radius=0.07\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit Imaging__\n", + "\n", + "Plot individual fit attributes with `aplt.plot_array()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=fit.data, title=\"Data\")\n", + "aplt.plot_array(array=fit.noise_map, title=\"Noise Map\")\n", + "aplt.plot_array(array=fit.signal_to_noise_map, title=\"Signal-to-Noise Map\")\n", + "aplt.plot_array(array=fit.model_data, title=\"Model Image\")\n", + "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", + "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", + "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi-Squared Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Plane Images__\n", + "\n", + "For a double Einstein ring (3-plane system), per-plane images are accessed via\n", + "`model_images_of_planes_list`, which has one entry per plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=fit.model_images_of_planes_list[0], title=\"Plane 0 Model Image\")\n", + "aplt.plot_array(array=fit.model_images_of_planes_list[1], title=\"Plane 1 Model Image\")\n", + "aplt.plot_array(array=fit.model_images_of_planes_list[2], title=\"Plane 2 Model Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Full Subplot__\n", + "\n", + "A multi-panel subplot overview is produced with `aplt.subplot_fit_imaging()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixelized Source Reconstruction__\n", + "\n", + "Now set up a double Einstein ring fit using pixelized source reconstructions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " pixelization=al.Pixelization(\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=(24, 24)),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " pixelization=al.Pixelization(\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=(24, 24)),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the fit overview." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixelized source reconstructions are accessed via the `inversion` property of the fit.\n", + "\n", + "For a double Einstein ring there are two reconstructions (one per source plane), indexed\n", + "by their position in `fit.inversion.reconstruction_dict`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=fit.model_images_of_planes_list[1],\n", + " title=\"Plane 1 Model Image (Pixelized)\",\n", + ")\n", + "aplt.plot_array(\n", + " array=fit.model_images_of_planes_list[2],\n", + " title=\"Plane 2 Model Image (Pixelized)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Inversion__\n", + "\n", + "The inversion is computed directly from a `Tracer` using `TracerToInversion`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(\n", + " tracer=tracer,\n", + " dataset=dataset,\n", + ")\n", + "\n", + "inversion = tracer_to_inversion.inversion" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the reconstructed source for each pixelization index." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=fit.model_images_of_planes_list[1],\n", + " title=\"Inversion Reconstruction (Plane 1)\",\n", + ")\n", + "aplt.plot_array(\n", + " array=fit.model_images_of_planes_list[2],\n", + " title=\"Inversion Reconstruction (Plane 2)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/plot/advanced/plotters_pixelization.ipynb b/notebooks/guides/plot/advanced/plotters_pixelization.ipynb index b780a5a53..32a60b950 100644 --- a/notebooks/guides/plot/advanced/plotters_pixelization.ipynb +++ b/notebooks/guides/plot/advanced/plotters_pixelization.ipynb @@ -1,446 +1,483 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plots: Pixelization\n", - "===================\n", - "\n", - "This example illustrates the plotting API for pixelized source reconstructions.\n", - "\n", - "The new API uses:\n", - "\n", - " - `aplt.plot_array()` \u2014 plot any 2D array (including source reconstructions).\n", - " - `aplt.plot_grid()` \u2014 plot a grid of coordinates.\n", - " - `aplt.subplot_fit_imaging()` \u2014 multi-panel fit overview.\n", - " - `aplt.subplot_fit_interferometer()` \u2014 interferometer fit overview.\n", - "\n", - "Inversion and mapper quantities are accessed via `fit.inversion` and plotted with `aplt.plot_array()`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "Refer to `plots/start_here.ipynb` for an introduction to the new plotting API.\n", - "\n", - "__Contents__\n", - "\n", - "- **Setup:** General setup for the analysis.\n", - "- **Fit Imaging:** Plot the multi-panel fit overview with `aplt.subplot_fit_imaging()`.\n", - "- **Inversion:** The `inversion` property contains the linear algebra, mesh calculations and other key quantities.\n", - "- **Mapper Grids:** The mapper maps pixels from the image-plane to the source-plane pixelization.\n", - "- **Mapper Galaxy Dict:** The mapper galaxy dict maps each mapper to its corresponding galaxy.\n", - "- **Fit Interferometer:** A fit to an interferometer dataset with a pixelized source is plotted with." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Setup__\n", - "\n", - "Set up the dataset and a fit with a pixelized source reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=(24, 24)),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit Imaging__\n", - "\n", - "Plot the multi-panel fit overview with `aplt.subplot_fit_imaging()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot individual fit attributes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=fit.data, title=\"Data\")\n", - "aplt.plot_array(array=fit.model_data, title=\"Model Image\")\n", - "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", - "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", - "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi-Squared Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixelized source reconstruction is accessed via `fit.model_images_of_planes_list[1]`,\n", - "which is the reconstructed image of the source plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=fit.model_images_of_planes_list[1],\n", - " title=\"Source Plane Reconstruction\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Inversion__\n", - "\n", - "The `inversion` property contains the linear algebra, mesh calculations and other key quantities\n", - "used to reconstruct the source galaxy.\n", - "\n", - "The reconstruction is accessed via `fit.inversion.reconstruction`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "\n", - "aplt.plot_array(\n", - " array=fit.model_images_of_planes_list[1],\n", - " title=\"Inversion Reconstruction\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "An inversion can also be computed directly from a `Tracer` using `TracerToInversion`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(\n", - " tracer=tracer,\n", - " dataset=dataset,\n", - ")\n", - "\n", - "inversion = tracer_to_inversion.inversion\n", - "\n", - "aplt.plot_array(\n", - " array=fit.model_images_of_planes_list[1],\n", - " title=\"Inversion Reconstruction (via TracerToInversion)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapper Grids__\n", - "\n", - "The mapper maps pixels from the image-plane to the source-plane pixelization.\n", - "\n", - "We can extract the image-plane and source-plane mesh grids and plot them as overlays." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = inversion.cls_list_from(cls=al.Mapper)[0]\n", - "\n", - "image_plane_mesh_grid = mapper.mask.derive_grid.unmasked\n", - "\n", - "aplt.plot_array(\n", - " array=fit.data,\n", - " title=\"Data with Image-Plane Mesh Grid\",\n", - " positions=image_plane_mesh_grid,\n", - ")\n", - "\n", - "source_plane_mesh_grid = tracer.traced_grid_2d_list_from(grid=image_plane_mesh_grid)[-1]\n", - "\n", - "aplt.plot_grid(\n", - " grid=source_plane_mesh_grid,\n", - " title=\"Source-Plane Mesh Grid\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapper Galaxy Dict__\n", - "\n", - "The mapper galaxy dict maps each mapper to its corresponding galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper_galaxy_dict = tracer_to_inversion.mapper_galaxy_dict\n", - "\n", - "mapper = list(mapper_galaxy_dict)[0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the image-plane mesh grid and source-plane mesh grid together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_plane_mesh_grid = mapper.mask.derive_grid.unmasked\n", - "source_plane_mesh_grid = tracer.traced_grid_2d_list_from(grid=image_plane_mesh_grid)[-1]\n", - "\n", - "aplt.plot_array(\n", - " array=fit.data,\n", - " title=\"Data with Mesh Grid Overlay\",\n", - " positions=image_plane_mesh_grid,\n", - ")\n", - "\n", - "aplt.plot_grid(\n", - " grid=source_plane_mesh_grid,\n", - " title=\"Source-Plane Mesh Grid\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit Interferometer__\n", - "\n", - "A fit to an interferometer dataset with a pixelized source is plotted with\n", - "`aplt.subplot_fit_interferometer()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(200, 200), pixel_scales=0.05, radius=3.0\n", - ")\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerDFT,\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=(24, 24)),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the dirty model image (the model image in real space for an interferometer fit)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=fit.dirty_model_image,\n", - " title=\"Dirty Model Image (Interferometer)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plots: Pixelization\n", + "===================\n", + "\n", + "This example illustrates the plotting API for pixelized source reconstructions.\n", + "\n", + "The new API uses:\n", + "\n", + " - `aplt.plot_array()` \u2014 plot any 2D array (including source reconstructions).\n", + " - `aplt.plot_grid()` \u2014 plot a grid of coordinates.\n", + " - `aplt.subplot_fit_imaging()` \u2014 multi-panel fit overview.\n", + " - `aplt.subplot_fit_interferometer()` \u2014 interferometer fit overview.\n", + "\n", + "Inversion and mapper quantities are accessed via `fit.inversion` and plotted with `aplt.plot_array()`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "Refer to `plots/start_here.ipynb` for an introduction to the new plotting API.\n", + "\n", + "__Contents__\n", + "\n", + "- **Setup:** General setup for the analysis.\n", + "- **Fit Imaging:** Plot the multi-panel fit overview with `aplt.subplot_fit_imaging()`.\n", + "- **Inversion:** The `inversion` property contains the linear algebra, mesh calculations and other key quantities.\n", + "- **Mapper Grids:** The mapper maps pixels from the image-plane to the source-plane pixelization.\n", + "- **Mapper Galaxy Dict:** The mapper galaxy dict maps each mapper to its corresponding galaxy.\n", + "- **Fit Interferometer:** A fit to an interferometer dataset with a pixelized source is plotted with." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Setup__\n", + "\n", + "Set up the dataset and a fit with a pixelized source reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=(24, 24)),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit Imaging__\n", + "\n", + "Plot the multi-panel fit overview with `aplt.subplot_fit_imaging()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot individual fit attributes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=fit.data, title=\"Data\")\n", + "aplt.plot_array(array=fit.model_data, title=\"Model Image\")\n", + "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", + "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", + "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi-Squared Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixelized source reconstruction is accessed via `fit.model_images_of_planes_list[1]`,\n", + "which is the reconstructed image of the source plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=fit.model_images_of_planes_list[1],\n", + " title=\"Source Plane Reconstruction\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Inversion__\n", + "\n", + "The `inversion` property contains the linear algebra, mesh calculations and other key quantities\n", + "used to reconstruct the source galaxy.\n", + "\n", + "The reconstruction is accessed via `fit.inversion.reconstruction`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "\n", + "aplt.plot_array(\n", + " array=fit.model_images_of_planes_list[1],\n", + " title=\"Inversion Reconstruction\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "An inversion can also be computed directly from a `Tracer` using `TracerToInversion`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(\n", + " tracer=tracer,\n", + " dataset=dataset,\n", + ")\n", + "\n", + "inversion = tracer_to_inversion.inversion\n", + "\n", + "aplt.plot_array(\n", + " array=fit.model_images_of_planes_list[1],\n", + " title=\"Inversion Reconstruction (via TracerToInversion)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapper Grids__\n", + "\n", + "The mapper maps pixels from the image-plane to the source-plane pixelization.\n", + "\n", + "We can extract the image-plane and source-plane mesh grids and plot them as overlays." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = inversion.cls_list_from(cls=al.Mapper)[0]\n", + "\n", + "image_plane_mesh_grid = mapper.mask.derive_grid.unmasked\n", + "\n", + "aplt.plot_array(\n", + " array=fit.data,\n", + " title=\"Data with Image-Plane Mesh Grid\",\n", + " positions=image_plane_mesh_grid,\n", + ")\n", + "\n", + "source_plane_mesh_grid = tracer.traced_grid_2d_list_from(grid=image_plane_mesh_grid)[-1]\n", + "\n", + "aplt.plot_grid(\n", + " grid=source_plane_mesh_grid,\n", + " title=\"Source-Plane Mesh Grid\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapper Galaxy Dict__\n", + "\n", + "The mapper galaxy dict maps each mapper to its corresponding galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper_galaxy_dict = tracer_to_inversion.mapper_galaxy_dict\n", + "\n", + "mapper = list(mapper_galaxy_dict)[0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the image-plane mesh grid and source-plane mesh grid together." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_plane_mesh_grid = mapper.mask.derive_grid.unmasked\n", + "source_plane_mesh_grid = tracer.traced_grid_2d_list_from(grid=image_plane_mesh_grid)[-1]\n", + "\n", + "aplt.plot_array(\n", + " array=fit.data,\n", + " title=\"Data with Mesh Grid Overlay\",\n", + " positions=image_plane_mesh_grid,\n", + ")\n", + "\n", + "aplt.plot_grid(\n", + " grid=source_plane_mesh_grid,\n", + " title=\"Source-Plane Mesh Grid\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit Interferometer__\n", + "\n", + "A fit to an interferometer dataset with a pixelized source is plotted with\n", + "`aplt.subplot_fit_interferometer()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(200, 200), pixel_scales=0.05, radius=3.0\n", + ")\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerDFT,\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=(24, 24)),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the dirty model image (the model image in real space for an interferometer fit)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=fit.dirty_model_image,\n", + " title=\"Dirty Model Image (Interferometer)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/plot/examples/mat_plot.ipynb b/notebooks/guides/plot/examples/mat_plot.ipynb index 6dd14ce0e..0be4d0ded 100644 --- a/notebooks/guides/plot/examples/mat_plot.ipynb +++ b/notebooks/guides/plot/examples/mat_plot.ipynb @@ -1,301 +1,338 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plots: Customization\n", - "====================\n", - "\n", - "This example illustrates how to customize the appearance of figures in the new plotting API.\n", - "\n", - "In the old API, customization was done via a `MatPlot2D` object passed to a `*Plotter` class.\n", - "Both `MatPlot2D` and all `*Plotter` classes have been removed.\n", - "\n", - "In the new API, customization is done by passing keyword arguments directly to `aplt.plot_array()`,\n", - "`aplt.plot_grid()`, and `aplt.subplot_*()` functions.\n", - "\n", - "The main customization kwargs are:\n", - "\n", - " - `title`: Figure title string.\n", - " - `colormap`: Matplotlib colormap name (e.g. \"jet\", \"gray\", \"hot\").\n", - " - `use_log10`: Plot colormap in log10 scale.\n", - " - `output_path`: Directory path for saving the figure.\n", - " - `output_format`: File format, e.g. \"png\" or \"pdf\".\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "Refer to `plots/start_here.ipynb` for a general introduction to the new plotting API.\n", - "\n", - "__Contents__\n", - "\n", - "- **Setup:** General setup for the analysis.\n", - "- **Output:** To save a figure to disk, pass `output_path` (a directory) and `output_format`.\n", - "- **Title:** The figure title is set with the `title=` kwarg.\n", - "- **Colormap:** The colormap is set with the `colormap=` kwarg.\n", - "- **Log10:** Many lensing quantities (images, convergence, potential) span many orders of magnitude and are.\n", - "- **Config Defaults:** All default values (colormaps, tick sizes, label fonts, etc.) are configured via the config files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Setup__\n", - "\n", - "Load a dataset to illustrate customization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\") / \"imaging\" / \"slacs1430+4105\"\n", - "data_path = dataset_path / \"data.fits\"\n", - "data = al.Array2D.from_fits(file_path=data_path, hdu=0, pixel_scales=0.03)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "To save a figure to disk, pass `output_path` (a directory) and `output_format`.\n", - "\n", - "The file is saved as `{output_path}/{title}.{output_format}` by default." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=data,\n", - " title=\"example\",\n", - " output_path=Path(\"notebooks\") / \"plot\" / \"plots\",\n", - " output_format=\"png\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Multiple formats can be specified as a list to save the same figure in multiple formats." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=data,\n", - " title=\"example\",\n", - " output_path=Path(\"notebooks\") / \"plot\" / \"plots\",\n", - " output_format=[\"png\", \"pdf\"],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To display the figure on screen (instead of saving it), omit `output_path`.\n", - "\n", - "This is also the default behaviour when no `output_path` is provided." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Title__\n", - "\n", - "The figure title is set with the `title=` kwarg." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"This is the title\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colormap__\n", - "\n", - "The colormap is set with the `colormap=` kwarg.\n", - "\n", - "Any valid matplotlib colormap name can be used." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"Jet Colormap\", colormap=\"jet\")\n", - "aplt.plot_array(array=data, title=\"Hot Colormap\", colormap=\"hot\")\n", - "aplt.plot_array(array=data, title=\"Gray Colormap\", colormap=\"gray\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log10__\n", - "\n", - "Many lensing quantities (images, convergence, potential) span many orders of magnitude and are\n", - "easier to interpret in log10 space.\n", - "\n", - "Pass `use_log10=True` to plot in log10 scale." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"Data (Log10)\", use_log10=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Log10 is particularly useful for convergence and potential maps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.6, ell_comps=(0.0, 0.0)),\n", - ")\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCoreSph(\n", - " centre=(0.0, 0.0), intensity=1.0, effective_radius=0.5, sersic_index=2.0\n", - " ),\n", - ")\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.convergence_2d_from(grid=grid),\n", - " title=\"Convergence (Log10)\",\n", - " use_log10=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Config Defaults__\n", - "\n", - "All default values (colormaps, tick sizes, label fonts, etc.) are configured via the config files:\n", - "\n", - " autolens_workspace/config/visualize/\n", - "\n", - "Key config files and entries:\n", - "\n", - " - `mat_wrap.yaml` -> Figure -> figure: -> figsize\n", - " - `mat_wrap.yaml` -> YLabel -> figure: -> fontsize\n", - " - `mat_wrap.yaml` -> XLabel -> figure: -> fontsize\n", - " - `mat_wrap.yaml` -> TickParams -> figure: -> labelsize\n", - " - `mat_wrap.yaml` -> YTicks -> figure: -> labelsize\n", - " - `mat_wrap.yaml` -> XTicks -> figure: -> labelsize\n", - "\n", - "When no explicit keyword is passed to a plotting function, the config value is used.\n", - "This allows project-wide defaults to be set without changing code." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "Finish.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plots: Customization\n", + "====================\n", + "\n", + "This example illustrates how to customize the appearance of figures in the new plotting API.\n", + "\n", + "In the old API, customization was done via a `MatPlot2D` object passed to a `*Plotter` class.\n", + "Both `MatPlot2D` and all `*Plotter` classes have been removed.\n", + "\n", + "In the new API, customization is done by passing keyword arguments directly to `aplt.plot_array()`,\n", + "`aplt.plot_grid()`, and `aplt.subplot_*()` functions.\n", + "\n", + "The main customization kwargs are:\n", + "\n", + " - `title`: Figure title string.\n", + " - `colormap`: Matplotlib colormap name (e.g. \"jet\", \"gray\", \"hot\").\n", + " - `use_log10`: Plot colormap in log10 scale.\n", + " - `output_path`: Directory path for saving the figure.\n", + " - `output_format`: File format, e.g. \"png\" or \"pdf\".\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "Refer to `plots/start_here.ipynb` for a general introduction to the new plotting API.\n", + "\n", + "__Contents__\n", + "\n", + "- **Setup:** General setup for the analysis.\n", + "- **Output:** To save a figure to disk, pass `output_path` (a directory) and `output_format`.\n", + "- **Title:** The figure title is set with the `title=` kwarg.\n", + "- **Colormap:** The colormap is set with the `colormap=` kwarg.\n", + "- **Log10:** Many lensing quantities (images, convergence, potential) span many orders of magnitude and are.\n", + "- **Config Defaults:** All default values (colormaps, tick sizes, label fonts, etc.) are configured via the config files." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Setup__\n", + "\n", + "Load a dataset to illustrate customization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / \"imaging\" / \"slacs1430+4105\"\n", + "data_path = dataset_path / \"data.fits\"\n", + "data = al.Array2D.from_fits(file_path=data_path, hdu=0, pixel_scales=0.03)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "To save a figure to disk, pass `output_path` (a directory) and `output_format`.\n", + "\n", + "The file is saved as `{output_path}/{title}.{output_format}` by default." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=data,\n", + " title=\"example\",\n", + " output_path=Path(\"notebooks\") / \"plot\" / \"plots\",\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Multiple formats can be specified as a list to save the same figure in multiple formats." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=data,\n", + " title=\"example\",\n", + " output_path=Path(\"notebooks\") / \"plot\" / \"plots\",\n", + " output_format=[\"png\", \"pdf\"],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To display the figure on screen (instead of saving it), omit `output_path`.\n", + "\n", + "This is also the default behaviour when no `output_path` is provided." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Title__\n", + "\n", + "The figure title is set with the `title=` kwarg." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"This is the title\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colormap__\n", + "\n", + "The colormap is set with the `colormap=` kwarg.\n", + "\n", + "Any valid matplotlib colormap name can be used." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"Jet Colormap\", colormap=\"jet\")\n", + "aplt.plot_array(array=data, title=\"Hot Colormap\", colormap=\"hot\")\n", + "aplt.plot_array(array=data, title=\"Gray Colormap\", colormap=\"gray\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log10__\n", + "\n", + "Many lensing quantities (images, convergence, potential) span many orders of magnitude and are\n", + "easier to interpret in log10 space.\n", + "\n", + "Pass `use_log10=True` to plot in log10 scale." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"Data (Log10)\", use_log10=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Log10 is particularly useful for convergence and potential maps." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.6, ell_comps=(0.0, 0.0)),\n", + ")\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCoreSph(\n", + " centre=(0.0, 0.0), intensity=1.0, effective_radius=0.5, sersic_index=2.0\n", + " ),\n", + ")\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.convergence_2d_from(grid=grid),\n", + " title=\"Convergence (Log10)\",\n", + " use_log10=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Config Defaults__\n", + "\n", + "All default values (colormaps, tick sizes, label fonts, etc.) are configured via the config files:\n", + "\n", + " autolens_workspace/config/visualize/\n", + "\n", + "Key config files and entries:\n", + "\n", + " - `mat_wrap.yaml` -> Figure -> figure: -> figsize\n", + " - `mat_wrap.yaml` -> YLabel -> figure: -> fontsize\n", + " - `mat_wrap.yaml` -> XLabel -> figure: -> fontsize\n", + " - `mat_wrap.yaml` -> TickParams -> figure: -> labelsize\n", + " - `mat_wrap.yaml` -> YTicks -> figure: -> labelsize\n", + " - `mat_wrap.yaml` -> XTicks -> figure: -> labelsize\n", + "\n", + "When no explicit keyword is passed to a plotting function, the config value is used.\n", + "This allows project-wide defaults to be set without changing code." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "Finish.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/plot/examples/plotters.ipynb b/notebooks/guides/plot/examples/plotters.ipynb index fc96e9598..c8ce37f55 100644 --- a/notebooks/guides/plot/examples/plotters.ipynb +++ b/notebooks/guides/plot/examples/plotters.ipynb @@ -1,662 +1,699 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plots: Plotters\n", - "===============\n", - "\n", - "This example illustrates the new plotting API for all key PyAutoLens objects.\n", - "\n", - "The old API used dedicated `*Plotter` classes (e.g. `Tracer`, `Imaging`,\n", - "`FitImaging`, `LightProfile`, `MassProfilePlotter`, etc.). These have all been removed.\n", - "\n", - "The new API uses:\n", - "\n", - " - `aplt.plot_array(array, title, ...)` \u2014 plot any 2D array.\n", - " - `aplt.plot_grid(grid, title, ...)` \u2014 plot a grid of coordinates.\n", - " - `aplt.subplot_imaging_dataset(dataset)` \u2014 multi-panel dataset overview.\n", - " - `aplt.subplot_tracer(tracer, grid)` \u2014 multi-panel tracer overview.\n", - " - `aplt.subplot_fit_imaging(fit)` \u2014 multi-panel fit overview.\n", - " - `aplt.subplot_interferometer_dirty_images(dataset)` \u2014 interferometer dataset overview.\n", - " - `aplt.subplot_fit_interferometer(fit)` \u2014 interferometer fit overview.\n", - " - `aplt.subplot_galaxies_images(tracer, grid)` \u2014 per-plane images.\n", - " - `aplt.subplot_fit_point(fit)` \u2014 point source fit overview.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "Refer to `plots/start_here.ipynb` for an introduction to the new plotting API.\n", - "\n", - "__Contents__\n", - "\n", - "- **Setup:** General setup for the analysis.\n", - "- **Array2D:** Any `Array2D` \u2014 images, convergence, noise-maps, etc.\n", - "- **Grid2D:** A `Grid2D` of (y,x) coordinates is plotted with `aplt.plot_grid()`.\n", - "- **Tracer:** Tracer quantities (image, convergence, potential, deflections, magnification) are computed via.\n", - "- **Imaging Dataset:** An `Imaging` dataset's data, noise-map and PSF are plotted individually with `aplt.plot_array()`.\n", - "- **Fit Imaging:** A fit's residuals, chi-squared, model image, etc.\n", - "- **Light Profile:** A light profile image is computed via `image_2d_from()` and plotted with `aplt.plot_array()`.\n", - "- **Mass Profile:** Mass profile quantities are computed and plotted individually.\n", - "- **Galaxy:** A galaxy's image and mass quantities are computed and plotted with `aplt.plot_array()`.\n", - "- **Interferometer:** Interferometer datasets and fits are plotted using their dedicated subplot functions.\n", - "\n", - "__Setup__\n", - "\n", - "Set up standard objects used throughout this example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Array2D__\n", - "\n", - "Any `Array2D` \u2014 images, convergence, noise-maps, etc. \u2014 is plotted with `aplt.plot_array()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "aplt.plot_array(array=dataset.noise_map, title=\"Noise Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid2D__\n", - "\n", - "A `Grid2D` of (y,x) coordinates is plotted with `aplt.plot_grid()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=grid, title=\"Uniform Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A ray-traced (lensed) grid can be computed and plotted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections = tracer.deflections_yx_2d_from(grid=grid)\n", - "lensed_grid = grid.grid_2d_via_deflection_grid_from(deflection_grid=deflections)\n", - "aplt.plot_grid(grid=lensed_grid, title=\"Lensed Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "Tracer quantities (image, convergence, potential, deflections, magnification) are computed\n", - "via method calls and plotted with `aplt.plot_array()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Tracer Image\")\n", - "aplt.plot_array(array=tracer.convergence_2d_from(grid=grid), title=\"Convergence\")\n", - "aplt.plot_array(array=tracer.potential_2d_from(grid=grid), title=\"Potential\")\n", - "\n", - "deflections_yx = tracer.deflections_yx_2d_from(grid=grid)\n", - "\n", - "import autoarray as aa\n", - "\n", - "aplt.plot_array(\n", - " array=aa.Array2D(values=deflections_yx.slim[:, 0], mask=grid.mask),\n", - " title=\"Deflections Y\",\n", - ")\n", - "aplt.plot_array(\n", - " array=aa.Array2D(values=deflections_yx.slim[:, 1], mask=grid.mask),\n", - " title=\"Deflections X\",\n", - ")\n", - "lens_calc = al.LensCalc.from_tracer(tracer=tracer)\n", - "aplt.plot_array(array=lens_calc.magnification_2d_from(grid=grid), title=\"Magnification\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A multi-panel subplot of the tracer is produced with `aplt.subplot_tracer()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_tracer(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A subplot of the per-plane images is produced with `aplt.subplot_galaxies_images()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_galaxies_images(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The source-plane image (plane index 1) is accessed via the image list." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=tracer.image_2d_list_from(grid=grid)[1],\n", - " title=\"Source Plane Image\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Imaging Dataset__\n", - "\n", - "An `Imaging` dataset's data, noise-map and PSF are plotted individually with `aplt.plot_array()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "aplt.plot_array(array=dataset.noise_map, title=\"Noise Map\")\n", - "aplt.plot_array(array=dataset.psf.kernel, title=\"PSF\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A multi-panel subplot of the dataset is produced with `aplt.subplot_imaging_dataset()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit Imaging__\n", - "\n", - "A fit's residuals, chi-squared, model image, etc. are accessed as attributes and plotted\n", - "with `aplt.plot_array()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=fit.data, title=\"Data\")\n", - "aplt.plot_array(array=fit.noise_map, title=\"Noise Map\")\n", - "aplt.plot_array(array=fit.signal_to_noise_map, title=\"Signal-to-Noise Map\")\n", - "aplt.plot_array(array=fit.model_data, title=\"Model Image\")\n", - "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", - "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", - "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi-Squared Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Per-plane model images are accessed via `model_images_of_planes_list`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=fit.model_images_of_planes_list[0], title=\"Plane 0 Model Image\")\n", - "aplt.plot_array(array=fit.model_images_of_planes_list[1], title=\"Plane 1 Model Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A multi-panel fit subplot is produced with `aplt.subplot_fit_imaging()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Light Profile__\n", - "\n", - "A light profile image is computed via `image_2d_from()` and plotted with `aplt.plot_array()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = tracer.galaxies[0].bulge\n", - "aplt.plot_array(array=bulge.image_2d_from(grid=grid), title=\"Bulge Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mass Profile__\n", - "\n", - "Mass profile quantities are computed and plotted individually." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = tracer.galaxies[0].mass\n", - "aplt.plot_array(array=mass.convergence_2d_from(grid=grid), title=\"Mass Convergence\")\n", - "aplt.plot_array(array=mass.potential_2d_from(grid=grid), title=\"Mass Potential\")\n", - "\n", - "mass_deflections = mass.deflections_yx_2d_from(grid=grid)\n", - "aplt.plot_array(\n", - " array=aa.Array2D(values=mass_deflections.slim[:, 0], mask=grid.mask),\n", - " title=\"Mass Deflections Y\",\n", - ")\n", - "aplt.plot_array(\n", - " array=aa.Array2D(values=mass_deflections.slim[:, 1], mask=grid.mask),\n", - " title=\"Mass Deflections X\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy__\n", - "\n", - "A galaxy's image and mass quantities are computed and plotted with `aplt.plot_array()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxy = tracer.galaxies[0]\n", - "aplt.plot_array(array=galaxy.image_2d_from(grid=grid), title=\"Galaxy Image\")\n", - "aplt.plot_array(array=galaxy.convergence_2d_from(grid=grid), title=\"Galaxy Convergence\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__1D Profiles__\n", - "\n", - "1D radial profiles are computed using a projected 2D grid and plotted with matplotlib directly.\n", - "\n", - "There is no 1D plotting function in the new API \u2014 use matplotlib." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_2d_projected = grid.grid_2d_radial_projected_from(\n", - " centre=galaxy.bulge.centre, angle=bulge.angle()\n", - ")\n", - "\n", - "image_1d = galaxy.bulge.image_2d_from(grid=grid_2d_projected)\n", - "\n", - "plt.plot(grid_2d_projected[:, 1], image_1d)\n", - "plt.xlabel(\"Radius (arcseconds)\")\n", - "plt.ylabel(\"Luminosity\")\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Using a radial grid of (y,x) coordinates along the x-axis plots the 1D radial profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "radii = np.arange(10000) * 0.01\n", - "grid_radial = al.Grid2DIrregular(values=[(0.0, r) for r in radii])\n", - "image_1d = bulge.image_2d_from(grid=grid_radial)\n", - "\n", - "plt.plot(radii, image_1d)\n", - "plt.xlabel(\"Radius (arcseconds)\")\n", - "plt.ylabel(\"Luminosity\")\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interferometer__\n", - "\n", - "Interferometer datasets and fits are plotted using their dedicated subplot functions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(200, 200), pixel_scales=0.05, radius=3.0\n", - ")\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerDFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.1, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=1.0,\n", - " sersic_index=2.5,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Dataset / Fit__\n", - "\n", - "A point source fit is plotted with `aplt.subplot_fit_point()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name\n", - "\n", - "if not (dataset_path / \"point_dataset_positions_only.json\").exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/point_source/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.from_json(\n", - " file_path=Path(dataset_path, \"point_dataset_positions_only.json\"),\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.8,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0, point_0=al.ps.PointFlux(centre=(0.0, 0.0), flux=0.8)\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "point_grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.2,\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=point_grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "fit = al.FitPointDataset(dataset=dataset, tracer=tracer, solver=solver)\n", - "\n", - "aplt.subplot_fit_point(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plots: Plotters\n", + "===============\n", + "\n", + "This example illustrates the new plotting API for all key PyAutoLens objects.\n", + "\n", + "The old API used dedicated `*Plotter` classes (e.g. `Tracer`, `Imaging`,\n", + "`FitImaging`, `LightProfile`, `MassProfilePlotter`, etc.). These have all been removed.\n", + "\n", + "The new API uses:\n", + "\n", + " - `aplt.plot_array(array, title, ...)` \u2014 plot any 2D array.\n", + " - `aplt.plot_grid(grid, title, ...)` \u2014 plot a grid of coordinates.\n", + " - `aplt.subplot_imaging_dataset(dataset)` \u2014 multi-panel dataset overview.\n", + " - `aplt.subplot_tracer(tracer, grid)` \u2014 multi-panel tracer overview.\n", + " - `aplt.subplot_fit_imaging(fit)` \u2014 multi-panel fit overview.\n", + " - `aplt.subplot_interferometer_dirty_images(dataset)` \u2014 interferometer dataset overview.\n", + " - `aplt.subplot_fit_interferometer(fit)` \u2014 interferometer fit overview.\n", + " - `aplt.subplot_galaxies_images(tracer, grid)` \u2014 per-plane images.\n", + " - `aplt.subplot_fit_point(fit)` \u2014 point source fit overview.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "Refer to `plots/start_here.ipynb` for an introduction to the new plotting API.\n", + "\n", + "__Contents__\n", + "\n", + "- **Setup:** General setup for the analysis.\n", + "- **Array2D:** Any `Array2D` \u2014 images, convergence, noise-maps, etc.\n", + "- **Grid2D:** A `Grid2D` of (y,x) coordinates is plotted with `aplt.plot_grid()`.\n", + "- **Tracer:** Tracer quantities (image, convergence, potential, deflections, magnification) are computed via.\n", + "- **Imaging Dataset:** An `Imaging` dataset's data, noise-map and PSF are plotted individually with `aplt.plot_array()`.\n", + "- **Fit Imaging:** A fit's residuals, chi-squared, model image, etc.\n", + "- **Light Profile:** A light profile image is computed via `image_2d_from()` and plotted with `aplt.plot_array()`.\n", + "- **Mass Profile:** Mass profile quantities are computed and plotted individually.\n", + "- **Galaxy:** A galaxy's image and mass quantities are computed and plotted with `aplt.plot_array()`.\n", + "- **Interferometer:** Interferometer datasets and fits are plotted using their dedicated subplot functions.\n", + "\n", + "__Setup__\n", + "\n", + "Set up standard objects used throughout this example." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Array2D__\n", + "\n", + "Any `Array2D` \u2014 images, convergence, noise-maps, etc. \u2014 is plotted with `aplt.plot_array()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "aplt.plot_array(array=dataset.noise_map, title=\"Noise Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid2D__\n", + "\n", + "A `Grid2D` of (y,x) coordinates is plotted with `aplt.plot_grid()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=grid, title=\"Uniform Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A ray-traced (lensed) grid can be computed and plotted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections = tracer.deflections_yx_2d_from(grid=grid)\n", + "lensed_grid = grid.grid_2d_via_deflection_grid_from(deflection_grid=deflections)\n", + "aplt.plot_grid(grid=lensed_grid, title=\"Lensed Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "Tracer quantities (image, convergence, potential, deflections, magnification) are computed\n", + "via method calls and plotted with `aplt.plot_array()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Tracer Image\")\n", + "aplt.plot_array(array=tracer.convergence_2d_from(grid=grid), title=\"Convergence\")\n", + "aplt.plot_array(array=tracer.potential_2d_from(grid=grid), title=\"Potential\")\n", + "\n", + "deflections_yx = tracer.deflections_yx_2d_from(grid=grid)\n", + "\n", + "import autoarray as aa\n", + "\n", + "aplt.plot_array(\n", + " array=aa.Array2D(values=deflections_yx.slim[:, 0], mask=grid.mask),\n", + " title=\"Deflections Y\",\n", + ")\n", + "aplt.plot_array(\n", + " array=aa.Array2D(values=deflections_yx.slim[:, 1], mask=grid.mask),\n", + " title=\"Deflections X\",\n", + ")\n", + "lens_calc = al.LensCalc.from_tracer(tracer=tracer)\n", + "aplt.plot_array(array=lens_calc.magnification_2d_from(grid=grid), title=\"Magnification\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A multi-panel subplot of the tracer is produced with `aplt.subplot_tracer()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_tracer(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A subplot of the per-plane images is produced with `aplt.subplot_galaxies_images()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_galaxies_images(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The source-plane image (plane index 1) is accessed via the image list." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=tracer.image_2d_list_from(grid=grid)[1],\n", + " title=\"Source Plane Image\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imaging Dataset__\n", + "\n", + "An `Imaging` dataset's data, noise-map and PSF are plotted individually with `aplt.plot_array()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "aplt.plot_array(array=dataset.noise_map, title=\"Noise Map\")\n", + "aplt.plot_array(array=dataset.psf.kernel, title=\"PSF\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A multi-panel subplot of the dataset is produced with `aplt.subplot_imaging_dataset()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit Imaging__\n", + "\n", + "A fit's residuals, chi-squared, model image, etc. are accessed as attributes and plotted\n", + "with `aplt.plot_array()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=fit.data, title=\"Data\")\n", + "aplt.plot_array(array=fit.noise_map, title=\"Noise Map\")\n", + "aplt.plot_array(array=fit.signal_to_noise_map, title=\"Signal-to-Noise Map\")\n", + "aplt.plot_array(array=fit.model_data, title=\"Model Image\")\n", + "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", + "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", + "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi-Squared Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Per-plane model images are accessed via `model_images_of_planes_list`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=fit.model_images_of_planes_list[0], title=\"Plane 0 Model Image\")\n", + "aplt.plot_array(array=fit.model_images_of_planes_list[1], title=\"Plane 1 Model Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A multi-panel fit subplot is produced with `aplt.subplot_fit_imaging()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Light Profile__\n", + "\n", + "A light profile image is computed via `image_2d_from()` and plotted with `aplt.plot_array()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = tracer.galaxies[0].bulge\n", + "aplt.plot_array(array=bulge.image_2d_from(grid=grid), title=\"Bulge Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mass Profile__\n", + "\n", + "Mass profile quantities are computed and plotted individually." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = tracer.galaxies[0].mass\n", + "aplt.plot_array(array=mass.convergence_2d_from(grid=grid), title=\"Mass Convergence\")\n", + "aplt.plot_array(array=mass.potential_2d_from(grid=grid), title=\"Mass Potential\")\n", + "\n", + "mass_deflections = mass.deflections_yx_2d_from(grid=grid)\n", + "aplt.plot_array(\n", + " array=aa.Array2D(values=mass_deflections.slim[:, 0], mask=grid.mask),\n", + " title=\"Mass Deflections Y\",\n", + ")\n", + "aplt.plot_array(\n", + " array=aa.Array2D(values=mass_deflections.slim[:, 1], mask=grid.mask),\n", + " title=\"Mass Deflections X\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy__\n", + "\n", + "A galaxy's image and mass quantities are computed and plotted with `aplt.plot_array()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxy = tracer.galaxies[0]\n", + "aplt.plot_array(array=galaxy.image_2d_from(grid=grid), title=\"Galaxy Image\")\n", + "aplt.plot_array(array=galaxy.convergence_2d_from(grid=grid), title=\"Galaxy Convergence\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__1D Profiles__\n", + "\n", + "1D radial profiles are computed using a projected 2D grid and plotted with matplotlib directly.\n", + "\n", + "There is no 1D plotting function in the new API \u2014 use matplotlib." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_2d_projected = grid.grid_2d_radial_projected_from(\n", + " centre=galaxy.bulge.centre, angle=bulge.angle()\n", + ")\n", + "\n", + "image_1d = galaxy.bulge.image_2d_from(grid=grid_2d_projected)\n", + "\n", + "plt.plot(grid_2d_projected[:, 1], image_1d)\n", + "plt.xlabel(\"Radius (arcseconds)\")\n", + "plt.ylabel(\"Luminosity\")\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Using a radial grid of (y,x) coordinates along the x-axis plots the 1D radial profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "radii = np.arange(10000) * 0.01\n", + "grid_radial = al.Grid2DIrregular(values=[(0.0, r) for r in radii])\n", + "image_1d = bulge.image_2d_from(grid=grid_radial)\n", + "\n", + "plt.plot(radii, image_1d)\n", + "plt.xlabel(\"Radius (arcseconds)\")\n", + "plt.ylabel(\"Luminosity\")\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interferometer__\n", + "\n", + "Interferometer datasets and fits are plotted using their dedicated subplot functions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(200, 200), pixel_scales=0.05, radius=3.0\n", + ")\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerDFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.1, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=1.0,\n", + " sersic_index=2.5,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Dataset / Fit__\n", + "\n", + "A point source fit is plotted with `aplt.subplot_fit_point()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name\n", + "\n", + "if not (dataset_path / \"point_dataset_positions_only.json\").exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/point_source/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.from_json(\n", + " file_path=Path(dataset_path, \"point_dataset_positions_only.json\"),\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.8,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0, point_0=al.ps.PointFlux(centre=(0.0, 0.0), flux=0.8)\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "point_grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.2,\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=point_grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "fit = al.FitPointDataset(dataset=dataset, tracer=tracer, solver=solver)\n", + "\n", + "aplt.subplot_fit_point(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/plot/examples/searches.ipynb b/notebooks/guides/plot/examples/searches.ipynb index 0b944ed41..fd8f381a5 100644 --- a/notebooks/guides/plot/examples/searches.ipynb +++ b/notebooks/guides/plot/examples/searches.ipynb @@ -1,1192 +1,1229 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plots: Searches\n", - "===============\n", - "\n", - "This example illustrates the API for plotting the results of different non-linear searches.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "You should refer to the `plots/start_here.ipynb` notebook first for a description of how visuals work and the default\n", - "behaviour of plotting visuals.\n", - "\n", - "__Contents__\n", - "\n", - "- **Setup:** General setup for the analysis.\n", - "- **Notation:** Plot are labeled with short hand parameter names (e.g.\n", - "- **DynestyPlotter:** We set up the Dynesty non-linear search and perform the fit to get the samples we will plot below.\n", - "- **Plots:** All plots use dynesty's inbuilt plotting library and the model.\n", - "- **EmceePlotter:** We now use the MCMC plotting functions to visualize emcee's results.\n", - "- **Search Specific Visualization:** The internal sampler can be used to plot the results of the non-linear search.\n", - "- **ZeusPlotter:** We now use the MCMC plotting functions to visualize Zeus's results.\n", - "- **GetDist:** This example illustrates how to plot visualization summarizing the results of model-fit using any.\n", - "- **Parameter Names:** Note that in order to customize the figure, we will use the `samples.model.parameter_names` list.\n", - "- **GetDist Plotter:** To make plots we use a GetDist plotter object, which can be customized to change the appearance of.\n", - "- **GetDist Subplots:** Using the plotter we can make different plots, for example a triangle plot showing the 1D and 2D.\n", - "- **GetDist Single Plots:** We can make plots of specific 1D or 2D PDFs, using the single plotter object.\n", - "- **Output:** A figure can be output using standard matplotlib functionality.\n", - "- **GetDist Other Plots:** There are many more ways to visualize PDFs possible with GetDist, checkout the official.\n", - "- **Plotting Multiple Samples:** Finally, we can plot the results of multiple different non-linear searches on the same plot, using.\n", - "\n", - "__Setup__\n", - "\n", - "To illustrate plotting, we require standard objects like a dataset and model which we will perform quick model-fits to\n", - "for illustration." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "dataset_name = \"simple__no_lens_light\"\n", - "\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Notation__\n", - "\n", - "Plot are labeled with short hand parameter names (e.g. the `centre` parameters are plotted using an `x`). \n", - "\n", - "The mappings of every parameter to its shorthand symbol for plots is specified in the `config/notation.yaml` file \n", - "and can be customized.\n", - "\n", - "Each label also has a superscript corresponding to the model component the parameter originates from. For example,\n", - "Gaussians are given the superscript `g`. This can also be customized in the `config/notation.yaml` file.\n", - "\n", - "__DynestyPlotter__\n", - "\n", - "We set up the Dynesty non-linear search and perform the fit to get the samples we will plot below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.DynestyStatic(\n", - " path_prefix=Path(\"plot\"),\n", - " name=\"DynestyPlotter\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now pass the samples to a `DynestyPlotter` which will allow us to use the corner plotting function of the\n", - "public library anesthetic\n", - "\n", - "The Dynesty readthedocs describes fully all of the methods used below \n", - "\n", - " - https://dynesty-sampler.readthedocs.io/en/latest/quickstart.html\n", - " - https://dynesty-sampler.readthedocs.io/en/latest/api.html#module-Dynesty.plotting\n", - " \n", - "In all the examples below, we use the `kwargs` of this function to pass in any of the input parameters that are \n", - "described in the API docs.\n", - "\n", - "Dynesty plotters use `_kwargs` dictionaries to pass visualization settings to matplotlib lib. For example, below,\n", - "we:\n", - "\n", - " - Set the fontsize of the x and y labels by passing `label_kwargs={\"fontsize\": 16}`.\n", - " - Set the fontsize of the title by passing `title_kwargs={\"fontsize\": \"10\"}`.\n", - " \n", - "There are other `_kwargs` inputs we pass as None, you should check out the Dynesty docs if you need to customize your\n", - "figure." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# %%\n", - "'''\n", - "The `corner_anesthetic` function produces a triangle of 1D and 2D PDF's of every parameter using the library `anesthetic`.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `corner_cornerpy` function produces a triangle of 1D and 2D PDF's of every parameter using the library `corner.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_cornerpy(\n", - " samples=result.samples,\n", - " dims=None,\n", - " span=None,\n", - " quantiles=[0.025, 0.5, 0.975],\n", - " color=\"black\",\n", - " smooth=0.02,\n", - " quantiles_2d=None,\n", - " hist_kwargs=None,\n", - " hist2d_kwargs=None,\n", - " label_kwargs={\"fontsize\": 16},\n", - " show_titles=True,\n", - " title_fmt=\".2f\",\n", - " title_kwargs={\"fontsize\": \"10\"},\n", - " truths=None,\n", - " truth_color=\"red\",\n", - " truth_kwargs=None,\n", - " max_n_ticks=5,\n", - " top_ticks=False,\n", - " use_math_text=False,\n", - " verbose=False,\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The internal sampler can be used to plot the results of the non-linear search. \n", - "\n", - "We do this using the `search_internal` attribute which contains the sampler in its native form.\n", - "\n", - "The first time you run a search, the `search_internal` attribute will be available because it is passed ot the\n", - "result via memory. \n", - "\n", - "If you rerun the fit on a completed result, it will not be available in memory, and therefore\n", - "will be loaded from the `files/search_internal` folder. The `search_internal` entry of the `output.yaml` must be true \n", - "for this to be possible." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_internal = result.search_internal" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Plots__\n", - "\n", - "All plots use dynesty's inbuilt plotting library and the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " from dynesty import plotting as dyplot\n", - "\n", - " model = result.model\n", - "\n", - " \"\"\"\n", - " The boundplot plots the bounding distribution used to propose either (1) live points at a given iteration or (2) a \n", - " specific dead point during the course of a run, projected onto the two dimensions specified by `dims`.\n", - " \"\"\"\n", - " dyplot.boundplot(\n", - " results=search_internal.results,\n", - " labels=model.parameter_labels_with_superscripts_latex,\n", - " dims=(2, 2),\n", - " it=-1, # -1 is the final iteration of the dynesty samples, change this to plot a different iteration\n", - " idx=None,\n", - " prior_transform=None,\n", - " periodic=None,\n", - " reflective=None,\n", - " ndraws=5000,\n", - " color=\"gray\",\n", - " plot_kwargs=None,\n", - " label_kwargs={\"fontsize\": 16},\n", - " max_n_ticks=5,\n", - " use_math_text=False,\n", - " show_live=False,\n", - " live_color=\"darkviolet\",\n", - " live_kwargs=None,\n", - " span=None,\n", - " fig=None,\n", - " )\n", - "\n", - " plt.show()\n", - " plt.close()\n", - "\n", - " \"\"\"\n", - " The cornerbound plots the bounding distribution used to propose either (1) live points at a given iteration or (2) a \n", - " specific dead point during the course of a run, projected onto all pairs of dimensions.\n", - " \"\"\"\n", - " try:\n", - " dyplot.cornerbound(\n", - " results=search_internal.results,\n", - " labels=model.parameter_labels_with_superscripts_latex,\n", - " it=-1, # -1 is the final iteration of the dynesty samples, change this to plot a different iteration\n", - " idx=None,\n", - " dims=None,\n", - " prior_transform=None,\n", - " periodic=None,\n", - " reflective=None,\n", - " ndraws=5000,\n", - " color=\"gray\",\n", - " plot_kwargs=None,\n", - " label_kwargs={\"fontsize\": 16},\n", - " max_n_ticks=5,\n", - " use_math_text=False,\n", - " show_live=False,\n", - " live_color=\"darkviolet\",\n", - " live_kwargs=None,\n", - " span=None,\n", - " fig=None,\n", - " )\n", - "\n", - " plt.show()\n", - " plt.close()\n", - "\n", - " except ValueError:\n", - " pass\n", - "\n", - " \"\"\"\n", - " The cornerplot plots a corner plot of the 1-D and 2-D marginalized posteriors.\n", - " \"\"\"\n", - "\n", - " try:\n", - " dyplot.cornerplot(\n", - " results=search_internal.results,\n", - " labels=model.parameter_labels_with_superscripts_latex,\n", - " dims=None,\n", - " span=None,\n", - " quantiles=[0.025, 0.5, 0.975],\n", - " color=\"black\",\n", - " smooth=0.02,\n", - " quantiles_2d=None,\n", - " hist_kwargs=None,\n", - " hist2d_kwargs=None,\n", - " label_kwargs={\"fontsize\": 16},\n", - " show_titles=True,\n", - " title_fmt=\".2f\",\n", - " title_kwargs={\"fontsize\": \"10\"},\n", - " truths=None,\n", - " truth_color=\"red\",\n", - " truth_kwargs=None,\n", - " max_n_ticks=5,\n", - " top_ticks=False,\n", - " use_math_text=False,\n", - " verbose=False,\n", - " )\n", - "\n", - " plt.show()\n", - " plt.close()\n", - "\n", - " except ValueError:\n", - " pass\n", - "\n", - " \"\"\"\n", - " The cornerpoints plots a (sub-)corner plot of (weighted) samples.\n", - " \"\"\"\n", - " dyplot.cornerpoints(\n", - " results=search_internal.results,\n", - " labels=model.parameter_labels_with_superscripts_latex,\n", - " dims=None,\n", - " thin=1,\n", - " span=None,\n", - " cmap=\"plasma\",\n", - " color=None,\n", - " kde=True,\n", - " nkde=1000,\n", - " plot_kwargs=None,\n", - " label_kwargs={\"fontsize\": 16},\n", - " truths=None,\n", - " truth_color=\"red\",\n", - " truth_kwargs=None,\n", - " max_n_ticks=5,\n", - " use_math_text=False,\n", - " fig=None,\n", - " )\n", - "\n", - " plt.show()\n", - " plt.close()\n", - "\n", - " \"\"\"\n", - " The runplot plots live points, ln(likelihood), ln(weight), and ln(evidence) as a function of ln(prior volume).\n", - " \"\"\"\n", - " dyplot.runplot(\n", - " results=search_internal.results,\n", - " span=None,\n", - " logplot=False,\n", - " kde=True,\n", - " nkde=1000,\n", - " color=\"blue\",\n", - " plot_kwargs=None,\n", - " label_kwargs={\"fontsize\": 16},\n", - " lnz_error=True,\n", - " lnz_truth=None,\n", - " truth_color=\"red\",\n", - " truth_kwargs=None,\n", - " max_x_ticks=8,\n", - " max_y_ticks=3,\n", - " use_math_text=True,\n", - " mark_final_live=True,\n", - " fig=None,\n", - " )\n", - "\n", - " plt.show()\n", - " plt.close()\n", - "\n", - " \"\"\"\n", - " The traceplot plots traces and marginalized posteriors for each parameter.\n", - " \"\"\"\n", - " dyplot.traceplot(\n", - " results=search_internal.results,\n", - " span=None,\n", - " quantiles=[0.025, 0.5, 0.975],\n", - " smooth=0.02,\n", - " thin=1,\n", - " dims=None,\n", - " post_color=\"blue\",\n", - " post_kwargs=None,\n", - " kde=True,\n", - " nkde=1000,\n", - " trace_cmap=\"plasma\",\n", - " trace_color=None,\n", - " trace_kwargs=None,\n", - " connect=False,\n", - " connect_highlight=10,\n", - " connect_color=\"red\",\n", - " connect_kwargs=None,\n", - " max_n_ticks=5,\n", - " use_math_text=False,\n", - " label_kwargs={\"fontsize\": 16},\n", - " show_titles=True,\n", - " title_fmt=\".2f\",\n", - " title_kwargs={\"fontsize\": \"10\"},\n", - " truths=None,\n", - " truth_color=\"red\",\n", - " truth_kwargs=None,\n", - " verbose=False,\n", - " fig=None,\n", - " )\n", - "\n", - " plt.show()\n", - " plt.close()\n", - "except (AttributeError, TypeError):\n", - " pass # search_internal unavailable in test mode" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__EmceePlotter__\n", - "\n", - "We now use the MCMC plotting functions to visualize emcee's results.\n", - "\n", - "The emcee readthedocs describes fully all of the methods used below\n", - "\n", - " - https://emcee.readthedocs.io/en/stable/user/sampler/\n", - "\n", - " The `corner_cornerpy` function wraps the library `corner.py` to make corner plots of the PDF:\n", - "\n", - "- https://corner.readthedocs.io/en/latest/index.html\n", - "\n", - "In all the examples below, we use the `kwargs` of this function to pass in any of the input parameters that are\n", - "described in the API docs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "The `corner_cornerpy` function produces a triangle of 1D and 2D PDF's of every parameter in the model fit.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_cornerpy(\n", - " samples=result.samples,\n", - " bins=20,\n", - " range=None,\n", - " color=\"k\",\n", - " hist_bin_factor=1,\n", - " smooth=None,\n", - " smooth1d=None,\n", - " label_kwargs=None,\n", - " titles=None,\n", - " show_titles=False,\n", - " title_fmt=\".2f\",\n", - " title_kwargs=None,\n", - " truths=None,\n", - " truth_color=\"#4682b4\",\n", - " scale_hist=False,\n", - " quantiles=None,\n", - " verbose=False,\n", - " fig=None,\n", - " max_n_ticks=5,\n", - " top_ticks=False,\n", - " use_math_text=False,\n", - " reverse=False,\n", - " labelpad=0.0,\n", - " hist_kwargs=None,\n", - " group=\"posterior\",\n", - " var_names=None,\n", - " filter_vars=None,\n", - " coords=None,\n", - " divergences=False,\n", - " divergences_kwargs=None,\n", - " labeller=None,\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Specific Visualization__\n", - "\n", - "The internal sampler can be used to plot the results of the non-linear search. \n", - "\n", - "We do this using the `search_internal` attribute which contains the sampler in its native form.\n", - "\n", - "The first time you run a search, the `search_internal` attribute will be available because it is passed ot the\n", - "result via memory. \n", - "\n", - "If you rerun the fit on a completed result, it will not be available in memory, and therefore\n", - "will be loaded from the `files/search_internal` folder. The `search_internal` entry of the `output.yaml` must be true \n", - "for this to be possible." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_internal = result.search_internal" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The method below shows a 2D projection of the walker trajectories." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " fig, axes = plt.subplots(result.model.prior_count, figsize=(10, 7))\n", - "\n", - " for i in range(result.model.prior_count):\n", - " for walker_index in range(search_internal.get_log_prob().shape[1]):\n", - " ax = axes[i]\n", - " ax.plot(\n", - " search_internal.get_chain()[:, walker_index, i],\n", - " search_internal.get_log_prob()[:, walker_index],\n", - " alpha=0.3,\n", - " )\n", - "\n", - " ax.set_ylabel(\"Log Likelihood\")\n", - " ax.set_xlabel(result.model.parameter_labels_with_superscripts_latex[i])\n", - "\n", - " plt.show()\n", - "\n", - " \"\"\"\n", - " This method shows the likelihood as a series of steps.\n", - " \"\"\"\n", - "\n", - " fig, axes = plt.subplots(1, figsize=(10, 7))\n", - "\n", - " for walker_index in range(search_internal.get_log_prob().shape[1]):\n", - " axes.plot(search_internal.get_log_prob()[:, walker_index], alpha=0.3)\n", - "\n", - " axes.set_ylabel(\"Log Likelihood\")\n", - " axes.set_xlabel(\"step number\")\n", - "\n", - " plt.show()\n", - "\n", - " \"\"\"\n", - " This method shows the parameter values of every walker at every step.\n", - " \"\"\"\n", - " fig, axes = plt.subplots(\n", - " result.samples.model.prior_count, figsize=(10, 7), sharex=True\n", - " )\n", - "\n", - " for i in range(result.samples.model.prior_count):\n", - " ax = axes[i]\n", - " ax.plot(search_internal.get_chain()[:, :, i], alpha=0.3)\n", - " ax.set_ylabel(result.model.parameter_labels_with_superscripts_latex[i])\n", - "\n", - " axes[-1].set_xlabel(\"step number\")\n", - "\n", - " plt.show()\n", - "except AttributeError:\n", - " pass # MCMC-specific methods not available for non-MCMC searches" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__ZeusPlotter__\n", - "\n", - "We now use the MCMC plotting functions to visualize Zeus's results.\n", - "\n", - "The zeus readthedocs describes fully all of the methods used below\n", - "\n", - " - https://zeus-mcmc.readthedocs.io/en/latest/api/plotting.html#cornerplot\n", - " - https://zeus-mcmc.readthedocs.io/en/latest/notebooks/normal_distribution.html\n", - "\n", - " The `corner_cornerpy` function wraps the library `corner.py` to make corner plots of the PDF:\n", - "\n", - "- https://corner.readthedocs.io/en/latest/index.html\n", - "\n", - "In all the examples below, we use the `kwargs` of this function to pass in any of the input parameters that are\n", - "described in the API docs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "The `corner_cornerpy` function produces a triangle of 1D and 2D PDF's of every parameter in the model fit.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_cornerpy(\n", - " samples=result.samples,\n", - " weight_list=None,\n", - " levels=None,\n", - " span=None,\n", - " quantiles=[0.025, 0.5, 0.975],\n", - " truth=None,\n", - " color=None,\n", - " alpha=0.5,\n", - " linewidth=1.5,\n", - " fill=True,\n", - " fontsize=10,\n", - " show_titles=True,\n", - " title_fmt=\".2f\",\n", - " title_fontsize=12,\n", - " cut=3,\n", - " fig=None,\n", - " size=(10, 10),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Specific Visualization__\n", - "\n", - "The internal sampler can be used to plot the results of the non-linear search. \n", - "\n", - "We do this using the `search_internal` attribute which contains the sampler in its native form.\n", - "\n", - "For zeus, the `search_internal` attribute is only available if the zeus sampler results are output to hard-disk\n", - "via hdf5. The `search_internal` entry of the `output.yaml` must be true for this to be the case." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_internal = result.search_internal" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Plots__\n", - "\n", - "The method below shows a 2D projection of the walker trajectories." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " fig, axes = plt.subplots(result.model.prior_count, figsize=(10, 7))\n", - "\n", - " for i in range(result.model.prior_count):\n", - " for walker_index in range(search_internal.get_log_prob().shape[1]):\n", - " ax = axes[i]\n", - " ax.plot(\n", - " search_internal.get_chain()[:, walker_index, i],\n", - " search_internal.get_log_prob()[:, walker_index],\n", - " alpha=0.3,\n", - " )\n", - "\n", - " ax.set_ylabel(\"Log Likelihood\")\n", - " ax.set_xlabel(result.model.parameter_labels_with_superscripts_latex[i])\n", - "\n", - " plt.show()\n", - "\n", - " \"\"\"\n", - " This method shows the likelihood as a series of steps.\n", - " \"\"\"\n", - "\n", - " fig, axes = plt.subplots(1, figsize=(10, 7))\n", - "\n", - " for walker_index in range(search_internal.get_log_prob().shape[1]):\n", - " axes.plot(search_internal.get_log_prob()[:, walker_index], alpha=0.3)\n", - "\n", - " axes.set_ylabel(\"Log Likelihood\")\n", - " axes.set_xlabel(\"step number\")\n", - "\n", - " plt.show()\n", - "\n", - " \"\"\"\n", - " This method shows the parameter values of every walker at every step.\n", - " \"\"\"\n", - " fig, axes = plt.subplots(\n", - " result.samples.model.prior_count, figsize=(10, 7), sharex=True\n", - " )\n", - "\n", - " for i in range(result.samples.model.prior_count):\n", - " ax = axes[i]\n", - " ax.plot(search_internal.get_chain()[:, :, i], alpha=0.3)\n", - " ax.set_ylabel(result.model.parameter_labels_with_superscripts_latex[i])\n", - "\n", - " axes[-1].set_xlabel(\"step number\")\n", - "\n", - " plt.show()\n", - "except AttributeError:\n", - " pass # MCMC-specific methods not available for non-MCMC searches\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__GetDist__\n", - "\n", - "This example illustrates how to plot visualization summarizing the results of model-fit using any non-linear search\n", - "using GetDist:\n", - "\n", - " - https://getdist.readthedocs.io/en/latest/\n", - "\n", - "GetDist is an optional library which creates 1D and 2D plots of probability distribution functions (PDF)s. Its\n", - "visualization tools has more than the in-built visualization tools of many non-linear searches (e.g. Nautilus /\n", - "emcee) and can often produce better looking plots.\n", - "\n", - "GetDist was developed for the analysis of Cosmological datasets.\n", - "\n", - "Because GetDist is an optional library, you will likely have to install it manually via the command:\n", - "\n", - "`pip install getdist`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from getdist import MCSamples\n", - "from getdist import plots\n", - "import numpy as np" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "GetDist uses a `model.paramnames` file to load the name of every parameter in the model-fit and pair it with the\n", - "latex symbol used to represent it in plots.\n", - "\n", - "This file is not created by default, but can be output by the `search.paths` object as shown below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search.paths._save_parameter_names_file(model=model)\n", - "search.paths.zip_remove()\n", - "search.paths._zip()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "GetDist uses an `MCSamples` object to store the samples of a non-linear search.\n", - "\n", - "We create this object via a conversion from **PyAutoFit** `Samples`, as well as using the `names`\n", - "and `labels` of parameters in the `Samples` object.\n", - "\n", - "The input `sampler=\"nested\"` is input because we used a nested sampling, `Nautilus`. For MCMC this should be\n", - "replaced with \"mcmc\"." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "gd_samples = MCSamples(\n", - " samples=np.asarray(samples.parameter_lists),\n", - " loglikes=np.asarray(samples.log_likelihood_list),\n", - " weights=np.asarray(samples.weight_list),\n", - " names=samples.model.model_component_and_parameter_names,\n", - " labels=samples.model.parameter_labels_with_superscripts,\n", - " sampler=\"nested\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parameter Names__\n", - "\n", - "Note that in order to customize the figure, we will use the `samples.model.parameter_names` list." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(samples.model.model_component_and_parameter_names)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__GetDist Plotter__\n", - "\n", - "To make plots we use a GetDist plotter object, which can be customized to change the appearance of the plots." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "gd_plotter = plots.get_subplot_plotter(width_inch=12)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__GetDist Subplots__\n", - "\n", - "Using the plotter we can make different plots, for example a triangle plot showing the 1D and 2D PDFs of every \n", - "parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " gd_plotter.triangle_plot(roots=gd_samples, filled=True)\n", - " plt.show()\n", - " plt.close()\n", - "except Exception:\n", - " pass" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A triangle plot with specific parameters can be plotted by using the `params` input, whereby we specify the specific\n", - "parameter names to plot." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " gd_plotter.triangle_plot(\n", - " roots=gd_samples,\n", - " filled=True,\n", - " params=[\n", - " \"galaxies_lens_mass_einstein_radius\",\n", - " \"galaxies_lens_mass_ell_comps_0\",\n", - " ],\n", - " )\n", - "except Exception:\n", - " pass\n", - "\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Rectangle plots can be used to show specific 2D combinations of parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " gd_plotter.rectangle_plot(\n", - " roots=gd_samples,\n", - " yparams=[\"galaxies_lens_mass_einstein_radius\"],\n", - " xparams=[\n", - " \"galaxies_lens_mass_ell_comps_0\",\n", - " \"galaxies_lens_mass_ell_comps_1\",\n", - " ],\n", - " )\n", - "except Exception:\n", - " pass\n", - "\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__GetDist Single Plots__\n", - "\n", - "We can make plots of specific 1D or 2D PDFs, using the single plotter object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "gd_plotter = plots.get_single_plotter()\n", - "\n", - "try:\n", - " gd_plotter.plot_1d(roots=gd_samples, param=\"galaxies_lens_mass_einstein_radius\")\n", - "except Exception:\n", - " pass\n", - "\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "gd_plotter = plots.get_single_plotter()\n", - "\n", - "try:\n", - " gd_plotter.plot_2d(\n", - " roots=gd_samples,\n", - " param1=\"galaxies_lens_mass_einstein_radius\",\n", - " param2=\"galaxies_lens_mass_ell_comps_0\",\n", - " )\n", - "except Exception:\n", - " pass\n", - "\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also make a 3D plot, where the 2D PDF is plotted colored by the value of a third parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "gd_plotter = plots.get_single_plotter()\n", - "\n", - "try:\n", - " gd_plotter.plot_3d(\n", - " roots=gd_samples,\n", - " params=[\n", - " \"galaxies_lens_mass_einstein_radius\",\n", - " \"galaxies_lens_mass_ell_comps_0\",\n", - " \"galaxies_lens_mass_ell_comps_1\",\n", - " ],\n", - " )\n", - "except Exception:\n", - " pass\n", - "\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "A figure can be output using standard matplotlib functionality." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "gd_plotter = plots.get_single_plotter()\n", - "\n", - "try:\n", - " gd_plotter.plot_3d(roots=gd_samples, params=[\"centre\", \"sigma\", \"normalization\"])\n", - "except Exception:\n", - " pass\n", - "\n", - "output_path = Path(\"output\")\n", - "\n", - "plt.savefig(Path(output_path, \"getdist.png\"))\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__GetDist Other Plots__\n", - "\n", - "There are many more ways to visualize PDFs possible with GetDist, checkout the official documentation for them all!\n", - "\n", - " - https://getdist.readthedocs.io/en/latest/\n", - " - https://getdist.readthedocs.io/en/latest/plots.html\n", - "\n", - "__Plotting Multiple Samples__\n", - "\n", - "Finally, we can plot the results of multiple different non-linear searches on the same plot, using all\n", - "of the functions above.\n", - "\n", - "Lets quickly make a second set of `Nautilus` results and plot them on the same figure above with the results\n", - "of the first search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "Nautilus = af.Nautilus(path_prefix=\"plot\", name=\"GetDist_2\")\n", - "\n", - "result_extra = Nautilus.fit(model=model, analysis=analysis)\n", - "\n", - "samples_extra = result_extra.samples\n", - "\n", - "gd_samples_extra = MCSamples(\n", - " samples=np.asarray(samples_extra.parameter_lists),\n", - " loglikes=np.asarray(samples_extra.log_likelihood_list),\n", - " weights=np.asarray(samples_extra.weight_list),\n", - " names=samples_extra.model.model_component_and_parameter_names,\n", - " labels=samples.model.parameter_labels_with_superscripts,\n", - " sampler=\"nested\",\n", - ")\n", - "\n", - "gd_plotter = plots.get_subplot_plotter(width_inch=12)\n", - "\n", - "try:\n", - " gd_plotter.triangle_plot(roots=[gd_samples, gd_samples_extra], filled=True)\n", - "except Exception:\n", - " pass\n", - "\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that the models do not need to be the same to make the plots above.\n", - "\n", - "GetDist will clever use the `names` of the parameters to combine the parameters into customizeable PDF plots." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plots: Searches\n", + "===============\n", + "\n", + "This example illustrates the API for plotting the results of different non-linear searches.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "You should refer to the `plots/start_here.ipynb` notebook first for a description of how visuals work and the default\n", + "behaviour of plotting visuals.\n", + "\n", + "__Contents__\n", + "\n", + "- **Setup:** General setup for the analysis.\n", + "- **Notation:** Plot are labeled with short hand parameter names (e.g.\n", + "- **DynestyPlotter:** We set up the Dynesty non-linear search and perform the fit to get the samples we will plot below.\n", + "- **Plots:** All plots use dynesty's inbuilt plotting library and the model.\n", + "- **EmceePlotter:** We now use the MCMC plotting functions to visualize emcee's results.\n", + "- **Search Specific Visualization:** The internal sampler can be used to plot the results of the non-linear search.\n", + "- **ZeusPlotter:** We now use the MCMC plotting functions to visualize Zeus's results.\n", + "- **GetDist:** This example illustrates how to plot visualization summarizing the results of model-fit using any.\n", + "- **Parameter Names:** Note that in order to customize the figure, we will use the `samples.model.parameter_names` list.\n", + "- **GetDist Plotter:** To make plots we use a GetDist plotter object, which can be customized to change the appearance of.\n", + "- **GetDist Subplots:** Using the plotter we can make different plots, for example a triangle plot showing the 1D and 2D.\n", + "- **GetDist Single Plots:** We can make plots of specific 1D or 2D PDFs, using the single plotter object.\n", + "- **Output:** A figure can be output using standard matplotlib functionality.\n", + "- **GetDist Other Plots:** There are many more ways to visualize PDFs possible with GetDist, checkout the official.\n", + "- **Plotting Multiple Samples:** Finally, we can plot the results of multiple different non-linear searches on the same plot, using.\n", + "\n", + "__Setup__\n", + "\n", + "To illustrate plotting, we require standard objects like a dataset and model which we will perform quick model-fits to\n", + "for illustration." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "dataset_name = \"simple__no_lens_light\"\n", + "\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Notation__\n", + "\n", + "Plot are labeled with short hand parameter names (e.g. the `centre` parameters are plotted using an `x`). \n", + "\n", + "The mappings of every parameter to its shorthand symbol for plots is specified in the `config/notation.yaml` file \n", + "and can be customized.\n", + "\n", + "Each label also has a superscript corresponding to the model component the parameter originates from. For example,\n", + "Gaussians are given the superscript `g`. This can also be customized in the `config/notation.yaml` file.\n", + "\n", + "__DynestyPlotter__\n", + "\n", + "We set up the Dynesty non-linear search and perform the fit to get the samples we will plot below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.DynestyStatic(\n", + " path_prefix=Path(\"plot\"),\n", + " name=\"DynestyPlotter\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now pass the samples to a `DynestyPlotter` which will allow us to use the corner plotting function of the\n", + "public library anesthetic\n", + "\n", + "The Dynesty readthedocs describes fully all of the methods used below \n", + "\n", + " - https://dynesty-sampler.readthedocs.io/en/latest/quickstart.html\n", + " - https://dynesty-sampler.readthedocs.io/en/latest/api.html#module-Dynesty.plotting\n", + " \n", + "In all the examples below, we use the `kwargs` of this function to pass in any of the input parameters that are \n", + "described in the API docs.\n", + "\n", + "Dynesty plotters use `_kwargs` dictionaries to pass visualization settings to matplotlib lib. For example, below,\n", + "we:\n", + "\n", + " - Set the fontsize of the x and y labels by passing `label_kwargs={\"fontsize\": 16}`.\n", + " - Set the fontsize of the title by passing `title_kwargs={\"fontsize\": \"10\"}`.\n", + " \n", + "There are other `_kwargs` inputs we pass as None, you should check out the Dynesty docs if you need to customize your\n", + "figure." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# %%\n", + "'''\n", + "The `corner_anesthetic` function produces a triangle of 1D and 2D PDF's of every parameter using the library `anesthetic`.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `corner_cornerpy` function produces a triangle of 1D and 2D PDF's of every parameter using the library `corner.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_cornerpy(\n", + " samples=result.samples,\n", + " dims=None,\n", + " span=None,\n", + " quantiles=[0.025, 0.5, 0.975],\n", + " color=\"black\",\n", + " smooth=0.02,\n", + " quantiles_2d=None,\n", + " hist_kwargs=None,\n", + " hist2d_kwargs=None,\n", + " label_kwargs={\"fontsize\": 16},\n", + " show_titles=True,\n", + " title_fmt=\".2f\",\n", + " title_kwargs={\"fontsize\": \"10\"},\n", + " truths=None,\n", + " truth_color=\"red\",\n", + " truth_kwargs=None,\n", + " max_n_ticks=5,\n", + " top_ticks=False,\n", + " use_math_text=False,\n", + " verbose=False,\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The internal sampler can be used to plot the results of the non-linear search. \n", + "\n", + "We do this using the `search_internal` attribute which contains the sampler in its native form.\n", + "\n", + "The first time you run a search, the `search_internal` attribute will be available because it is passed ot the\n", + "result via memory. \n", + "\n", + "If you rerun the fit on a completed result, it will not be available in memory, and therefore\n", + "will be loaded from the `files/search_internal` folder. The `search_internal` entry of the `output.yaml` must be true \n", + "for this to be possible." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_internal = result.search_internal" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Plots__\n", + "\n", + "All plots use dynesty's inbuilt plotting library and the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " from dynesty import plotting as dyplot\n", + "\n", + " model = result.model\n", + "\n", + " \"\"\"\n", + " The boundplot plots the bounding distribution used to propose either (1) live points at a given iteration or (2) a \n", + " specific dead point during the course of a run, projected onto the two dimensions specified by `dims`.\n", + " \"\"\"\n", + " dyplot.boundplot(\n", + " results=search_internal.results,\n", + " labels=model.parameter_labels_with_superscripts_latex,\n", + " dims=(2, 2),\n", + " it=-1, # -1 is the final iteration of the dynesty samples, change this to plot a different iteration\n", + " idx=None,\n", + " prior_transform=None,\n", + " periodic=None,\n", + " reflective=None,\n", + " ndraws=5000,\n", + " color=\"gray\",\n", + " plot_kwargs=None,\n", + " label_kwargs={\"fontsize\": 16},\n", + " max_n_ticks=5,\n", + " use_math_text=False,\n", + " show_live=False,\n", + " live_color=\"darkviolet\",\n", + " live_kwargs=None,\n", + " span=None,\n", + " fig=None,\n", + " )\n", + "\n", + " plt.show()\n", + " plt.close()\n", + "\n", + " \"\"\"\n", + " The cornerbound plots the bounding distribution used to propose either (1) live points at a given iteration or (2) a \n", + " specific dead point during the course of a run, projected onto all pairs of dimensions.\n", + " \"\"\"\n", + " try:\n", + " dyplot.cornerbound(\n", + " results=search_internal.results,\n", + " labels=model.parameter_labels_with_superscripts_latex,\n", + " it=-1, # -1 is the final iteration of the dynesty samples, change this to plot a different iteration\n", + " idx=None,\n", + " dims=None,\n", + " prior_transform=None,\n", + " periodic=None,\n", + " reflective=None,\n", + " ndraws=5000,\n", + " color=\"gray\",\n", + " plot_kwargs=None,\n", + " label_kwargs={\"fontsize\": 16},\n", + " max_n_ticks=5,\n", + " use_math_text=False,\n", + " show_live=False,\n", + " live_color=\"darkviolet\",\n", + " live_kwargs=None,\n", + " span=None,\n", + " fig=None,\n", + " )\n", + "\n", + " plt.show()\n", + " plt.close()\n", + "\n", + " except ValueError:\n", + " pass\n", + "\n", + " \"\"\"\n", + " The cornerplot plots a corner plot of the 1-D and 2-D marginalized posteriors.\n", + " \"\"\"\n", + "\n", + " try:\n", + " dyplot.cornerplot(\n", + " results=search_internal.results,\n", + " labels=model.parameter_labels_with_superscripts_latex,\n", + " dims=None,\n", + " span=None,\n", + " quantiles=[0.025, 0.5, 0.975],\n", + " color=\"black\",\n", + " smooth=0.02,\n", + " quantiles_2d=None,\n", + " hist_kwargs=None,\n", + " hist2d_kwargs=None,\n", + " label_kwargs={\"fontsize\": 16},\n", + " show_titles=True,\n", + " title_fmt=\".2f\",\n", + " title_kwargs={\"fontsize\": \"10\"},\n", + " truths=None,\n", + " truth_color=\"red\",\n", + " truth_kwargs=None,\n", + " max_n_ticks=5,\n", + " top_ticks=False,\n", + " use_math_text=False,\n", + " verbose=False,\n", + " )\n", + "\n", + " plt.show()\n", + " plt.close()\n", + "\n", + " except ValueError:\n", + " pass\n", + "\n", + " \"\"\"\n", + " The cornerpoints plots a (sub-)corner plot of (weighted) samples.\n", + " \"\"\"\n", + " dyplot.cornerpoints(\n", + " results=search_internal.results,\n", + " labels=model.parameter_labels_with_superscripts_latex,\n", + " dims=None,\n", + " thin=1,\n", + " span=None,\n", + " cmap=\"plasma\",\n", + " color=None,\n", + " kde=True,\n", + " nkde=1000,\n", + " plot_kwargs=None,\n", + " label_kwargs={\"fontsize\": 16},\n", + " truths=None,\n", + " truth_color=\"red\",\n", + " truth_kwargs=None,\n", + " max_n_ticks=5,\n", + " use_math_text=False,\n", + " fig=None,\n", + " )\n", + "\n", + " plt.show()\n", + " plt.close()\n", + "\n", + " \"\"\"\n", + " The runplot plots live points, ln(likelihood), ln(weight), and ln(evidence) as a function of ln(prior volume).\n", + " \"\"\"\n", + " dyplot.runplot(\n", + " results=search_internal.results,\n", + " span=None,\n", + " logplot=False,\n", + " kde=True,\n", + " nkde=1000,\n", + " color=\"blue\",\n", + " plot_kwargs=None,\n", + " label_kwargs={\"fontsize\": 16},\n", + " lnz_error=True,\n", + " lnz_truth=None,\n", + " truth_color=\"red\",\n", + " truth_kwargs=None,\n", + " max_x_ticks=8,\n", + " max_y_ticks=3,\n", + " use_math_text=True,\n", + " mark_final_live=True,\n", + " fig=None,\n", + " )\n", + "\n", + " plt.show()\n", + " plt.close()\n", + "\n", + " \"\"\"\n", + " The traceplot plots traces and marginalized posteriors for each parameter.\n", + " \"\"\"\n", + " dyplot.traceplot(\n", + " results=search_internal.results,\n", + " span=None,\n", + " quantiles=[0.025, 0.5, 0.975],\n", + " smooth=0.02,\n", + " thin=1,\n", + " dims=None,\n", + " post_color=\"blue\",\n", + " post_kwargs=None,\n", + " kde=True,\n", + " nkde=1000,\n", + " trace_cmap=\"plasma\",\n", + " trace_color=None,\n", + " trace_kwargs=None,\n", + " connect=False,\n", + " connect_highlight=10,\n", + " connect_color=\"red\",\n", + " connect_kwargs=None,\n", + " max_n_ticks=5,\n", + " use_math_text=False,\n", + " label_kwargs={\"fontsize\": 16},\n", + " show_titles=True,\n", + " title_fmt=\".2f\",\n", + " title_kwargs={\"fontsize\": \"10\"},\n", + " truths=None,\n", + " truth_color=\"red\",\n", + " truth_kwargs=None,\n", + " verbose=False,\n", + " fig=None,\n", + " )\n", + "\n", + " plt.show()\n", + " plt.close()\n", + "except (AttributeError, TypeError):\n", + " pass # search_internal unavailable in test mode" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__EmceePlotter__\n", + "\n", + "We now use the MCMC plotting functions to visualize emcee's results.\n", + "\n", + "The emcee readthedocs describes fully all of the methods used below\n", + "\n", + " - https://emcee.readthedocs.io/en/stable/user/sampler/\n", + "\n", + " The `corner_cornerpy` function wraps the library `corner.py` to make corner plots of the PDF:\n", + "\n", + "- https://corner.readthedocs.io/en/latest/index.html\n", + "\n", + "In all the examples below, we use the `kwargs` of this function to pass in any of the input parameters that are\n", + "described in the API docs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "The `corner_cornerpy` function produces a triangle of 1D and 2D PDF's of every parameter in the model fit.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_cornerpy(\n", + " samples=result.samples,\n", + " bins=20,\n", + " range=None,\n", + " color=\"k\",\n", + " hist_bin_factor=1,\n", + " smooth=None,\n", + " smooth1d=None,\n", + " label_kwargs=None,\n", + " titles=None,\n", + " show_titles=False,\n", + " title_fmt=\".2f\",\n", + " title_kwargs=None,\n", + " truths=None,\n", + " truth_color=\"#4682b4\",\n", + " scale_hist=False,\n", + " quantiles=None,\n", + " verbose=False,\n", + " fig=None,\n", + " max_n_ticks=5,\n", + " top_ticks=False,\n", + " use_math_text=False,\n", + " reverse=False,\n", + " labelpad=0.0,\n", + " hist_kwargs=None,\n", + " group=\"posterior\",\n", + " var_names=None,\n", + " filter_vars=None,\n", + " coords=None,\n", + " divergences=False,\n", + " divergences_kwargs=None,\n", + " labeller=None,\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Specific Visualization__\n", + "\n", + "The internal sampler can be used to plot the results of the non-linear search. \n", + "\n", + "We do this using the `search_internal` attribute which contains the sampler in its native form.\n", + "\n", + "The first time you run a search, the `search_internal` attribute will be available because it is passed ot the\n", + "result via memory. \n", + "\n", + "If you rerun the fit on a completed result, it will not be available in memory, and therefore\n", + "will be loaded from the `files/search_internal` folder. The `search_internal` entry of the `output.yaml` must be true \n", + "for this to be possible." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_internal = result.search_internal" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The method below shows a 2D projection of the walker trajectories." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " fig, axes = plt.subplots(result.model.prior_count, figsize=(10, 7))\n", + "\n", + " for i in range(result.model.prior_count):\n", + " for walker_index in range(search_internal.get_log_prob().shape[1]):\n", + " ax = axes[i]\n", + " ax.plot(\n", + " search_internal.get_chain()[:, walker_index, i],\n", + " search_internal.get_log_prob()[:, walker_index],\n", + " alpha=0.3,\n", + " )\n", + "\n", + " ax.set_ylabel(\"Log Likelihood\")\n", + " ax.set_xlabel(result.model.parameter_labels_with_superscripts_latex[i])\n", + "\n", + " plt.show()\n", + "\n", + " \"\"\"\n", + " This method shows the likelihood as a series of steps.\n", + " \"\"\"\n", + "\n", + " fig, axes = plt.subplots(1, figsize=(10, 7))\n", + "\n", + " for walker_index in range(search_internal.get_log_prob().shape[1]):\n", + " axes.plot(search_internal.get_log_prob()[:, walker_index], alpha=0.3)\n", + "\n", + " axes.set_ylabel(\"Log Likelihood\")\n", + " axes.set_xlabel(\"step number\")\n", + "\n", + " plt.show()\n", + "\n", + " \"\"\"\n", + " This method shows the parameter values of every walker at every step.\n", + " \"\"\"\n", + " fig, axes = plt.subplots(\n", + " result.samples.model.prior_count, figsize=(10, 7), sharex=True\n", + " )\n", + "\n", + " for i in range(result.samples.model.prior_count):\n", + " ax = axes[i]\n", + " ax.plot(search_internal.get_chain()[:, :, i], alpha=0.3)\n", + " ax.set_ylabel(result.model.parameter_labels_with_superscripts_latex[i])\n", + "\n", + " axes[-1].set_xlabel(\"step number\")\n", + "\n", + " plt.show()\n", + "except AttributeError:\n", + " pass # MCMC-specific methods not available for non-MCMC searches" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__ZeusPlotter__\n", + "\n", + "We now use the MCMC plotting functions to visualize Zeus's results.\n", + "\n", + "The zeus readthedocs describes fully all of the methods used below\n", + "\n", + " - https://zeus-mcmc.readthedocs.io/en/latest/api/plotting.html#cornerplot\n", + " - https://zeus-mcmc.readthedocs.io/en/latest/notebooks/normal_distribution.html\n", + "\n", + " The `corner_cornerpy` function wraps the library `corner.py` to make corner plots of the PDF:\n", + "\n", + "- https://corner.readthedocs.io/en/latest/index.html\n", + "\n", + "In all the examples below, we use the `kwargs` of this function to pass in any of the input parameters that are\n", + "described in the API docs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "The `corner_cornerpy` function produces a triangle of 1D and 2D PDF's of every parameter in the model fit.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_cornerpy(\n", + " samples=result.samples,\n", + " weight_list=None,\n", + " levels=None,\n", + " span=None,\n", + " quantiles=[0.025, 0.5, 0.975],\n", + " truth=None,\n", + " color=None,\n", + " alpha=0.5,\n", + " linewidth=1.5,\n", + " fill=True,\n", + " fontsize=10,\n", + " show_titles=True,\n", + " title_fmt=\".2f\",\n", + " title_fontsize=12,\n", + " cut=3,\n", + " fig=None,\n", + " size=(10, 10),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Specific Visualization__\n", + "\n", + "The internal sampler can be used to plot the results of the non-linear search. \n", + "\n", + "We do this using the `search_internal` attribute which contains the sampler in its native form.\n", + "\n", + "For zeus, the `search_internal` attribute is only available if the zeus sampler results are output to hard-disk\n", + "via hdf5. The `search_internal` entry of the `output.yaml` must be true for this to be the case." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_internal = result.search_internal" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Plots__\n", + "\n", + "The method below shows a 2D projection of the walker trajectories." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " fig, axes = plt.subplots(result.model.prior_count, figsize=(10, 7))\n", + "\n", + " for i in range(result.model.prior_count):\n", + " for walker_index in range(search_internal.get_log_prob().shape[1]):\n", + " ax = axes[i]\n", + " ax.plot(\n", + " search_internal.get_chain()[:, walker_index, i],\n", + " search_internal.get_log_prob()[:, walker_index],\n", + " alpha=0.3,\n", + " )\n", + "\n", + " ax.set_ylabel(\"Log Likelihood\")\n", + " ax.set_xlabel(result.model.parameter_labels_with_superscripts_latex[i])\n", + "\n", + " plt.show()\n", + "\n", + " \"\"\"\n", + " This method shows the likelihood as a series of steps.\n", + " \"\"\"\n", + "\n", + " fig, axes = plt.subplots(1, figsize=(10, 7))\n", + "\n", + " for walker_index in range(search_internal.get_log_prob().shape[1]):\n", + " axes.plot(search_internal.get_log_prob()[:, walker_index], alpha=0.3)\n", + "\n", + " axes.set_ylabel(\"Log Likelihood\")\n", + " axes.set_xlabel(\"step number\")\n", + "\n", + " plt.show()\n", + "\n", + " \"\"\"\n", + " This method shows the parameter values of every walker at every step.\n", + " \"\"\"\n", + " fig, axes = plt.subplots(\n", + " result.samples.model.prior_count, figsize=(10, 7), sharex=True\n", + " )\n", + "\n", + " for i in range(result.samples.model.prior_count):\n", + " ax = axes[i]\n", + " ax.plot(search_internal.get_chain()[:, :, i], alpha=0.3)\n", + " ax.set_ylabel(result.model.parameter_labels_with_superscripts_latex[i])\n", + "\n", + " axes[-1].set_xlabel(\"step number\")\n", + "\n", + " plt.show()\n", + "except AttributeError:\n", + " pass # MCMC-specific methods not available for non-MCMC searches\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__GetDist__\n", + "\n", + "This example illustrates how to plot visualization summarizing the results of model-fit using any non-linear search\n", + "using GetDist:\n", + "\n", + " - https://getdist.readthedocs.io/en/latest/\n", + "\n", + "GetDist is an optional library which creates 1D and 2D plots of probability distribution functions (PDF)s. Its\n", + "visualization tools has more than the in-built visualization tools of many non-linear searches (e.g. Nautilus /\n", + "emcee) and can often produce better looking plots.\n", + "\n", + "GetDist was developed for the analysis of Cosmological datasets.\n", + "\n", + "Because GetDist is an optional library, you will likely have to install it manually via the command:\n", + "\n", + "`pip install getdist`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from getdist import MCSamples\n", + "from getdist import plots\n", + "import numpy as np" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "GetDist uses a `model.paramnames` file to load the name of every parameter in the model-fit and pair it with the\n", + "latex symbol used to represent it in plots.\n", + "\n", + "This file is not created by default, but can be output by the `search.paths` object as shown below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search.paths._save_parameter_names_file(model=model)\n", + "search.paths.zip_remove()\n", + "search.paths._zip()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "GetDist uses an `MCSamples` object to store the samples of a non-linear search.\n", + "\n", + "We create this object via a conversion from **PyAutoFit** `Samples`, as well as using the `names`\n", + "and `labels` of parameters in the `Samples` object.\n", + "\n", + "The input `sampler=\"nested\"` is input because we used a nested sampling, `Nautilus`. For MCMC this should be\n", + "replaced with \"mcmc\"." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "gd_samples = MCSamples(\n", + " samples=np.asarray(samples.parameter_lists),\n", + " loglikes=np.asarray(samples.log_likelihood_list),\n", + " weights=np.asarray(samples.weight_list),\n", + " names=samples.model.model_component_and_parameter_names,\n", + " labels=samples.model.parameter_labels_with_superscripts,\n", + " sampler=\"nested\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parameter Names__\n", + "\n", + "Note that in order to customize the figure, we will use the `samples.model.parameter_names` list." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(samples.model.model_component_and_parameter_names)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__GetDist Plotter__\n", + "\n", + "To make plots we use a GetDist plotter object, which can be customized to change the appearance of the plots." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "gd_plotter = plots.get_subplot_plotter(width_inch=12)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__GetDist Subplots__\n", + "\n", + "Using the plotter we can make different plots, for example a triangle plot showing the 1D and 2D PDFs of every \n", + "parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " gd_plotter.triangle_plot(roots=gd_samples, filled=True)\n", + " plt.show()\n", + " plt.close()\n", + "except Exception:\n", + " pass" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A triangle plot with specific parameters can be plotted by using the `params` input, whereby we specify the specific\n", + "parameter names to plot." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " gd_plotter.triangle_plot(\n", + " roots=gd_samples,\n", + " filled=True,\n", + " params=[\n", + " \"galaxies_lens_mass_einstein_radius\",\n", + " \"galaxies_lens_mass_ell_comps_0\",\n", + " ],\n", + " )\n", + "except Exception:\n", + " pass\n", + "\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Rectangle plots can be used to show specific 2D combinations of parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " gd_plotter.rectangle_plot(\n", + " roots=gd_samples,\n", + " yparams=[\"galaxies_lens_mass_einstein_radius\"],\n", + " xparams=[\n", + " \"galaxies_lens_mass_ell_comps_0\",\n", + " \"galaxies_lens_mass_ell_comps_1\",\n", + " ],\n", + " )\n", + "except Exception:\n", + " pass\n", + "\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__GetDist Single Plots__\n", + "\n", + "We can make plots of specific 1D or 2D PDFs, using the single plotter object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "gd_plotter = plots.get_single_plotter()\n", + "\n", + "try:\n", + " gd_plotter.plot_1d(roots=gd_samples, param=\"galaxies_lens_mass_einstein_radius\")\n", + "except Exception:\n", + " pass\n", + "\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "gd_plotter = plots.get_single_plotter()\n", + "\n", + "try:\n", + " gd_plotter.plot_2d(\n", + " roots=gd_samples,\n", + " param1=\"galaxies_lens_mass_einstein_radius\",\n", + " param2=\"galaxies_lens_mass_ell_comps_0\",\n", + " )\n", + "except Exception:\n", + " pass\n", + "\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also make a 3D plot, where the 2D PDF is plotted colored by the value of a third parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "gd_plotter = plots.get_single_plotter()\n", + "\n", + "try:\n", + " gd_plotter.plot_3d(\n", + " roots=gd_samples,\n", + " params=[\n", + " \"galaxies_lens_mass_einstein_radius\",\n", + " \"galaxies_lens_mass_ell_comps_0\",\n", + " \"galaxies_lens_mass_ell_comps_1\",\n", + " ],\n", + " )\n", + "except Exception:\n", + " pass\n", + "\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "A figure can be output using standard matplotlib functionality." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "gd_plotter = plots.get_single_plotter()\n", + "\n", + "try:\n", + " gd_plotter.plot_3d(roots=gd_samples, params=[\"centre\", \"sigma\", \"normalization\"])\n", + "except Exception:\n", + " pass\n", + "\n", + "output_path = Path(\"output\")\n", + "\n", + "plt.savefig(Path(output_path, \"getdist.png\"))\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__GetDist Other Plots__\n", + "\n", + "There are many more ways to visualize PDFs possible with GetDist, checkout the official documentation for them all!\n", + "\n", + " - https://getdist.readthedocs.io/en/latest/\n", + " - https://getdist.readthedocs.io/en/latest/plots.html\n", + "\n", + "__Plotting Multiple Samples__\n", + "\n", + "Finally, we can plot the results of multiple different non-linear searches on the same plot, using all\n", + "of the functions above.\n", + "\n", + "Lets quickly make a second set of `Nautilus` results and plot them on the same figure above with the results\n", + "of the first search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "Nautilus = af.Nautilus(path_prefix=\"plot\", name=\"GetDist_2\")\n", + "\n", + "result_extra = Nautilus.fit(model=model, analysis=analysis)\n", + "\n", + "samples_extra = result_extra.samples\n", + "\n", + "gd_samples_extra = MCSamples(\n", + " samples=np.asarray(samples_extra.parameter_lists),\n", + " loglikes=np.asarray(samples_extra.log_likelihood_list),\n", + " weights=np.asarray(samples_extra.weight_list),\n", + " names=samples_extra.model.model_component_and_parameter_names,\n", + " labels=samples.model.parameter_labels_with_superscripts,\n", + " sampler=\"nested\",\n", + ")\n", + "\n", + "gd_plotter = plots.get_subplot_plotter(width_inch=12)\n", + "\n", + "try:\n", + " gd_plotter.triangle_plot(roots=[gd_samples, gd_samples_extra], filled=True)\n", + "except Exception:\n", + " pass\n", + "\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that the models do not need to be the same to make the plots above.\n", + "\n", + "GetDist will clever use the `names` of the parameters to combine the parameters into customizeable PDF plots." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/plot/examples/visuals.ipynb b/notebooks/guides/plot/examples/visuals.ipynb index ab75f1dd8..fb9f8ea9f 100644 --- a/notebooks/guides/plot/examples/visuals.ipynb +++ b/notebooks/guides/plot/examples/visuals.ipynb @@ -1,436 +1,473 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plots: Visuals (Overlays)\n", - "=========================\n", - "\n", - "This example illustrates how to add overlays to plots using the new API.\n", - "\n", - "Overlays are specified via two keyword arguments on `aplt.plot_array()` and `aplt.plot_grid()`:\n", - "\n", - " - `lines=`: A list of `Grid2DIrregular` objects drawn as lines (e.g. critical curves, caustics).\n", - " - `positions=`: A `Grid2DIrregular` object drawn as scatter points (e.g. image positions).\n", - "\n", - "The old `Visuals2D` and `MatPlot2D` objects that configured overlays have been removed.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "Refer to `plots/start_here.ipynb` for a general introduction to the new plotting API.\n", - "\n", - "__Contents__\n", - "\n", - "- **Setup:** General setup for the analysis.\n", - "- **Critical Curves:** Critical curves are plotted as lines over the image using the `lines=` keyword argument.\n", - "- **Multiple Critical Curves:** If a `Tracer` has multiple lens galaxies it may have multiple tangential and radial critical curves.\n", - "- **Caustics:** Caustics are the critical curves mapped to the source plane.\n", - "- **Image Positions:** The multiple image positions of a lensed source can be plotted using `positions=`.\n", - "- **Light Profile Centres:** The centres of light profiles can be extracted and plotted as positions over an image.\n", - "- **Mass Profile Centres:** Mass profile centres can be extracted and overlaid in the same way.\n", - "- **Combined Overlays:** `lines=` and `positions=` can be used together on the same plot." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Setup__\n", - "\n", - "Create the standard objects used to illustrate overlays." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.6, ell_comps=(0.2, 0.2)),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCoreSph(\n", - " centre=(0.1, 0.1), intensity=0.3, effective_radius=1.0, sersic_index=2.5\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "lens_calc = al.LensCalc.from_tracer(tracer=tracer)\n", - "\n", - "lens_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(-1.0, 0.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(-1.0, 0.0), einstein_radius=0.8, ell_comps=(0.2, 0.2)\n", - " ),\n", - ")\n", - "\n", - "source_galaxy_1 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCoreSph(\n", - " centre=(0.2, 0.2), intensity=0.3, effective_radius=1.0, sersic_index=2.5\n", - " ),\n", - ")\n", - "\n", - "tracer_x2 = al.Tracer(\n", - " galaxies=[lens_galaxy, lens_galaxy_1, source_galaxy, source_galaxy_1]\n", - ")\n", - "\n", - "lens_calc_x2 = al.LensCalc.from_tracer(tracer=tracer_x2)\n", - "\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / \"slacs1430+4105\"\n", - "data_path = dataset_path / \"data.fits\"\n", - "data = al.Array2D.from_fits(file_path=data_path, hdu=0, pixel_scales=0.03)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Critical Curves__\n", - "\n", - "Critical curves are plotted as lines over the image using the `lines=` keyword argument.\n", - "\n", - "`tangential_critical_curve_list_from` returns a list of `Grid2DIrregular` objects, one per\n", - "tangential critical curve. Pass this list directly to `lines=`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tangential_critical_curve_list = lens_calc.tangential_critical_curve_list_from(\n", - " grid=grid\n", - ")\n", - "\n", - "image = tracer.image_2d_from(grid=grid)\n", - "\n", - "aplt.plot_array(\n", - " array=image,\n", - " title=\"Image with Tangential Critical Curves\",\n", - " lines=tangential_critical_curve_list,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Radial critical curves can be overlaid in the same way. Combine both lists with `+` to\n", - "overlay tangential and radial critical curves together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "radial_critical_curve_list = lens_calc.radial_critical_curve_list_from(grid=grid)\n", - "\n", - "aplt.plot_array(\n", - " array=image,\n", - " title=\"Image with All Critical Curves\",\n", - " lines=tangential_critical_curve_list + radial_critical_curve_list,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multiple Critical Curves__\n", - "\n", - "If a `Tracer` has multiple lens galaxies it may have multiple tangential and radial critical\n", - "curves. These are all contained in the returned lists and plotted together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tangential_critical_curve_list = lens_calc_x2.tangential_critical_curve_list_from(\n", - " grid=grid\n", - ")\n", - "radial_critical_curve_list = lens_calc_x2.radial_critical_curve_list_from(grid=grid)\n", - "\n", - "image_x2 = tracer_x2.image_2d_from(grid=grid)\n", - "\n", - "aplt.plot_array(\n", - " array=image_x2,\n", - " title=\"Two-Galaxy System Critical Curves\",\n", - " lines=tangential_critical_curve_list + radial_critical_curve_list,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Caustics__\n", - "\n", - "Caustics are the critical curves mapped to the source plane. They are plotted over the\n", - "source-plane image using `lines=`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tangential_caustic_list = lens_calc.tangential_caustic_list_from(grid=grid)\n", - "radial_caustic_list = lens_calc.radial_caustic_list_from(grid=grid)\n", - "\n", - "source_image = tracer.image_2d_list_from(grid=grid)[1]\n", - "\n", - "aplt.plot_array(\n", - " array=source_image,\n", - " title=\"Source Plane with Tangential Caustics\",\n", - " lines=tangential_caustic_list,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=source_image,\n", - " title=\"Source Plane with All Caustics\",\n", - " lines=tangential_caustic_list + radial_caustic_list,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image Positions__\n", - "\n", - "The multiple image positions of a lensed source can be plotted using `positions=`.\n", - "\n", - "`positions=` accepts an `al.Grid2DIrregular` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "multiple_images = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=image,\n", - " title=\"Image with Multiple Images\",\n", - " positions=multiple_images,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Arbitrary (y,x) coordinates can also be plotted as positions, for example to mark\n", - "interesting regions on an image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(values=[(1.0, 1.0), (2.0, 2.0), (3.0, 3.0)])\n", - "\n", - "aplt.plot_array(\n", - " array=data,\n", - " title=\"Data with Positions\",\n", - " positions=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Light Profile Centres__\n", - "\n", - "The centres of light profiles can be extracted and plotted as positions over an image.\n", - "\n", - "We extract image-plane centres from the first (lens) galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "light_profile_centres = tracer.galaxies[0].extract_attribute(\n", - " cls=al.LightProfile, attr_name=\"centre\"\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=image,\n", - " title=\"Image with Light Profile Centres\",\n", - " positions=light_profile_centres,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Source-plane centres can be extracted from the last galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_profile_centres = tracer.galaxies[-1].extract_attribute(\n", - " cls=al.LightProfile, attr_name=\"centre\"\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=source_image,\n", - " title=\"Source Plane with Light Profile Centres\",\n", - " positions=source_profile_centres,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mass Profile Centres__\n", - "\n", - "Mass profile centres can be extracted and overlaid in the same way." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass_profile_centres = tracer.extract_attribute(\n", - " cls=al.mp.MassProfile, attr_name=\"centre\"\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=image,\n", - " title=\"Image with Mass Profile Centres\",\n", - " positions=mass_profile_centres,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Combined Overlays__\n", - "\n", - "`lines=` and `positions=` can be used together on the same plot." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tangential_critical_curve_list = lens_calc.tangential_critical_curve_list_from(\n", - " grid=grid\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=image,\n", - " title=\"Image with Critical Curves and Multiple Images\",\n", - " lines=tangential_critical_curve_list,\n", - " positions=multiple_images,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plots: Visuals (Overlays)\n", + "=========================\n", + "\n", + "This example illustrates how to add overlays to plots using the new API.\n", + "\n", + "Overlays are specified via two keyword arguments on `aplt.plot_array()` and `aplt.plot_grid()`:\n", + "\n", + " - `lines=`: A list of `Grid2DIrregular` objects drawn as lines (e.g. critical curves, caustics).\n", + " - `positions=`: A `Grid2DIrregular` object drawn as scatter points (e.g. image positions).\n", + "\n", + "The old `Visuals2D` and `MatPlot2D` objects that configured overlays have been removed.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "Refer to `plots/start_here.ipynb` for a general introduction to the new plotting API.\n", + "\n", + "__Contents__\n", + "\n", + "- **Setup:** General setup for the analysis.\n", + "- **Critical Curves:** Critical curves are plotted as lines over the image using the `lines=` keyword argument.\n", + "- **Multiple Critical Curves:** If a `Tracer` has multiple lens galaxies it may have multiple tangential and radial critical curves.\n", + "- **Caustics:** Caustics are the critical curves mapped to the source plane.\n", + "- **Image Positions:** The multiple image positions of a lensed source can be plotted using `positions=`.\n", + "- **Light Profile Centres:** The centres of light profiles can be extracted and plotted as positions over an image.\n", + "- **Mass Profile Centres:** Mass profile centres can be extracted and overlaid in the same way.\n", + "- **Combined Overlays:** `lines=` and `positions=` can be used together on the same plot." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Setup__\n", + "\n", + "Create the standard objects used to illustrate overlays." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.6, ell_comps=(0.2, 0.2)),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCoreSph(\n", + " centre=(0.1, 0.1), intensity=0.3, effective_radius=1.0, sersic_index=2.5\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "lens_calc = al.LensCalc.from_tracer(tracer=tracer)\n", + "\n", + "lens_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(-1.0, 0.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(-1.0, 0.0), einstein_radius=0.8, ell_comps=(0.2, 0.2)\n", + " ),\n", + ")\n", + "\n", + "source_galaxy_1 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCoreSph(\n", + " centre=(0.2, 0.2), intensity=0.3, effective_radius=1.0, sersic_index=2.5\n", + " ),\n", + ")\n", + "\n", + "tracer_x2 = al.Tracer(\n", + " galaxies=[lens_galaxy, lens_galaxy_1, source_galaxy, source_galaxy_1]\n", + ")\n", + "\n", + "lens_calc_x2 = al.LensCalc.from_tracer(tracer=tracer_x2)\n", + "\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / \"slacs1430+4105\"\n", + "data_path = dataset_path / \"data.fits\"\n", + "data = al.Array2D.from_fits(file_path=data_path, hdu=0, pixel_scales=0.03)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Critical Curves__\n", + "\n", + "Critical curves are plotted as lines over the image using the `lines=` keyword argument.\n", + "\n", + "`tangential_critical_curve_list_from` returns a list of `Grid2DIrregular` objects, one per\n", + "tangential critical curve. Pass this list directly to `lines=`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tangential_critical_curve_list = lens_calc.tangential_critical_curve_list_from(\n", + " grid=grid\n", + ")\n", + "\n", + "image = tracer.image_2d_from(grid=grid)\n", + "\n", + "aplt.plot_array(\n", + " array=image,\n", + " title=\"Image with Tangential Critical Curves\",\n", + " lines=tangential_critical_curve_list,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Radial critical curves can be overlaid in the same way. Combine both lists with `+` to\n", + "overlay tangential and radial critical curves together." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "radial_critical_curve_list = lens_calc.radial_critical_curve_list_from(grid=grid)\n", + "\n", + "aplt.plot_array(\n", + " array=image,\n", + " title=\"Image with All Critical Curves\",\n", + " lines=tangential_critical_curve_list + radial_critical_curve_list,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multiple Critical Curves__\n", + "\n", + "If a `Tracer` has multiple lens galaxies it may have multiple tangential and radial critical\n", + "curves. These are all contained in the returned lists and plotted together." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tangential_critical_curve_list = lens_calc_x2.tangential_critical_curve_list_from(\n", + " grid=grid\n", + ")\n", + "radial_critical_curve_list = lens_calc_x2.radial_critical_curve_list_from(grid=grid)\n", + "\n", + "image_x2 = tracer_x2.image_2d_from(grid=grid)\n", + "\n", + "aplt.plot_array(\n", + " array=image_x2,\n", + " title=\"Two-Galaxy System Critical Curves\",\n", + " lines=tangential_critical_curve_list + radial_critical_curve_list,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Caustics__\n", + "\n", + "Caustics are the critical curves mapped to the source plane. They are plotted over the\n", + "source-plane image using `lines=`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tangential_caustic_list = lens_calc.tangential_caustic_list_from(grid=grid)\n", + "radial_caustic_list = lens_calc.radial_caustic_list_from(grid=grid)\n", + "\n", + "source_image = tracer.image_2d_list_from(grid=grid)[1]\n", + "\n", + "aplt.plot_array(\n", + " array=source_image,\n", + " title=\"Source Plane with Tangential Caustics\",\n", + " lines=tangential_caustic_list,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=source_image,\n", + " title=\"Source Plane with All Caustics\",\n", + " lines=tangential_caustic_list + radial_caustic_list,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image Positions__\n", + "\n", + "The multiple image positions of a lensed source can be plotted using `positions=`.\n", + "\n", + "`positions=` accepts an `al.Grid2DIrregular` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "multiple_images = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=image,\n", + " title=\"Image with Multiple Images\",\n", + " positions=multiple_images,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Arbitrary (y,x) coordinates can also be plotted as positions, for example to mark\n", + "interesting regions on an image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(values=[(1.0, 1.0), (2.0, 2.0), (3.0, 3.0)])\n", + "\n", + "aplt.plot_array(\n", + " array=data,\n", + " title=\"Data with Positions\",\n", + " positions=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Light Profile Centres__\n", + "\n", + "The centres of light profiles can be extracted and plotted as positions over an image.\n", + "\n", + "We extract image-plane centres from the first (lens) galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "light_profile_centres = tracer.galaxies[0].extract_attribute(\n", + " cls=al.LightProfile, attr_name=\"centre\"\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=image,\n", + " title=\"Image with Light Profile Centres\",\n", + " positions=light_profile_centres,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Source-plane centres can be extracted from the last galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_profile_centres = tracer.galaxies[-1].extract_attribute(\n", + " cls=al.LightProfile, attr_name=\"centre\"\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=source_image,\n", + " title=\"Source Plane with Light Profile Centres\",\n", + " positions=source_profile_centres,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mass Profile Centres__\n", + "\n", + "Mass profile centres can be extracted and overlaid in the same way." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass_profile_centres = tracer.extract_attribute(\n", + " cls=al.mp.MassProfile, attr_name=\"centre\"\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=image,\n", + " title=\"Image with Mass Profile Centres\",\n", + " positions=mass_profile_centres,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Combined Overlays__\n", + "\n", + "`lines=` and `positions=` can be used together on the same plot." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tangential_critical_curve_list = lens_calc.tangential_critical_curve_list_from(\n", + " grid=grid\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=image,\n", + " title=\"Image with Critical Curves and Multiple Images\",\n", + " lines=tangential_critical_curve_list,\n", + " positions=multiple_images,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/plot/start_here.ipynb b/notebooks/guides/plot/start_here.ipynb index f8ad94b67..e0212798a 100644 --- a/notebooks/guides/plot/start_here.ipynb +++ b/notebooks/guides/plot/start_here.ipynb @@ -1,374 +1,411 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plots: Start Here\n", - "=================\n", - "\n", - "This example introduces the new plotting API in PyAutoLens.\n", - "\n", - "The old API (removed) used `*Plotter` classes (e.g. `Imaging`, `Tracer`) together with\n", - "`MatPlot2D` and `Visuals2D` helper objects. These have all been removed.\n", - "\n", - "The new API uses standalone functions:\n", - "\n", - " - `aplt.plot_array()` \u2014 plot any 2D array.\n", - " - `aplt.plot_grid()` \u2014 plot a 2D grid of (y,x) coordinates.\n", - " - `aplt.subplot_tracer()`, `aplt.subplot_fit_imaging()`, etc. \u2014 multi-panel subplots for standard objects.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Customization:** Each plotting function accepts direct keyword arguments for customization.\n", - "- **Config Defaults:** All default plotting values are configured via config files in.\n", - "- **Overlays:** Overlays are added to plots using the `lines=` and `positions=` keyword arguments." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load an example imaging dataset and set up objects used throughout this example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "data_path = dataset_path / \"data.fits\"\n", - "data = al.Array2D.from_fits(file_path=data_path, hdu=0, pixel_scales=0.1)\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.6, ell_comps=(0.2, 0.2)),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCoreSph(\n", - " centre=(0.1, 0.1), intensity=0.3, effective_radius=1.0, sersic_index=2.5\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__plot_array__\n", - "\n", - "The fundamental plotting function is `aplt.plot_array()`, which displays any 2D `Array2D`.\n", - "\n", - "We can plot the raw data array loaded from a .fits file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also plot quantities computed from a tracer, such as its image, convergence and potential." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Tracer Image\")\n", - "aplt.plot_array(array=tracer.convergence_2d_from(grid=grid), title=\"Convergence\")\n", - "aplt.plot_array(array=tracer.potential_2d_from(grid=grid), title=\"Potential\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__plot_grid__\n", - "\n", - "The `aplt.plot_grid()` function displays a 2D grid of (y,x) coordinates.\n", - "\n", - "This is useful for visualizing image-plane and source-plane grids." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=grid, title=\"Uniform Grid\")\n", - "\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=grid)[1]\n", - "aplt.plot_grid(grid=traced_grid, title=\"Source-Plane Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Customization__\n", - "\n", - "Each plotting function accepts direct keyword arguments for customization:\n", - "\n", - " - `title`: The figure title string.\n", - " - `colormap`: The matplotlib colormap name (e.g. \"jet\", \"hot\", \"gray\").\n", - " - `use_log10`: If True, the colormap is plotted in log10 scale.\n", - " - `output_path`: Directory path to save the figure on disk.\n", - " - `output_format`: Format of the saved file, e.g. \"png\" or \"pdf\".\n", - "\n", - "These replace the old `MatPlot2D` object entirely \u2014 there is no `MatPlot2D` anymore." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Tracer Image (Log10)\",\n", - " use_log10=True,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Tracer Image (Jet Colormap)\",\n", - " colormap=\"jet\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To save a figure to disk, pass `output_path` and `output_format`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=data,\n", - " title=\"Data Saved to Disk\",\n", - " output_path=Path(\"output\"),\n", - " output_format=\"png\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Config Defaults__\n", - "\n", - "All default plotting values are configured via config files in:\n", - "\n", - " autolens_workspace/config/visualize/\n", - "\n", - "When no explicit keyword is passed to a plotting function the config value is used, allowing\n", - "the default appearance to be controlled project-wide without changing code." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Overlays__\n", - "\n", - "Overlays are added to plots using the `lines=` and `positions=` keyword arguments:\n", - "\n", - " - `lines=`: A list of `Grid2DIrregular` objects drawn as lines (e.g. critical curves, caustics).\n", - " - `positions=`: An `Grid2DIrregular` object drawn as scatter points (e.g. image positions).\n", - "\n", - "These replace the old `Visuals2D` object entirely \u2014 there is no `Visuals2D` anymore.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_calc = al.LensCalc.from_tracer(tracer=tracer)\n", - "tangential_critical_curve_list = lens_calc.tangential_critical_curve_list_from(\n", - " grid=grid\n", - ")\n", - "tangential_caustic_list = lens_calc.tangential_caustic_list_from(grid=grid)\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Image with Critical Curves\",\n", - " lines=tangential_critical_curve_list,\n", - ")\n", - "\n", - "source_image = tracer.image_2d_list_from(grid=grid)[1]\n", - "aplt.plot_array(\n", - " array=source_image,\n", - " title=\"Source Plane with Caustics\",\n", - " lines=tangential_caustic_list,\n", - ")\n", - "\n", - "positions = al.Grid2DIrregular(values=[(1.0, 1.0), (2.0, 2.0), (-1.0, 0.5)])\n", - "aplt.plot_array(\n", - " array=data,\n", - " title=\"Data with Positions\",\n", - " positions=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__subplot_* Functions__\n", - "\n", - "For standard objects (datasets, tracers, fits), dedicated subplot functions produce\n", - "multi-panel overviews automatically.\n", - "\n", - "These replace all the old `*Plotter` class `.subplot_*()` method calls." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - "aplt.subplot_tracer(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_galaxies_images(tracer=tracer, grid=grid)\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=3.0,\n", - ")\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The search plotting functions (`aplt.corner_anesthetic`, `aplt.corner_cornerpy`, etc.) provide\n", - "visualization of non-linear search results -- see `scripts/guides/plot/examples/searches.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plots: Start Here\n", + "=================\n", + "\n", + "This example introduces the new plotting API in PyAutoLens.\n", + "\n", + "The old API (removed) used `*Plotter` classes (e.g. `Imaging`, `Tracer`) together with\n", + "`MatPlot2D` and `Visuals2D` helper objects. These have all been removed.\n", + "\n", + "The new API uses standalone functions:\n", + "\n", + " - `aplt.plot_array()` \u2014 plot any 2D array.\n", + " - `aplt.plot_grid()` \u2014 plot a 2D grid of (y,x) coordinates.\n", + " - `aplt.subplot_tracer()`, `aplt.subplot_fit_imaging()`, etc. \u2014 multi-panel subplots for standard objects.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Customization:** Each plotting function accepts direct keyword arguments for customization.\n", + "- **Config Defaults:** All default plotting values are configured via config files in.\n", + "- **Overlays:** Overlays are added to plots using the `lines=` and `positions=` keyword arguments." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load an example imaging dataset and set up objects used throughout this example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "data_path = dataset_path / \"data.fits\"\n", + "data = al.Array2D.from_fits(file_path=data_path, hdu=0, pixel_scales=0.1)\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.6, ell_comps=(0.2, 0.2)),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCoreSph(\n", + " centre=(0.1, 0.1), intensity=0.3, effective_radius=1.0, sersic_index=2.5\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__plot_array__\n", + "\n", + "The fundamental plotting function is `aplt.plot_array()`, which displays any 2D `Array2D`.\n", + "\n", + "We can plot the raw data array loaded from a .fits file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also plot quantities computed from a tracer, such as its image, convergence and potential." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Tracer Image\")\n", + "aplt.plot_array(array=tracer.convergence_2d_from(grid=grid), title=\"Convergence\")\n", + "aplt.plot_array(array=tracer.potential_2d_from(grid=grid), title=\"Potential\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__plot_grid__\n", + "\n", + "The `aplt.plot_grid()` function displays a 2D grid of (y,x) coordinates.\n", + "\n", + "This is useful for visualizing image-plane and source-plane grids." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=grid, title=\"Uniform Grid\")\n", + "\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=grid)[1]\n", + "aplt.plot_grid(grid=traced_grid, title=\"Source-Plane Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Customization__\n", + "\n", + "Each plotting function accepts direct keyword arguments for customization:\n", + "\n", + " - `title`: The figure title string.\n", + " - `colormap`: The matplotlib colormap name (e.g. \"jet\", \"hot\", \"gray\").\n", + " - `use_log10`: If True, the colormap is plotted in log10 scale.\n", + " - `output_path`: Directory path to save the figure on disk.\n", + " - `output_format`: Format of the saved file, e.g. \"png\" or \"pdf\".\n", + "\n", + "These replace the old `MatPlot2D` object entirely \u2014 there is no `MatPlot2D` anymore." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Tracer Image (Log10)\",\n", + " use_log10=True,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Tracer Image (Jet Colormap)\",\n", + " colormap=\"jet\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To save a figure to disk, pass `output_path` and `output_format`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=data,\n", + " title=\"Data Saved to Disk\",\n", + " output_path=Path(\"output\"),\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Config Defaults__\n", + "\n", + "All default plotting values are configured via config files in:\n", + "\n", + " autolens_workspace/config/visualize/\n", + "\n", + "When no explicit keyword is passed to a plotting function the config value is used, allowing\n", + "the default appearance to be controlled project-wide without changing code." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Overlays__\n", + "\n", + "Overlays are added to plots using the `lines=` and `positions=` keyword arguments:\n", + "\n", + " - `lines=`: A list of `Grid2DIrregular` objects drawn as lines (e.g. critical curves, caustics).\n", + " - `positions=`: An `Grid2DIrregular` object drawn as scatter points (e.g. image positions).\n", + "\n", + "These replace the old `Visuals2D` object entirely \u2014 there is no `Visuals2D` anymore.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_calc = al.LensCalc.from_tracer(tracer=tracer)\n", + "tangential_critical_curve_list = lens_calc.tangential_critical_curve_list_from(\n", + " grid=grid\n", + ")\n", + "tangential_caustic_list = lens_calc.tangential_caustic_list_from(grid=grid)\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Image with Critical Curves\",\n", + " lines=tangential_critical_curve_list,\n", + ")\n", + "\n", + "source_image = tracer.image_2d_list_from(grid=grid)[1]\n", + "aplt.plot_array(\n", + " array=source_image,\n", + " title=\"Source Plane with Caustics\",\n", + " lines=tangential_caustic_list,\n", + ")\n", + "\n", + "positions = al.Grid2DIrregular(values=[(1.0, 1.0), (2.0, 2.0), (-1.0, 0.5)])\n", + "aplt.plot_array(\n", + " array=data,\n", + " title=\"Data with Positions\",\n", + " positions=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__subplot_* Functions__\n", + "\n", + "For standard objects (datasets, tracers, fits), dedicated subplot functions produce\n", + "multi-panel overviews automatically.\n", + "\n", + "These replace all the old `*Plotter` class `.subplot_*()` method calls." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + "aplt.subplot_tracer(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_galaxies_images(tracer=tracer, grid=grid)\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=3.0,\n", + ")\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The search plotting functions (`aplt.corner_anesthetic`, `aplt.corner_cornerpy`, etc.) provide\n", + "visualization of non-linear search results -- see `scripts/guides/plot/examples/searches.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/point_source_pairing.ipynb b/notebooks/guides/point_source_pairing.ipynb new file mode 100644 index 000000000..3db488f71 --- /dev/null +++ b/notebooks/guides/point_source_pairing.ipynb @@ -0,0 +1,313 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Guide: Point-Source Pairing, Over-Prediction and Under-Prediction\n", + "=================================================================\n", + "\n", + "When fitting multiple-image positions \u2014 the bread and butter of group- and cluster-scale lens\n", + "modeling \u2014 the model tracer will not, in general, predict exactly the images you observed. A wrong\n", + "(or merely uncertain) mass model predicts *extra* images that were never detected, or fails to\n", + "produce an observed image at all. What the likelihood does in those two situations decides which\n", + "models a sampler rewards, and historically lensing codes have handled it with quiet conventions\n", + "rather than explicit choices.\n", + "\n", + "This guide documents PyAutoLens's choices: the three image-plane pairing schemes, the\n", + "over/under-prediction policies, the solver settings that interact with them at cluster scale, and\n", + "the source-plane vs image-plane chi-squared trade-off. It is the reference the cluster examples\n", + "(``scripts/cluster/``, including the Lenstool walkthrough in ``scripts/cluster/lenstool/``) point\n", + "at for likelihood choices.\n", + "\n", + "__The two failure modes__\n", + "\n", + "**Under-prediction** (n_model < n_observed): the model cannot produce an observed image. This is\n", + "always physically damning \u2014 you *saw* the image \u2014 so every scheme must penalize it hard. A\n", + "likelihood that quietly drops unmatched observed images actively rewards mass models that lens\n", + "less, and samplers will find and exploit that reward.\n", + "\n", + "**Over-prediction** (n_model > n_observed): the model predicts images you did not detect. This is\n", + "sometimes damning (a bright predicted image where the data shows blank sky) and sometimes entirely\n", + "fine \u2014 most lens models predict a strongly *demagnified* central image that real observations\n", + "cannot detect. The observational convention (shared by Lenstool practice) is therefore: extra\n", + "images below the detection limit are tolerated; bright extra images count against the model.\n", + "\n", + "__The three pairing schemes__\n", + "\n", + "- ``FitPositionsImagePairRepeat`` (the model-fit default): every observed position pairs to its\n", + " *nearest* model position, repeats allowed. Under-prediction is penalized by construction (an\n", + " unmatched observed image pays its distance to the nearest surviving image; if the solver returns\n", + " no images at all, a large finite floor applies). Over-prediction is governed by the\n", + " ``unmatched_model_policy`` described below.\n", + "\n", + "- ``FitPositionsImagePair``: Hungarian (linear-sum-assignment) pairing without repeats. Unmatched\n", + " observed positions (under-prediction) now contribute their distance to the nearest model\n", + " position \u2014 this scheme previously *dropped* them, which rewarded under-predicting models and is\n", + " why its docstring long carried a do-not-use warning. Repeats-forbidden pairing is mainly useful\n", + " when images are well separated and you want strict one-to-one bookkeeping.\n", + "\n", + "- ``FitPositionsImagePairAll``: a mixture likelihood \u2014 each observed position marginalizes over\n", + " every model position, and the 1/n_permutations normalization acts as an Occam factor that\n", + " mildly penalizes extra images. Statistically the most principled, and differentiable end to end\n", + " (it is the scheme the JAX point-source likelihood tests exercise); its penalties are implicit\n", + " rather than tunable.\n", + "\n", + "__The over-prediction policy (FitPositionsImagePairRepeat)__\n", + "\n", + "The ``unmatched_model_policy`` class attribute selects what happens to model images no observed\n", + "position paired to:\n", + "\n", + "- ``\"magnification_filter\"`` (default): model images with absolute magnification below\n", + " ``magnification_threshold`` (default 0.1) are exempt \u2014 the demagnified-central convention \u2014\n", + " and every *other* unmatched model image adds its distance to the nearest observed position as\n", + " a residual (normalized by the mean position noise).\n", + "- ``\"penalize\"``: as above with no magnification exemption.\n", + "- ``\"ignore\"``: extra images cost nothing \u2014 the historical behaviour, now an explicit opt-in.\n", + "\n", + "Switching policy uses the class-attribute pattern (no constructor plumbing):\n", + "\n", + " class FitStrict(al.FitPositionsImagePairRepeat):\n", + " unmatched_model_policy = \"penalize\"\n", + "\n", + " analysis = al.AnalysisPoint(dataset=dataset, solver=solver, fit_positions_cls=FitStrict)\n", + "\n", + "The ``n_unmatched_model_positions`` property reports how many extras the policy counted \u2014 worth\n", + "inspecting on any max-likelihood fit before trusting it.\n", + "\n", + "__Solver settings that masquerade as physics__\n", + "\n", + "The ``PointSolver`` tiles the image plane in triangles and refines toward the source position; a\n", + "too-coarse starting grid can *miss* a genuine image entirely. That looks exactly like model\n", + "under-prediction \u2014 but it is a numerical artifact, and it will bias the sampler for numerical\n", + "rather than physical reasons. Rules of thumb at cluster scale:\n", + "\n", + "- The starting grid must resolve the smallest image separation you care about: member-galaxy-scale\n", + " perturbations produce image pairs separated by ~1\", so grids much coarser than that will merge\n", + " or miss them.\n", + "- ``pixel_scale_precision`` sets the refinement floor; the cluster profiling scripts\n", + " (``autolens_profiling/likelihood_breakdown/cluster/image_plane.py``) time the cost of tightening\n", + " it (solve ~0.3 s/call at a 200x200 @ 0.7\" grid and 0.01\" precision, with a ~10 s one-off JAX\n", + " compile per source plane).\n", + "- If a fit reports under-prediction, re-solve the max-likelihood model at double resolution\n", + " before believing it: if the missing image appears, it was the grid.\n", + "\n", + "__Source-plane vs image-plane chi-squared__\n", + "\n", + "The source-plane chi-squared (``FitPositionsSource``, Lenstool's default) ray-traces observed\n", + "images backwards and never solves the lens equation \u2014 it is ~100x cheaper per evaluation (3 ms vs\n", + "0.3 s on the standard cluster model, per the profiling breakdowns) and pairing is trivial because\n", + "every observed image maps to one source. Its costs: it cannot see over-prediction *at all* (no\n", + "forward solve, so extra images never exist), magnification weighting only approximates the\n", + "image-plane noise mapping, and the magnification amplification gives it a documented precision\n", + "floor on high-magnification systems (see ``autolens_workspace_test/scripts/cluster/\n", + "likelihood_sanity.py``) \u2014 penalty terms of the image-plane policies sit far above that floor, but\n", + "sub-percent mass perturbations do not.\n", + "\n", + "The pragmatic workflow at cluster scale: **search with the source-plane chi-squared, validate with\n", + "the image-plane chi-squared** \u2014 run the image-plane fit (and inspect ``n_unmatched_model_positions``\n", + "plus the per-system image counts) on the max-likelihood model before publishing, exactly as the\n", + "Lenstool-users example does.\n", + "\n", + "__Demonstration__\n", + "\n", + "The code below builds a toy under- and over-predicting fit so the policies are visible in numbers,\n", + "using a mock solver so it runs in seconds." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "import numpy as np\n", + "\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "An isothermal lens with a point source: two bright observed images near the Einstein radius. The\n", + "mock solver lets us hand the fit whatever \"model\" images we want, isolating the pairing behaviour\n", + "from the lens equation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0),\n", + ")\n", + "source = al.Galaxy(redshift=1.0, point_0=al.ps.Point(centre=(0.0, 0.0)))\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "data = al.Grid2DIrregular([(0.0, 1.05), (0.0, -0.95)])\n", + "noise_map = al.ArrayIrregular([0.5, 0.5])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Case 1 \u2014 perfect prediction plus a demagnified central image__\n", + "\n", + "The model predicts both observed images exactly, plus a third image at 0.01\" from the lens centre\n", + "where an isothermal profile demagnifies to |mu| ~ 0.01. Under the default policy the central image\n", + "is exempt: chi-squared stays zero." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_data = al.Grid2DIrregular([(0.0, 1.05), (0.0, -0.95), (0.0, 0.01)])\n", + "solver = al.m.MockPointSolver(model_positions=model_data)\n", + "\n", + "fit = al.FitPositionsImagePairRepeat(\n", + " name=\"point_0\", data=data, noise_map=noise_map, tracer=tracer, solver=solver\n", + ")\n", + "print(\"Case 1 \u2014 extra demagnified central image (default policy):\")\n", + "print(f\" n_unmatched_model_positions = {int(fit.n_unmatched_model_positions)}\")\n", + "print(f\" chi_squared = {float(fit.chi_squared):.4f} (exempt below |mu| = 0.1)\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Case 2 \u2014 a bright unobserved image__\n", + "\n", + "Move the extra image out to 3\" \u2014 magnification order unity, no exemption. The model now pays for\n", + "predicting an image the data does not show." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_data = al.Grid2DIrregular([(0.0, 1.05), (0.0, -0.95), (0.0, 3.0)])\n", + "solver = al.m.MockPointSolver(model_positions=model_data)\n", + "\n", + "fit = al.FitPositionsImagePairRepeat(\n", + " name=\"point_0\", data=data, noise_map=noise_map, tracer=tracer, solver=solver\n", + ")\n", + "print(\"Case 2 \u2014 bright unobserved image:\")\n", + "print(f\" n_unmatched_model_positions = {int(fit.n_unmatched_model_positions)}\")\n", + "print(f\" chi_squared = {float(fit.chi_squared):.4f}\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Case 3 \u2014 under-prediction__\n", + "\n", + "The model produces only one of the two observed images. The unmatched observed image pays its full\n", + "distance to the surviving image \u2014 under-prediction is never free, under any scheme or policy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_data = al.Grid2DIrregular([(0.0, 1.05)])\n", + "solver = al.m.MockPointSolver(model_positions=model_data)\n", + "\n", + "fit = al.FitPositionsImagePairRepeat(\n", + " name=\"point_0\", data=data, noise_map=noise_map, tracer=tracer, solver=solver\n", + ")\n", + "print(\"Case 3 \u2014 missing image:\")\n", + "print(f\" residuals = {[round(float(r), 3) for r in np.asarray(fit.residual_map)]}\")\n", + "print(f\" chi_squared = {float(fit.chi_squared):.4f}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished. For the production-scale picture \u2014 real solver, multi-plane cluster tracer, timings \u2014\n", + "see ``scripts/cluster/likelihood_function.py`` and the profiling breakdowns referenced above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/guides/profiles/light.ipynb b/notebooks/guides/profiles/light.ipynb index e6c908d68..71e8a3ecc 100644 --- a/notebooks/guides/profiles/light.ipynb +++ b/notebooks/guides/profiles/light.ipynb @@ -1,882 +1,919 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Light Profiles\n", - "==============\n", - "\n", - "This guide is the single-page tour of every light profile available in **PyAutoLens** (all of\n", - "which are re-exported from **PyAutoGalaxy**): how to construct each one, how to evaluate its\n", - "image on a grid, how to compose it into a model, and how to pull an instance back out of that\n", - "model.\n", - "\n", - "Strong-lensing systems typically have **two** populations of light: emission from the lens\n", - "galaxy in the foreground and emission from the lensed source in the background. In this\n", - "guide both are demonstrated, with the multi-plane combination handled by the `Tracer` \u2014\n", - "**PyAutoLens**'s container for galaxies grouped by redshift plane.\n", - "\n", - "The guide is deliberately broad rather than deep \u2014 for each family it shows the *shape* of the\n", - "API and points you at the relevant `features/` package for the workflow details.\n", - "\n", - "__Contents__\n", - "\n", - "- **Overview & Docs URL:** Where the canonical API reference lives.\n", - "- **All Light Profiles (Survey):** A high-level run-through of every profile in `al.lp.*` and\n", - " the related namespaces, without yet evaluating any images.\n", - "- **Detailed Example: Sersic Image:** Build a `Grid2D`, instantiate `al.lp.Sersic`, evaluate\n", - " `image_2d_from`, plot it.\n", - "- **Linear Light Profiles:** One-line API for `al.lp_linear.*` \u2014 intensity solved by inversion.\n", - "- **Operated Light Profiles:** One-line API for `al.lp_operated.*` \u2014 emission post PSF.\n", - "- **Basis:** The grouping object that lets many profiles behave as a single composite \u2014 the\n", - " building block of Multi-Gaussian Expansion (MGE) and shapelet decompositions.\n", - "- **Light Profile in a Model:** Wrap a profile in `af.Model`, compose it into an\n", - " `af.Collection` via a `Galaxy` and a `Tracer`, inspect the model info.\n", - "- **Model Instance from Light Profile:** Realise an instance from the model's prior medians,\n", - " drop it into a `Tracer`, and evaluate `image_2d_from` on the tracer.\n", - "- **Multipole Light Profiles:** The newer `SersicMultipole` and `GaussianMultipole`, with the\n", - " m=3 / m=4 Fourier perturbation on the eccentric radius explained and plotted.\n", - "- **Remaining Profiles Walkthrough:** Compact `image_2d_from` block for every standard profile\n", - " not yet shown, emphasising the API is the same as the Sersic example above.\n", - "\n", - "__Units__\n", - "\n", - "In this guide, all quantities use the internal unit coordinates of **PyAutoLens**: spatial\n", - "coordinates in arc-seconds, luminosities in electrons per second, and mass quantities (e.g.\n", - "convergence) are dimensionless.\n", - "\n", - "The `guides/units_and_cosmology.ipynb` guide illustrates how to convert these to physical\n", - "quantities (kiloparsecs, magnitudes, solar masses).\n", - "\n", - "__Data Structures__\n", - "\n", - "Images returned by `image_2d_from` are wrapped in the `Array2D` data structure with `slim`\n", - "and `native` views. The `guides/data_structures.py` guide covers this in detail; here we\n", - "only use the default `slim` 1D representation when printing values.\n", - "\n", - "__Docs URL__\n", - "\n", - "The published API reference for these classes lives at:\n", - "\n", - " https://pyautolens.readthedocs.io/en/latest/api/light.html\n", - "\n", - "The autosummary on that page is the authoritative list of every public light-profile class.\n", - "This guide mirrors it section-by-section, so a class shown here as `al.lp.SersicCore` is\n", - "documented there under the `Standard [ag.lp]` autosummary, and so on for `al.lp_linear`,\n", - "`al.lp_operated`, `al.lp_basis`. Note that the API reference uses the `ag.*` namespace\n", - "labels because the classes are defined in PyAutoGalaxy and re-exported here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "To evaluate the image of any light profile we need a 2D Cartesian grid of (y,x) coordinates.\n", - "We build a 100x100 grid here at a 0.05\" pixel scale \u2014 used by every section below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.05,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__All Light Profiles (Survey)__\n", - "\n", - "**PyAutoLens** groups light profiles into five namespaces, each with a clear purpose:\n", - "\n", - "- `al.lp.*` \u2014 *Standard* parametric profiles. `intensity` is a free model parameter.\n", - "- `al.lp_linear.*` \u2014 *Linear* profiles. `intensity` is removed from the model and instead\n", - " solved analytically via a linear matrix inversion during each likelihood evaluation.\n", - "- `al.lp_operated.*` \u2014 *Operated* profiles representing emission that has already had an\n", - " instrument operation (e.g. PSF convolution) applied to it; `operated_only` on the fit\n", - " classes controls inclusion.\n", - "- `al.lp_basis.Basis` \u2014 A grouping object that bundles multiple light profiles into a single\n", - " composite profile (e.g. an MGE built from many Gaussians).\n", - "- `al.lp_snr.*` \u2014 Standard profiles parameterised by *signal-to-noise ratio* rather than\n", - " intensity; useful when simulating a dataset with a target SNR. Not covered further in this\n", - " guide, but shares the API of the Standard profiles.\n", - "\n", - "Below we construct each standard profile with default parameters. No `image_2d_from` is\n", - "evaluated yet \u2014 that comes in the next section. The goal here is purely a catalogue of what\n", - "is available, in the same order they appear in the API reference." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Sersic family\n", - "sersic = al.lp.Sersic()\n", - "sersic_sph = al.lp.SersicSph()\n", - "sersic_core = al.lp.SersicCore()\n", - "sersic_core_sph = al.lp.SersicCoreSph()\n", - "sersic_multipole = al.lp.SersicMultipole()\n", - "\n", - "# Exponential family (Sersic with sersic_index fixed to 1)\n", - "exponential = al.lp.Exponential()\n", - "exponential_sph = al.lp.ExponentialSph()\n", - "exponential_core = al.lp.ExponentialCore()\n", - "exponential_core_sph = al.lp.ExponentialCoreSph()\n", - "\n", - "# de Vaucouleurs (Sersic with sersic_index fixed to 4)\n", - "dev_vaucouleurs = al.lp.DevVaucouleurs()\n", - "dev_vaucouleurs_sph = al.lp.DevVaucouleursSph()\n", - "\n", - "# Gaussian / Moffat / Multipole-Gaussian\n", - "gaussian = al.lp.Gaussian()\n", - "gaussian_sph = al.lp.GaussianSph()\n", - "gaussian_multipole = al.lp.GaussianMultipole()\n", - "moffat = al.lp.Moffat()\n", - "moffat_sph = al.lp.MoffatSph()\n", - "\n", - "# Specialised: Chameleon (NFW-like double-isothermal) and Elson-Free-Fall (King-like)\n", - "chameleon = al.lp.Chameleon()\n", - "chameleon_sph = al.lp.ChameleonSph()\n", - "eff = al.lp.ElsonFreeFall()\n", - "eff_sph = al.lp.ElsonFreeFallSph()\n", - "\n", - "# Shapelets \u2014 n_y, n_x (Cartesian) or n, m (Polar) pick the basis index\n", - "shapelet_cartesian = al.lp.ShapeletCartesian(n_y=0, n_x=0)\n", - "shapelet_polar = al.lp.ShapeletPolar(n=0, m=0)\n", - "shapelet_exponential = al.lp.ShapeletExponential(n=0, m=0)\n", - "\n", - "# Basis \u2014 bundles a list of light profiles into a single composite\n", - "basis = al.lp_basis.Basis(profile_list=[al.lp_linear.Gaussian(sigma=0.5)])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two things worth knowing about this list before we move on:\n", - "\n", - "1. Every elliptical profile (e.g. `Sersic`, `Gaussian`) has a spherical sibling whose name\n", - " ends in `Sph` (e.g. `SersicSph`, `GaussianSph`). The spherical variant fixes the\n", - " ellipticity components `ell_comps` to `(0, 0)`, which is useful when you want to model a\n", - " round galaxy and avoid two redundant parameters in the non-linear search.\n", - "2. The `Multipole` variants (`SersicMultipole`, `GaussianMultipole`) only exist as\n", - " *elliptical* profiles \u2014 the m=3 / m=4 perturbations are angular distortions and are not\n", - " meaningful without an underlying elliptical reference frame.\n", - "\n", - "We now move on to seeing what these profiles actually produce when evaluated on a grid.\n", - "\n", - "__Detailed Example: Sersic Image__\n", - "\n", - "The `Sersic` profile is the canonical galaxy light profile, controlled by:\n", - "\n", - "- `centre` \u2014 the (y, x) arc-second coordinate of the profile's centre.\n", - "- `ell_comps` \u2014 the two ellipticity components `(e1, e2)`. Use\n", - " `al.convert.ell_comps_from(axis_ratio=..., angle=...)` to convert from human-friendly\n", - " axis ratio and position angle.\n", - "- `intensity` \u2014 overall brightness normalisation.\n", - "- `effective_radius` \u2014 the half-light radius (arc-seconds).\n", - "- `sersic_index` \u2014 the Sersic concentration. `n=1` reduces to an exponential disc and\n", - " `n=4` reduces to a de Vaucouleurs profile.\n", - "\n", - "Build a Sersic and evaluate its image on our grid:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sersic = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")\n", - "\n", - "image = sersic.image_2d_from(grid=grid)\n", - "\n", - "aplt.plot_array(array=image, title=\"Sersic Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The returned `image` is an `Array2D` \u2014 the `slim` view is a 1D numpy array of length\n", - "`total_pixels`, and `native` gives a 2D `(shape_native_y, shape_native_x)` array.\n", - "\n", - "This same `image_2d_from(grid=grid)` call exists on every light profile in this guide,\n", - "returning an image of identical shape and units. Every section below is a small variation\n", - "on this one \u2014 the API is uniform.\n", - "\n", - "__Linear Light Profiles__\n", - "\n", - "For a non-linear search, the `intensity` parameter of a standard light profile is a free\n", - "parameter sampled by the fitter. This works fine for one or two profiles, but adds a free\n", - "dimension to the parameter space for every extra profile you bolt on.\n", - "\n", - "Linear light profiles solve this by removing `intensity` from the model entirely and instead\n", - "recovering it analytically via a linear matrix inversion at each likelihood evaluation. This\n", - "keeps the non-linear parameter space small even when you combine many profiles, and is the\n", - "default in our modern lens-modeling examples.\n", - "\n", - "The API is identical to the standard profile, just without `intensity`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "linear_sersic = al.lp_linear.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Every standard profile in `al.lp.*` has an `al.lp_linear.*` counterpart, **including the\n", - "newer `SersicMultipole` and `GaussianMultipole`** \u2014 `al.lp_linear.SersicMultipole` and\n", - "`al.lp_linear.GaussianMultipole` both exist and behave the same way (the multipole comps\n", - "are non-linear parameters; only the overall intensity is solved by inversion).\n", - "\n", - "The full workflow (likelihood function, fits, modeling) is documented in:\n", - "\n", - " scripts/imaging/features/linear_light_profiles/\n", - "\n", - "That folder contains `fit.py`, `modeling.py`, and `likelihood_function.py` showing how to\n", - "build lens models with linear profiles and what the likelihood looks like under the hood.\n", - "\n", - "__Operated Light Profiles__\n", - "\n", - "Some emission components \u2014 chiefly the unresolved bright cores of AGN \u2014 are already PSF-\n", - "convolved by the time you receive the image. Standard profiles get PSF-convolved during the\n", - "fit, so applying a PSF a second time double-blurs them.\n", - "\n", - "Operated light profiles tell the fit \"this profile's emission has already had the PSF\n", - "operation applied; do not blur it again\". The fit classes expose an `operated_only` flag\n", - "that controls whether these profiles are included or excluded from a given image computation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "operated_gaussian = al.lp_operated.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=0.3,\n", - " sigma=0.05,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Three operated profiles are available: `al.lp_operated.Gaussian`, `al.lp_operated.Moffat`,\n", - "and `al.lp_operated.Sersic`.\n", - "\n", - "The full workflow (simulating with operated profiles, modeling with them) is documented in:\n", - "\n", - " scripts/imaging/features/operated_light_profile/\n", - "\n", - "That folder contains `simulator.py` and `modeling.py`.\n", - "\n", - "__Basis__\n", - "\n", - "A `Basis` is not a profile in its own right but a *grouping* of profiles that behave as a\n", - "single composite. The classic application is the Multi-Gaussian Expansion (MGE), where a\n", - "galaxy's light is decomposed into a sum of many concentric Gaussians at fixed centres and\n", - "ellipticities but with increasing widths \u2014 together they reproduce arbitrary radial profiles\n", - "the standard parametric forms cannot capture. In strong-lens modelling MGEs are routinely\n", - "used for the lens-galaxy light to soak up complex morphology without inflating the parameter\n", - "space.\n", - "\n", - "The `Basis` constructor takes a `profile_list` of any light or mass profiles. Below we\n", - "build a four-Gaussian MGE with shared centre and ellipticity, sigmas that span an order of\n", - "magnitude, and explicit decreasing `intensity` so the innermost Gaussian dominates the\n", - "core and the wider ones add the outer envelope. Standard `al.lp.Gaussian` profiles are\n", - "used here so the demo image is meaningful; in an actual lens fit you would swap these for\n", - "`al.lp_linear.Gaussian` (see the note after the plot)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "basis = al.lp_basis.Basis(\n", - " profile_list=[\n", - " al.lp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " sigma=0.05,\n", - " ),\n", - " al.lp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=0.5,\n", - " sigma=0.15,\n", - " ),\n", - " al.lp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=0.25,\n", - " sigma=0.4,\n", - " ),\n", - " al.lp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=0.1,\n", - " sigma=1.0,\n", - " ),\n", - " ]\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=basis.image_2d_from(grid=grid),\n", - " title=\"Basis Image (4-Gaussian MGE)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two things make `Basis` powerful:\n", - "\n", - "- It slots into a `Galaxy` exactly like a `Sersic` would \u2014 once wrapped, the rest of the\n", - " modelling code doesn't have to know it's looking at four Gaussians under the hood.\n", - "- When the constituents are `LightProfileLinear` instances (e.g. `al.lp_linear.Gaussian`)\n", - " rather than the standard `al.lp.Gaussian` used in the demo above, all of their\n", - " `intensity` values are solved together in a **single combined inversion** at each\n", - " likelihood evaluation. This means an MGE built from, say, 30 Gaussians adds only the\n", - " shared geometric parameters to the non-linear search rather than 30 extra intensities.\n", - " The demo above uses standard Gaussians purely so the image is non-zero on a static\n", - " plot \u2014 in a real lens fit you would build the basis from `al.lp_linear.Gaussian` and\n", - " let the inversion solve the intensities.\n", - "\n", - "The full MGE workflow \u2014 choosing how many Gaussians to use, how to space their `sigma`\n", - "values, and how the inversion plays with regularisation \u2014 is documented in:\n", - "\n", - " scripts/imaging/features/multi_gaussian_expansion/\n", - "\n", - "Shapelet decompositions follow the same `Basis` pattern, using `al.lp.ShapeletPolar` /\n", - "`al.lp.ShapeletCartesian` / `al.lp.ShapeletExponential` (and their linear counterparts).\n", - "The full shapelets workflow is documented in:\n", - "\n", - " scripts/imaging/features/shapelets/\n", - "\n", - "__Light Profile in a Model__\n", - "\n", - "So far we have been instantiating profiles with concrete parameter values. When fitting a\n", - "real strong-lens dataset we instead build a *model* of the profile and let the non-linear\n", - "search find the best-fit parameters. This is what `af.Model` is for.\n", - "\n", - "For a lens system the model typically contains two `Galaxy` objects on different planes: a\n", - "foreground lens (with light + mass) and a background source (with light only). In this\n", - "section we focus on the *light* side of that picture and show how a Sersic light profile\n", - "plugs in to both." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_bulge_model = af.Model(al.lp.Sersic)\n", - "source_bulge_model = af.Model(al.lp.Sersic)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `af.Model` wraps the profile class. Every constructor argument that has a numerical\n", - "default now becomes a *prior* \u2014 by default the priors are `UniformPriors` covering a sensible\n", - "range for each parameter (see the autogalaxy config for the configured ranges).\n", - "\n", - "You can override individual priors before fitting:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_bulge_model.sersic_index = af.UniformPrior(lower_limit=0.5, upper_limit=8.0)\n", - "source_bulge_model.effective_radius = af.UniformPrior(\n", - " lower_limit=0.01, upper_limit=10.0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A model profile by itself is not a complete model \u2014 for a strong lens we wrap each profile\n", - "in a `Galaxy` at its own redshift and assemble them in an `af.Collection`. The `Tracer` is\n", - "not part of the *model* spec itself; it is what `AnalysisImaging.log_likelihood_function`\n", - "builds out of the realised instance at fit time." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy_model = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge_model)\n", - "source_galaxy_model = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge_model)\n", - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens_galaxy_model, source=source_galaxy_model)\n", - ")\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing `model.info` prints the full priors-and-defaults summary \u2014 useful before kicking off\n", - "a long fit to confirm the model looks the way you expect.\n", - "\n", - "The model API is the same for **every** light profile in this guide \u2014 swap `al.lp.Sersic`\n", - "for `al.lp.SersicMultipole`, `al.lp_linear.Gaussian`, `al.lp_basis.Basis`, etc., and the\n", - "rest of the snippet is unchanged. Multipole comps and Basis constituent lists are wired\n", - "into the prior machinery automatically.\n", - "\n", - "Full lens-modelling end-to-end examples live in `scripts/imaging/modeling/start_here.py`\n", - "and the topic-specific guides under `scripts/imaging/features/`.\n", - "\n", - "__Model Instance from Light Profile__\n", - "\n", - "A model is a description of *possible* profiles. To get an actual profile back out \u2014 for\n", - "example to plot what the prior medians look like before running a fit \u2014 call\n", - "`instance_from_prior_medians()`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_bulge_instance = lens_bulge_model.instance_from_prior_medians()\n", - "print(type(lens_bulge_instance)) # autogalaxy.profiles.light.standard.sersic.Sersic\n", - "\n", - "image = lens_bulge_instance.image_2d_from(grid=grid)\n", - "aplt.plot_array(array=image, title=\"Lens Bulge Instance from Prior Medians\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The instance returned from `instance_from_prior_medians()` is a real `al.lp.Sersic` \u2014 the\n", - "same class we constructed by hand at the top of the guide \u2014 and supports the full API\n", - "including `image_2d_from`.\n", - "\n", - "The same flow works at the full-model level. We realise an instance of the lens-and-source\n", - "collection and drop the resulting galaxies into a `Tracer`, which is the **PyAutoLens**\n", - "object that handles ray-tracing across redshift planes and combining the lens-galaxy light\n", - "with the lensed source-galaxy light into a single image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_instance = model.instance_from_prior_medians()\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[model_instance.galaxies.lens, model_instance.galaxies.source]\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Tracer Image (Lens Light + Lensed Source Light)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Tracer` takes care of two things at once: it sums the lens-galaxy light (no deflection)\n", - "with the source-galaxy light *after* ray-tracing the grid through the lens-plane deflections,\n", - "producing the combined observed image you would see at the telescope.\n", - "\n", - "When the lens galaxy has a mass profile, the source would appear as a lensed arc. Our\n", - "example builds the model with no mass profile, so the source contributes its un-deflected\n", - "emission \u2014 useful for sanity-checking the *light* side of the model alone. The `tracer.py`\n", - "guide in this directory shows the full lensing flow with mass profiles included.\n", - "\n", - "You can still pull the individual light-profile instances back out of the same object:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_bulge_instance = model_instance.galaxies.lens.bulge\n", - "source_bulge_instance = model_instance.galaxies.source.bulge\n", - "print(type(lens_bulge_instance)) # autogalaxy.profiles.light.standard.sersic.Sersic\n", - "print(type(source_bulge_instance)) # autogalaxy.profiles.light.standard.sersic.Sersic" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "After a fit completes, `result.max_log_likelihood_tracer` returns the same shape of object,\n", - "with the prior medians replaced by the fitted parameter values. See\n", - "`scripts/guides/results/start_here.py` for the full results-introspection guide.\n", - "\n", - "__Multipole Light Profiles__\n", - "\n", - "`SersicMultipole` and `GaussianMultipole` are recent additions that bolt m=3 and m=4 Fourier\n", - "angular perturbations onto the eccentric radius of a base profile. The perturbed radius is\n", - "\n", - " r' = r * (1 + c3 cos(3 theta) + s3 sin(3 theta)\n", - " + c4 cos(4 theta) + s4 sin(4 theta))\n", - "\n", - "where `theta` is the polar angle in the profile's elliptical reference frame, and the\n", - "`multipole_3_comps = (c3, s3)` and `multipole_4_comps = (c4, s4)` parameters control the\n", - "amplitude of each perturbation.\n", - "\n", - "When both `multipole_*_comps` are `(0.0, 0.0)` (the defaults), the profile reduces exactly\n", - "to the base profile. This is by design \u2014 you can swap a `Sersic` for a `SersicMultipole`\n", - "in any model without changing its predictions, and the multipole comps simply add four\n", - "extra free parameters that can capture boxy / discy / lopsided morphologies in lens-galaxy\n", - "light or in resolved source morphology. Plugging one into the `af.Model` / `af.Collection`\n", - "/ `Galaxy` / `Tracer` pattern shown above works exactly as it did for the plain `Sersic` \u2014\n", - "the multipole comps are picked up as priors automatically.\n", - "\n", - "Build a `SersicMultipole` with non-zero multipole components, alongside the unperturbed\n", - "Sersic that produced our reference image earlier in the guide:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sersic_multipole = al.lp.SersicMultipole(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " multipole_3_comps=(0.05, 0.00),\n", - " multipole_4_comps=(0.00, 0.04),\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=sersic_multipole.image_2d_from(grid=grid),\n", - " title=\"SersicMultipole Image (m=3 + m=4 perturbation)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For comparison, here is the unperturbed Sersic image \u2014 the two should look almost\n", - "identical with the perturbation showing as a subtle azimuthal modulation:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=sersic.image_2d_from(grid=grid),\n", - " title=\"Sersic Image (no multipole perturbation)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `GaussianMultipole` profile applies the same perturbation to a Gaussian base \u2014 useful\n", - "when you want a multipole component inside a Multi-Gaussian Expansion (the `Basis` section\n", - "above shows how to bundle Gaussians together):" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "gaussian_multipole = al.lp.GaussianMultipole(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " sigma=0.4,\n", - " multipole_3_comps=(0.05, 0.00),\n", - " multipole_4_comps=(0.00, 0.04),\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=gaussian_multipole.image_2d_from(grid=grid),\n", - " title=\"GaussianMultipole Image\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two practical notes on the multipole variants:\n", - "\n", - "- There is **no spherical (`*Sph`) variant** of either multipole. The perturbation is an\n", - " angular distortion measured in the elliptical reference frame, so it only makes sense for\n", - " an elliptical profile (a spherical profile has no preferred angle).\n", - "- Both multipoles exist as **linear variants** too: `al.lp_linear.SersicMultipole` and\n", - " `al.lp_linear.GaussianMultipole`. In the linear form the multipole comps remain non-\n", - " linear parameters but the overall intensity is solved by inversion, just like for the\n", - " ordinary linear profiles.\n", - "\n", - "__Remaining Profiles Walkthrough__\n", - "\n", - "We have shown the full `image_2d_from` \u2192 `af.Model` \u2192 `instance` \u2192 `Tracer` flow for the\n", - "`Sersic` profile. Every remaining standard profile uses the **same API** \u2014 the only thing\n", - "that changes is which parameters appear in the constructor.\n", - "\n", - "The compact tour below builds each remaining profile with sensible parameter values and\n", - "plots its image, so you can see what each looks like. When you want to use any of these in\n", - "a lens model, repeat the `af.Model(...)` / `af.Collection(...)` / `Tracer(...)` pattern\n", - "from the previous section." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(\n", - " array=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " radius_break=0.05,\n", - " gamma=0.2,\n", - " alpha=3.0,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"SersicCore Image\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.lp.Exponential(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", - " intensity=0.5,\n", - " effective_radius=1.6,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"Exponential Image\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.lp.ExponentialCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", - " intensity=0.5,\n", - " effective_radius=1.6,\n", - " radius_break=0.05,\n", - " gamma=0.2,\n", - " alpha=3.0,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"ExponentialCore Image\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.lp.DevVaucouleurs(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"DevVaucouleurs Image\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.lp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " sigma=0.4,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"Gaussian Image\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.lp.Moffat(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " alpha=0.4,\n", - " beta=2.5,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"Moffat Image\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.lp.Chameleon(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " core_radius_0=0.05,\n", - " core_radius_1=0.3,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"Chameleon Image\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.lp.ElsonFreeFall(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " eta=2.0,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"ElsonFreeFall Image\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The spherical variants (`SersicSph`, `GaussianSph`, etc.) are constructed identically with\n", - "the `ell_comps` argument removed \u2014 they look like a rotationally symmetric version of the\n", - "corresponding elliptical plot.\n", - "\n", - "The shapelet profiles are normally used inside a `Basis` rather than individually, but for\n", - "completeness here is the lowest-order Cartesian shapelet on its own:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=al.lp.ShapeletCartesian(\n", - " n_y=0,\n", - " n_x=0,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=1.0,\n", - " beta=0.2,\n", - " ).image_2d_from(grid=grid),\n", - " title=\"ShapeletCartesian (n_y=0, n_x=0) Image\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And that completes the tour. If you arrived here from the API reference and now want to use\n", - "any of these profiles in an actual lens fit, the next step is\n", - "`scripts/imaging/modeling/start_here.py`, which sets up an `AnalysisImaging` and runs a\n", - "non-linear search end-to-end on a strong-lens dataset. The `scripts/imaging/features/`\n", - "subpackages handle the family-specific workflows referenced throughout this guide:\n", - "\n", - "- `linear_light_profiles/` \u2014 using `al.lp_linear.*` in a fit.\n", - "- `operated_light_profile/` \u2014 using `al.lp_operated.*` in a fit.\n", - "- `multi_gaussian_expansion/` \u2014 building and fitting an MGE-style `Basis`.\n", - "- `shapelets/` \u2014 building and fitting a shapelet-style `Basis`.\n", - "\n", - "For the mass-profile counterpart of this guide, see `scripts/guides/profiles/mass.py`\n", - "(when available). For a full walk-through of how light and mass profiles are combined in\n", - "the `Tracer` to produce lensed images, see `scripts/guides/tracer.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Light Profiles\n", + "==============\n", + "\n", + "This guide is the single-page tour of every light profile available in **PyAutoLens** (all of\n", + "which are re-exported from **PyAutoGalaxy**): how to construct each one, how to evaluate its\n", + "image on a grid, how to compose it into a model, and how to pull an instance back out of that\n", + "model.\n", + "\n", + "Strong-lensing systems typically have **two** populations of light: emission from the lens\n", + "galaxy in the foreground and emission from the lensed source in the background. In this\n", + "guide both are demonstrated, with the multi-plane combination handled by the `Tracer` \u2014\n", + "**PyAutoLens**'s container for galaxies grouped by redshift plane.\n", + "\n", + "The guide is deliberately broad rather than deep \u2014 for each family it shows the *shape* of the\n", + "API and points you at the relevant `features/` package for the workflow details.\n", + "\n", + "__Contents__\n", + "\n", + "- **Overview & Docs URL:** Where the canonical API reference lives.\n", + "- **All Light Profiles (Survey):** A high-level run-through of every profile in `al.lp.*` and\n", + " the related namespaces, without yet evaluating any images.\n", + "- **Detailed Example: Sersic Image:** Build a `Grid2D`, instantiate `al.lp.Sersic`, evaluate\n", + " `image_2d_from`, plot it.\n", + "- **Linear Light Profiles:** One-line API for `al.lp_linear.*` \u2014 intensity solved by inversion.\n", + "- **Operated Light Profiles:** One-line API for `al.lp_operated.*` \u2014 emission post PSF.\n", + "- **Basis:** The grouping object that lets many profiles behave as a single composite \u2014 the\n", + " building block of Multi-Gaussian Expansion (MGE) and shapelet decompositions.\n", + "- **Light Profile in a Model:** Wrap a profile in `af.Model`, compose it into an\n", + " `af.Collection` via a `Galaxy` and a `Tracer`, inspect the model info.\n", + "- **Model Instance from Light Profile:** Realise an instance from the model's prior medians,\n", + " drop it into a `Tracer`, and evaluate `image_2d_from` on the tracer.\n", + "- **Multipole Light Profiles:** The newer `SersicMultipole` and `GaussianMultipole`, with the\n", + " m=3 / m=4 Fourier perturbation on the eccentric radius explained and plotted.\n", + "- **Remaining Profiles Walkthrough:** Compact `image_2d_from` block for every standard profile\n", + " not yet shown, emphasising the API is the same as the Sersic example above.\n", + "\n", + "__Units__\n", + "\n", + "In this guide, all quantities use the internal unit coordinates of **PyAutoLens**: spatial\n", + "coordinates in arc-seconds, luminosities in electrons per second, and mass quantities (e.g.\n", + "convergence) are dimensionless.\n", + "\n", + "The `guides/units_and_cosmology.ipynb` guide illustrates how to convert these to physical\n", + "quantities (kiloparsecs, magnitudes, solar masses).\n", + "\n", + "__Data Structures__\n", + "\n", + "Images returned by `image_2d_from` are wrapped in the `Array2D` data structure with `slim`\n", + "and `native` views. The `guides/data_structures.py` guide covers this in detail; here we\n", + "only use the default `slim` 1D representation when printing values.\n", + "\n", + "__Docs URL__\n", + "\n", + "The published API reference for these classes lives at:\n", + "\n", + " https://pyautolens.readthedocs.io/en/latest/api/light.html\n", + "\n", + "The autosummary on that page is the authoritative list of every public light-profile class.\n", + "This guide mirrors it section-by-section, so a class shown here as `al.lp.SersicCore` is\n", + "documented there under the `Standard [ag.lp]` autosummary, and so on for `al.lp_linear`,\n", + "`al.lp_operated`, `al.lp_basis`. Note that the API reference uses the `ag.*` namespace\n", + "labels because the classes are defined in PyAutoGalaxy and re-exported here." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "To evaluate the image of any light profile we need a 2D Cartesian grid of (y,x) coordinates.\n", + "We build a 100x100 grid here at a 0.05\" pixel scale \u2014 used by every section below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.05,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__All Light Profiles (Survey)__\n", + "\n", + "**PyAutoLens** groups light profiles into five namespaces, each with a clear purpose:\n", + "\n", + "- `al.lp.*` \u2014 *Standard* parametric profiles. `intensity` is a free model parameter.\n", + "- `al.lp_linear.*` \u2014 *Linear* profiles. `intensity` is removed from the model and instead\n", + " solved analytically via a linear matrix inversion during each likelihood evaluation.\n", + "- `al.lp_operated.*` \u2014 *Operated* profiles representing emission that has already had an\n", + " instrument operation (e.g. PSF convolution) applied to it; `operated_only` on the fit\n", + " classes controls inclusion.\n", + "- `al.lp_basis.Basis` \u2014 A grouping object that bundles multiple light profiles into a single\n", + " composite profile (e.g. an MGE built from many Gaussians).\n", + "- `al.lp_snr.*` \u2014 Standard profiles parameterised by *signal-to-noise ratio* rather than\n", + " intensity; useful when simulating a dataset with a target SNR. Not covered further in this\n", + " guide, but shares the API of the Standard profiles.\n", + "\n", + "Below we construct each standard profile with default parameters. No `image_2d_from` is\n", + "evaluated yet \u2014 that comes in the next section. The goal here is purely a catalogue of what\n", + "is available, in the same order they appear in the API reference." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Sersic family\n", + "sersic = al.lp.Sersic()\n", + "sersic_sph = al.lp.SersicSph()\n", + "sersic_core = al.lp.SersicCore()\n", + "sersic_core_sph = al.lp.SersicCoreSph()\n", + "sersic_multipole = al.lp.SersicMultipole()\n", + "\n", + "# Exponential family (Sersic with sersic_index fixed to 1)\n", + "exponential = al.lp.Exponential()\n", + "exponential_sph = al.lp.ExponentialSph()\n", + "exponential_core = al.lp.ExponentialCore()\n", + "exponential_core_sph = al.lp.ExponentialCoreSph()\n", + "\n", + "# de Vaucouleurs (Sersic with sersic_index fixed to 4)\n", + "dev_vaucouleurs = al.lp.DevVaucouleurs()\n", + "dev_vaucouleurs_sph = al.lp.DevVaucouleursSph()\n", + "\n", + "# Gaussian / Moffat / Multipole-Gaussian\n", + "gaussian = al.lp.Gaussian()\n", + "gaussian_sph = al.lp.GaussianSph()\n", + "gaussian_multipole = al.lp.GaussianMultipole()\n", + "moffat = al.lp.Moffat()\n", + "moffat_sph = al.lp.MoffatSph()\n", + "\n", + "# Specialised: Chameleon (NFW-like double-isothermal) and Elson-Free-Fall (King-like)\n", + "chameleon = al.lp.Chameleon()\n", + "chameleon_sph = al.lp.ChameleonSph()\n", + "eff = al.lp.ElsonFreeFall()\n", + "eff_sph = al.lp.ElsonFreeFallSph()\n", + "\n", + "# Shapelets \u2014 n_y, n_x (Cartesian) or n, m (Polar) pick the basis index\n", + "shapelet_cartesian = al.lp.ShapeletCartesian(n_y=0, n_x=0)\n", + "shapelet_polar = al.lp.ShapeletPolar(n=0, m=0)\n", + "shapelet_exponential = al.lp.ShapeletExponential(n=0, m=0)\n", + "\n", + "# Basis \u2014 bundles a list of light profiles into a single composite\n", + "basis = al.lp_basis.Basis(profile_list=[al.lp_linear.Gaussian(sigma=0.5)])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Two things worth knowing about this list before we move on:\n", + "\n", + "1. Every elliptical profile (e.g. `Sersic`, `Gaussian`) has a spherical sibling whose name\n", + " ends in `Sph` (e.g. `SersicSph`, `GaussianSph`). The spherical variant fixes the\n", + " ellipticity components `ell_comps` to `(0, 0)`, which is useful when you want to model a\n", + " round galaxy and avoid two redundant parameters in the non-linear search.\n", + "2. The `Multipole` variants (`SersicMultipole`, `GaussianMultipole`) only exist as\n", + " *elliptical* profiles \u2014 the m=3 / m=4 perturbations are angular distortions and are not\n", + " meaningful without an underlying elliptical reference frame.\n", + "\n", + "We now move on to seeing what these profiles actually produce when evaluated on a grid.\n", + "\n", + "__Detailed Example: Sersic Image__\n", + "\n", + "The `Sersic` profile is the canonical galaxy light profile, controlled by:\n", + "\n", + "- `centre` \u2014 the (y, x) arc-second coordinate of the profile's centre.\n", + "- `ell_comps` \u2014 the two ellipticity components `(e1, e2)`. Use\n", + " `al.convert.ell_comps_from(axis_ratio=..., angle=...)` to convert from human-friendly\n", + " axis ratio and position angle.\n", + "- `intensity` \u2014 overall brightness normalisation.\n", + "- `effective_radius` \u2014 the half-light radius (arc-seconds).\n", + "- `sersic_index` \u2014 the Sersic concentration. `n=1` reduces to an exponential disc and\n", + " `n=4` reduces to a de Vaucouleurs profile.\n", + "\n", + "Build a Sersic and evaluate its image on our grid:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sersic = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")\n", + "\n", + "image = sersic.image_2d_from(grid=grid)\n", + "\n", + "aplt.plot_array(array=image, title=\"Sersic Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The returned `image` is an `Array2D` \u2014 the `slim` view is a 1D numpy array of length\n", + "`total_pixels`, and `native` gives a 2D `(shape_native_y, shape_native_x)` array.\n", + "\n", + "This same `image_2d_from(grid=grid)` call exists on every light profile in this guide,\n", + "returning an image of identical shape and units. Every section below is a small variation\n", + "on this one \u2014 the API is uniform.\n", + "\n", + "__Linear Light Profiles__\n", + "\n", + "For a non-linear search, the `intensity` parameter of a standard light profile is a free\n", + "parameter sampled by the fitter. This works fine for one or two profiles, but adds a free\n", + "dimension to the parameter space for every extra profile you bolt on.\n", + "\n", + "Linear light profiles solve this by removing `intensity` from the model entirely and instead\n", + "recovering it analytically via a linear matrix inversion at each likelihood evaluation. This\n", + "keeps the non-linear parameter space small even when you combine many profiles, and is the\n", + "default in our modern lens-modeling examples.\n", + "\n", + "The API is identical to the standard profile, just without `intensity`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "linear_sersic = al.lp_linear.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Every standard profile in `al.lp.*` has an `al.lp_linear.*` counterpart, **including the\n", + "newer `SersicMultipole` and `GaussianMultipole`** \u2014 `al.lp_linear.SersicMultipole` and\n", + "`al.lp_linear.GaussianMultipole` both exist and behave the same way (the multipole comps\n", + "are non-linear parameters; only the overall intensity is solved by inversion).\n", + "\n", + "The full workflow (likelihood function, fits, modeling) is documented in:\n", + "\n", + " scripts/imaging/features/linear_light_profiles/\n", + "\n", + "That folder contains `fit.py`, `modeling.py`, and `likelihood_function.py` showing how to\n", + "build lens models with linear profiles and what the likelihood looks like under the hood.\n", + "\n", + "__Operated Light Profiles__\n", + "\n", + "Some emission components \u2014 chiefly the unresolved bright cores of AGN \u2014 are already PSF-\n", + "convolved by the time you receive the image. Standard profiles get PSF-convolved during the\n", + "fit, so applying a PSF a second time double-blurs them.\n", + "\n", + "Operated light profiles tell the fit \"this profile's emission has already had the PSF\n", + "operation applied; do not blur it again\". The fit classes expose an `operated_only` flag\n", + "that controls whether these profiles are included or excluded from a given image computation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "operated_gaussian = al.lp_operated.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=0.3,\n", + " sigma=0.05,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Three operated profiles are available: `al.lp_operated.Gaussian`, `al.lp_operated.Moffat`,\n", + "and `al.lp_operated.Sersic`.\n", + "\n", + "The full workflow (simulating with operated profiles, modeling with them) is documented in:\n", + "\n", + " scripts/imaging/features/operated_light_profile/\n", + "\n", + "That folder contains `simulator.py` and `modeling.py`.\n", + "\n", + "__Basis__\n", + "\n", + "A `Basis` is not a profile in its own right but a *grouping* of profiles that behave as a\n", + "single composite. The classic application is the Multi-Gaussian Expansion (MGE), where a\n", + "galaxy's light is decomposed into a sum of many concentric Gaussians at fixed centres and\n", + "ellipticities but with increasing widths \u2014 together they reproduce arbitrary radial profiles\n", + "the standard parametric forms cannot capture. In strong-lens modelling MGEs are routinely\n", + "used for the lens-galaxy light to soak up complex morphology without inflating the parameter\n", + "space.\n", + "\n", + "The `Basis` constructor takes a `profile_list` of any light or mass profiles. Below we\n", + "build a four-Gaussian MGE with shared centre and ellipticity, sigmas that span an order of\n", + "magnitude, and explicit decreasing `intensity` so the innermost Gaussian dominates the\n", + "core and the wider ones add the outer envelope. Standard `al.lp.Gaussian` profiles are\n", + "used here so the demo image is meaningful; in an actual lens fit you would swap these for\n", + "`al.lp_linear.Gaussian` (see the note after the plot)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "basis = al.lp_basis.Basis(\n", + " profile_list=[\n", + " al.lp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " sigma=0.05,\n", + " ),\n", + " al.lp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=0.5,\n", + " sigma=0.15,\n", + " ),\n", + " al.lp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=0.25,\n", + " sigma=0.4,\n", + " ),\n", + " al.lp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=0.1,\n", + " sigma=1.0,\n", + " ),\n", + " ]\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=basis.image_2d_from(grid=grid),\n", + " title=\"Basis Image (4-Gaussian MGE)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Two things make `Basis` powerful:\n", + "\n", + "- It slots into a `Galaxy` exactly like a `Sersic` would \u2014 once wrapped, the rest of the\n", + " modelling code doesn't have to know it's looking at four Gaussians under the hood.\n", + "- When the constituents are `LightProfileLinear` instances (e.g. `al.lp_linear.Gaussian`)\n", + " rather than the standard `al.lp.Gaussian` used in the demo above, all of their\n", + " `intensity` values are solved together in a **single combined inversion** at each\n", + " likelihood evaluation. This means an MGE built from, say, 30 Gaussians adds only the\n", + " shared geometric parameters to the non-linear search rather than 30 extra intensities.\n", + " The demo above uses standard Gaussians purely so the image is non-zero on a static\n", + " plot \u2014 in a real lens fit you would build the basis from `al.lp_linear.Gaussian` and\n", + " let the inversion solve the intensities.\n", + "\n", + "The full MGE workflow \u2014 choosing how many Gaussians to use, how to space their `sigma`\n", + "values, and how the inversion plays with regularisation \u2014 is documented in:\n", + "\n", + " scripts/imaging/features/multi_gaussian_expansion/\n", + "\n", + "Shapelet decompositions follow the same `Basis` pattern, using `al.lp.ShapeletPolar` /\n", + "`al.lp.ShapeletCartesian` / `al.lp.ShapeletExponential` (and their linear counterparts).\n", + "The full shapelets workflow is documented in:\n", + "\n", + " scripts/imaging/features/shapelets/\n", + "\n", + "__Light Profile in a Model__\n", + "\n", + "So far we have been instantiating profiles with concrete parameter values. When fitting a\n", + "real strong-lens dataset we instead build a *model* of the profile and let the non-linear\n", + "search find the best-fit parameters. This is what `af.Model` is for.\n", + "\n", + "For a lens system the model typically contains two `Galaxy` objects on different planes: a\n", + "foreground lens (with light + mass) and a background source (with light only). In this\n", + "section we focus on the *light* side of that picture and show how a Sersic light profile\n", + "plugs in to both." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_bulge_model = af.Model(al.lp.Sersic)\n", + "source_bulge_model = af.Model(al.lp.Sersic)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `af.Model` wraps the profile class. Every constructor argument that has a numerical\n", + "default now becomes a *prior* \u2014 by default the priors are `UniformPriors` covering a sensible\n", + "range for each parameter (see the autogalaxy config for the configured ranges).\n", + "\n", + "You can override individual priors before fitting:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_bulge_model.sersic_index = af.UniformPrior(lower_limit=0.5, upper_limit=8.0)\n", + "source_bulge_model.effective_radius = af.UniformPrior(\n", + " lower_limit=0.01, upper_limit=10.0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A model profile by itself is not a complete model \u2014 for a strong lens we wrap each profile\n", + "in a `Galaxy` at its own redshift and assemble them in an `af.Collection`. The `Tracer` is\n", + "not part of the *model* spec itself; it is what `AnalysisImaging.log_likelihood_function`\n", + "builds out of the realised instance at fit time." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy_model = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge_model)\n", + "source_galaxy_model = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge_model)\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens_galaxy_model, source=source_galaxy_model)\n", + ")\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing `model.info` prints the full priors-and-defaults summary \u2014 useful before kicking off\n", + "a long fit to confirm the model looks the way you expect.\n", + "\n", + "The model API is the same for **every** light profile in this guide \u2014 swap `al.lp.Sersic`\n", + "for `al.lp.SersicMultipole`, `al.lp_linear.Gaussian`, `al.lp_basis.Basis`, etc., and the\n", + "rest of the snippet is unchanged. Multipole comps and Basis constituent lists are wired\n", + "into the prior machinery automatically.\n", + "\n", + "Full lens-modelling end-to-end examples live in `scripts/imaging/modeling/start_here.py`\n", + "and the topic-specific guides under `scripts/imaging/features/`.\n", + "\n", + "__Model Instance from Light Profile__\n", + "\n", + "A model is a description of *possible* profiles. To get an actual profile back out \u2014 for\n", + "example to plot what the prior medians look like before running a fit \u2014 call\n", + "`instance_from_prior_medians()`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_bulge_instance = lens_bulge_model.instance_from_prior_medians()\n", + "print(type(lens_bulge_instance)) # autogalaxy.profiles.light.standard.sersic.Sersic\n", + "\n", + "image = lens_bulge_instance.image_2d_from(grid=grid)\n", + "aplt.plot_array(array=image, title=\"Lens Bulge Instance from Prior Medians\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The instance returned from `instance_from_prior_medians()` is a real `al.lp.Sersic` \u2014 the\n", + "same class we constructed by hand at the top of the guide \u2014 and supports the full API\n", + "including `image_2d_from`.\n", + "\n", + "The same flow works at the full-model level. We realise an instance of the lens-and-source\n", + "collection and drop the resulting galaxies into a `Tracer`, which is the **PyAutoLens**\n", + "object that handles ray-tracing across redshift planes and combining the lens-galaxy light\n", + "with the lensed source-galaxy light into a single image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_instance = model.instance_from_prior_medians()\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[model_instance.galaxies.lens, model_instance.galaxies.source]\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Tracer Image (Lens Light + Lensed Source Light)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Tracer` takes care of two things at once: it sums the lens-galaxy light (no deflection)\n", + "with the source-galaxy light *after* ray-tracing the grid through the lens-plane deflections,\n", + "producing the combined observed image you would see at the telescope.\n", + "\n", + "When the lens galaxy has a mass profile, the source would appear as a lensed arc. Our\n", + "example builds the model with no mass profile, so the source contributes its un-deflected\n", + "emission \u2014 useful for sanity-checking the *light* side of the model alone. The `tracer.py`\n", + "guide in this directory shows the full lensing flow with mass profiles included.\n", + "\n", + "You can still pull the individual light-profile instances back out of the same object:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_bulge_instance = model_instance.galaxies.lens.bulge\n", + "source_bulge_instance = model_instance.galaxies.source.bulge\n", + "print(type(lens_bulge_instance)) # autogalaxy.profiles.light.standard.sersic.Sersic\n", + "print(type(source_bulge_instance)) # autogalaxy.profiles.light.standard.sersic.Sersic" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "After a fit completes, `result.max_log_likelihood_tracer` returns the same shape of object,\n", + "with the prior medians replaced by the fitted parameter values. See\n", + "`scripts/guides/results/start_here.py` for the full results-introspection guide.\n", + "\n", + "__Multipole Light Profiles__\n", + "\n", + "`SersicMultipole` and `GaussianMultipole` are recent additions that bolt m=3 and m=4 Fourier\n", + "angular perturbations onto the eccentric radius of a base profile. The perturbed radius is\n", + "\n", + " r' = r * (1 + c3 cos(3 theta) + s3 sin(3 theta)\n", + " + c4 cos(4 theta) + s4 sin(4 theta))\n", + "\n", + "where `theta` is the polar angle in the profile's elliptical reference frame, and the\n", + "`multipole_3_comps = (c3, s3)` and `multipole_4_comps = (c4, s4)` parameters control the\n", + "amplitude of each perturbation.\n", + "\n", + "When both `multipole_*_comps` are `(0.0, 0.0)` (the defaults), the profile reduces exactly\n", + "to the base profile. This is by design \u2014 you can swap a `Sersic` for a `SersicMultipole`\n", + "in any model without changing its predictions, and the multipole comps simply add four\n", + "extra free parameters that can capture boxy / discy / lopsided morphologies in lens-galaxy\n", + "light or in resolved source morphology. Plugging one into the `af.Model` / `af.Collection`\n", + "/ `Galaxy` / `Tracer` pattern shown above works exactly as it did for the plain `Sersic` \u2014\n", + "the multipole comps are picked up as priors automatically.\n", + "\n", + "Build a `SersicMultipole` with non-zero multipole components, alongside the unperturbed\n", + "Sersic that produced our reference image earlier in the guide:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sersic_multipole = al.lp.SersicMultipole(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " multipole_3_comps=(0.05, 0.00),\n", + " multipole_4_comps=(0.00, 0.04),\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=sersic_multipole.image_2d_from(grid=grid),\n", + " title=\"SersicMultipole Image (m=3 + m=4 perturbation)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For comparison, here is the unperturbed Sersic image \u2014 the two should look almost\n", + "identical with the perturbation showing as a subtle azimuthal modulation:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=sersic.image_2d_from(grid=grid),\n", + " title=\"Sersic Image (no multipole perturbation)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `GaussianMultipole` profile applies the same perturbation to a Gaussian base \u2014 useful\n", + "when you want a multipole component inside a Multi-Gaussian Expansion (the `Basis` section\n", + "above shows how to bundle Gaussians together):" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "gaussian_multipole = al.lp.GaussianMultipole(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " sigma=0.4,\n", + " multipole_3_comps=(0.05, 0.00),\n", + " multipole_4_comps=(0.00, 0.04),\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=gaussian_multipole.image_2d_from(grid=grid),\n", + " title=\"GaussianMultipole Image\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Two practical notes on the multipole variants:\n", + "\n", + "- There is **no spherical (`*Sph`) variant** of either multipole. The perturbation is an\n", + " angular distortion measured in the elliptical reference frame, so it only makes sense for\n", + " an elliptical profile (a spherical profile has no preferred angle).\n", + "- Both multipoles exist as **linear variants** too: `al.lp_linear.SersicMultipole` and\n", + " `al.lp_linear.GaussianMultipole`. In the linear form the multipole comps remain non-\n", + " linear parameters but the overall intensity is solved by inversion, just like for the\n", + " ordinary linear profiles.\n", + "\n", + "__Remaining Profiles Walkthrough__\n", + "\n", + "We have shown the full `image_2d_from` \u2192 `af.Model` \u2192 `instance` \u2192 `Tracer` flow for the\n", + "`Sersic` profile. Every remaining standard profile uses the **same API** \u2014 the only thing\n", + "that changes is which parameters appear in the constructor.\n", + "\n", + "The compact tour below builds each remaining profile with sensible parameter values and\n", + "plots its image, so you can see what each looks like. When you want to use any of these in\n", + "a lens model, repeat the `af.Model(...)` / `af.Collection(...)` / `Tracer(...)` pattern\n", + "from the previous section." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(\n", + " array=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " radius_break=0.05,\n", + " gamma=0.2,\n", + " alpha=3.0,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"SersicCore Image\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.lp.Exponential(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", + " intensity=0.5,\n", + " effective_radius=1.6,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"Exponential Image\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.lp.ExponentialCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", + " intensity=0.5,\n", + " effective_radius=1.6,\n", + " radius_break=0.05,\n", + " gamma=0.2,\n", + " alpha=3.0,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"ExponentialCore Image\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.lp.DevVaucouleurs(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"DevVaucouleurs Image\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.lp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " sigma=0.4,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"Gaussian Image\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.lp.Moffat(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " alpha=0.4,\n", + " beta=2.5,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"Moffat Image\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.lp.Chameleon(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " core_radius_0=0.05,\n", + " core_radius_1=0.3,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"Chameleon Image\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.lp.ElsonFreeFall(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " eta=2.0,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"ElsonFreeFall Image\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The spherical variants (`SersicSph`, `GaussianSph`, etc.) are constructed identically with\n", + "the `ell_comps` argument removed \u2014 they look like a rotationally symmetric version of the\n", + "corresponding elliptical plot.\n", + "\n", + "The shapelet profiles are normally used inside a `Basis` rather than individually, but for\n", + "completeness here is the lowest-order Cartesian shapelet on its own:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=al.lp.ShapeletCartesian(\n", + " n_y=0,\n", + " n_x=0,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=1.0,\n", + " beta=0.2,\n", + " ).image_2d_from(grid=grid),\n", + " title=\"ShapeletCartesian (n_y=0, n_x=0) Image\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "And that completes the tour. If you arrived here from the API reference and now want to use\n", + "any of these profiles in an actual lens fit, the next step is\n", + "`scripts/imaging/modeling/start_here.py`, which sets up an `AnalysisImaging` and runs a\n", + "non-linear search end-to-end on a strong-lens dataset. The `scripts/imaging/features/`\n", + "subpackages handle the family-specific workflows referenced throughout this guide:\n", + "\n", + "- `linear_light_profiles/` \u2014 using `al.lp_linear.*` in a fit.\n", + "- `operated_light_profile/` \u2014 using `al.lp_operated.*` in a fit.\n", + "- `multi_gaussian_expansion/` \u2014 building and fitting an MGE-style `Basis`.\n", + "- `shapelets/` \u2014 building and fitting a shapelet-style `Basis`.\n", + "\n", + "For the mass-profile counterpart of this guide, see `scripts/guides/profiles/mass.py`\n", + "(when available). For a full walk-through of how light and mass profiles are combined in\n", + "the `Tracer` to produce lensed images, see `scripts/guides/tracer.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/profiles/light_and_mass_profiles.ipynb b/notebooks/guides/profiles/light_and_mass_profiles.ipynb index 31721ff31..b9f564ecb 100644 --- a/notebooks/guides/profiles/light_and_mass_profiles.ipynb +++ b/notebooks/guides/profiles/light_and_mass_profiles.ipynb @@ -1,965 +1,1002 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Light-and-Mass Profiles\n", - "=======================\n", - "\n", - "This guide is the third and final entry in the `scripts/guides/profiles/` trilogy. It\n", - "covers the *stellar*, *dark-matter*, and *combined light-and-mass* profiles \u2014 the ones used\n", - "to decompose a lens galaxy into its constituent matter components rather than to model the\n", - "total mass with a single parametric profile.\n", - "\n", - "It pairs with:\n", - "\n", - "- `scripts/guides/profiles/light.py` \u2014 pure light profiles in `al.lp.*` and friends.\n", - "- `scripts/guides/profiles/mass.py` \u2014 parametric lensing mass profiles in `al.mp.*`\n", - " (Total, Mass Sheets, Multipoles, Point Mass).\n", - "\n", - "Where `mass.py` shows the *total* mass distribution of a lens galaxy parameterised by a\n", - "single power-law / isothermal / dPIE profile, this guide shows the *decomposed* picture:\n", - "stellar mass (Sersic, Chameleon, ...) plus dark-matter halo (NFW family). The `al.lmp.*`\n", - "and `al.lmp_linear.*` namespaces tie a galaxy's *visible* light to its *stellar* mass via\n", - "a single shared `mass_to_light_ratio` parameter, so one model object produces both an\n", - "`image_2d_from` (light) and a `convergence_2d_from` (mass) consistently.\n", - "\n", - "__Contents__\n", - "\n", - "- **Overview & Docs URL:** The three-guide layout and where the canonical API reference lives.\n", - "- **All Profiles (Survey):** A high-level catalogue of every stellar / dark / lmp /\n", - " lmp_linear class.\n", - "- **Stellar Mass Detailed Example:** `al.mp.Sersic` \u2014 pure stellar mass. Convergence and\n", - " the lensed source image via `Tracer`.\n", - "- **Dark Mass Detailed Example:** `al.mp.NFW` \u2014 the standard cuspy dark-matter halo.\n", - " Convergence in linear and log10 scales.\n", - "- **NFW Variants:** `gNFW`, `cNFW`, `NFWTruncated`, plus the MCR (mass-concentration), Virial-\n", - " mass, and Scatter reparameterisations. One-liner construction and when to reach for each.\n", - "- **Combined Light + Mass Profiles (`al.lmp`):** The headline feature \u2014 one Sersic-shaped\n", - " object emits BOTH `image_2d_from` (light) AND `convergence_2d_from` (mass) via a shared\n", - " `mass_to_light_ratio`.\n", - "- **Linear Combined Light + Mass (`al.lmp_linear`):** The inversion-aware variant of the\n", - " lmp family; the intensity-via-inversion semantics from `al.lp_linear` carry over.\n", - "- **Composing a Decomposed Bulge+Halo Model:** `af.Model` with a stellar Sersic mass + NFW\n", - " halo on the lens galaxy. The standard decomposed-lens recipe.\n", - "- **Model Instance from Decomposed Model:** `instance_from_prior_medians()` -> `Tracer` ->\n", - " lensed image.\n", - "- **Remaining Profiles Walkthrough:** Compact `convergence_2d_from` / `image_2d_from` block\n", - " per remaining stellar / dark / lmp / lmp_linear profile.\n", - "- **Back-References:** Pointer to `light.py` and `mass.py`.\n", - "\n", - "__Units__\n", - "\n", - "Spatial coordinates are in arc-seconds; convergence is dimensionless; intensity is in\n", - "electrons per second; `mass_to_light_ratio` is dimensionless (in the units of the workspace\n", - "config). The `guides/units_and_cosmology.ipynb` guide covers physical-unit conversion.\n", - "\n", - "__Data Structures__\n", - "\n", - "`convergence_2d_from` and `image_2d_from` both return `Array2D`. Plotting through\n", - "`aplt.plot_array` is direct. See `mass.py` for the deflection-magnitude pattern when you\n", - "need to inspect the vector quantities \u2014 this guide focuses on convergence and image.\n", - "\n", - "__Docs URL__\n", - "\n", - "The published API reference for these classes lives at:\n", - "\n", - " https://pyautolens.readthedocs.io/en/latest/api/mass.html\n", - "\n", - "The autosummary on that page is the authoritative list of every public mass-profile class\n", - "(stellar and dark are in the same page under their respective sections). Light-and-mass\n", - "combined profiles in `al.lmp.*` and `al.lmp_linear.*` are not yet documented on the\n", - "reference page \u2014 refer to the source under\n", - "`autogalaxy/profiles/light_and_mass_profiles.py` for the full list." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "To evaluate any quantity on a profile we need a 2D Cartesian grid of (y,x) coordinates.\n", - "We build a 100x100 grid here at a 0.05\" pixel scale \u2014 used by every section below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.05,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__All Profiles (Survey)__\n", - "\n", - "**PyAutoLens** organises matter-decomposition profiles into four namespaces:\n", - "\n", - "- `al.mp.*` *stellar* mass \u2014 Sersic-like profiles parameterised by light shape plus a\n", - " `mass_to_light_ratio` that converts intensity into surface mass density.\n", - "- `al.mp.*` *dark-matter* mass \u2014 NFW family halos parameterised by `kappa_s` /\n", - " `scale_radius`, or alternatively by halo mass + concentration via the MCR and Virial\n", - " variants.\n", - "- `al.lmp.*` *combined light-and-mass* profiles \u2014 one object that emits both\n", - " `image_2d_from` (light) and `convergence_2d_from` (mass) via a shared\n", - " `mass_to_light_ratio`.\n", - "- `al.lmp_linear.*` \u2014 the inversion-aware version of `al.lmp.*`. Same classes, intensity\n", - " solved by inversion at fit time.\n", - "\n", - "Below we construct one profile from each family with sensible defaults. No quantity is\n", - "evaluated yet \u2014 this is the catalogue." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Stellar mass \u2014 pure mass parameterised by light shape * mass_to_light_ratio\n", - "stellar_sersic = al.mp.Sersic()\n", - "stellar_chameleon = al.mp.Chameleon()\n", - "stellar_gaussian = al.mp.Gaussian()\n", - "\n", - "# Dark mass \u2014 NFW family\n", - "dark_nfw = al.mp.NFW()\n", - "dark_gnfw = al.mp.gNFW()\n", - "dark_cnfw = al.mp.cNFW()\n", - "dark_nfw_truncated = al.mp.NFWTruncatedSph()\n", - "dark_nfw_mcr = al.mp.NFWMCRLudlow()\n", - "dark_nfw_virial = al.mp.NFWVirialMassConcSph()\n", - "\n", - "# Combined light-and-mass \u2014 one object, two outputs\n", - "lmp_sersic = al.lmp.Sersic()\n", - "lmp_chameleon = al.lmp.Chameleon()\n", - "lmp_gaussian = al.lmp.Gaussian()\n", - "lmp_gaussian_gradient = al.lmp.GaussianGradient()\n", - "\n", - "# Linear combined light-and-mass \u2014 intensity solved by inversion at fit time\n", - "lmp_linear_sersic = al.lmp_linear.Sersic()\n", - "lmp_linear_gaussian = al.lmp_linear.Gaussian()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two things worth knowing about this list before we move on:\n", - "\n", - "1. The *stellar* and *combined* profiles use the **same parametric forms** as the light\n", - " profiles (Sersic, Chameleon, Gaussian, ...). Their `convergence_2d_from` is the\n", - " profile's light shape multiplied by `mass_to_light_ratio` \u2014 this is a fundamentally\n", - " different parameterisation from the dark-matter halos, which are governed by halo-mass\n", - " physics (concentration-mass relations, virial parameters).\n", - "2. The dark-matter NFW family has *many* variants that are physically equivalent but\n", - " parameterised differently \u2014 they trade `kappa_s` / `scale_radius` for `mass_at_200` +\n", - " concentration (via MCR), or for `virial_mass` + concentration (Virial). Picking the\n", - " right parameterisation matters for prior choice and degeneracy management; the NFW\n", - " Variants section below is a one-page tour of the menagerie.\n", - "\n", - "__Stellar Mass Detailed Example__\n", - "\n", - "`al.mp.Sersic` is the stellar mass profile. It is parameterised by exactly the same\n", - "light-shape arguments as `al.lp.Sersic` (`intensity`, `effective_radius`, `sersic_index`)\n", - "plus a `mass_to_light_ratio` that converts the surface brightness profile into a surface\n", - "mass density. Crucially it does NOT emit light \u2014 `al.mp.Sersic` only carries mass.\n", - "\n", - "Build a stellar Sersic mass profile and plot its convergence:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "stellar_sersic = al.mp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " mass_to_light_ratio=2.0,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=stellar_sersic.convergence_2d_from(grid=grid),\n", - " title=\"Stellar Sersic Mass Convergence\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The convergence map looks exactly like a Sersic image because, geometrically, it is one\n", - "\u2014 the same Sersic profile, normalised by `mass_to_light_ratio`.\n", - "\n", - "To see what this stellar mass does to a background source, drop it into a `Tracer` with a\n", - "source-plane galaxy carrying a small Sersic light:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(redshift=0.5, mass=stellar_sersic)\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.05, 0.05),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=0.1,\n", - " sersic_index=1.5,\n", - " ),\n", - ")\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Tracer Image (Stellar Sersic Lens + Sersic Source)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that the lens galaxy carries *only* mass \u2014 no light is emitted from the lens plane in\n", - "the tracer image. To add lens-galaxy light you would attach a separate `al.lp.*` light\n", - "profile to the same `Galaxy`, OR use a `al.lmp.*` profile (covered later in this guide)\n", - "which carries both light and mass in a single object.\n", - "\n", - "__Dark Mass Detailed Example__\n", - "\n", - "`al.mp.NFW` is the canonical dark-matter halo \u2014 the elliptical Navarro-Frenk-White profile\n", - "parameterised by:\n", - "\n", - "- `kappa_s` \u2014 the dimensionless characteristic convergence at the scale radius.\n", - "- `scale_radius` \u2014 the radius (arc-seconds) at which the density slope transitions from\n", - " ~r^-1 (inner cusp) to ~r^-3 (outer fall-off).\n", - "\n", - "This is the \"natural\" lensing parameterisation; physically equivalent reparameterisations\n", - "in terms of halo mass + concentration are covered in the NFW Variants section below.\n", - "\n", - "Build an `al.mp.NFW` and plot its convergence on linear and log10 scales:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dark_nfw = al.mp.NFW(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", - " kappa_s=0.1,\n", - " scale_radius=2.0,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=dark_nfw.convergence_2d_from(grid=grid),\n", - " title=\"NFW Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=dark_nfw.convergence_2d_from(grid=grid),\n", - " title=\"NFW Convergence (log10 \u2014 central cusp visible)\",\n", - " use_log10=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The log10 view shows the characteristic central cusp that distinguishes NFW from the cored\n", - "profiles (`Isothermal`, `dPIEMass`). In a real lens model the cusp is what makes dark-\n", - "matter mass distinguishable from a steep stellar profile.\n", - "\n", - "__NFW Variants__\n", - "\n", - "The NFW family has several reparameterisations. All produce identical convergence maps\n", - "for matched halo parameters \u2014 they differ only in which parameters the *user* specifies\n", - "and which are derived. Picking the right one is mostly about prior choice and parameter\n", - "degeneracies.\n", - "\n", - "- **Generalised NFW (`al.mp.gNFW`)** \u2014 adds a free `inner_slope` parameter. Reduces to\n", - " `NFW` when `inner_slope=1.0`. Useful when you want to test whether the inner cusp is\n", - " steeper or shallower than the NFW prediction.\n", - "- **Cored NFW (`al.mp.cNFW`)** \u2014 adds a `core_radius` parameter that flattens the central\n", - " cusp. Useful for self-interacting dark matter scenarios.\n", - "- **Truncated NFW (`al.mp.NFWTruncatedSph`)** \u2014 adds a `truncation_radius` beyond which\n", - " the density is forced to zero. Useful for satellite halos with stripped outer envelopes.\n", - "- **MCR variants (`NFWMCR*` / `cNFWMCR*` / `gNFWMCR*`)** \u2014 replace `kappa_s` /\n", - " `scale_radius` with `mass_at_200` and use a fixed concentration-mass relation (Duffy\n", - " 2008 or Ludlow 2016) to derive the concentration. Removes one free parameter at the\n", - " cost of assuming the concentration follows the relation.\n", - "- **Scatter variants (`NFWMCRScatter*`)** \u2014 the Ludlow MCR plus a `scatter_sigma` knob\n", - " that adds Gaussian scatter on the log-concentration. Useful when you want to *fit* the\n", - " concentration but anchor it to the Ludlow relation with finite tolerance.\n", - "- **Virial-mass variants (`*VirialMassConc*`)** \u2014 parameterise the halo by its virial mass\n", - " and concentration directly, instead of the `mass_at_200` convention. Choose whichever\n", - " matches the units used in your cosmology pipeline.\n", - "\n", - "One-liner construction of the major variants:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "nfw_gnfw = al.mp.gNFW(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", - " kappa_s=0.1,\n", - " inner_slope=1.2,\n", - " scale_radius=2.0,\n", - ")\n", - "\n", - "nfw_cnfw = al.mp.cNFW(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", - " kappa_s=0.1,\n", - " scale_radius=2.0,\n", - " core_radius=0.2,\n", - ")\n", - "\n", - "nfw_truncated = al.mp.NFWTruncatedSph(\n", - " centre=(0.0, 0.0),\n", - " kappa_s=0.1,\n", - " scale_radius=2.0,\n", - " truncation_radius=4.0,\n", - ")\n", - "\n", - "nfw_mcr = al.mp.NFWMCRLudlowSph(\n", - " centre=(0.0, 0.0),\n", - " mass_at_200=1.0e12,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - ")\n", - "\n", - "nfw_mcr_scatter = al.mp.NFWMCRScatterLudlowSph(\n", - " centre=(0.0, 0.0),\n", - " mass_at_200=1.0e12,\n", - " scatter_sigma=0.1,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - ")\n", - "\n", - "nfw_virial = al.mp.NFWVirialMassConcSph(\n", - " centre=(0.0, 0.0),\n", - " virial_mass=1.0e12,\n", - " concentration=10.0,\n", - " virial_overdens=200.0,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For comparison, here is the cored variant's convergence \u2014 note the flat central plateau:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=nfw_cnfw.convergence_2d_from(grid=grid),\n", - " title=\"cNFW Convergence (cored central region)\",\n", - " use_log10=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And the MCR variant produced from a halo mass rather than `kappa_s` / `scale_radius`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=nfw_mcr.convergence_2d_from(grid=grid),\n", - " title=\"NFWMCRLudlowSph Convergence (mass=1e12, concentration via Ludlow)\",\n", - " use_log10=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Combined Light + Mass Profiles (`al.lmp`)__\n", - "\n", - "The `al.lmp.*` namespace is the headline feature of this guide. Each `al.lmp` class\n", - "represents a galaxy component that emits BOTH light AND mass \u2014 a single object with one\n", - "set of parameters governing both `image_2d_from` and `convergence_2d_from` via a shared\n", - "`mass_to_light_ratio`.\n", - "\n", - "This is the standard way to model a lens galaxy's bulge: one `al.lmp.Sersic` carries the\n", - "visible Sersic light and contributes the stellar mass component to the deflection, tied\n", - "together by `mass_to_light_ratio`. Pair it with a separate `al.mp.NFW` (or one of the\n", - "dark variants above) for the dark-matter halo and you have the canonical bulge+halo\n", - "decomposition.\n", - "\n", - "Build one `al.lmp.Sersic` and plot both its image AND its convergence:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lmp_sersic = al.lmp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " mass_to_light_ratio=2.0,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=lmp_sersic.image_2d_from(grid=grid),\n", - " title=\"lmp.Sersic Image (light side)\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=lmp_sersic.convergence_2d_from(grid=grid),\n", - " title=\"lmp.Sersic Convergence (mass side, mass_to_light_ratio=2.0)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The two maps are *the same Sersic profile*, with the mass map scaled by\n", - "`mass_to_light_ratio`. This is the geometric guarantee of the `al.lmp.*` family \u2014 the\n", - "stellar light and stellar mass cannot diverge from each other except through this one\n", - "parameter.\n", - "\n", - "`al.lmp.SersicGradient` and `al.lmp.GaussianGradient` extend this with a *radial gradient*\n", - "in the mass-to-light ratio (one extra parameter), useful when the inner regions of a\n", - "galaxy have a systematically different M/L from the outskirts.\n", - "\n", - "__Linear Combined Light + Mass (`al.lmp_linear`)__\n", - "\n", - "`al.lmp_linear.*` mirrors `al.lmp.*` class-for-class. The only difference is that the\n", - "`intensity` parameter is solved analytically via the linear inversion at each likelihood\n", - "evaluation (the same mechanism that `al.lp_linear.*` uses for pure light profiles). This\n", - "means a lens model with a `al.lmp_linear.Sersic` bulge has one fewer free non-linear\n", - "parameter than the corresponding `al.lmp.Sersic` model \u2014 the bulge intensity is no longer\n", - "sampled by the search.\n", - "\n", - "The construction API is identical to `al.lmp.*`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lmp_linear_sersic = al.lmp_linear.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " mass_to_light_ratio=2.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The full workflow for linear profiles in lens models \u2014 including how the inversion\n", - "combines light and mass contributions \u2014 is documented in:\n", - "\n", - " scripts/imaging/features/linear_light_profiles/\n", - "\n", - "That folder is the canonical reference; this guide stops at the construction API.\n", - "\n", - "__Composing a Decomposed Bulge+Halo Model__\n", - "\n", - "The standard decomposed lens-mass model attaches:\n", - "\n", - "- one `al.mp.Sersic` (or `al.lmp.Sersic`) as the lens-galaxy stellar component, AND\n", - "- one `al.mp.NFW` (or another dark variant) as the lens-galaxy dark-matter halo,\n", - "\n", - "so the total convergence is the sum of stellar + dark contributions. Compose this via\n", - "`af.Model` exactly as you would for a single-profile lens \u2014 `Galaxy` accepts arbitrary\n", - "kwargs for its mass components." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_bulge_mass_model = af.Model(al.mp.Sersic)\n", - "lens_dark_model = af.Model(al.mp.NFW)\n", - "source_bulge_model = af.Model(al.lp.Sersic)\n", - "\n", - "lens_galaxy_model = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=lens_bulge_mass_model,\n", - " dark=lens_dark_model,\n", - ")\n", - "source_galaxy_model = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge_model)\n", - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens_galaxy_model, source=source_galaxy_model)\n", - ")\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing `model.info` shows the priors-and-defaults summary. Notice the lens galaxy has\n", - "*two* mass components (`bulge` and `dark`) whose convergences will be summed during the\n", - "fit. This is how a decomposed lens model differs from the single-profile `Isothermal`\n", - "lens shown in `mass.py` \u2014 the model is more flexible but also higher-dimensional.\n", - "\n", - "Swapping `al.mp.Sersic` for `al.lmp.Sersic` on the lens bulge would additionally model the\n", - "lens-galaxy *light* through the same object \u2014 useful when the lens galaxy is visible in\n", - "the data and you want to tie its light to its stellar mass.\n", - "\n", - "Full end-to-end lens fits with this model live under `scripts/imaging/features/` and the\n", - "SLaM pipeline guides; this section just shows the model spec.\n", - "\n", - "__Model Instance from Decomposed Model__\n", - "\n", - "To realise an instance from the model's prior medians and visualise it, call\n", - "`instance_from_prior_medians()` on the collection and feed the result into a `Tracer`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_instance = model.instance_from_prior_medians()\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[model_instance.galaxies.lens, model_instance.galaxies.source]\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Tracer Image from Decomposed Model (Stellar + NFW)\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.convergence_2d_from(grid=grid),\n", - " title=\"Tracer Total Convergence (Stellar + NFW summed)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The total convergence map shows the summed stellar + dark contributions \u2014 exactly the\n", - "lensing mass distribution that the search would optimise.\n", - "\n", - "After a fit completes, `result.max_log_likelihood_tracer` returns the same shape of object\n", - "with the prior medians replaced by the fitted parameters. See\n", - "`scripts/guides/results/start_here.py` for the full results-introspection guide.\n", - "\n", - "__Remaining Profiles Walkthrough__\n", - "\n", - "We have shown the full flow for `al.mp.Sersic` (stellar), `al.mp.NFW` (dark), and\n", - "`al.lmp.Sersic` (combined). Every remaining stellar / dark / lmp / lmp_linear profile uses\n", - "the same API \u2014 only the constructor parameters change.\n", - "\n", - "The compact tour below builds each remaining profile with sensible defaults. Stellar and\n", - "dark profiles show `convergence_2d_from`; `lmp` profiles show both `image_2d_from` and\n", - "`convergence_2d_from` so the dual-output property is visible per class.\n", - "\n", - "Stellar mass:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=al.mp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=4.0,\n", - " radius_break=0.05,\n", - " gamma=0.25,\n", - " alpha=3.0,\n", - " mass_to_light_ratio=2.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"Stellar SersicCore Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.SersicGradient(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " mass_to_light_ratio=2.0,\n", - " mass_to_light_gradient=0.3,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"Stellar SersicGradient Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.DevVaucouleurs(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " mass_to_light_ratio=2.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"Stellar DevVaucouleurs Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.Exponential(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", - " intensity=0.5,\n", - " effective_radius=1.6,\n", - " mass_to_light_ratio=2.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"Stellar Exponential Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " sigma=0.4,\n", - " mass_to_light_ratio=2.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"Stellar Gaussian Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.GaussianGradient(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " sigma=0.4,\n", - " mass_to_light_ratio_base=2.0,\n", - " mass_to_light_gradient=0.2,\n", - " mass_to_light_radius=0.5,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"Stellar GaussianGradient Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.Chameleon(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " core_radius_0=0.05,\n", - " core_radius_1=0.3,\n", - " mass_to_light_ratio=2.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"Stellar Chameleon Convergence\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Dark mass \u2014 the variant menagerie, one block each:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=al.mp.gNFW(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", - " kappa_s=0.1,\n", - " inner_slope=1.2,\n", - " scale_radius=2.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"gNFW Convergence (inner_slope=1.2)\",\n", - " use_log10=True,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.NFWTruncatedSph(\n", - " centre=(0.0, 0.0),\n", - " kappa_s=0.1,\n", - " scale_radius=2.0,\n", - " truncation_radius=4.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"NFWTruncatedSph Convergence\",\n", - " use_log10=True,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.NFWMCRDuffySph(\n", - " centre=(0.0, 0.0),\n", - " mass_at_200=1.0e12,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"NFWMCRDuffySph Convergence (Duffy 2008 MCR)\",\n", - " use_log10=True,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.NFWMCRScatterLudlowSph(\n", - " centre=(0.0, 0.0),\n", - " mass_at_200=1.0e12,\n", - " scatter_sigma=0.1,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"NFWMCRScatterLudlowSph Convergence (Ludlow + sigma=0.1)\",\n", - " use_log10=True,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.NFWVirialMassConcSph(\n", - " centre=(0.0, 0.0),\n", - " virial_mass=1.0e12,\n", - " concentration=10.0,\n", - " virial_overdens=200.0,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"NFWVirialMassConcSph Convergence\",\n", - " use_log10=True,\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.cNFWMCRLudlowSph(\n", - " centre=(0.0, 0.0),\n", - " mass_at_200=1.0e12,\n", - " f_c=0.05,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"cNFWMCRLudlowSph Convergence (cored + Ludlow MCR)\",\n", - " use_log10=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Combined light-and-mass \u2014 every remaining `al.lmp` profile, with both image and convergence\n", - "shown side-by-side via stacked `plot_array` calls:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lmp_chameleon = al.lmp.Chameleon(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " core_radius_0=0.05,\n", - " core_radius_1=0.3,\n", - " mass_to_light_ratio=2.0,\n", - ")\n", - "aplt.plot_array(\n", - " array=lmp_chameleon.image_2d_from(grid=grid), title=\"lmp.Chameleon Image\"\n", - ")\n", - "aplt.plot_array(\n", - " array=lmp_chameleon.convergence_2d_from(grid=grid),\n", - " title=\"lmp.Chameleon Convergence\",\n", - ")\n", - "\n", - "lmp_devvaucouleurs = al.lmp.DevVaucouleurs(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " mass_to_light_ratio=2.0,\n", - ")\n", - "aplt.plot_array(\n", - " array=lmp_devvaucouleurs.image_2d_from(grid=grid),\n", - " title=\"lmp.DevVaucouleurs Image\",\n", - ")\n", - "aplt.plot_array(\n", - " array=lmp_devvaucouleurs.convergence_2d_from(grid=grid),\n", - " title=\"lmp.DevVaucouleurs Convergence\",\n", - ")\n", - "\n", - "lmp_exponential = al.lmp.Exponential(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", - " intensity=0.5,\n", - " effective_radius=1.6,\n", - " mass_to_light_ratio=2.0,\n", - ")\n", - "aplt.plot_array(\n", - " array=lmp_exponential.image_2d_from(grid=grid),\n", - " title=\"lmp.Exponential Image\",\n", - ")\n", - "aplt.plot_array(\n", - " array=lmp_exponential.convergence_2d_from(grid=grid),\n", - " title=\"lmp.Exponential Convergence\",\n", - ")\n", - "\n", - "lmp_gaussian = al.lmp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " sigma=0.4,\n", - " mass_to_light_ratio=2.0,\n", - ")\n", - "aplt.plot_array(array=lmp_gaussian.image_2d_from(grid=grid), title=\"lmp.Gaussian Image\")\n", - "aplt.plot_array(\n", - " array=lmp_gaussian.convergence_2d_from(grid=grid),\n", - " title=\"lmp.Gaussian Convergence\",\n", - ")\n", - "\n", - "lmp_sersic_gradient = al.lmp.SersicGradient(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " mass_to_light_ratio=2.0,\n", - " mass_to_light_gradient=0.3,\n", - ")\n", - "aplt.plot_array(\n", - " array=lmp_sersic_gradient.image_2d_from(grid=grid),\n", - " title=\"lmp.SersicGradient Image\",\n", - ")\n", - "aplt.plot_array(\n", - " array=lmp_sersic_gradient.convergence_2d_from(grid=grid),\n", - " title=\"lmp.SersicGradient Convergence\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The spherical variants (`SersicSph`, `DevVaucouleursSph`, `ExponentialSph`, `ChameleonSph`,\n", - "etc., across all three namespaces) are constructed identically with the `ell_comps`\n", - "argument removed. Each looks like a rotationally symmetric version of its elliptical\n", - "counterpart.\n", - "\n", - "`al.lmp_linear.*` shares the construction API with `al.lmp.*` \u2014 instantiate any\n", - "`al.lmp_linear.X` exactly as you would `al.lmp.X`. Their distinct behaviour only shows up\n", - "inside a model fit, when the inversion solves for `intensity` rather than treating it as\n", - "a free non-linear parameter.\n", - "\n", - "__Back-References__\n", - "\n", - "This guide completes the three-guide tour of the `scripts/guides/profiles/` folder:\n", - "\n", - "- `scripts/guides/profiles/light.py` \u2014 pure light profiles, the `al.lp.*` /\n", - " `al.lp_linear.*` / `al.lp_operated.*` / `al.lp_basis.*` namespaces.\n", - "- `scripts/guides/profiles/mass.py` \u2014 parametric lensing mass profiles, the *Total* /\n", - " *Mass Sheets* / *Multipoles* / *Point Mass* families of `al.mp.*`.\n", - "- `scripts/guides/profiles/light_and_mass_profiles.py` \u2014 this guide; stellar / dark / lmp /\n", - " lmp_linear matter decomposition.\n", - "\n", - "For the next step into actual lens modelling, see `scripts/imaging/modeling/start_here.py`\n", - "(strong-lens fit end-to-end) and the topic-specific feature packages under\n", - "`scripts/imaging/features/`. The SLaM pipeline guide (`scripts/guides/modeling/slam_start_here.py`)\n", - "shows the production decomposed-lens workflow built on these profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Light-and-Mass Profiles\n", + "=======================\n", + "\n", + "This guide is the third and final entry in the `scripts/guides/profiles/` trilogy. It\n", + "covers the *stellar*, *dark-matter*, and *combined light-and-mass* profiles \u2014 the ones used\n", + "to decompose a lens galaxy into its constituent matter components rather than to model the\n", + "total mass with a single parametric profile.\n", + "\n", + "It pairs with:\n", + "\n", + "- `scripts/guides/profiles/light.py` \u2014 pure light profiles in `al.lp.*` and friends.\n", + "- `scripts/guides/profiles/mass.py` \u2014 parametric lensing mass profiles in `al.mp.*`\n", + " (Total, Mass Sheets, Multipoles, Point Mass).\n", + "\n", + "Where `mass.py` shows the *total* mass distribution of a lens galaxy parameterised by a\n", + "single power-law / isothermal / dPIE profile, this guide shows the *decomposed* picture:\n", + "stellar mass (Sersic, Chameleon, ...) plus dark-matter halo (NFW family). The `al.lmp.*`\n", + "and `al.lmp_linear.*` namespaces tie a galaxy's *visible* light to its *stellar* mass via\n", + "a single shared `mass_to_light_ratio` parameter, so one model object produces both an\n", + "`image_2d_from` (light) and a `convergence_2d_from` (mass) consistently.\n", + "\n", + "__Contents__\n", + "\n", + "- **Overview & Docs URL:** The three-guide layout and where the canonical API reference lives.\n", + "- **All Profiles (Survey):** A high-level catalogue of every stellar / dark / lmp /\n", + " lmp_linear class.\n", + "- **Stellar Mass Detailed Example:** `al.mp.Sersic` \u2014 pure stellar mass. Convergence and\n", + " the lensed source image via `Tracer`.\n", + "- **Dark Mass Detailed Example:** `al.mp.NFW` \u2014 the standard cuspy dark-matter halo.\n", + " Convergence in linear and log10 scales.\n", + "- **NFW Variants:** `gNFW`, `cNFW`, `NFWTruncated`, plus the MCR (mass-concentration), Virial-\n", + " mass, and Scatter reparameterisations. One-liner construction and when to reach for each.\n", + "- **Combined Light + Mass Profiles (`al.lmp`):** The headline feature \u2014 one Sersic-shaped\n", + " object emits BOTH `image_2d_from` (light) AND `convergence_2d_from` (mass) via a shared\n", + " `mass_to_light_ratio`.\n", + "- **Linear Combined Light + Mass (`al.lmp_linear`):** The inversion-aware variant of the\n", + " lmp family; the intensity-via-inversion semantics from `al.lp_linear` carry over.\n", + "- **Composing a Decomposed Bulge+Halo Model:** `af.Model` with a stellar Sersic mass + NFW\n", + " halo on the lens galaxy. The standard decomposed-lens recipe.\n", + "- **Model Instance from Decomposed Model:** `instance_from_prior_medians()` -> `Tracer` ->\n", + " lensed image.\n", + "- **Remaining Profiles Walkthrough:** Compact `convergence_2d_from` / `image_2d_from` block\n", + " per remaining stellar / dark / lmp / lmp_linear profile.\n", + "- **Back-References:** Pointer to `light.py` and `mass.py`.\n", + "\n", + "__Units__\n", + "\n", + "Spatial coordinates are in arc-seconds; convergence is dimensionless; intensity is in\n", + "electrons per second; `mass_to_light_ratio` is dimensionless (in the units of the workspace\n", + "config). The `guides/units_and_cosmology.ipynb` guide covers physical-unit conversion.\n", + "\n", + "__Data Structures__\n", + "\n", + "`convergence_2d_from` and `image_2d_from` both return `Array2D`. Plotting through\n", + "`aplt.plot_array` is direct. See `mass.py` for the deflection-magnitude pattern when you\n", + "need to inspect the vector quantities \u2014 this guide focuses on convergence and image.\n", + "\n", + "__Docs URL__\n", + "\n", + "The published API reference for these classes lives at:\n", + "\n", + " https://pyautolens.readthedocs.io/en/latest/api/mass.html\n", + "\n", + "The autosummary on that page is the authoritative list of every public mass-profile class\n", + "(stellar and dark are in the same page under their respective sections). Light-and-mass\n", + "combined profiles in `al.lmp.*` and `al.lmp_linear.*` are not yet documented on the\n", + "reference page \u2014 refer to the source under\n", + "`autogalaxy/profiles/light_and_mass_profiles.py` for the full list." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "To evaluate any quantity on a profile we need a 2D Cartesian grid of (y,x) coordinates.\n", + "We build a 100x100 grid here at a 0.05\" pixel scale \u2014 used by every section below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.05,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__All Profiles (Survey)__\n", + "\n", + "**PyAutoLens** organises matter-decomposition profiles into four namespaces:\n", + "\n", + "- `al.mp.*` *stellar* mass \u2014 Sersic-like profiles parameterised by light shape plus a\n", + " `mass_to_light_ratio` that converts intensity into surface mass density.\n", + "- `al.mp.*` *dark-matter* mass \u2014 NFW family halos parameterised by `kappa_s` /\n", + " `scale_radius`, or alternatively by halo mass + concentration via the MCR and Virial\n", + " variants.\n", + "- `al.lmp.*` *combined light-and-mass* profiles \u2014 one object that emits both\n", + " `image_2d_from` (light) and `convergence_2d_from` (mass) via a shared\n", + " `mass_to_light_ratio`.\n", + "- `al.lmp_linear.*` \u2014 the inversion-aware version of `al.lmp.*`. Same classes, intensity\n", + " solved by inversion at fit time.\n", + "\n", + "Below we construct one profile from each family with sensible defaults. No quantity is\n", + "evaluated yet \u2014 this is the catalogue." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Stellar mass \u2014 pure mass parameterised by light shape * mass_to_light_ratio\n", + "stellar_sersic = al.mp.Sersic()\n", + "stellar_chameleon = al.mp.Chameleon()\n", + "stellar_gaussian = al.mp.Gaussian()\n", + "\n", + "# Dark mass \u2014 NFW family\n", + "dark_nfw = al.mp.NFW()\n", + "dark_gnfw = al.mp.gNFW()\n", + "dark_cnfw = al.mp.cNFW()\n", + "dark_nfw_truncated = al.mp.NFWTruncatedSph()\n", + "dark_nfw_mcr = al.mp.NFWMCRLudlow()\n", + "dark_nfw_virial = al.mp.NFWVirialMassConcSph()\n", + "\n", + "# Combined light-and-mass \u2014 one object, two outputs\n", + "lmp_sersic = al.lmp.Sersic()\n", + "lmp_chameleon = al.lmp.Chameleon()\n", + "lmp_gaussian = al.lmp.Gaussian()\n", + "lmp_gaussian_gradient = al.lmp.GaussianGradient()\n", + "\n", + "# Linear combined light-and-mass \u2014 intensity solved by inversion at fit time\n", + "lmp_linear_sersic = al.lmp_linear.Sersic()\n", + "lmp_linear_gaussian = al.lmp_linear.Gaussian()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Two things worth knowing about this list before we move on:\n", + "\n", + "1. The *stellar* and *combined* profiles use the **same parametric forms** as the light\n", + " profiles (Sersic, Chameleon, Gaussian, ...). Their `convergence_2d_from` is the\n", + " profile's light shape multiplied by `mass_to_light_ratio` \u2014 this is a fundamentally\n", + " different parameterisation from the dark-matter halos, which are governed by halo-mass\n", + " physics (concentration-mass relations, virial parameters).\n", + "2. The dark-matter NFW family has *many* variants that are physically equivalent but\n", + " parameterised differently \u2014 they trade `kappa_s` / `scale_radius` for `mass_at_200` +\n", + " concentration (via MCR), or for `virial_mass` + concentration (Virial). Picking the\n", + " right parameterisation matters for prior choice and degeneracy management; the NFW\n", + " Variants section below is a one-page tour of the menagerie.\n", + "\n", + "__Stellar Mass Detailed Example__\n", + "\n", + "`al.mp.Sersic` is the stellar mass profile. It is parameterised by exactly the same\n", + "light-shape arguments as `al.lp.Sersic` (`intensity`, `effective_radius`, `sersic_index`)\n", + "plus a `mass_to_light_ratio` that converts the surface brightness profile into a surface\n", + "mass density. Crucially it does NOT emit light \u2014 `al.mp.Sersic` only carries mass.\n", + "\n", + "Build a stellar Sersic mass profile and plot its convergence:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "stellar_sersic = al.mp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " mass_to_light_ratio=2.0,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=stellar_sersic.convergence_2d_from(grid=grid),\n", + " title=\"Stellar Sersic Mass Convergence\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The convergence map looks exactly like a Sersic image because, geometrically, it is one\n", + "\u2014 the same Sersic profile, normalised by `mass_to_light_ratio`.\n", + "\n", + "To see what this stellar mass does to a background source, drop it into a `Tracer` with a\n", + "source-plane galaxy carrying a small Sersic light:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(redshift=0.5, mass=stellar_sersic)\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.05, 0.05),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=0.1,\n", + " sersic_index=1.5,\n", + " ),\n", + ")\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Tracer Image (Stellar Sersic Lens + Sersic Source)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that the lens galaxy carries *only* mass \u2014 no light is emitted from the lens plane in\n", + "the tracer image. To add lens-galaxy light you would attach a separate `al.lp.*` light\n", + "profile to the same `Galaxy`, OR use a `al.lmp.*` profile (covered later in this guide)\n", + "which carries both light and mass in a single object.\n", + "\n", + "__Dark Mass Detailed Example__\n", + "\n", + "`al.mp.NFW` is the canonical dark-matter halo \u2014 the elliptical Navarro-Frenk-White profile\n", + "parameterised by:\n", + "\n", + "- `kappa_s` \u2014 the dimensionless characteristic convergence at the scale radius.\n", + "- `scale_radius` \u2014 the radius (arc-seconds) at which the density slope transitions from\n", + " ~r^-1 (inner cusp) to ~r^-3 (outer fall-off).\n", + "\n", + "This is the \"natural\" lensing parameterisation; physically equivalent reparameterisations\n", + "in terms of halo mass + concentration are covered in the NFW Variants section below.\n", + "\n", + "Build an `al.mp.NFW` and plot its convergence on linear and log10 scales:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dark_nfw = al.mp.NFW(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", + " kappa_s=0.1,\n", + " scale_radius=2.0,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=dark_nfw.convergence_2d_from(grid=grid),\n", + " title=\"NFW Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=dark_nfw.convergence_2d_from(grid=grid),\n", + " title=\"NFW Convergence (log10 \u2014 central cusp visible)\",\n", + " use_log10=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The log10 view shows the characteristic central cusp that distinguishes NFW from the cored\n", + "profiles (`Isothermal`, `dPIEMass`). In a real lens model the cusp is what makes dark-\n", + "matter mass distinguishable from a steep stellar profile.\n", + "\n", + "__NFW Variants__\n", + "\n", + "The NFW family has several reparameterisations. All produce identical convergence maps\n", + "for matched halo parameters \u2014 they differ only in which parameters the *user* specifies\n", + "and which are derived. Picking the right one is mostly about prior choice and parameter\n", + "degeneracies.\n", + "\n", + "- **Generalised NFW (`al.mp.gNFW`)** \u2014 adds a free `inner_slope` parameter. Reduces to\n", + " `NFW` when `inner_slope=1.0`. Useful when you want to test whether the inner cusp is\n", + " steeper or shallower than the NFW prediction.\n", + "- **Cored NFW (`al.mp.cNFW`)** \u2014 adds a `core_radius` parameter that flattens the central\n", + " cusp. Useful for self-interacting dark matter scenarios.\n", + "- **Truncated NFW (`al.mp.NFWTruncatedSph`)** \u2014 adds a `truncation_radius` beyond which\n", + " the density is forced to zero. Useful for satellite halos with stripped outer envelopes.\n", + "- **MCR variants (`NFWMCR*` / `cNFWMCR*` / `gNFWMCR*`)** \u2014 replace `kappa_s` /\n", + " `scale_radius` with `mass_at_200` and use a fixed concentration-mass relation (Duffy\n", + " 2008 or Ludlow 2016) to derive the concentration. Removes one free parameter at the\n", + " cost of assuming the concentration follows the relation.\n", + "- **Scatter variants (`NFWMCRScatter*`)** \u2014 the Ludlow MCR plus a `scatter_sigma` knob\n", + " that adds Gaussian scatter on the log-concentration. Useful when you want to *fit* the\n", + " concentration but anchor it to the Ludlow relation with finite tolerance.\n", + "- **Virial-mass variants (`*VirialMassConc*`)** \u2014 parameterise the halo by its virial mass\n", + " and concentration directly, instead of the `mass_at_200` convention. Choose whichever\n", + " matches the units used in your cosmology pipeline.\n", + "\n", + "One-liner construction of the major variants:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "nfw_gnfw = al.mp.gNFW(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", + " kappa_s=0.1,\n", + " inner_slope=1.2,\n", + " scale_radius=2.0,\n", + ")\n", + "\n", + "nfw_cnfw = al.mp.cNFW(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", + " kappa_s=0.1,\n", + " scale_radius=2.0,\n", + " core_radius=0.2,\n", + ")\n", + "\n", + "nfw_truncated = al.mp.NFWTruncatedSph(\n", + " centre=(0.0, 0.0),\n", + " kappa_s=0.1,\n", + " scale_radius=2.0,\n", + " truncation_radius=4.0,\n", + ")\n", + "\n", + "nfw_mcr = al.mp.NFWMCRLudlowSph(\n", + " centre=(0.0, 0.0),\n", + " mass_at_200=1.0e12,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + ")\n", + "\n", + "nfw_mcr_scatter = al.mp.NFWMCRScatterLudlowSph(\n", + " centre=(0.0, 0.0),\n", + " mass_at_200=1.0e12,\n", + " scatter_sigma=0.1,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + ")\n", + "\n", + "nfw_virial = al.mp.NFWVirialMassConcSph(\n", + " centre=(0.0, 0.0),\n", + " virial_mass=1.0e12,\n", + " concentration=10.0,\n", + " virial_overdens=200.0,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For comparison, here is the cored variant's convergence \u2014 note the flat central plateau:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=nfw_cnfw.convergence_2d_from(grid=grid),\n", + " title=\"cNFW Convergence (cored central region)\",\n", + " use_log10=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "And the MCR variant produced from a halo mass rather than `kappa_s` / `scale_radius`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=nfw_mcr.convergence_2d_from(grid=grid),\n", + " title=\"NFWMCRLudlowSph Convergence (mass=1e12, concentration via Ludlow)\",\n", + " use_log10=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Combined Light + Mass Profiles (`al.lmp`)__\n", + "\n", + "The `al.lmp.*` namespace is the headline feature of this guide. Each `al.lmp` class\n", + "represents a galaxy component that emits BOTH light AND mass \u2014 a single object with one\n", + "set of parameters governing both `image_2d_from` and `convergence_2d_from` via a shared\n", + "`mass_to_light_ratio`.\n", + "\n", + "This is the standard way to model a lens galaxy's bulge: one `al.lmp.Sersic` carries the\n", + "visible Sersic light and contributes the stellar mass component to the deflection, tied\n", + "together by `mass_to_light_ratio`. Pair it with a separate `al.mp.NFW` (or one of the\n", + "dark variants above) for the dark-matter halo and you have the canonical bulge+halo\n", + "decomposition.\n", + "\n", + "Build one `al.lmp.Sersic` and plot both its image AND its convergence:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lmp_sersic = al.lmp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " mass_to_light_ratio=2.0,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=lmp_sersic.image_2d_from(grid=grid),\n", + " title=\"lmp.Sersic Image (light side)\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=lmp_sersic.convergence_2d_from(grid=grid),\n", + " title=\"lmp.Sersic Convergence (mass side, mass_to_light_ratio=2.0)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The two maps are *the same Sersic profile*, with the mass map scaled by\n", + "`mass_to_light_ratio`. This is the geometric guarantee of the `al.lmp.*` family \u2014 the\n", + "stellar light and stellar mass cannot diverge from each other except through this one\n", + "parameter.\n", + "\n", + "`al.lmp.SersicGradient` and `al.lmp.GaussianGradient` extend this with a *radial gradient*\n", + "in the mass-to-light ratio (one extra parameter), useful when the inner regions of a\n", + "galaxy have a systematically different M/L from the outskirts.\n", + "\n", + "__Linear Combined Light + Mass (`al.lmp_linear`)__\n", + "\n", + "`al.lmp_linear.*` mirrors `al.lmp.*` class-for-class. The only difference is that the\n", + "`intensity` parameter is solved analytically via the linear inversion at each likelihood\n", + "evaluation (the same mechanism that `al.lp_linear.*` uses for pure light profiles). This\n", + "means a lens model with a `al.lmp_linear.Sersic` bulge has one fewer free non-linear\n", + "parameter than the corresponding `al.lmp.Sersic` model \u2014 the bulge intensity is no longer\n", + "sampled by the search.\n", + "\n", + "The construction API is identical to `al.lmp.*`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lmp_linear_sersic = al.lmp_linear.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " mass_to_light_ratio=2.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The full workflow for linear profiles in lens models \u2014 including how the inversion\n", + "combines light and mass contributions \u2014 is documented in:\n", + "\n", + " scripts/imaging/features/linear_light_profiles/\n", + "\n", + "That folder is the canonical reference; this guide stops at the construction API.\n", + "\n", + "__Composing a Decomposed Bulge+Halo Model__\n", + "\n", + "The standard decomposed lens-mass model attaches:\n", + "\n", + "- one `al.mp.Sersic` (or `al.lmp.Sersic`) as the lens-galaxy stellar component, AND\n", + "- one `al.mp.NFW` (or another dark variant) as the lens-galaxy dark-matter halo,\n", + "\n", + "so the total convergence is the sum of stellar + dark contributions. Compose this via\n", + "`af.Model` exactly as you would for a single-profile lens \u2014 `Galaxy` accepts arbitrary\n", + "kwargs for its mass components." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_bulge_mass_model = af.Model(al.mp.Sersic)\n", + "lens_dark_model = af.Model(al.mp.NFW)\n", + "source_bulge_model = af.Model(al.lp.Sersic)\n", + "\n", + "lens_galaxy_model = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=lens_bulge_mass_model,\n", + " dark=lens_dark_model,\n", + ")\n", + "source_galaxy_model = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge_model)\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens_galaxy_model, source=source_galaxy_model)\n", + ")\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing `model.info` shows the priors-and-defaults summary. Notice the lens galaxy has\n", + "*two* mass components (`bulge` and `dark`) whose convergences will be summed during the\n", + "fit. This is how a decomposed lens model differs from the single-profile `Isothermal`\n", + "lens shown in `mass.py` \u2014 the model is more flexible but also higher-dimensional.\n", + "\n", + "Swapping `al.mp.Sersic` for `al.lmp.Sersic` on the lens bulge would additionally model the\n", + "lens-galaxy *light* through the same object \u2014 useful when the lens galaxy is visible in\n", + "the data and you want to tie its light to its stellar mass.\n", + "\n", + "Full end-to-end lens fits with this model live under `scripts/imaging/features/` and the\n", + "SLaM pipeline guides; this section just shows the model spec.\n", + "\n", + "__Model Instance from Decomposed Model__\n", + "\n", + "To realise an instance from the model's prior medians and visualise it, call\n", + "`instance_from_prior_medians()` on the collection and feed the result into a `Tracer`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_instance = model.instance_from_prior_medians()\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[model_instance.galaxies.lens, model_instance.galaxies.source]\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Tracer Image from Decomposed Model (Stellar + NFW)\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.convergence_2d_from(grid=grid),\n", + " title=\"Tracer Total Convergence (Stellar + NFW summed)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The total convergence map shows the summed stellar + dark contributions \u2014 exactly the\n", + "lensing mass distribution that the search would optimise.\n", + "\n", + "After a fit completes, `result.max_log_likelihood_tracer` returns the same shape of object\n", + "with the prior medians replaced by the fitted parameters. See\n", + "`scripts/guides/results/start_here.py` for the full results-introspection guide.\n", + "\n", + "__Remaining Profiles Walkthrough__\n", + "\n", + "We have shown the full flow for `al.mp.Sersic` (stellar), `al.mp.NFW` (dark), and\n", + "`al.lmp.Sersic` (combined). Every remaining stellar / dark / lmp / lmp_linear profile uses\n", + "the same API \u2014 only the constructor parameters change.\n", + "\n", + "The compact tour below builds each remaining profile with sensible defaults. Stellar and\n", + "dark profiles show `convergence_2d_from`; `lmp` profiles show both `image_2d_from` and\n", + "`convergence_2d_from` so the dual-output property is visible per class.\n", + "\n", + "Stellar mass:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=al.mp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=4.0,\n", + " radius_break=0.05,\n", + " gamma=0.25,\n", + " alpha=3.0,\n", + " mass_to_light_ratio=2.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"Stellar SersicCore Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.SersicGradient(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " mass_to_light_ratio=2.0,\n", + " mass_to_light_gradient=0.3,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"Stellar SersicGradient Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.DevVaucouleurs(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " mass_to_light_ratio=2.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"Stellar DevVaucouleurs Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.Exponential(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", + " intensity=0.5,\n", + " effective_radius=1.6,\n", + " mass_to_light_ratio=2.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"Stellar Exponential Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " sigma=0.4,\n", + " mass_to_light_ratio=2.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"Stellar Gaussian Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.GaussianGradient(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " sigma=0.4,\n", + " mass_to_light_ratio_base=2.0,\n", + " mass_to_light_gradient=0.2,\n", + " mass_to_light_radius=0.5,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"Stellar GaussianGradient Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.Chameleon(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " core_radius_0=0.05,\n", + " core_radius_1=0.3,\n", + " mass_to_light_ratio=2.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"Stellar Chameleon Convergence\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Dark mass \u2014 the variant menagerie, one block each:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=al.mp.gNFW(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", + " kappa_s=0.1,\n", + " inner_slope=1.2,\n", + " scale_radius=2.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"gNFW Convergence (inner_slope=1.2)\",\n", + " use_log10=True,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.NFWTruncatedSph(\n", + " centre=(0.0, 0.0),\n", + " kappa_s=0.1,\n", + " scale_radius=2.0,\n", + " truncation_radius=4.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"NFWTruncatedSph Convergence\",\n", + " use_log10=True,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.NFWMCRDuffySph(\n", + " centre=(0.0, 0.0),\n", + " mass_at_200=1.0e12,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"NFWMCRDuffySph Convergence (Duffy 2008 MCR)\",\n", + " use_log10=True,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.NFWMCRScatterLudlowSph(\n", + " centre=(0.0, 0.0),\n", + " mass_at_200=1.0e12,\n", + " scatter_sigma=0.1,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"NFWMCRScatterLudlowSph Convergence (Ludlow + sigma=0.1)\",\n", + " use_log10=True,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.NFWVirialMassConcSph(\n", + " centre=(0.0, 0.0),\n", + " virial_mass=1.0e12,\n", + " concentration=10.0,\n", + " virial_overdens=200.0,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"NFWVirialMassConcSph Convergence\",\n", + " use_log10=True,\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.cNFWMCRLudlowSph(\n", + " centre=(0.0, 0.0),\n", + " mass_at_200=1.0e12,\n", + " f_c=0.05,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"cNFWMCRLudlowSph Convergence (cored + Ludlow MCR)\",\n", + " use_log10=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Combined light-and-mass \u2014 every remaining `al.lmp` profile, with both image and convergence\n", + "shown side-by-side via stacked `plot_array` calls:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lmp_chameleon = al.lmp.Chameleon(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " core_radius_0=0.05,\n", + " core_radius_1=0.3,\n", + " mass_to_light_ratio=2.0,\n", + ")\n", + "aplt.plot_array(\n", + " array=lmp_chameleon.image_2d_from(grid=grid), title=\"lmp.Chameleon Image\"\n", + ")\n", + "aplt.plot_array(\n", + " array=lmp_chameleon.convergence_2d_from(grid=grid),\n", + " title=\"lmp.Chameleon Convergence\",\n", + ")\n", + "\n", + "lmp_devvaucouleurs = al.lmp.DevVaucouleurs(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " mass_to_light_ratio=2.0,\n", + ")\n", + "aplt.plot_array(\n", + " array=lmp_devvaucouleurs.image_2d_from(grid=grid),\n", + " title=\"lmp.DevVaucouleurs Image\",\n", + ")\n", + "aplt.plot_array(\n", + " array=lmp_devvaucouleurs.convergence_2d_from(grid=grid),\n", + " title=\"lmp.DevVaucouleurs Convergence\",\n", + ")\n", + "\n", + "lmp_exponential = al.lmp.Exponential(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", + " intensity=0.5,\n", + " effective_radius=1.6,\n", + " mass_to_light_ratio=2.0,\n", + ")\n", + "aplt.plot_array(\n", + " array=lmp_exponential.image_2d_from(grid=grid),\n", + " title=\"lmp.Exponential Image\",\n", + ")\n", + "aplt.plot_array(\n", + " array=lmp_exponential.convergence_2d_from(grid=grid),\n", + " title=\"lmp.Exponential Convergence\",\n", + ")\n", + "\n", + "lmp_gaussian = al.lmp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " sigma=0.4,\n", + " mass_to_light_ratio=2.0,\n", + ")\n", + "aplt.plot_array(array=lmp_gaussian.image_2d_from(grid=grid), title=\"lmp.Gaussian Image\")\n", + "aplt.plot_array(\n", + " array=lmp_gaussian.convergence_2d_from(grid=grid),\n", + " title=\"lmp.Gaussian Convergence\",\n", + ")\n", + "\n", + "lmp_sersic_gradient = al.lmp.SersicGradient(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " mass_to_light_ratio=2.0,\n", + " mass_to_light_gradient=0.3,\n", + ")\n", + "aplt.plot_array(\n", + " array=lmp_sersic_gradient.image_2d_from(grid=grid),\n", + " title=\"lmp.SersicGradient Image\",\n", + ")\n", + "aplt.plot_array(\n", + " array=lmp_sersic_gradient.convergence_2d_from(grid=grid),\n", + " title=\"lmp.SersicGradient Convergence\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The spherical variants (`SersicSph`, `DevVaucouleursSph`, `ExponentialSph`, `ChameleonSph`,\n", + "etc., across all three namespaces) are constructed identically with the `ell_comps`\n", + "argument removed. Each looks like a rotationally symmetric version of its elliptical\n", + "counterpart.\n", + "\n", + "`al.lmp_linear.*` shares the construction API with `al.lmp.*` \u2014 instantiate any\n", + "`al.lmp_linear.X` exactly as you would `al.lmp.X`. Their distinct behaviour only shows up\n", + "inside a model fit, when the inversion solves for `intensity` rather than treating it as\n", + "a free non-linear parameter.\n", + "\n", + "__Back-References__\n", + "\n", + "This guide completes the three-guide tour of the `scripts/guides/profiles/` folder:\n", + "\n", + "- `scripts/guides/profiles/light.py` \u2014 pure light profiles, the `al.lp.*` /\n", + " `al.lp_linear.*` / `al.lp_operated.*` / `al.lp_basis.*` namespaces.\n", + "- `scripts/guides/profiles/mass.py` \u2014 parametric lensing mass profiles, the *Total* /\n", + " *Mass Sheets* / *Multipoles* / *Point Mass* families of `al.mp.*`.\n", + "- `scripts/guides/profiles/light_and_mass_profiles.py` \u2014 this guide; stellar / dark / lmp /\n", + " lmp_linear matter decomposition.\n", + "\n", + "For the next step into actual lens modelling, see `scripts/imaging/modeling/start_here.py`\n", + "(strong-lens fit end-to-end) and the topic-specific feature packages under\n", + "`scripts/imaging/features/`. The SLaM pipeline guide (`scripts/guides/modeling/slam_start_here.py`)\n", + "shows the production decomposed-lens workflow built on these profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/profiles/mass.ipynb b/notebooks/guides/profiles/mass.ipynb index 213eabcbc..8d5665395 100644 --- a/notebooks/guides/profiles/mass.ipynb +++ b/notebooks/guides/profiles/mass.ipynb @@ -1,844 +1,881 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Mass Profiles\n", - "=============\n", - "\n", - "This guide is the single-page tour of every lensing mass profile available in **PyAutoLens**\n", - "(re-exported from **PyAutoGalaxy**): how to construct each one, how to evaluate its\n", - "convergence and deflections on a grid, how to compose it into a model, and how to pull an\n", - "instance back out of that model.\n", - "\n", - "It is the companion to `scripts/guides/profiles/light.py` \u2014 the section flow, prose style\n", - "and \"detailed example then walkthrough\" rhythm intentionally mirror that guide so the two\n", - "can be read as a pair. Where `light.py` shows `image_2d_from`, this guide shows\n", - "`convergence_2d_from` and `deflections_yx_2d_from`; mass profiles do not produce images of\n", - "their own (a mass profile only deflects light from a *source* through ray-tracing \u2014 that is\n", - "what the `Tracer` is for, and we use it for the detailed example).\n", - "\n", - "This guide covers the *lensing* mass profiles \u2014 Total, Mass Sheets, Multipoles, and Point\n", - "Mass. The *stellar* (`al.mp.Sersic`, `al.mp.Chameleon`, ...) and *dark-matter* (NFW family)\n", - "mass profiles, along with the combined light+mass profiles in `al.lmp.*` / `al.lmp_linear.*`,\n", - "get their own dedicated guide at `scripts/guides/profiles/light_and_mass_profiles.py` where\n", - "the stellar-plus-dark decomposition story is told properly.\n", - "\n", - "__Contents__\n", - "\n", - "- **Overview & Docs URL:** Where the canonical API reference lives.\n", - "- **All Mass Profiles (Survey):** A high-level run-through of every profile in `al.mp.*`\n", - " (Total / Mass Sheets / Multipoles / Point Mass), without yet evaluating any quantities.\n", - "- **Detailed Example: Isothermal:** Build a `Grid2D`, instantiate `al.mp.Isothermal`, plot\n", - " convergence, potential, deflection-magnitude, and the lensed source image produced when\n", - " the isothermal mass is dropped into a `Tracer`.\n", - "- **Mass Sheets:** `ExternalShear`, `MassSheet`, `ExternalPotential` \u2014 global perturbations\n", - " rather than parametric matter distributions.\n", - "- **Point Mass:** `PointMass`, `SMBH`, `SMBHBinary` \u2014 delta-function-like mass for\n", - " microlensing and supermassive black hole lensing.\n", - "- **Mass Profile in a Model:** Wrap a profile in `af.Model`, compose lens + source via\n", - " `af.Collection`, inspect the model info.\n", - "- **Model Instance from Mass Profile:** Realise an instance from the model's prior medians\n", - " and drop it into a `Tracer` to produce a lensed image.\n", - "- **Multipole Mass Profile:** `PowerLawMultipole` \u2014 m=3 / m=4 Fourier perturbation on the\n", - " power-law convergence, the lensing counterpart of `lp.SersicMultipole`.\n", - "- **Remaining Profiles Walkthrough:** Compact `convergence_2d_from` block for every total\n", - " profile not yet shown.\n", - "- **Follow-Up:** Pointer to `light_and_mass_profiles.py` for stellar / dark / lmp coverage.\n", - "\n", - "__Units__\n", - "\n", - "Spatial coordinates are in arc-seconds, mass quantities are dimensionless (convergence) or\n", - "in their natural lensing units (potential, deflection angles in arc-seconds). The\n", - "`guides/units_and_cosmology.ipynb` guide covers conversion to physical units.\n", - "\n", - "__Data Structures__\n", - "\n", - "`convergence_2d_from` and `potential_2d_from` return `Array2D`. `deflections_yx_2d_from`\n", - "returns `VectorYX2D` (a 2D vector field). `aplt.plot_array` accepts `Array2D` directly;\n", - "to plot a vector field we either pull out the y / x components individually or compute the\n", - "magnitude and wrap it in an `Array2D`. This guide uses the magnitude approach for the\n", - "detailed example.\n", - "\n", - "__Docs URL__\n", - "\n", - "The published API reference for these classes lives at:\n", - "\n", - " https://pyautolens.readthedocs.io/en/latest/api/mass.html\n", - "\n", - "The autosummary on that page is the authoritative list of every public mass-profile class.\n", - "Note that the reference uses the `ag.mp` namespace label because the classes are defined in\n", - "PyAutoGalaxy and re-exported here as `al.mp`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "\n", - "import autofit as af\n", - "import autoarray as aa\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "To evaluate any quantity on a mass profile we need a 2D Cartesian grid of (y,x) coordinates.\n", - "We build a 100x100 grid here at a 0.05\" pixel scale \u2014 used by every section below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.05,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__All Mass Profiles (Survey)__\n", - "\n", - "**PyAutoLens** groups mass profiles into four families relevant to lensing:\n", - "\n", - "- `al.mp.*` total profiles \u2014 the canonical parametric lens mass distributions\n", - " (`Isothermal`, `PowerLaw`, their cored / broken / dPIE variants).\n", - "- `al.mp.*` mass sheets \u2014 `ExternalShear`, `MassSheet`, `ExternalPotential`. Global\n", - " perturbations representing line-of-sight contributions.\n", - "- `al.mp.*` multipoles \u2014 `PowerLawMultipole`, the m=3 / m=4 angular extension of the power-\n", - " law family.\n", - "- `al.mp.*` point masses \u2014 `PointMass`, `SMBH`, `SMBHBinary` \u2014 delta-function-like sources\n", - " for microlensing and central black holes.\n", - "\n", - "The *stellar* and *dark-matter* mass profiles (e.g. `al.mp.Sersic`, the NFW family) and the\n", - "combined `al.lmp.*` / `al.lmp_linear.*` light-and-mass namespaces are covered in\n", - "`scripts/guides/profiles/light_and_mass_profiles.py`. They share the same API but tell a\n", - "different story (stellar-plus-dark decomposition).\n", - "\n", - "Below we construct each lensing profile with default parameters. No quantity is evaluated\n", - "yet \u2014 that comes in the next section. The goal here is purely a catalogue of what is\n", - "available." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Power-law family (Isothermal is PowerLaw with slope = 2)\n", - "isothermal = al.mp.Isothermal()\n", - "isothermal_sph = al.mp.IsothermalSph()\n", - "isothermal_core = al.mp.IsothermalCore()\n", - "isothermal_core_sph = al.mp.IsothermalCoreSph()\n", - "power_law = al.mp.PowerLaw()\n", - "power_law_sph = al.mp.PowerLawSph()\n", - "power_law_core = al.mp.PowerLawCore()\n", - "power_law_core_sph = al.mp.PowerLawCoreSph()\n", - "power_law_broken = al.mp.PowerLawBroken()\n", - "power_law_broken_sph = al.mp.PowerLawBrokenSph()\n", - "\n", - "# Pseudo-isothermal family (mass and potential parameterisations)\n", - "# Note: dPIEMass with default ell_comps=(0,0) triggers a divide-by-zero in the complex-plane\n", - "# formula; we use a small ellipticity here so the survey constructions succeed cleanly.\n", - "dpie_mass = al.mp.dPIEMass(ell_comps=(0.05, 0.0))\n", - "dpie_mass_sph = al.mp.dPIEMassSph()\n", - "pie_mass = al.mp.PIEMass(ell_comps=(0.05, 0.0))\n", - "dpie_potential = al.mp.dPIEPotential()\n", - "dpie_potential_sph = al.mp.dPIEPotentialSph()\n", - "\n", - "# Mass sheets \u2014 global perturbations\n", - "external_shear = al.mp.ExternalShear()\n", - "mass_sheet = al.mp.MassSheet()\n", - "external_potential = al.mp.ExternalPotential()\n", - "\n", - "# Multipole \u2014 power-law angular extension\n", - "power_law_multipole = al.mp.PowerLawMultipole()\n", - "\n", - "# Point masses \u2014 microlensing and supermassive black holes\n", - "point_mass = al.mp.PointMass()\n", - "smbh = al.mp.SMBH()\n", - "smbh_binary = al.mp.SMBHBinary()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two things worth knowing about this list before we move on:\n", - "\n", - "1. Every elliptical lensing profile (e.g. `Isothermal`, `PowerLaw`, `PowerLawCore`) has a\n", - " spherical sibling whose name ends in `Sph` (e.g. `IsothermalSph`). The spherical variant\n", - " fixes the ellipticity components `ell_comps` to `(0, 0)`, which is useful when you want\n", - " to model a circular halo and avoid two redundant parameters in the non-linear search.\n", - "2. The `PowerLawMultipole` variant only exists as an *elliptical* profile \u2014 the m=3 / m=4\n", - " perturbations are angular distortions and are not meaningful without an underlying\n", - " reference frame, exactly like the light multipoles in `light.py`.\n", - "\n", - "We now move on to seeing what these profiles actually produce when evaluated on a grid.\n", - "\n", - "__Detailed Example: Isothermal__\n", - "\n", - "The `Isothermal` profile is the canonical strong-lens mass profile \u2014 equivalent to the\n", - "elliptical power-law with `slope = 2.0`. Its three parameters are:\n", - "\n", - "- `centre` \u2014 the (y, x) arc-second coordinate of the profile centre.\n", - "- `ell_comps` \u2014 the two ellipticity components `(e1, e2)`. Use\n", - " `al.convert.ell_comps_from(axis_ratio=..., angle=...)` to convert from human-friendly\n", - " axis ratio and position angle.\n", - "- `einstein_radius` \u2014 the Einstein radius in arc-seconds. This sets the strength of the\n", - " lens and the size of the Einstein ring for an axisymmetric source on-axis.\n", - "\n", - "Build an isothermal and evaluate every standard lensing quantity on our grid:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "isothermal = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " einstein_radius=1.6,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "First the *convergence* \u2014 the projected surface mass density in units of the critical\n", - "surface density. This is the simplest visualisation of a mass profile and the most\n", - "analogous to the `image_2d_from` plot we used throughout `light.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=isothermal.convergence_2d_from(grid=grid),\n", - " title=\"Isothermal Convergence\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next the *lensing potential* \u2014 the 2D field whose gradient gives the deflection angles.\n", - "We plot it in log10 to make the central cusp visible." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=isothermal.potential_2d_from(grid=grid),\n", - " title=\"Isothermal Potential (log10)\",\n", - " use_log10=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The *deflection angles* are a vector field rather than a scalar, so `deflections_yx_2d_from`\n", - "returns a `VectorYX2D` rather than an `Array2D`. `aplt.plot_array` does not accept vectors\n", - "directly; the conventional pattern in this workspace is to either plot the y / x components\n", - "separately or compute the magnitude. Below we use the magnitude \u2014 a single map showing\n", - "\"how much light is bent here\", which is the most useful single-figure summary." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections = isothermal.deflections_yx_2d_from(grid=grid)\n", - "deflection_magnitude = aa.Array2D(\n", - " values=np.hypot(deflections.slim[:, 0], deflections.slim[:, 1]),\n", - " mask=grid.mask,\n", - ")\n", - "aplt.plot_array(\n", - " array=deflection_magnitude,\n", - " title=\"Isothermal Deflection Magnitude\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "None of these maps show what the isothermal mass actually *does* to a background source \u2014\n", - "for that we need a `Tracer`. A `Tracer` groups galaxies by redshift plane, ray-traces the\n", - "grid through every plane, and combines the resulting emission into a single observed image.\n", - "\n", - "Below we build a two-plane lens system: a lens galaxy at `z=0.5` carrying our isothermal\n", - "mass profile, and a source galaxy at `z=1.0` with a small Sersic light profile. The\n", - "`tracer.image_2d_from` call produces the lensed image you would see at the telescope." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(redshift=0.5, mass=isothermal)\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.05, 0.05),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=0.1,\n", - " sersic_index=1.5,\n", - " ),\n", - ")\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Tracer Image (Isothermal Lens + Sersic Source)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The arc / ring you see is the source's Sersic light, deflected by the isothermal mass into\n", - "the characteristic strong-lens morphology. The same `image_2d_from(grid=grid)` call exists\n", - "on every `Tracer` regardless of which mass profile is on the lens galaxy \u2014 every section\n", - "below is a small variation on this one, swapping the lens-galaxy mass for a different\n", - "profile.\n", - "\n", - "__Mass Sheets__\n", - "\n", - "Mass sheets are *global* perturbations rather than localised mass distributions. They show\n", - "up in strong-lens models to capture line-of-sight contributions: nearby group / cluster\n", - "members, large-scale structure, and any other mass that is not part of the primary lens but\n", - "still distorts the light path.\n", - "\n", - "Three mass-sheet profiles are available:\n", - "\n", - "- `al.mp.ExternalShear(gamma_1, gamma_2)` \u2014 a constant shear with two components. The most\n", - " common line-of-sight correction; you will see this routinely added to lens-mass models.\n", - "- `al.mp.MassSheet(centre, kappa)` \u2014 a uniform convergence (positive or negative). Useful\n", - " for representing diffuse line-of-sight mass.\n", - "- `al.mp.ExternalPotential(centre, gamma_1, gamma_2, tau_1, tau_2, delta_1, delta_2)` \u2014 the\n", - " Powell et al. 2022 generalisation of external shear that adds the next-order line-of-sight\n", - " terms. Reduces to `ExternalShear` when `tau_*` and `delta_*` are zero." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "external_shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.03)\n", - "deflections = external_shear.deflections_yx_2d_from(grid=grid)\n", - "aplt.plot_array(\n", - " array=aa.Array2D(\n", - " values=np.hypot(deflections.slim[:, 0], deflections.slim[:, 1]),\n", - " mask=grid.mask,\n", - " ),\n", - " title=\"ExternalShear Deflection Magnitude\",\n", - ")\n", - "\n", - "mass_sheet = al.mp.MassSheet(centre=(0.0, 0.0), kappa=0.1)\n", - "aplt.plot_array(\n", - " array=mass_sheet.convergence_2d_from(grid=grid),\n", - " title=\"MassSheet Convergence (uniform kappa=0.1)\",\n", - ")\n", - "\n", - "external_potential = al.mp.ExternalPotential(\n", - " centre=(0.0, 0.0),\n", - " gamma_1=0.04,\n", - " gamma_2=0.02,\n", - " tau_1=0.005,\n", - " tau_2=0.0,\n", - " delta_1=0.0,\n", - " delta_2=0.003,\n", - ")\n", - "deflections = external_potential.deflections_yx_2d_from(grid=grid)\n", - "aplt.plot_array(\n", - " array=aa.Array2D(\n", - " values=np.hypot(deflections.slim[:, 0], deflections.slim[:, 1]),\n", - " mask=grid.mask,\n", - " ),\n", - " title=\"ExternalPotential Deflection Magnitude\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A mass sheet alone produces no observable lensing \u2014 its deflections are uniform (constant)\n", - "across the image and so are absorbed into the global astrometric solution. Mass sheets only\n", - "become visible *in combination* with another mass profile; their job is to perturb the\n", - "isothermal or power-law deflections at the few-percent level.\n", - "\n", - "__Point Mass__\n", - "\n", - "Point masses are delta-function-like mass distributions. Three flavours exist:\n", - "\n", - "- `al.mp.PointMass(centre, einstein_radius)` \u2014 a single point mass parameterised by its\n", - " Einstein radius. Used for microlensing.\n", - "- `al.mp.SMBH(centre, mass, redshift_object, redshift_source)` \u2014 a supermassive black hole\n", - " parameterised by mass (in solar masses); internally converts to a `PointMass` Einstein\n", - " radius using the input redshifts.\n", - "- `al.mp.SMBHBinary(centre, separation, angle_binary, mass, mass_ratio, ...)` \u2014 two SMBHs\n", - " separated by a tunable angle and separation. Used for binary SMBH lensing.\n", - "\n", - "`PointMass.convergence_2d_from` returns a raw numpy array rather than the `Array2D` that\n", - "elliptical profiles return \u2014 a library quirk. We wrap the point mass in a `Galaxy` and plot\n", - "the galaxy's convergence map instead, which goes through the standard `Array2D` path and is\n", - "also the natural usage pattern for a lensing model anyway." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "point_mass = al.mp.PointMass(centre=(0.0, 0.0), einstein_radius=0.3)\n", - "point_galaxy = al.Galaxy(redshift=0.5, mass=point_mass)\n", - "\n", - "aplt.plot_array(\n", - " array=point_galaxy.convergence_2d_from(grid=grid),\n", - " title=\"PointMass Convergence (delta-function-like)\",\n", - ")\n", - "\n", - "smbh = al.mp.SMBH(\n", - " centre=(0.0, 0.0),\n", - " mass=1.0e9,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - ")\n", - "smbh_galaxy = al.Galaxy(redshift=0.5, mass=smbh)\n", - "aplt.plot_array(\n", - " array=smbh_galaxy.convergence_2d_from(grid=grid),\n", - " title=\"SMBH Convergence (M = 1e9 M_sun)\",\n", - ")\n", - "\n", - "smbh_binary = al.mp.SMBHBinary(\n", - " centre=(0.0, 0.0),\n", - " separation=0.4,\n", - " angle_binary=45.0,\n", - " mass=2.0e9,\n", - " mass_ratio=0.5,\n", - " redshift_object=0.5,\n", - " redshift_source=1.0,\n", - ")\n", - "smbh_binary_galaxy = al.Galaxy(redshift=0.5, mass=smbh_binary)\n", - "aplt.plot_array(\n", - " array=smbh_binary_galaxy.convergence_2d_from(grid=grid),\n", - " title=\"SMBHBinary Convergence\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mass Profile in a Model__\n", - "\n", - "So far we have been instantiating mass profiles with concrete parameter values. When\n", - "fitting a real strong-lens dataset we instead build a *model* of the mass profile and let\n", - "the non-linear search find the best-fit parameters. This is what `af.Model` is for.\n", - "\n", - "For a strong-lens system the model typically contains two `Galaxy` objects on different\n", - "planes: a foreground lens (light + mass) and a background source (light only). In this\n", - "section we focus on the *mass* side of that picture and show how an isothermal mass profile\n", - "plugs in to a lens galaxy that has only mass (no light) \u2014 the simplest case." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_mass_model = af.Model(al.mp.Isothermal)\n", - "source_bulge_model = af.Model(al.lp.Sersic)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `af.Model` wraps the profile class. Every constructor argument with a numerical default\n", - "becomes a *prior* \u2014 by default the priors are sensible distributions for each parameter\n", - "(see the autogalaxy config for the configured ranges).\n", - "\n", - "You can override priors before fitting:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_mass_model.einstein_radius = af.UniformPrior(lower_limit=0.5, upper_limit=3.0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Wrap each profile in a `Galaxy` at its own redshift and assemble them in an `af.Collection`.\n", - "The `Tracer` is not part of the *model* spec itself; it is what `AnalysisImaging` builds out\n", - "of the realised instance at fit time." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy_model = af.Model(al.Galaxy, redshift=0.5, mass=lens_mass_model)\n", - "source_galaxy_model = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge_model)\n", - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens_galaxy_model, source=source_galaxy_model)\n", - ")\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing `model.info` shows the full priors-and-defaults summary \u2014 useful before kicking\n", - "off a long fit to confirm the model is shaped the way you expect.\n", - "\n", - "The model API is the same for **every** mass profile in this guide \u2014 swap\n", - "`al.mp.Isothermal` for `al.mp.PowerLaw`, `al.mp.dPIEMass`, `al.mp.PowerLawMultipole`, etc.,\n", - "and the rest of the snippet is unchanged.\n", - "\n", - "Full lens-modelling end-to-end examples live in `scripts/imaging/modeling/start_here.py`\n", - "and the topic-specific guides under `scripts/imaging/features/`.\n", - "\n", - "__Model Instance from Mass Profile__\n", - "\n", - "A model is a description of *possible* profiles. To get an actual profile back out \u2014 for\n", - "example to plot what the prior medians look like before running a fit \u2014 call\n", - "`instance_from_prior_medians()`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_mass_instance = lens_mass_model.instance_from_prior_medians()\n", - "print(type(lens_mass_instance)) # autogalaxy.profiles.mass.total.isothermal.Isothermal\n", - "\n", - "aplt.plot_array(\n", - " array=lens_mass_instance.convergence_2d_from(grid=grid),\n", - " title=\"Isothermal Instance Convergence\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The instance returned is a real `al.mp.Isothermal` \u2014 the same class we constructed by hand\n", - "in the detailed example above \u2014 and supports the full mass-profile API.\n", - "\n", - "The same flow works at the full-model level. We realise an instance of the lens-and-source\n", - "collection and drop the resulting galaxies into a `Tracer`, which combines the lens mass\n", - "deflections with the source light to produce the lensed image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_instance = model.instance_from_prior_medians()\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[model_instance.galaxies.lens, model_instance.galaxies.source]\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Tracer Image from Model Instance\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "After a fit completes, `result.max_log_likelihood_tracer` returns the same shape of object,\n", - "with the prior medians replaced by the fitted parameter values. See\n", - "`scripts/guides/results/start_here.py` for the full results-introspection guide.\n", - "\n", - "__Multipole Mass Profile__\n", - "\n", - "`PowerLawMultipole` is the lensing counterpart of `lp.SersicMultipole` \u2014 an m=3 / m=4\n", - "Fourier angular perturbation, here on the eccentric radius of the power-law convergence\n", - "field rather than the Sersic intensity field. The signature is\n", - "\n", - " al.mp.PowerLawMultipole(\n", - " m=4, # multipole order (3 or 4)\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.0,\n", - " slope=2.0,\n", - " multipole_comps=(c, s), # cosine / sine components of the m-th term\n", - " )\n", - "\n", - "A `PowerLawMultipole` produces only the *perturbation* \u2014 not the underlying power-law mass\n", - "distribution. The standard usage is therefore to add the multipole alongside a `PowerLaw`\n", - "or `Isothermal` profile sharing the same `centre`, `einstein_radius`, and `slope`, so that\n", - "the combined convergence reads as a power law with boxy / discy / lopsided distortions.\n", - "\n", - "Build a `PowerLawMultipole` with non-trivial multipole components and plot its convergence\n", - "alongside the underlying `PowerLaw` for comparison:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "power_law_base = al.mp.PowerLaw(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " einstein_radius=1.6,\n", - " slope=2.0,\n", - ")\n", - "\n", - "power_law_multipole_m4 = al.mp.PowerLawMultipole(\n", - " m=4,\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " slope=2.0,\n", - " multipole_comps=(0.05, 0.02),\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=power_law_base.convergence_2d_from(grid=grid),\n", - " title=\"PowerLaw Convergence (base profile)\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=power_law_multipole_m4.convergence_2d_from(grid=grid),\n", - " title=\"PowerLawMultipole Convergence (m=4 perturbation only)\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two practical notes on the multipole:\n", - "\n", - "- The multipole profile carries the *perturbation*, not the base mass distribution. In a\n", - " model you typically wrap both `PowerLaw` and `PowerLawMultipole` on the same lens galaxy\n", - " and link their shared parameters (`centre`, `einstein_radius`, `slope`).\n", - "- There is **no spherical (`*Sph`) variant** of `PowerLawMultipole`. The perturbation is\n", - " an angular distortion measured in the lens reference frame, so it only makes sense for an\n", - " elliptical (or quasi-elliptical) host.\n", - "\n", - "Plugging `PowerLawMultipole` into the `af.Model` / `af.Collection` / `Galaxy` / `Tracer`\n", - "pattern shown above works exactly as it did for the plain `Isothermal` \u2014 the multipole\n", - "components are picked up as priors automatically.\n", - "\n", - "__Remaining Profiles Walkthrough__\n", - "\n", - "We have shown the full `convergence_2d_from` \u2192 `af.Model` \u2192 `instance` \u2192 `Tracer` flow for\n", - "the `Isothermal` profile. Every remaining lensing mass profile uses the **same API** \u2014 the\n", - "only thing that changes is which parameters appear in the constructor.\n", - "\n", - "The compact tour below builds each remaining total profile with sensible parameter values\n", - "and plots its convergence. When you want to use any of these in a lens model, repeat the\n", - "`af.Model(...)` / `af.Collection(...)` / `Tracer(...)` pattern from the previous section." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(\n", - " array=al.mp.PowerLaw(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " einstein_radius=1.6,\n", - " slope=2.2,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"PowerLaw Convergence (slope=2.2)\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.PowerLawCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " einstein_radius=1.6,\n", - " slope=2.0,\n", - " core_radius=0.05,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"PowerLawCore Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.PowerLawBroken(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " einstein_radius=1.6,\n", - " inner_slope=1.5,\n", - " outer_slope=2.5,\n", - " break_radius=0.5,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"PowerLawBroken Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.IsothermalCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " einstein_radius=1.6,\n", - " core_radius=0.05,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"IsothermalCore Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.dPIEMass(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ra=0.1,\n", - " b0=0.5,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"dPIEMass Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.PIEMass(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ra=0.1,\n", - " b0=0.5,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"PIEMass Convergence\",\n", - ")\n", - "\n", - "aplt.plot_array(\n", - " array=al.mp.dPIEPotential(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", - " ra=0.1,\n", - " rs=2.0,\n", - " b0=1.0,\n", - " ).convergence_2d_from(grid=grid),\n", - " title=\"dPIEPotential Convergence\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The spherical variants (`IsothermalSph`, `PowerLawSph`, `PowerLawCoreSph`,\n", - "`PowerLawBrokenSph`, `dPIEMassSph`) are constructed identically with the `ell_comps`\n", - "argument removed. Each looks like a rotationally symmetric version of its elliptical\n", - "counterpart.\n", - "\n", - "__Follow-Up: Stellar, Dark Matter, and Combined Light+Mass Profiles__\n", - "\n", - "This guide deliberately stops at the parametric lensing mass profiles. The remaining mass\n", - "families \u2014 *stellar* mass profiles (e.g. `al.mp.Sersic`, `al.mp.Chameleon`,\n", - "`al.mp.GaussianGradient`), the *dark-matter* NFW family (`al.mp.NFW`, `al.mp.gNFW`,\n", - "`al.mp.cNFW`, all their MCR / virial / scatter variants), and the combined *light-and-mass*\n", - "namespaces `al.lmp.*` and `al.lmp_linear.*` \u2014 are the subject of a separate companion\n", - "guide:\n", - "\n", - " scripts/guides/profiles/light_and_mass_profiles.py\n", - "\n", - "That guide tells the stellar-plus-dark decomposition story: how a Sersic mass component\n", - "representing the stellar matter combines with an NFW component representing the dark matter\n", - "halo, how the `lmp` profiles tie a Sersic *light* to a Sersic *mass* through a shared\n", - "mass-to-light ratio, and how those compositions feed into the standard model / instance /\n", - "Tracer flow shown here.\n", - "\n", - "If you arrived at this guide from the API reference and now want to use any of these mass\n", - "profiles in an actual lens fit, the next step is `scripts/imaging/modeling/start_here.py`,\n", - "which sets up an `AnalysisImaging` and runs a non-linear search end-to-end on a strong-lens\n", - "dataset. For a deeper walk-through of how mass profiles combine with light profiles in the\n", - "`Tracer` to produce lensed images, see `scripts/guides/tracer.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Mass Profiles\n", + "=============\n", + "\n", + "This guide is the single-page tour of every lensing mass profile available in **PyAutoLens**\n", + "(re-exported from **PyAutoGalaxy**): how to construct each one, how to evaluate its\n", + "convergence and deflections on a grid, how to compose it into a model, and how to pull an\n", + "instance back out of that model.\n", + "\n", + "It is the companion to `scripts/guides/profiles/light.py` \u2014 the section flow, prose style\n", + "and \"detailed example then walkthrough\" rhythm intentionally mirror that guide so the two\n", + "can be read as a pair. Where `light.py` shows `image_2d_from`, this guide shows\n", + "`convergence_2d_from` and `deflections_yx_2d_from`; mass profiles do not produce images of\n", + "their own (a mass profile only deflects light from a *source* through ray-tracing \u2014 that is\n", + "what the `Tracer` is for, and we use it for the detailed example).\n", + "\n", + "This guide covers the *lensing* mass profiles \u2014 Total, Mass Sheets, Multipoles, and Point\n", + "Mass. The *stellar* (`al.mp.Sersic`, `al.mp.Chameleon`, ...) and *dark-matter* (NFW family)\n", + "mass profiles, along with the combined light+mass profiles in `al.lmp.*` / `al.lmp_linear.*`,\n", + "get their own dedicated guide at `scripts/guides/profiles/light_and_mass_profiles.py` where\n", + "the stellar-plus-dark decomposition story is told properly.\n", + "\n", + "__Contents__\n", + "\n", + "- **Overview & Docs URL:** Where the canonical API reference lives.\n", + "- **All Mass Profiles (Survey):** A high-level run-through of every profile in `al.mp.*`\n", + " (Total / Mass Sheets / Multipoles / Point Mass), without yet evaluating any quantities.\n", + "- **Detailed Example: Isothermal:** Build a `Grid2D`, instantiate `al.mp.Isothermal`, plot\n", + " convergence, potential, deflection-magnitude, and the lensed source image produced when\n", + " the isothermal mass is dropped into a `Tracer`.\n", + "- **Mass Sheets:** `ExternalShear`, `MassSheet`, `ExternalPotential` \u2014 global perturbations\n", + " rather than parametric matter distributions.\n", + "- **Point Mass:** `PointMass`, `SMBH`, `SMBHBinary` \u2014 delta-function-like mass for\n", + " microlensing and supermassive black hole lensing.\n", + "- **Mass Profile in a Model:** Wrap a profile in `af.Model`, compose lens + source via\n", + " `af.Collection`, inspect the model info.\n", + "- **Model Instance from Mass Profile:** Realise an instance from the model's prior medians\n", + " and drop it into a `Tracer` to produce a lensed image.\n", + "- **Multipole Mass Profile:** `PowerLawMultipole` \u2014 m=3 / m=4 Fourier perturbation on the\n", + " power-law convergence, the lensing counterpart of `lp.SersicMultipole`.\n", + "- **Remaining Profiles Walkthrough:** Compact `convergence_2d_from` block for every total\n", + " profile not yet shown.\n", + "- **Follow-Up:** Pointer to `light_and_mass_profiles.py` for stellar / dark / lmp coverage.\n", + "\n", + "__Units__\n", + "\n", + "Spatial coordinates are in arc-seconds, mass quantities are dimensionless (convergence) or\n", + "in their natural lensing units (potential, deflection angles in arc-seconds). The\n", + "`guides/units_and_cosmology.ipynb` guide covers conversion to physical units.\n", + "\n", + "__Data Structures__\n", + "\n", + "`convergence_2d_from` and `potential_2d_from` return `Array2D`. `deflections_yx_2d_from`\n", + "returns `VectorYX2D` (a 2D vector field). `aplt.plot_array` accepts `Array2D` directly;\n", + "to plot a vector field we either pull out the y / x components individually or compute the\n", + "magnitude and wrap it in an `Array2D`. This guide uses the magnitude approach for the\n", + "detailed example.\n", + "\n", + "__Docs URL__\n", + "\n", + "The published API reference for these classes lives at:\n", + "\n", + " https://pyautolens.readthedocs.io/en/latest/api/mass.html\n", + "\n", + "The autosummary on that page is the authoritative list of every public mass-profile class.\n", + "Note that the reference uses the `ag.mp` namespace label because the classes are defined in\n", + "PyAutoGalaxy and re-exported here as `al.mp`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "\n", + "import autofit as af\n", + "import autoarray as aa\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "To evaluate any quantity on a mass profile we need a 2D Cartesian grid of (y,x) coordinates.\n", + "We build a 100x100 grid here at a 0.05\" pixel scale \u2014 used by every section below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.05,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__All Mass Profiles (Survey)__\n", + "\n", + "**PyAutoLens** groups mass profiles into four families relevant to lensing:\n", + "\n", + "- `al.mp.*` total profiles \u2014 the canonical parametric lens mass distributions\n", + " (`Isothermal`, `PowerLaw`, their cored / broken / dPIE variants).\n", + "- `al.mp.*` mass sheets \u2014 `ExternalShear`, `MassSheet`, `ExternalPotential`. Global\n", + " perturbations representing line-of-sight contributions.\n", + "- `al.mp.*` multipoles \u2014 `PowerLawMultipole`, the m=3 / m=4 angular extension of the power-\n", + " law family.\n", + "- `al.mp.*` point masses \u2014 `PointMass`, `SMBH`, `SMBHBinary` \u2014 delta-function-like sources\n", + " for microlensing and central black holes.\n", + "\n", + "The *stellar* and *dark-matter* mass profiles (e.g. `al.mp.Sersic`, the NFW family) and the\n", + "combined `al.lmp.*` / `al.lmp_linear.*` light-and-mass namespaces are covered in\n", + "`scripts/guides/profiles/light_and_mass_profiles.py`. They share the same API but tell a\n", + "different story (stellar-plus-dark decomposition).\n", + "\n", + "Below we construct each lensing profile with default parameters. No quantity is evaluated\n", + "yet \u2014 that comes in the next section. The goal here is purely a catalogue of what is\n", + "available." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Power-law family (Isothermal is PowerLaw with slope = 2)\n", + "isothermal = al.mp.Isothermal()\n", + "isothermal_sph = al.mp.IsothermalSph()\n", + "isothermal_core = al.mp.IsothermalCore()\n", + "isothermal_core_sph = al.mp.IsothermalCoreSph()\n", + "power_law = al.mp.PowerLaw()\n", + "power_law_sph = al.mp.PowerLawSph()\n", + "power_law_core = al.mp.PowerLawCore()\n", + "power_law_core_sph = al.mp.PowerLawCoreSph()\n", + "power_law_broken = al.mp.PowerLawBroken()\n", + "power_law_broken_sph = al.mp.PowerLawBrokenSph()\n", + "\n", + "# Pseudo-isothermal family (mass and potential parameterisations)\n", + "# Note: dPIEMass with default ell_comps=(0,0) triggers a divide-by-zero in the complex-plane\n", + "# formula; we use a small ellipticity here so the survey constructions succeed cleanly.\n", + "dpie_mass = al.mp.dPIEMass(ell_comps=(0.05, 0.0))\n", + "dpie_mass_sph = al.mp.dPIEMassSph()\n", + "pie_mass = al.mp.PIEMass(ell_comps=(0.05, 0.0))\n", + "dpie_potential = al.mp.dPIEPotential()\n", + "dpie_potential_sph = al.mp.dPIEPotentialSph()\n", + "\n", + "# Mass sheets \u2014 global perturbations\n", + "external_shear = al.mp.ExternalShear()\n", + "mass_sheet = al.mp.MassSheet()\n", + "external_potential = al.mp.ExternalPotential()\n", + "\n", + "# Multipole \u2014 power-law angular extension\n", + "power_law_multipole = al.mp.PowerLawMultipole()\n", + "\n", + "# Point masses \u2014 microlensing and supermassive black holes\n", + "point_mass = al.mp.PointMass()\n", + "smbh = al.mp.SMBH()\n", + "smbh_binary = al.mp.SMBHBinary()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Two things worth knowing about this list before we move on:\n", + "\n", + "1. Every elliptical lensing profile (e.g. `Isothermal`, `PowerLaw`, `PowerLawCore`) has a\n", + " spherical sibling whose name ends in `Sph` (e.g. `IsothermalSph`). The spherical variant\n", + " fixes the ellipticity components `ell_comps` to `(0, 0)`, which is useful when you want\n", + " to model a circular halo and avoid two redundant parameters in the non-linear search.\n", + "2. The `PowerLawMultipole` variant only exists as an *elliptical* profile \u2014 the m=3 / m=4\n", + " perturbations are angular distortions and are not meaningful without an underlying\n", + " reference frame, exactly like the light multipoles in `light.py`.\n", + "\n", + "We now move on to seeing what these profiles actually produce when evaluated on a grid.\n", + "\n", + "__Detailed Example: Isothermal__\n", + "\n", + "The `Isothermal` profile is the canonical strong-lens mass profile \u2014 equivalent to the\n", + "elliptical power-law with `slope = 2.0`. Its three parameters are:\n", + "\n", + "- `centre` \u2014 the (y, x) arc-second coordinate of the profile centre.\n", + "- `ell_comps` \u2014 the two ellipticity components `(e1, e2)`. Use\n", + " `al.convert.ell_comps_from(axis_ratio=..., angle=...)` to convert from human-friendly\n", + " axis ratio and position angle.\n", + "- `einstein_radius` \u2014 the Einstein radius in arc-seconds. This sets the strength of the\n", + " lens and the size of the Einstein ring for an axisymmetric source on-axis.\n", + "\n", + "Build an isothermal and evaluate every standard lensing quantity on our grid:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "isothermal = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " einstein_radius=1.6,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "First the *convergence* \u2014 the projected surface mass density in units of the critical\n", + "surface density. This is the simplest visualisation of a mass profile and the most\n", + "analogous to the `image_2d_from` plot we used throughout `light.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=isothermal.convergence_2d_from(grid=grid),\n", + " title=\"Isothermal Convergence\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next the *lensing potential* \u2014 the 2D field whose gradient gives the deflection angles.\n", + "We plot it in log10 to make the central cusp visible." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=isothermal.potential_2d_from(grid=grid),\n", + " title=\"Isothermal Potential (log10)\",\n", + " use_log10=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The *deflection angles* are a vector field rather than a scalar, so `deflections_yx_2d_from`\n", + "returns a `VectorYX2D` rather than an `Array2D`. `aplt.plot_array` does not accept vectors\n", + "directly; the conventional pattern in this workspace is to either plot the y / x components\n", + "separately or compute the magnitude. Below we use the magnitude \u2014 a single map showing\n", + "\"how much light is bent here\", which is the most useful single-figure summary." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections = isothermal.deflections_yx_2d_from(grid=grid)\n", + "deflection_magnitude = aa.Array2D(\n", + " values=np.hypot(deflections.slim[:, 0], deflections.slim[:, 1]),\n", + " mask=grid.mask,\n", + ")\n", + "aplt.plot_array(\n", + " array=deflection_magnitude,\n", + " title=\"Isothermal Deflection Magnitude\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "None of these maps show what the isothermal mass actually *does* to a background source \u2014\n", + "for that we need a `Tracer`. A `Tracer` groups galaxies by redshift plane, ray-traces the\n", + "grid through every plane, and combines the resulting emission into a single observed image.\n", + "\n", + "Below we build a two-plane lens system: a lens galaxy at `z=0.5` carrying our isothermal\n", + "mass profile, and a source galaxy at `z=1.0` with a small Sersic light profile. The\n", + "`tracer.image_2d_from` call produces the lensed image you would see at the telescope." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(redshift=0.5, mass=isothermal)\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.05, 0.05),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=0.1,\n", + " sersic_index=1.5,\n", + " ),\n", + ")\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Tracer Image (Isothermal Lens + Sersic Source)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The arc / ring you see is the source's Sersic light, deflected by the isothermal mass into\n", + "the characteristic strong-lens morphology. The same `image_2d_from(grid=grid)` call exists\n", + "on every `Tracer` regardless of which mass profile is on the lens galaxy \u2014 every section\n", + "below is a small variation on this one, swapping the lens-galaxy mass for a different\n", + "profile.\n", + "\n", + "__Mass Sheets__\n", + "\n", + "Mass sheets are *global* perturbations rather than localised mass distributions. They show\n", + "up in strong-lens models to capture line-of-sight contributions: nearby group / cluster\n", + "members, large-scale structure, and any other mass that is not part of the primary lens but\n", + "still distorts the light path.\n", + "\n", + "Three mass-sheet profiles are available:\n", + "\n", + "- `al.mp.ExternalShear(gamma_1, gamma_2)` \u2014 a constant shear with two components. The most\n", + " common line-of-sight correction; you will see this routinely added to lens-mass models.\n", + "- `al.mp.MassSheet(centre, kappa)` \u2014 a uniform convergence (positive or negative). Useful\n", + " for representing diffuse line-of-sight mass.\n", + "- `al.mp.ExternalPotential(centre, gamma_1, gamma_2, tau_1, tau_2, delta_1, delta_2)` \u2014 the\n", + " Powell et al. 2022 generalisation of external shear that adds the next-order line-of-sight\n", + " terms. Reduces to `ExternalShear` when `tau_*` and `delta_*` are zero." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "external_shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.03)\n", + "deflections = external_shear.deflections_yx_2d_from(grid=grid)\n", + "aplt.plot_array(\n", + " array=aa.Array2D(\n", + " values=np.hypot(deflections.slim[:, 0], deflections.slim[:, 1]),\n", + " mask=grid.mask,\n", + " ),\n", + " title=\"ExternalShear Deflection Magnitude\",\n", + ")\n", + "\n", + "mass_sheet = al.mp.MassSheet(centre=(0.0, 0.0), kappa=0.1)\n", + "aplt.plot_array(\n", + " array=mass_sheet.convergence_2d_from(grid=grid),\n", + " title=\"MassSheet Convergence (uniform kappa=0.1)\",\n", + ")\n", + "\n", + "external_potential = al.mp.ExternalPotential(\n", + " centre=(0.0, 0.0),\n", + " gamma_1=0.04,\n", + " gamma_2=0.02,\n", + " tau_1=0.005,\n", + " tau_2=0.0,\n", + " delta_1=0.0,\n", + " delta_2=0.003,\n", + ")\n", + "deflections = external_potential.deflections_yx_2d_from(grid=grid)\n", + "aplt.plot_array(\n", + " array=aa.Array2D(\n", + " values=np.hypot(deflections.slim[:, 0], deflections.slim[:, 1]),\n", + " mask=grid.mask,\n", + " ),\n", + " title=\"ExternalPotential Deflection Magnitude\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A mass sheet alone produces no observable lensing \u2014 its deflections are uniform (constant)\n", + "across the image and so are absorbed into the global astrometric solution. Mass sheets only\n", + "become visible *in combination* with another mass profile; their job is to perturb the\n", + "isothermal or power-law deflections at the few-percent level.\n", + "\n", + "__Point Mass__\n", + "\n", + "Point masses are delta-function-like mass distributions. Three flavours exist:\n", + "\n", + "- `al.mp.PointMass(centre, einstein_radius)` \u2014 a single point mass parameterised by its\n", + " Einstein radius. Used for microlensing.\n", + "- `al.mp.SMBH(centre, mass, redshift_object, redshift_source)` \u2014 a supermassive black hole\n", + " parameterised by mass (in solar masses); internally converts to a `PointMass` Einstein\n", + " radius using the input redshifts.\n", + "- `al.mp.SMBHBinary(centre, separation, angle_binary, mass, mass_ratio, ...)` \u2014 two SMBHs\n", + " separated by a tunable angle and separation. Used for binary SMBH lensing.\n", + "\n", + "`PointMass.convergence_2d_from` returns a raw numpy array rather than the `Array2D` that\n", + "elliptical profiles return \u2014 a library quirk. We wrap the point mass in a `Galaxy` and plot\n", + "the galaxy's convergence map instead, which goes through the standard `Array2D` path and is\n", + "also the natural usage pattern for a lensing model anyway." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "point_mass = al.mp.PointMass(centre=(0.0, 0.0), einstein_radius=0.3)\n", + "point_galaxy = al.Galaxy(redshift=0.5, mass=point_mass)\n", + "\n", + "aplt.plot_array(\n", + " array=point_galaxy.convergence_2d_from(grid=grid),\n", + " title=\"PointMass Convergence (delta-function-like)\",\n", + ")\n", + "\n", + "smbh = al.mp.SMBH(\n", + " centre=(0.0, 0.0),\n", + " mass=1.0e9,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + ")\n", + "smbh_galaxy = al.Galaxy(redshift=0.5, mass=smbh)\n", + "aplt.plot_array(\n", + " array=smbh_galaxy.convergence_2d_from(grid=grid),\n", + " title=\"SMBH Convergence (M = 1e9 M_sun)\",\n", + ")\n", + "\n", + "smbh_binary = al.mp.SMBHBinary(\n", + " centre=(0.0, 0.0),\n", + " separation=0.4,\n", + " angle_binary=45.0,\n", + " mass=2.0e9,\n", + " mass_ratio=0.5,\n", + " redshift_object=0.5,\n", + " redshift_source=1.0,\n", + ")\n", + "smbh_binary_galaxy = al.Galaxy(redshift=0.5, mass=smbh_binary)\n", + "aplt.plot_array(\n", + " array=smbh_binary_galaxy.convergence_2d_from(grid=grid),\n", + " title=\"SMBHBinary Convergence\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mass Profile in a Model__\n", + "\n", + "So far we have been instantiating mass profiles with concrete parameter values. When\n", + "fitting a real strong-lens dataset we instead build a *model* of the mass profile and let\n", + "the non-linear search find the best-fit parameters. This is what `af.Model` is for.\n", + "\n", + "For a strong-lens system the model typically contains two `Galaxy` objects on different\n", + "planes: a foreground lens (light + mass) and a background source (light only). In this\n", + "section we focus on the *mass* side of that picture and show how an isothermal mass profile\n", + "plugs in to a lens galaxy that has only mass (no light) \u2014 the simplest case." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_mass_model = af.Model(al.mp.Isothermal)\n", + "source_bulge_model = af.Model(al.lp.Sersic)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `af.Model` wraps the profile class. Every constructor argument with a numerical default\n", + "becomes a *prior* \u2014 by default the priors are sensible distributions for each parameter\n", + "(see the autogalaxy config for the configured ranges).\n", + "\n", + "You can override priors before fitting:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_mass_model.einstein_radius = af.UniformPrior(lower_limit=0.5, upper_limit=3.0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Wrap each profile in a `Galaxy` at its own redshift and assemble them in an `af.Collection`.\n", + "The `Tracer` is not part of the *model* spec itself; it is what `AnalysisImaging` builds out\n", + "of the realised instance at fit time." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy_model = af.Model(al.Galaxy, redshift=0.5, mass=lens_mass_model)\n", + "source_galaxy_model = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge_model)\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens_galaxy_model, source=source_galaxy_model)\n", + ")\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing `model.info` shows the full priors-and-defaults summary \u2014 useful before kicking\n", + "off a long fit to confirm the model is shaped the way you expect.\n", + "\n", + "The model API is the same for **every** mass profile in this guide \u2014 swap\n", + "`al.mp.Isothermal` for `al.mp.PowerLaw`, `al.mp.dPIEMass`, `al.mp.PowerLawMultipole`, etc.,\n", + "and the rest of the snippet is unchanged.\n", + "\n", + "Full lens-modelling end-to-end examples live in `scripts/imaging/modeling/start_here.py`\n", + "and the topic-specific guides under `scripts/imaging/features/`.\n", + "\n", + "__Model Instance from Mass Profile__\n", + "\n", + "A model is a description of *possible* profiles. To get an actual profile back out \u2014 for\n", + "example to plot what the prior medians look like before running a fit \u2014 call\n", + "`instance_from_prior_medians()`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_mass_instance = lens_mass_model.instance_from_prior_medians()\n", + "print(type(lens_mass_instance)) # autogalaxy.profiles.mass.total.isothermal.Isothermal\n", + "\n", + "aplt.plot_array(\n", + " array=lens_mass_instance.convergence_2d_from(grid=grid),\n", + " title=\"Isothermal Instance Convergence\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The instance returned is a real `al.mp.Isothermal` \u2014 the same class we constructed by hand\n", + "in the detailed example above \u2014 and supports the full mass-profile API.\n", + "\n", + "The same flow works at the full-model level. We realise an instance of the lens-and-source\n", + "collection and drop the resulting galaxies into a `Tracer`, which combines the lens mass\n", + "deflections with the source light to produce the lensed image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_instance = model.instance_from_prior_medians()\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[model_instance.galaxies.lens, model_instance.galaxies.source]\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Tracer Image from Model Instance\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "After a fit completes, `result.max_log_likelihood_tracer` returns the same shape of object,\n", + "with the prior medians replaced by the fitted parameter values. See\n", + "`scripts/guides/results/start_here.py` for the full results-introspection guide.\n", + "\n", + "__Multipole Mass Profile__\n", + "\n", + "`PowerLawMultipole` is the lensing counterpart of `lp.SersicMultipole` \u2014 an m=3 / m=4\n", + "Fourier angular perturbation, here on the eccentric radius of the power-law convergence\n", + "field rather than the Sersic intensity field. The signature is\n", + "\n", + " al.mp.PowerLawMultipole(\n", + " m=4, # multipole order (3 or 4)\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.0,\n", + " slope=2.0,\n", + " multipole_comps=(c, s), # cosine / sine components of the m-th term\n", + " )\n", + "\n", + "A `PowerLawMultipole` produces only the *perturbation* \u2014 not the underlying power-law mass\n", + "distribution. The standard usage is therefore to add the multipole alongside a `PowerLaw`\n", + "or `Isothermal` profile sharing the same `centre`, `einstein_radius`, and `slope`, so that\n", + "the combined convergence reads as a power law with boxy / discy / lopsided distortions.\n", + "\n", + "Build a `PowerLawMultipole` with non-trivial multipole components and plot its convergence\n", + "alongside the underlying `PowerLaw` for comparison:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "power_law_base = al.mp.PowerLaw(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " einstein_radius=1.6,\n", + " slope=2.0,\n", + ")\n", + "\n", + "power_law_multipole_m4 = al.mp.PowerLawMultipole(\n", + " m=4,\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " slope=2.0,\n", + " multipole_comps=(0.05, 0.02),\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=power_law_base.convergence_2d_from(grid=grid),\n", + " title=\"PowerLaw Convergence (base profile)\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=power_law_multipole_m4.convergence_2d_from(grid=grid),\n", + " title=\"PowerLawMultipole Convergence (m=4 perturbation only)\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Two practical notes on the multipole:\n", + "\n", + "- The multipole profile carries the *perturbation*, not the base mass distribution. In a\n", + " model you typically wrap both `PowerLaw` and `PowerLawMultipole` on the same lens galaxy\n", + " and link their shared parameters (`centre`, `einstein_radius`, `slope`).\n", + "- There is **no spherical (`*Sph`) variant** of `PowerLawMultipole`. The perturbation is\n", + " an angular distortion measured in the lens reference frame, so it only makes sense for an\n", + " elliptical (or quasi-elliptical) host.\n", + "\n", + "Plugging `PowerLawMultipole` into the `af.Model` / `af.Collection` / `Galaxy` / `Tracer`\n", + "pattern shown above works exactly as it did for the plain `Isothermal` \u2014 the multipole\n", + "components are picked up as priors automatically.\n", + "\n", + "__Remaining Profiles Walkthrough__\n", + "\n", + "We have shown the full `convergence_2d_from` \u2192 `af.Model` \u2192 `instance` \u2192 `Tracer` flow for\n", + "the `Isothermal` profile. Every remaining lensing mass profile uses the **same API** \u2014 the\n", + "only thing that changes is which parameters appear in the constructor.\n", + "\n", + "The compact tour below builds each remaining total profile with sensible parameter values\n", + "and plots its convergence. When you want to use any of these in a lens model, repeat the\n", + "`af.Model(...)` / `af.Collection(...)` / `Tracer(...)` pattern from the previous section." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(\n", + " array=al.mp.PowerLaw(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " einstein_radius=1.6,\n", + " slope=2.2,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"PowerLaw Convergence (slope=2.2)\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.PowerLawCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " einstein_radius=1.6,\n", + " slope=2.0,\n", + " core_radius=0.05,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"PowerLawCore Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.PowerLawBroken(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " einstein_radius=1.6,\n", + " inner_slope=1.5,\n", + " outer_slope=2.5,\n", + " break_radius=0.5,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"PowerLawBroken Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.IsothermalCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " einstein_radius=1.6,\n", + " core_radius=0.05,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"IsothermalCore Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.dPIEMass(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ra=0.1,\n", + " b0=0.5,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"dPIEMass Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.PIEMass(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ra=0.1,\n", + " b0=0.5,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"PIEMass Convergence\",\n", + ")\n", + "\n", + "aplt.plot_array(\n", + " array=al.mp.dPIEPotential(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.85, angle=45.0),\n", + " ra=0.1,\n", + " rs=2.0,\n", + " b0=1.0,\n", + " ).convergence_2d_from(grid=grid),\n", + " title=\"dPIEPotential Convergence\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The spherical variants (`IsothermalSph`, `PowerLawSph`, `PowerLawCoreSph`,\n", + "`PowerLawBrokenSph`, `dPIEMassSph`) are constructed identically with the `ell_comps`\n", + "argument removed. Each looks like a rotationally symmetric version of its elliptical\n", + "counterpart.\n", + "\n", + "__Follow-Up: Stellar, Dark Matter, and Combined Light+Mass Profiles__\n", + "\n", + "This guide deliberately stops at the parametric lensing mass profiles. The remaining mass\n", + "families \u2014 *stellar* mass profiles (e.g. `al.mp.Sersic`, `al.mp.Chameleon`,\n", + "`al.mp.GaussianGradient`), the *dark-matter* NFW family (`al.mp.NFW`, `al.mp.gNFW`,\n", + "`al.mp.cNFW`, all their MCR / virial / scatter variants), and the combined *light-and-mass*\n", + "namespaces `al.lmp.*` and `al.lmp_linear.*` \u2014 are the subject of a separate companion\n", + "guide:\n", + "\n", + " scripts/guides/profiles/light_and_mass_profiles.py\n", + "\n", + "That guide tells the stellar-plus-dark decomposition story: how a Sersic mass component\n", + "representing the stellar matter combines with an NFW component representing the dark matter\n", + "halo, how the `lmp` profiles tie a Sersic *light* to a Sersic *mass* through a shared\n", + "mass-to-light ratio, and how those compositions feed into the standard model / instance /\n", + "Tracer flow shown here.\n", + "\n", + "If you arrived at this guide from the API reference and now want to use any of these mass\n", + "profiles in an actual lens fit, the next step is `scripts/imaging/modeling/start_here.py`,\n", + "which sets up an `AnalysisImaging` and runs a non-linear search end-to-end on a strong-lens\n", + "dataset. For a deeper walk-through of how mass profiles combine with light profiles in the\n", + "`Tracer` to produce lensed images, see `scripts/guides/tracer.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/_quick_fit.ipynb b/notebooks/guides/results/_quick_fit.ipynb index e303e19ef..d111dce90 100644 --- a/notebooks/guides/results/_quick_fit.ipynb +++ b/notebooks/guides/results/_quick_fit.ipynb @@ -1,190 +1,227 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Quick Fit Helper\n", - "=========================\n", - "\n", - "Internal helper invoked via subprocess from the tutorials in this folder.\n", - "Produces two fast, capped Nautilus fits at ``output/results_folder/`` so the\n", - "aggregator and workflow examples have a populated results directory to read\n", - "from.\n", - "\n", - "Idempotent: exits immediately if ``output/results_folder/`` already contains\n", - "the two completed imaging fits, so concurrent or repeated invocations are\n", - "cheap.\n", - "\n", - "Not a tutorial. The model and dataset mirror those used in ``start_here.py``\n", - "(``simple__no_lens_light`` imaging, isothermal lens + MGE source), but the\n", - "search is hard-capped at ``n_like_max=300`` likelihood evaluations rather\n", - "than running to convergence. This produces a shallow but valid posterior\n", - "fast enough to fit inside the per-script CI timeout." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "import shutil\n", - "import sys\n", - "from pathlib import Path\n", - "\n", - "\n", - "def has_workflow_results(results_path):\n", - " return (\n", - " len(list(results_path.glob(\"**/image/dataset.fits\"))) >= 2\n", - " and len(list(results_path.glob(\"**/files/latent/latent_summary.json\"))) >= 2\n", - " and len(list(results_path.glob(\"**/image/fit.png\"))) >= 2\n", - " and len(list(results_path.glob(\"**/image/fit.fits\"))) >= 2\n", - " )\n", - "\n", - "\n", - "results_path = Path(\"output\") / \"results_folder\"\n", - "if has_workflow_results(results_path):\n", - " sys.exit(0)\n", - "\n", - "if results_path.exists():\n", - " shutil.rmtree(results_path)\n", - "\n", - "import os\n", - "\n", - "# The aggregator tutorials that invoke this helper read image/dataset.fits via\n", - "# fit.value(\"dataset\"). Visualization-skipping environment variables suppress\n", - "# the visualizer that writes that file, so neutralize them here.\n", - "os.environ.pop(\"PYAUTO_SKIP_VISUALIZATION\", None)\n", - "os.environ.pop(\"PYAUTO_SKIP_FIT_OUTPUT\", None)\n", - "os.environ.pop(\"PYAUTO_FAST_PLOTS\", None)\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "from autoconf import conf\n", - "\n", - "# This deliberately shallow helper must retain its exploratory samples because\n", - "# the results tutorials demonstrate indexed sample access.\n", - "conf.instance[\"output\"][\"samples_weight_threshold\"] = None\n", - "\n", - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - " ),\n", - " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=None),\n", - " ),\n", - ")\n", - "\n", - "\n", - "class LatentShear(al.Latent):\n", - " \"\"\"\n", - " Custom latent catalogue reporting the lens external-shear magnitude and\n", - " angle for the workflow CSV example.\n", - " \"\"\"\n", - "\n", - " @staticmethod\n", - " def keys(analysis):\n", - " return [\n", - " \"galaxies.lens.shear.magnitude\",\n", - " \"galaxies.lens.shear.angle\",\n", - " ]\n", - "\n", - " @staticmethod\n", - " def variables(analysis, parameters, model):\n", - " instance = model.instance_from_vector(vector=parameters)\n", - "\n", - " import jax.numpy as jnp\n", - "\n", - " magnitude, angle = al.convert.shear_magnitude_and_angle_from(\n", - " gamma_1=instance.galaxies.lens.shear.gamma_1,\n", - " gamma_2=instance.galaxies.lens.shear.gamma_2,\n", - " xp=jnp,\n", - " )\n", - "\n", - " return (magnitude, angle)\n", - "\n", - "\n", - "class AnalysisLatent(al.AnalysisImaging):\n", - " Latent = LatentShear\n", - "\n", - "\n", - "analysis = AnalysisLatent(dataset=dataset, use_jax=True)\n", - "\n", - "for i in range(2):\n", - " search = af.Nautilus(\n", - " path_prefix=Path(\"results_folder\"),\n", - " name=\"results\",\n", - " unique_tag=f\"{dataset_name}_{i}\",\n", - " n_live=100,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " n_like_max=300,\n", - " )\n", - "\n", - " search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Quick Fit Helper\n", + "=========================\n", + "\n", + "Internal helper invoked via subprocess from the tutorials in this folder.\n", + "Produces two fast, capped Nautilus fits at ``output/results_folder/`` so the\n", + "aggregator and workflow examples have a populated results directory to read\n", + "from.\n", + "\n", + "Idempotent: exits immediately if ``output/results_folder/`` already contains\n", + "the two completed imaging fits, so concurrent or repeated invocations are\n", + "cheap.\n", + "\n", + "Not a tutorial. The model and dataset mirror those used in ``start_here.py``\n", + "(``simple__no_lens_light`` imaging, isothermal lens + MGE source), but the\n", + "search is hard-capped at ``n_like_max=300`` likelihood evaluations rather\n", + "than running to convergence. This produces a shallow but valid posterior\n", + "fast enough to fit inside the per-script CI timeout." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "import shutil\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "\n", + "def has_workflow_results(results_path):\n", + " return (\n", + " len(list(results_path.glob(\"**/image/dataset.fits\"))) >= 2\n", + " and len(list(results_path.glob(\"**/files/latent/latent_summary.json\"))) >= 2\n", + " and len(list(results_path.glob(\"**/image/fit.png\"))) >= 2\n", + " and len(list(results_path.glob(\"**/image/fit.fits\"))) >= 2\n", + " )\n", + "\n", + "\n", + "results_path = Path(\"output\") / \"results_folder\"\n", + "if has_workflow_results(results_path):\n", + " sys.exit(0)\n", + "\n", + "if results_path.exists():\n", + " shutil.rmtree(results_path)\n", + "\n", + "import os\n", + "\n", + "# The aggregator tutorials that invoke this helper read image/dataset.fits via\n", + "# fit.value(\"dataset\"). Visualization-skipping environment variables suppress\n", + "# the visualizer that writes that file, so neutralize them here.\n", + "os.environ.pop(\"PYAUTO_SKIP_VISUALIZATION\", None)\n", + "os.environ.pop(\"PYAUTO_SKIP_FIT_OUTPUT\", None)\n", + "os.environ.pop(\"PYAUTO_FAST_PLOTS\", None)\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "from autoconf import conf\n", + "\n", + "# This deliberately shallow helper must retain its exploratory samples because\n", + "# the results tutorials demonstrate indexed sample access.\n", + "conf.instance[\"output\"][\"samples_weight_threshold\"] = None\n", + "\n", + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + " ),\n", + " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=None),\n", + " ),\n", + ")\n", + "\n", + "\n", + "class LatentShear(al.Latent):\n", + " \"\"\"\n", + " Custom latent catalogue reporting the lens external-shear magnitude and\n", + " angle for the workflow CSV example.\n", + " \"\"\"\n", + "\n", + " @staticmethod\n", + " def keys(analysis):\n", + " return [\n", + " \"galaxies.lens.shear.magnitude\",\n", + " \"galaxies.lens.shear.angle\",\n", + " ]\n", + "\n", + " @staticmethod\n", + " def variables(analysis, parameters, model):\n", + " instance = model.instance_from_vector(vector=parameters)\n", + "\n", + " import jax.numpy as jnp\n", + "\n", + " magnitude, angle = al.convert.shear_magnitude_and_angle_from(\n", + " gamma_1=instance.galaxies.lens.shear.gamma_1,\n", + " gamma_2=instance.galaxies.lens.shear.gamma_2,\n", + " xp=jnp,\n", + " )\n", + "\n", + " return (magnitude, angle)\n", + "\n", + "\n", + "class AnalysisLatent(al.AnalysisImaging):\n", + " Latent = LatentShear\n", + "\n", + "\n", + "analysis = AnalysisLatent(dataset=dataset, use_jax=True)\n", + "\n", + "for i in range(2):\n", + " search = af.Nautilus(\n", + " path_prefix=Path(\"results_folder\"),\n", + " name=\"results\",\n", + " unique_tag=f\"{dataset_name}_{i}\",\n", + " n_live=100,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " n_like_max=300,\n", + " )\n", + "\n", + " search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/aggregator/data_fitting.ipynb b/notebooks/guides/results/aggregator/data_fitting.ipynb index e15966e66..dbb393b30 100644 --- a/notebooks/guides/results/aggregator/data_fitting.ipynb +++ b/notebooks/guides/results/aggregator/data_fitting.ipynb @@ -1,351 +1,388 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Data Fitting\n", - "=====================\n", - "\n", - "In this tutorial, we use the aggregator to load models and data from a non-linear search and use them to perform\n", - "fits to the data.\n", - "\n", - "We show how to use these tools to inspect the maximum log likelihood model of a fit to the data, customize things\n", - "like its visualization and also inspect fits randomly drawm from the PDF.\n", - "\n", - "__Contents__\n", - "\n", - "- **Interferometer:** This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", - "- **Aggregator:** First, set up the aggregator as shown in `start_here.py`.\n", - "- **Fits via Aggregator:** Having performed a model-fit, we now want to interpret and visualize the results.\n", - "- **Modification:** The `FitImagingAgg` allow us to modify the fit settings.\n", - "- **Visualization Customization:** The benefit of inspecting fits using the aggregator, rather than the files outputs to the.\n", - "\n", - "__Interferometer__\n", - "\n", - "This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", - "\n", - "The only entries that needs changing are:\n", - "\n", - " - `ImagingAgg` -> `InterferometerAgg`.\n", - " - `FitImagingAgg` -> `FitInterferometerAgg`.\n", - " - `aplt.subplot_imaging_dataset` -> `aplt.subplot_interferometer_dirty_images`.\n", - " - `aplt.subplot_fit_imaging` -> `aplt.subplot_fit_interferometer`.\n", - "\n", - "Quantities specific to an interfometer, for example its uv-wavelengths real space mask, are accessed using the same API\n", - "(e.g. `values(\"dataset.uv_wavelengths\")` and `.values{\"dataset.real_space_mask\"))." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import os\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Aggregator__\n", - "\n", - "First, set up the aggregator as shown in `start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "results_path = Path(\"output\") / \"results_folder\"\n", - "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "agg = Aggregator.from_directory(\n", - " directory=results_path,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The masks we used to fit the lenses is accessible via the aggregator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_gen = agg.values(\"dataset.mask\")\n", - "print([mask for mask in mask_gen])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The info dictionary we passed is also available." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Info:\")\n", - "info_gen = agg.values(\"info\")\n", - "print([info for info in info_gen])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fits via Aggregator__\n", - "\n", - "Having performed a model-fit, we now want to interpret and visualize the results. In this example, we inspect \n", - "the `Imaging` objects that gave good fits to the data. \n", - "\n", - "Using the API shown in the `start_here.py` example this would require us to create a `Samples` object and manually \n", - "compose our own `Imaging` object. For large datasets, this would require us to use generators to ensure it is \n", - "memory-light, which are cumbersome to write.\n", - "\n", - "This example therefore uses the `ImagingAgg` object, which conveniently loads the `Imaging` objects of every fit via \n", - "generators for us. \n", - "\n", - "We get a dataset generator via the `al.agg.ImagingAgg` object, where this `dataset_gen` contains the maximum log\n", - "likelihood `Imaging `object of every model-fit.\n", - "\n", - "The `dataset_gen` returns a list of `Imaging` objects, as opposed to just a single `Imaging` object. This is because\n", - "only a single `Analysis` class was used in the model-fit, meaning there was only one `Imaging` dataset that was\n", - "fit. \n", - "\n", - "The `multi` package of the workspace illustrates model-fits which fit multiple datasets \n", - "simultaneously, (e.g. multi-wavelength imaging) by summing `Analysis` objects together, where the `dataset_list` \n", - "would contain multiple `Imaging` objects." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_agg = al.agg.ImagingAgg(aggregator=agg)\n", - "dataset_gen = dataset_agg.dataset_gen_from()\n", - "\n", - "for dataset_list in dataset_gen:\n", - " # Only one `Analysis` so take first and only dataset.\n", - " dataset = dataset_list[0]\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now use the aggregator to load a generator containing the fit of the maximum log likelihood model (and therefore \n", - "fit) to each dataset.\n", - "\n", - "Analogous to the `dataset_gen` above returning a list with one `Imaging` object, the `fit_gen` returns a list of\n", - "`FitImaging` objects, because only one `Analysis` was used to perform the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit_agg = al.agg.FitImagingAgg(aggregator=agg)\n", - "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", - "\n", - "for fit_list in fit_gen:\n", - " # Only one `Analysis` so take first and only dataset.\n", - " fit = fit_list[0]\n", - "\n", - " aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Modification__\n", - "\n", - "The `FitImagingAgg` allow us to modify the fit settings. \n", - "\n", - "However, we can change these settings such that the fit is performed differently. For example, what if I wanted to see \n", - "how the fit looks where the pixelization didn`t use a border? \n", - "\n", - "You can do this by passing the settings objects, which overwrite the ones used by the analysis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit_agg = al.agg.FitImagingAgg(\n", - " aggregator=agg,\n", - " settings=al.Settings(use_border_relocator=False),\n", - ")\n", - "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", - "\n", - "for fit_list in fit_gen:\n", - " # Only one `Analysis` so take first and only dataset.\n", - " fit = fit_list[0]\n", - "\n", - " aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization Customization__\n", - "\n", - "The benefit of inspecting fits using the aggregator, rather than the files outputs to the hard-disk, is that we can \n", - "customize the plots using the `plot_yx` and `plot_array`/`subplot_\\*` objects..\n", - "\n", - "We create a new function to apply as a generator to do this. However, we use a convenience method available \n", - "in the aggregator package to set up the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit_agg = al.agg.FitImagingAgg(aggregator=agg)\n", - "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", - "\n", - "for fit_list in fit_gen:\n", - " # Only one `Analysis` so take first and only dataset.\n", - " fit = fit_list[0]\n", - "\n", - " aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Making this plot for a paper? You can output it to hard disk." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit_agg = al.agg.FitImagingAgg(aggregator=agg)\n", - "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", - "\n", - "for fit_list in fit_gen:\n", - " # Only one `Analysis` so take first and only dataset.\n", - " fit = fit_list[0]\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors (Random draws from PDF)__\n", - "\n", - "In the `examples/models.py` example we showed how `Tracer objects could be randomly drawn form the Probability \n", - "Distribution Function, in order to quantity things such as errors.\n", - "\n", - "The same approach can be used with `FitImaging` objects, to investigate how the properties of the fit vary within\n", - "the errors (e.g. showing source reconstructions fot different fits consistent with the errors)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit_agg = al.agg.FitImagingAgg(aggregator=agg)\n", - "fit_gen = fit_agg.randomly_drawn_via_pdf_gen_from(total_samples=2)\n", - "\n", - "\n", - "for fit_list_gen in fit_gen: # 1 Dataset so just one fit\n", - " for (\n", - " fit_list\n", - " ) in (\n", - " fit_list_gen\n", - " ): # Iterate over each total_samples=2, each with one fits for 1 analysis.\n", - " # Only one `Analysis` so take first and only dataset.\n", - " fit = fit_list[0]\n", - "\n", - " aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Data Fitting\n", + "=====================\n", + "\n", + "In this tutorial, we use the aggregator to load models and data from a non-linear search and use them to perform\n", + "fits to the data.\n", + "\n", + "We show how to use these tools to inspect the maximum log likelihood model of a fit to the data, customize things\n", + "like its visualization and also inspect fits randomly drawm from the PDF.\n", + "\n", + "__Contents__\n", + "\n", + "- **Interferometer:** This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", + "- **Aggregator:** First, set up the aggregator as shown in `start_here.py`.\n", + "- **Fits via Aggregator:** Having performed a model-fit, we now want to interpret and visualize the results.\n", + "- **Modification:** The `FitImagingAgg` allow us to modify the fit settings.\n", + "- **Visualization Customization:** The benefit of inspecting fits using the aggregator, rather than the files outputs to the.\n", + "\n", + "__Interferometer__\n", + "\n", + "This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", + "\n", + "The only entries that needs changing are:\n", + "\n", + " - `ImagingAgg` -> `InterferometerAgg`.\n", + " - `FitImagingAgg` -> `FitInterferometerAgg`.\n", + " - `aplt.subplot_imaging_dataset` -> `aplt.subplot_interferometer_dirty_images`.\n", + " - `aplt.subplot_fit_imaging` -> `aplt.subplot_fit_interferometer`.\n", + "\n", + "Quantities specific to an interfometer, for example its uv-wavelengths real space mask, are accessed using the same API\n", + "(e.g. `values(\"dataset.uv_wavelengths\")` and `.values{\"dataset.real_space_mask\"))." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import os\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Aggregator__\n", + "\n", + "First, set up the aggregator as shown in `start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "results_path = Path(\"output\") / \"results_folder\"\n", + "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "agg = Aggregator.from_directory(\n", + " directory=results_path,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The masks we used to fit the lenses is accessible via the aggregator." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_gen = agg.values(\"dataset.mask\")\n", + "print([mask for mask in mask_gen])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The info dictionary we passed is also available." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Info:\")\n", + "info_gen = agg.values(\"info\")\n", + "print([info for info in info_gen])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fits via Aggregator__\n", + "\n", + "Having performed a model-fit, we now want to interpret and visualize the results. In this example, we inspect \n", + "the `Imaging` objects that gave good fits to the data. \n", + "\n", + "Using the API shown in the `start_here.py` example this would require us to create a `Samples` object and manually \n", + "compose our own `Imaging` object. For large datasets, this would require us to use generators to ensure it is \n", + "memory-light, which are cumbersome to write.\n", + "\n", + "This example therefore uses the `ImagingAgg` object, which conveniently loads the `Imaging` objects of every fit via \n", + "generators for us. \n", + "\n", + "We get a dataset generator via the `al.agg.ImagingAgg` object, where this `dataset_gen` contains the maximum log\n", + "likelihood `Imaging `object of every model-fit.\n", + "\n", + "The `dataset_gen` returns a list of `Imaging` objects, as opposed to just a single `Imaging` object. This is because\n", + "only a single `Analysis` class was used in the model-fit, meaning there was only one `Imaging` dataset that was\n", + "fit. \n", + "\n", + "The `multi` package of the workspace illustrates model-fits which fit multiple datasets \n", + "simultaneously, (e.g. multi-wavelength imaging) by summing `Analysis` objects together, where the `dataset_list` \n", + "would contain multiple `Imaging` objects." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_agg = al.agg.ImagingAgg(aggregator=agg)\n", + "dataset_gen = dataset_agg.dataset_gen_from()\n", + "\n", + "for dataset_list in dataset_gen:\n", + " # Only one `Analysis` so take first and only dataset.\n", + " dataset = dataset_list[0]\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now use the aggregator to load a generator containing the fit of the maximum log likelihood model (and therefore \n", + "fit) to each dataset.\n", + "\n", + "Analogous to the `dataset_gen` above returning a list with one `Imaging` object, the `fit_gen` returns a list of\n", + "`FitImaging` objects, because only one `Analysis` was used to perform the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_agg = al.agg.FitImagingAgg(aggregator=agg)\n", + "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", + "\n", + "for fit_list in fit_gen:\n", + " # Only one `Analysis` so take first and only dataset.\n", + " fit = fit_list[0]\n", + "\n", + " aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Modification__\n", + "\n", + "The `FitImagingAgg` allow us to modify the fit settings. \n", + "\n", + "However, we can change these settings such that the fit is performed differently. For example, what if I wanted to see \n", + "how the fit looks where the pixelization didn`t use a border? \n", + "\n", + "You can do this by passing the settings objects, which overwrite the ones used by the analysis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_agg = al.agg.FitImagingAgg(\n", + " aggregator=agg,\n", + " settings=al.Settings(use_border_relocator=False),\n", + ")\n", + "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", + "\n", + "for fit_list in fit_gen:\n", + " # Only one `Analysis` so take first and only dataset.\n", + " fit = fit_list[0]\n", + "\n", + " aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization Customization__\n", + "\n", + "The benefit of inspecting fits using the aggregator, rather than the files outputs to the hard-disk, is that we can \n", + "customize the plots using the `plot_yx` and `plot_array`/`subplot_\\*` objects..\n", + "\n", + "We create a new function to apply as a generator to do this. However, we use a convenience method available \n", + "in the aggregator package to set up the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_agg = al.agg.FitImagingAgg(aggregator=agg)\n", + "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", + "\n", + "for fit_list in fit_gen:\n", + " # Only one `Analysis` so take first and only dataset.\n", + " fit = fit_list[0]\n", + "\n", + " aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Making this plot for a paper? You can output it to hard disk." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_agg = al.agg.FitImagingAgg(aggregator=agg)\n", + "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", + "\n", + "for fit_list in fit_gen:\n", + " # Only one `Analysis` so take first and only dataset.\n", + " fit = fit_list[0]\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors (Random draws from PDF)__\n", + "\n", + "In the `examples/models.py` example we showed how `Tracer objects could be randomly drawn form the Probability \n", + "Distribution Function, in order to quantity things such as errors.\n", + "\n", + "The same approach can be used with `FitImaging` objects, to investigate how the properties of the fit vary within\n", + "the errors (e.g. showing source reconstructions fot different fits consistent with the errors)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_agg = al.agg.FitImagingAgg(aggregator=agg)\n", + "fit_gen = fit_agg.randomly_drawn_via_pdf_gen_from(total_samples=2)\n", + "\n", + "\n", + "for fit_list_gen in fit_gen: # 1 Dataset so just one fit\n", + " for (\n", + " fit_list\n", + " ) in (\n", + " fit_list_gen\n", + " ): # Iterate over each total_samples=2, each with one fits for 1 analysis.\n", + " # Only one `Analysis` so take first and only dataset.\n", + " fit = fit_list[0]\n", + "\n", + " aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/aggregator/galaxies_fits.ipynb b/notebooks/guides/results/aggregator/galaxies_fits.ipynb index 62003ae46..b407ac315 100644 --- a/notebooks/guides/results/aggregator/galaxies_fits.ipynb +++ b/notebooks/guides/results/aggregator/galaxies_fits.ipynb @@ -1,396 +1,433 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Galaxies and Fits\n", - "===========================\n", - "\n", - "This tutorial inspects an inferred model using galaxies inferred by the non-linear search.\n", - "This allows us to visualize and interpret its results.\n", - "\n", - "The galaxies and fit API is described fully in the guides:\n", - "\n", - " - `autolens_workspace/*/guides/tracer.ipynb`\n", - " - `autolens_workspace/*/guides/fit.ipynb`\n", - " - `autolens_workspace/*/guides/galaxies.ipynb`\n", - "\n", - "This result example only explains specific functionality for using a `Result` object to inspect galaxies or a fit\n", - "and therefore you should read these guides in detail first.\n", - "\n", - "__Contents__\n", - "\n", - "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", - "- **Data Structures:** Arrays inspected in this example use bespoke data structures for storing arrays, grids, vectors and.\n", - "- **Model Fit:** Perform the model-fit using the search and analysis.\n", - "- **Max Likelihood Tracer:** As seen elsewhere in the workspace, the result contains a `max_log_likelihood_tracer` which we can.\n", - "- **Refitting:** Using the API introduced in the first tutorial, we can also refit the data locally.\n", - "- **Samples API:** In the first results tutorial, we used `Samples` objects to inspect the results of a model.\n", - "- **Max Likelihood Fit:** As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit` which we can.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Units__\n", - "\n", - "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", - "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", - "\n", - "The guide `guides/units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", - "kiloparsecs, magnitudes and solar masses.\n", - "\n", - "__Data Structures__\n", - "\n", - "Arrays inspected in this example use bespoke data structures for storing arrays, grids,\n", - "vectors and other 1D and 2D quantities. These use the `slim` and `native` API to toggle between representing the\n", - "data in 1D numpy arrays or high dimension numpy arrays.\n", - "\n", - "This tutorial will only use the `slim` properties which show results in 1D numpy arrays of\n", - "shape [total_unmasked_pixels]. This is a slimmed-down representation of the data in 1D that contains only the\n", - "unmasked data points\n", - "\n", - "These are documented fully in the `autolens_workspace/*/guides/data_structures.ipynb` guide.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `results/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Quick Fit Auto-Trigger__\n", - "\n", - "If a previous fit has not been run yet, the shared helper ``_quick_fit.py`` is invoked to produce one.\n", - "The helper writes a capped Nautilus fit to ``output/results_folder/`` so this tutorial (and its siblings\n", - "in this folder) have results to work with. When that folder already exists the helper exits immediately,\n", - "so re-running this tutorial is cheap." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "results_path = Path(\"output\") / \"results_folder\"\n", - "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "To illustrate results, we need to perform a model-fit in order to create a `Result` object.\n", - "\n", - "The code below performs a model-fit using nautilus. The helper above already wrote a completed fit to\n", - "``output/results_folder/``, so the ``search.fit(...)`` call below resumes from that checkpoint and\n", - "returns the in-memory ``Result`` object without redoing the search.\n", - "\n", - "You should be familiar with modeling already, if not read the `modeling/start_here.py` script before reading this one." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", - " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=None),\n", - " ),\n", - ")\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"results_folder\"),\n", - " name=\"results\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " n_like_max=300,\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Max Likelihood Tracer__\n", - "\n", - "As seen elsewhere in the workspace, the result contains a `max_log_likelihood_tracer` which we can visualize." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "aplt.subplot_tracer(tracer=tracer, grid=mask.derive_grid.all_false)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Refitting__\n", - "\n", - "Using the API introduced in the first tutorial, we can also refit the data locally. \n", - "\n", - "This allows us to inspect how the tracer changes for models with similar log likelihoods. We create and plot\n", - "the tracer of the tenth-last accepted model by Nautilus." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "instance = samples.from_sample_index(sample_index=-10)\n", - "\n", - "# Input to FitImaging to solve for linear light profile intensities, see `start_here.py` for details.\n", - "tracer = al.Tracer(galaxies=instance.galaxies)\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", - "\n", - "aplt.subplot_tracer(tracer=tracer, grid=mask.derive_grid.all_false)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Samples API__\n", - "\n", - "In the first results tutorial, we used `Samples` objects to inspect the results of a model.\n", - "\n", - "We saw how these samples created instances, which include a `galaxies` property that mains the API of the `Model`\n", - "creates above (e.g. `galaxies.source.bulge`). \n", - "\n", - "We can also use this instance to extract individual components of the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "ml_instance = samples.max_log_likelihood()\n", - "\n", - "# Input to FitImaging to solve for linear light profile intensities, see `start_here.py` for details.\n", - "tracer = al.Tracer(galaxies=instance.galaxies)\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", - "\n", - "bulge = tracer.galaxies[-1].bulge\n", - "\n", - "bulge_image_2d = bulge.image_2d_from(grid=dataset.grid)\n", - "print(bulge_image_2d.slim[0])\n", - "\n", - "aplt.plot_array(array=bulge.image_2d_from(grid=dataset.grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In fact, if we create a `Tracer` from an instance (which is how `result.max_log_likelihood_tracer` is created) we\n", - "can choose whether to access its attributes using each API: " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "print(tracer.galaxies[-1].bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Max Likelihood Fit__\n", - "\n", - "As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit` which we can visualize." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = result.max_log_likelihood_fit\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Refitting__\n", - "\n", - "Using the API introduced in the first tutorial, we can also refit the data locally. \n", - "\n", - "This allows us to inspect how the fit changes for models with similar log likelihoods. Below, we refit and plot\n", - "the fit of the tenth-last accepted model by Nautilus." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "instance = samples.from_sample_index(sample_index=-10)\n", - "\n", - "tracer = al.Tracer(galaxies=instance.galaxies)\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "We have learnt how to extract individual planes, galaxies, light and mass profiles from the tracer that results from\n", - "a model-fit and use these objects to compute specific quantities of each component." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Galaxies and Fits\n", + "===========================\n", + "\n", + "This tutorial inspects an inferred model using galaxies inferred by the non-linear search.\n", + "This allows us to visualize and interpret its results.\n", + "\n", + "The galaxies and fit API is described fully in the guides:\n", + "\n", + " - `autolens_workspace/*/guides/tracer.ipynb`\n", + " - `autolens_workspace/*/guides/fit.ipynb`\n", + " - `autolens_workspace/*/guides/galaxies.ipynb`\n", + "\n", + "This result example only explains specific functionality for using a `Result` object to inspect galaxies or a fit\n", + "and therefore you should read these guides in detail first.\n", + "\n", + "__Contents__\n", + "\n", + "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", + "- **Data Structures:** Arrays inspected in this example use bespoke data structures for storing arrays, grids, vectors and.\n", + "- **Model Fit:** Perform the model-fit using the search and analysis.\n", + "- **Max Likelihood Tracer:** As seen elsewhere in the workspace, the result contains a `max_log_likelihood_tracer` which we can.\n", + "- **Refitting:** Using the API introduced in the first tutorial, we can also refit the data locally.\n", + "- **Samples API:** In the first results tutorial, we used `Samples` objects to inspect the results of a model.\n", + "- **Max Likelihood Fit:** As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit` which we can.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Units__\n", + "\n", + "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", + "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", + "\n", + "The guide `guides/units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", + "kiloparsecs, magnitudes and solar masses.\n", + "\n", + "__Data Structures__\n", + "\n", + "Arrays inspected in this example use bespoke data structures for storing arrays, grids,\n", + "vectors and other 1D and 2D quantities. These use the `slim` and `native` API to toggle between representing the\n", + "data in 1D numpy arrays or high dimension numpy arrays.\n", + "\n", + "This tutorial will only use the `slim` properties which show results in 1D numpy arrays of\n", + "shape [total_unmasked_pixels]. This is a slimmed-down representation of the data in 1D that contains only the\n", + "unmasked data points\n", + "\n", + "These are documented fully in the `autolens_workspace/*/guides/data_structures.ipynb` guide.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `results/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Quick Fit Auto-Trigger__\n", + "\n", + "If a previous fit has not been run yet, the shared helper ``_quick_fit.py`` is invoked to produce one.\n", + "The helper writes a capped Nautilus fit to ``output/results_folder/`` so this tutorial (and its siblings\n", + "in this folder) have results to work with. When that folder already exists the helper exits immediately,\n", + "so re-running this tutorial is cheap." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "results_path = Path(\"output\") / \"results_folder\"\n", + "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "To illustrate results, we need to perform a model-fit in order to create a `Result` object.\n", + "\n", + "The code below performs a model-fit using nautilus. The helper above already wrote a completed fit to\n", + "``output/results_folder/``, so the ``search.fit(...)`` call below resumes from that checkpoint and\n", + "returns the in-memory ``Result`` object without redoing the search.\n", + "\n", + "You should be familiar with modeling already, if not read the `modeling/start_here.py` script before reading this one." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", + " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=None),\n", + " ),\n", + ")\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"results_folder\"),\n", + " name=\"results\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " n_like_max=300,\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Max Likelihood Tracer__\n", + "\n", + "As seen elsewhere in the workspace, the result contains a `max_log_likelihood_tracer` which we can visualize." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "aplt.subplot_tracer(tracer=tracer, grid=mask.derive_grid.all_false)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Refitting__\n", + "\n", + "Using the API introduced in the first tutorial, we can also refit the data locally. \n", + "\n", + "This allows us to inspect how the tracer changes for models with similar log likelihoods. We create and plot\n", + "the tracer of the tenth-last accepted model by Nautilus." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "instance = samples.from_sample_index(sample_index=-10)\n", + "\n", + "# Input to FitImaging to solve for linear light profile intensities, see `start_here.py` for details.\n", + "tracer = al.Tracer(galaxies=instance.galaxies)\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", + "\n", + "aplt.subplot_tracer(tracer=tracer, grid=mask.derive_grid.all_false)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Samples API__\n", + "\n", + "In the first results tutorial, we used `Samples` objects to inspect the results of a model.\n", + "\n", + "We saw how these samples created instances, which include a `galaxies` property that mains the API of the `Model`\n", + "creates above (e.g. `galaxies.source.bulge`). \n", + "\n", + "We can also use this instance to extract individual components of the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "ml_instance = samples.max_log_likelihood()\n", + "\n", + "# Input to FitImaging to solve for linear light profile intensities, see `start_here.py` for details.\n", + "tracer = al.Tracer(galaxies=instance.galaxies)\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", + "\n", + "bulge = tracer.galaxies[-1].bulge\n", + "\n", + "bulge_image_2d = bulge.image_2d_from(grid=dataset.grid)\n", + "print(bulge_image_2d.slim[0])\n", + "\n", + "aplt.plot_array(array=bulge.image_2d_from(grid=dataset.grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In fact, if we create a `Tracer` from an instance (which is how `result.max_log_likelihood_tracer` is created) we\n", + "can choose whether to access its attributes using each API: " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "print(tracer.galaxies[-1].bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Max Likelihood Fit__\n", + "\n", + "As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit` which we can visualize." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = result.max_log_likelihood_fit\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Refitting__\n", + "\n", + "Using the API introduced in the first tutorial, we can also refit the data locally. \n", + "\n", + "This allows us to inspect how the fit changes for models with similar log likelihoods. Below, we refit and plot\n", + "the fit of the tenth-last accepted model by Nautilus." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "instance = samples.from_sample_index(sample_index=-10)\n", + "\n", + "tracer = al.Tracer(galaxies=instance.galaxies)\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "We have learnt how to extract individual planes, galaxies, light and mass profiles from the tracer that results from\n", + "a model-fit and use these objects to compute specific quantities of each component." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/aggregator/interferometer.ipynb b/notebooks/guides/results/aggregator/interferometer.ipynb index a83b04fb9..0355ea4e4 100644 --- a/notebooks/guides/results/aggregator/interferometer.ipynb +++ b/notebooks/guides/results/aggregator/interferometer.ipynb @@ -1,89 +1,126 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This script still needs writing, I have kept some notes on questions asked by users which may help you..." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - " > For interferometric data, which units PyAutoLens uses for brightness? I think they are in Jy/arcsec^2 (?) since \n", - " I have computed the magnification from from my original image in Jy/beam I just wanted to be sure that the conversions \n", - " I have assumed are fine.\n", - "\n", - "This is correct, the units of brightness are Jy/arcsec^2\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "> -When converting the reconstructed source image to a .fits file, what is the best image shape to assume in the \n", - " interpolation ? By now I am using the same shape as the native one used when defining the real space mask at the \n", - " beginning, should it be fine? And also, what are the brightness units here?\n", - "\n", - "The reconstruction is essentially a devonvolved image which if you sum up all pixel you get the total flux of the \n", - "source. If if was a regular grid in each pixel the units are Jy/pixel or Jy/arcsec^2\n", - "\n", - "The shape of the grid is really you're choice, so long as it does not extend spatially beyond the extent of the source\n", - "reconstruction (as there is no reconstructed source values here in order to enable an accurate interpolation). \n", - "\n", - "I would use a shape which covers the whole source you are reconstructing (with a bit of padding), and gives visually\n", - "appealing data you can use for whatever science you're interested in.\n", - "\n", - "I would imagine the shape of the real space mask is much larger than you really need, but probably doesn't do much\n", - "harm either.\n", - "\n", - "One caveat may be magnifications. It could be that as you make the shape of the interpolation grid bigger, the\n", - "magnification changes. I'm not sure on this but would advise you experiment to see if the magnifications \n", - "seem \"stable\" (assuming you are estimate a magnification at some point)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This script still needs writing, I have kept some notes on questions asked by users which may help you..." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + " > For interferometric data, which units PyAutoLens uses for brightness? I think they are in Jy/arcsec^2 (?) since \n", + " I have computed the magnification from from my original image in Jy/beam I just wanted to be sure that the conversions \n", + " I have assumed are fine.\n", + "\n", + "This is correct, the units of brightness are Jy/arcsec^2\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> -When converting the reconstructed source image to a .fits file, what is the best image shape to assume in the \n", + " interpolation ? By now I am using the same shape as the native one used when defining the real space mask at the \n", + " beginning, should it be fine? And also, what are the brightness units here?\n", + "\n", + "The reconstruction is essentially a devonvolved image which if you sum up all pixel you get the total flux of the \n", + "source. If if was a regular grid in each pixel the units are Jy/pixel or Jy/arcsec^2\n", + "\n", + "The shape of the grid is really you're choice, so long as it does not extend spatially beyond the extent of the source\n", + "reconstruction (as there is no reconstructed source values here in order to enable an accurate interpolation). \n", + "\n", + "I would use a shape which covers the whole source you are reconstructing (with a bit of padding), and gives visually\n", + "appealing data you can use for whatever science you're interested in.\n", + "\n", + "I would imagine the shape of the real space mask is much larger than you really need, but probably doesn't do much\n", + "harm either.\n", + "\n", + "One caveat may be magnifications. It could be that as you make the shape of the interpolation grid bigger, the\n", + "magnification changes. I'm not sure on this but would advise you experiment to see if the magnifications \n", + "seem \"stable\" (assuming you are estimate a magnification at some point)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/aggregator/models.ipynb b/notebooks/guides/results/aggregator/models.ipynb index deaeabdf0..8d77511f5 100644 --- a/notebooks/guides/results/aggregator/models.ipynb +++ b/notebooks/guides/results/aggregator/models.ipynb @@ -1,349 +1,386 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Models\n", - "===============\n", - "\n", - "In this tutorial, we use the aggregator to load models and `Tracer`'s from a non-linear search. This allows us to\n", - "visualize and interpret its results.\n", - "\n", - "We then show how the aggregator also allows us to load many `Tracer`'s correspond to many samples of the non-linear\n", - "search. This allows us to compute the errors on quantities that the `Tracer` contains, but were not sampled directly\n", - "by the non-linear search.\n", - "\n", - "__Contents__\n", - "\n", - "- **Aggregator:** First, set up the aggregator as shown in `start_here.py`.\n", - "- **Tracer via Aggregator:** Having performed a model-fit, we now want to interpret and visualize the results.\n", - "- **Einstein Mass Example:** Each tracer has the information we need to compute the Einstein mass of a model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import os\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Aggregator__\n", - "\n", - "First, set up the aggregator as shown in `start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "results_path = Path(\"output\") / \"results_folder\"\n", - "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "agg = Aggregator.from_directory(\n", - " directory=results_path,\n", - ")\n", - "#" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer via Aggregator__\n", - "\n", - "Having performed a model-fit, we now want to interpret and visualize the results. In this example, we want to inspect\n", - "the `Tracer` objects that gave good fits to the data. \n", - "\n", - "Using the API shown in the `start_here.py` example this would require us to create a `Samples` object and manually \n", - "compose our own `Tracer` object. For large datasets, this would require us to use generators to ensure it is \n", - "memory-light, which are cumbersome to write.\n", - "\n", - "This example therefore uses the `TracerAgg` object, which conveniently loads the `Tracer` objects of every fit via \n", - "generators for us. \n", - "\n", - "We get a tracer generator via the `al.agg.TracerAgg` object, where this `tracer_gen` contains the maximum log\n", - "likelihood tracer of every model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_agg = al.agg.TracerAgg(aggregator=agg)\n", - "tracer_gen = tracer_agg.max_log_likelihood_gen_from()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now iterate over our tracer generator to make the plots we desire.\n", - "\n", - "The `tracer_gen` returns a list of `Tracer` objects, as opposed to just a single `Tracer`object. This is because\n", - "only a single `Analysis` class was used in the model-fit, meaning there was only one `Tracer` dataset that was\n", - "fit. \n", - "\n", - "The `multi` package of the workspace illustrates model-fits which fit multiple datasets \n", - "simultaneously, (e.g. multi-wavelength imaging) by summing `Analysis` objects together, where the `tracer_list` \n", - "would contain multiple `Tracer` objects.\n", - "\n", - "The parameters of galaxies in the `Tracer` may vary across the datasets (e.g. different light profile intensities \n", - "for different wavelengths), which would be reflected in the tracer list." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_agg = al.agg.ImagingAgg(aggregator=agg)\n", - "dataset_gen = dataset_agg.dataset_gen_from()\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.1)\n", - "\n", - "for dataset_list, tracer_list in zip(dataset_gen, tracer_gen):\n", - " # Only one `Analysis` so take first and only dataset.\n", - " dataset = dataset_list[0]\n", - "\n", - " # Only one `Analysis` so take first and only tracer.\n", - " tracer = tracer_list[0]\n", - "\n", - " # Input to FitImaging to solve for linear light profile intensities, see `start_here.py` for details.\n", - " fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - " tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", - "\n", - " aplt.plot_array(array=tracer.convergence_2d_from(grid=grid), title=\"Convergence\")\n", - " aplt.plot_array(array=tracer.potential_2d_from(grid=grid), title=\"Potential\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Einstein Mass Example__\n", - "\n", - "Each tracer has the information we need to compute the Einstein mass of a model. Therefore, lets print \n", - "the Einstein mass of each of our most-likely lens galaxies.\n", - "\n", - "The model instance uses the model defined by a pipeline. In this pipeline, we called the lens galaxy `lens`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_agg = al.agg.TracerAgg(aggregator=agg)\n", - "tracer_gen = tracer_agg.max_log_likelihood_gen_from()\n", - "\n", - "print(\"Maximum Log Likelihood Lens Einstein Masses:\")\n", - "\n", - "for dataset_list, tracer_list in zip(dataset_gen, tracer_gen):\n", - " # Only one `Analysis` so take first and only dataset.\n", - " dataset = dataset_list[0]\n", - "\n", - " # Only one `Analysis` so take first and only tracer.\n", - " tracer = tracer_list[0]\n", - "\n", - " einstein_mass = tracer.galaxies[0].einstein_mass_angular_from(grid=grid)\n", - " print(\"Einstein Mass (angular units) = \", einstein_mass)\n", - "\n", - " cosmology = al.cosmo.Planck15()\n", - "\n", - " critical_surface_density = (\n", - " cosmology.critical_surface_density_between_redshifts_from(\n", - " redshift_0=tracer.galaxies[0].redshift,\n", - " redshift_1=tracer.galaxies[-1].redshift,\n", - " )\n", - " )\n", - "\n", - " einstein_mass_kpc = einstein_mass * critical_surface_density\n", - "\n", - " print(\"Einstein Mass (kpc) = \", einstein_mass_kpc)\n", - " print(\"Einstein Mass (kpc) = \", \"{:.4e}\".format(einstein_mass_kpc))\n", - "\n", - " print(einstein_mass)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors (PDF from samples)__\n", - "\n", - "In this example, we will compute the errors on the axis ratio of a model. Computing the errors on a quantity \n", - "like the trap `density` is simple, because it is sampled by the non-linear search. The errors are therefore accessible\n", - "via the `Samples`, by marginalizing over all over parameters via the 1D Probability Density Function (PDF).\n", - "\n", - "Computing the errors on the axis ratio is more tricky, because it is a derived quantity. It is a parameter or \n", - "measurement that we want to calculate but was not sampled directly by the non-linear search. The `TracerAgg` object \n", - "object has everything we need to compute the errors of derived quantities.\n", - "\n", - "Below, we compute the axis ratio of every model sampled by the non-linear search and use this determine the PDF \n", - "of the axis ratio. When combining each axis ratio we weight each value by its `weight`. For Nautilus, \n", - "the nested sampler used by the fit, this ensures models which gave a bad fit (and thus have a low weight) do not \n", - "contribute significantly to the axis ratio error estimate.\n", - "\n", - "We set `minimum_weight=`1e-4`, such that any sample with a weight below this value is discarded when computing the \n", - "error. This speeds up the error computation by only using a small fraction of the total number of samples. Computing\n", - "a delta ellipticity is cheap, and this is probably not necessary. However, certain quantities have a non-negligible\n", - "computational overhead is being calculated and setting a minimum weight can speed up the calculation without \n", - "significantly changing the inferred errors.\n", - "\n", - "Below, we use the `TracerAgg` to get the `Tracer` of every Nautilus sample in each model-fit. We extract from each \n", - "tracer the model's axis-ratio, store them in a list and find the value via the PDF and quantile method. This again\n", - "uses generators, ensuring minimal memory use. \n", - "\n", - "In order to use these samples in the function `quantile`, we also need the weight list of the sample weights. We \n", - "compute this using the `TracerAgg`'s function `weights_above_gen_from`, which computes generators of the weights of all \n", - "points above this minimum value. This again ensures memory use in minimal." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_agg = al.agg.TracerAgg(aggregator=agg)\n", - "tracer_list_gen = tracer_agg.all_above_weight_gen_from(minimum_weight=1e-4)\n", - "weight_list_gen = tracer_agg.weights_above_gen_from(minimum_weight=1e-4)\n", - "\n", - "for tracer_gen, weight_gen in zip(tracer_list_gen, weight_list_gen):\n", - " axis_ratio_list = []\n", - "\n", - " for tracer_list in tracer_gen:\n", - " # Only one `Analysis` so take first and only tracer.\n", - " tracer = tracer_list[0]\n", - "\n", - " axis_ratio = al.convert.axis_ratio_from(\n", - " ell_comps=tracer.galaxies[0].mass.ell_comps\n", - " )\n", - "\n", - " axis_ratio_list.append(axis_ratio)\n", - "\n", - " weight_list = [weight for weight in weight_gen]\n", - "\n", - " if len(axis_ratio_list) > 1:\n", - " median_axis_ratio, lower_axis_ratio, upper_axis_ratio = af.marginalize(\n", - " parameter_list=axis_ratio_list, sigma=3.0, weight_list=weight_list\n", - " )\n", - "\n", - " print(\n", - " f\"Axis-Ratio = {median_axis_ratio} ({upper_axis_ratio} {lower_axis_ratio}\"\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors (Random draws from PDF)__\n", - "\n", - "An alternative approach to estimating the errors on a derived quantity is to randomly draw samples from the PDF \n", - "of the non-linear search. For a sufficiently high number of random draws, this should be as accurate and precise\n", - "as the method above. However, it can be difficult to be certain how many random draws are necessary.\n", - "\n", - "The weights of each sample are used to make every random draw. Therefore, when we compute the axis-ratio and its errors\n", - "we no longer need to pass the `weight_list` to the `quantile` function." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_agg = al.agg.TracerAgg(aggregator=agg)\n", - "tracer_list_gen = tracer_agg.randomly_drawn_via_pdf_gen_from(total_samples=2)\n", - "\n", - "for tracer_gen in tracer_list_gen:\n", - " axis_ratio_list = []\n", - "\n", - " for tracer_list in tracer_gen:\n", - " # Only one `Analysis` so take first and only tracer.\n", - " tracer = tracer_list[0]\n", - "\n", - " axis_ratio = al.convert.axis_ratio_from(\n", - " ell_comps=tracer.galaxies[0].mass.ell_comps\n", - " )\n", - "\n", - " axis_ratio_list.append(axis_ratio)\n", - "\n", - " median_axis_ratio, lower_axis_ratio, upper_axis_ratio = af.marginalize(\n", - " parameter_list=axis_ratio_list, sigma=3.0\n", - " )\n", - "\n", - " print(f\"Axis-Ratio = {median_axis_ratio} ({upper_axis_ratio} {lower_axis_ratio}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Models\n", + "===============\n", + "\n", + "In this tutorial, we use the aggregator to load models and `Tracer`'s from a non-linear search. This allows us to\n", + "visualize and interpret its results.\n", + "\n", + "We then show how the aggregator also allows us to load many `Tracer`'s correspond to many samples of the non-linear\n", + "search. This allows us to compute the errors on quantities that the `Tracer` contains, but were not sampled directly\n", + "by the non-linear search.\n", + "\n", + "__Contents__\n", + "\n", + "- **Aggregator:** First, set up the aggregator as shown in `start_here.py`.\n", + "- **Tracer via Aggregator:** Having performed a model-fit, we now want to interpret and visualize the results.\n", + "- **Einstein Mass Example:** Each tracer has the information we need to compute the Einstein mass of a model." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import os\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Aggregator__\n", + "\n", + "First, set up the aggregator as shown in `start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "results_path = Path(\"output\") / \"results_folder\"\n", + "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "agg = Aggregator.from_directory(\n", + " directory=results_path,\n", + ")\n", + "#" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer via Aggregator__\n", + "\n", + "Having performed a model-fit, we now want to interpret and visualize the results. In this example, we want to inspect\n", + "the `Tracer` objects that gave good fits to the data. \n", + "\n", + "Using the API shown in the `start_here.py` example this would require us to create a `Samples` object and manually \n", + "compose our own `Tracer` object. For large datasets, this would require us to use generators to ensure it is \n", + "memory-light, which are cumbersome to write.\n", + "\n", + "This example therefore uses the `TracerAgg` object, which conveniently loads the `Tracer` objects of every fit via \n", + "generators for us. \n", + "\n", + "We get a tracer generator via the `al.agg.TracerAgg` object, where this `tracer_gen` contains the maximum log\n", + "likelihood tracer of every model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_agg = al.agg.TracerAgg(aggregator=agg)\n", + "tracer_gen = tracer_agg.max_log_likelihood_gen_from()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now iterate over our tracer generator to make the plots we desire.\n", + "\n", + "The `tracer_gen` returns a list of `Tracer` objects, as opposed to just a single `Tracer`object. This is because\n", + "only a single `Analysis` class was used in the model-fit, meaning there was only one `Tracer` dataset that was\n", + "fit. \n", + "\n", + "The `multi` package of the workspace illustrates model-fits which fit multiple datasets \n", + "simultaneously, (e.g. multi-wavelength imaging) by summing `Analysis` objects together, where the `tracer_list` \n", + "would contain multiple `Tracer` objects.\n", + "\n", + "The parameters of galaxies in the `Tracer` may vary across the datasets (e.g. different light profile intensities \n", + "for different wavelengths), which would be reflected in the tracer list." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_agg = al.agg.ImagingAgg(aggregator=agg)\n", + "dataset_gen = dataset_agg.dataset_gen_from()\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.1)\n", + "\n", + "for dataset_list, tracer_list in zip(dataset_gen, tracer_gen):\n", + " # Only one `Analysis` so take first and only dataset.\n", + " dataset = dataset_list[0]\n", + "\n", + " # Only one `Analysis` so take first and only tracer.\n", + " tracer = tracer_list[0]\n", + "\n", + " # Input to FitImaging to solve for linear light profile intensities, see `start_here.py` for details.\n", + " fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + " tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", + "\n", + " aplt.plot_array(array=tracer.convergence_2d_from(grid=grid), title=\"Convergence\")\n", + " aplt.plot_array(array=tracer.potential_2d_from(grid=grid), title=\"Potential\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Einstein Mass Example__\n", + "\n", + "Each tracer has the information we need to compute the Einstein mass of a model. Therefore, lets print \n", + "the Einstein mass of each of our most-likely lens galaxies.\n", + "\n", + "The model instance uses the model defined by a pipeline. In this pipeline, we called the lens galaxy `lens`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_agg = al.agg.TracerAgg(aggregator=agg)\n", + "tracer_gen = tracer_agg.max_log_likelihood_gen_from()\n", + "\n", + "print(\"Maximum Log Likelihood Lens Einstein Masses:\")\n", + "\n", + "for dataset_list, tracer_list in zip(dataset_gen, tracer_gen):\n", + " # Only one `Analysis` so take first and only dataset.\n", + " dataset = dataset_list[0]\n", + "\n", + " # Only one `Analysis` so take first and only tracer.\n", + " tracer = tracer_list[0]\n", + "\n", + " einstein_mass = tracer.galaxies[0].einstein_mass_angular_from(grid=grid)\n", + " print(\"Einstein Mass (angular units) = \", einstein_mass)\n", + "\n", + " cosmology = al.cosmo.Planck15()\n", + "\n", + " critical_surface_density = (\n", + " cosmology.critical_surface_density_between_redshifts_from(\n", + " redshift_0=tracer.galaxies[0].redshift,\n", + " redshift_1=tracer.galaxies[-1].redshift,\n", + " )\n", + " )\n", + "\n", + " einstein_mass_kpc = einstein_mass * critical_surface_density\n", + "\n", + " print(\"Einstein Mass (kpc) = \", einstein_mass_kpc)\n", + " print(\"Einstein Mass (kpc) = \", \"{:.4e}\".format(einstein_mass_kpc))\n", + "\n", + " print(einstein_mass)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors (PDF from samples)__\n", + "\n", + "In this example, we will compute the errors on the axis ratio of a model. Computing the errors on a quantity \n", + "like the trap `density` is simple, because it is sampled by the non-linear search. The errors are therefore accessible\n", + "via the `Samples`, by marginalizing over all over parameters via the 1D Probability Density Function (PDF).\n", + "\n", + "Computing the errors on the axis ratio is more tricky, because it is a derived quantity. It is a parameter or \n", + "measurement that we want to calculate but was not sampled directly by the non-linear search. The `TracerAgg` object \n", + "object has everything we need to compute the errors of derived quantities.\n", + "\n", + "Below, we compute the axis ratio of every model sampled by the non-linear search and use this determine the PDF \n", + "of the axis ratio. When combining each axis ratio we weight each value by its `weight`. For Nautilus, \n", + "the nested sampler used by the fit, this ensures models which gave a bad fit (and thus have a low weight) do not \n", + "contribute significantly to the axis ratio error estimate.\n", + "\n", + "We set `minimum_weight=`1e-4`, such that any sample with a weight below this value is discarded when computing the \n", + "error. This speeds up the error computation by only using a small fraction of the total number of samples. Computing\n", + "a delta ellipticity is cheap, and this is probably not necessary. However, certain quantities have a non-negligible\n", + "computational overhead is being calculated and setting a minimum weight can speed up the calculation without \n", + "significantly changing the inferred errors.\n", + "\n", + "Below, we use the `TracerAgg` to get the `Tracer` of every Nautilus sample in each model-fit. We extract from each \n", + "tracer the model's axis-ratio, store them in a list and find the value via the PDF and quantile method. This again\n", + "uses generators, ensuring minimal memory use. \n", + "\n", + "In order to use these samples in the function `quantile`, we also need the weight list of the sample weights. We \n", + "compute this using the `TracerAgg`'s function `weights_above_gen_from`, which computes generators of the weights of all \n", + "points above this minimum value. This again ensures memory use in minimal." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_agg = al.agg.TracerAgg(aggregator=agg)\n", + "tracer_list_gen = tracer_agg.all_above_weight_gen_from(minimum_weight=1e-4)\n", + "weight_list_gen = tracer_agg.weights_above_gen_from(minimum_weight=1e-4)\n", + "\n", + "for tracer_gen, weight_gen in zip(tracer_list_gen, weight_list_gen):\n", + " axis_ratio_list = []\n", + "\n", + " for tracer_list in tracer_gen:\n", + " # Only one `Analysis` so take first and only tracer.\n", + " tracer = tracer_list[0]\n", + "\n", + " axis_ratio = al.convert.axis_ratio_from(\n", + " ell_comps=tracer.galaxies[0].mass.ell_comps\n", + " )\n", + "\n", + " axis_ratio_list.append(axis_ratio)\n", + "\n", + " weight_list = [weight for weight in weight_gen]\n", + "\n", + " if len(axis_ratio_list) > 1:\n", + " median_axis_ratio, lower_axis_ratio, upper_axis_ratio = af.marginalize(\n", + " parameter_list=axis_ratio_list, sigma=3.0, weight_list=weight_list\n", + " )\n", + "\n", + " print(\n", + " f\"Axis-Ratio = {median_axis_ratio} ({upper_axis_ratio} {lower_axis_ratio}\"\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors (Random draws from PDF)__\n", + "\n", + "An alternative approach to estimating the errors on a derived quantity is to randomly draw samples from the PDF \n", + "of the non-linear search. For a sufficiently high number of random draws, this should be as accurate and precise\n", + "as the method above. However, it can be difficult to be certain how many random draws are necessary.\n", + "\n", + "The weights of each sample are used to make every random draw. Therefore, when we compute the axis-ratio and its errors\n", + "we no longer need to pass the `weight_list` to the `quantile` function." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_agg = al.agg.TracerAgg(aggregator=agg)\n", + "tracer_list_gen = tracer_agg.randomly_drawn_via_pdf_gen_from(total_samples=2)\n", + "\n", + "for tracer_gen in tracer_list_gen:\n", + " axis_ratio_list = []\n", + "\n", + " for tracer_list in tracer_gen:\n", + " # Only one `Analysis` so take first and only tracer.\n", + " tracer = tracer_list[0]\n", + "\n", + " axis_ratio = al.convert.axis_ratio_from(\n", + " ell_comps=tracer.galaxies[0].mass.ell_comps\n", + " )\n", + "\n", + " axis_ratio_list.append(axis_ratio)\n", + "\n", + " median_axis_ratio, lower_axis_ratio, upper_axis_ratio = af.marginalize(\n", + " parameter_list=axis_ratio_list, sigma=3.0\n", + " )\n", + "\n", + " print(f\"Axis-Ratio = {median_axis_ratio} ({upper_axis_ratio} {lower_axis_ratio}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/aggregator/queries.ipynb b/notebooks/guides/results/aggregator/queries.ipynb index 31fccb4a1..ebbfffb8b 100644 --- a/notebooks/guides/results/aggregator/queries.ipynb +++ b/notebooks/guides/results/aggregator/queries.ipynb @@ -1,280 +1,317 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Queries\n", - "================\n", - "\n", - "Suppose we have the results of many fits in the database and we only wanted to load and inspect a specific set\n", - "of model-fits (e.g. the results of `tutorial_1_introduction`). We can use the database's querying tools to only load\n", - "the results we are interested in.\n", - "\n", - "The database also supports advanced querying, so that specific model-fits (e.g., which fit a certain model or dataset)\n", - "can be loaded.\n", - "\n", - "__Contents__\n", - "\n", - "- **Aggregator:** First, set up the aggregator as shown in `start_here.py`.\n", - "- **Unique Tag:** We can use the `Aggregator` to query the database and return only specific fits that we are.\n", - "- **Search Name:** We can also use the `name` of the search used to fit to the model as a query.\n", - "- **Model Queries:** We can also query based on the model fitted.\n", - "- **Logic:** Advanced queries can be constructed using logic, for example we below we combine the two queries." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Aggregator__\n", - "\n", - "First, set up the aggregator as shown in `start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "results_path = Path(\"output\") / \"results_folder\"\n", - "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "agg = Aggregator.from_directory(\n", - " directory=results_path,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Unique Tag__\n", - "\n", - "We can use the `Aggregator` to query the database and return only specific fits that we are interested in. We first \n", - "do this using the `unique_tag` which we can query to load the results of a specific `dataset_name` string we \n", - "input into the model-fit's search. \n", - "\n", - "By querying using the string `lens_sersic` the model-fit to only the second dataset is returned:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "unique_tag = agg.search.unique_tag\n", - "agg_query = agg.query(unique_tag == \"simple__no_lens_light\")\n", - "samples_gen = agg_query.values(\"samples\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As expected, this list now has only 1 `SamplesNest` corresponding to the second dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Directory Filtered DynestySampler Samples: \\n\")\n", - "print(\"Total Samples Objects via unique tag = \", len(list(samples_gen)), \"\\n\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we query using an incorrect dataset name we get no results:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "unique_tag = agg.search.unique_tag\n", - "agg_query = agg.query(unique_tag == \"incorrect_name\")\n", - "samples_gen = agg_query.values(\"samples\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Name__\n", - "\n", - "We can also use the `name` of the search used to fit to the model as a query. \n", - "\n", - "In this example, all three fits used the same search, which had the `name` `database_example`. Thus, using it as a \n", - "query in this example is somewhat pointless. However, querying based on the search name is very useful for model-fits\n", - "which use search chaining (see chapter 3 **HowToLens**), where the results of a particular fit in the chain can be\n", - "instantly loaded.\n", - "\n", - "As expected, this query contains all 3 results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "name = agg.search.name\n", - "agg_query = agg.query(name == \"database_example\")\n", - "print(\"Total Queried Results via search name = \", len(agg_query), \"\\n\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Queries__\n", - "\n", - "We can also query based on the model fitted. \n", - "\n", - "For example, we can load all results which fitted an `Isothermal` model-component, which in this simple \n", - "example is all 3 model-fits.\n", - "\n", - "The ability to query via the model is extremely powerful. It enables a user to fit many models to large samples \n", - "of lenses efficiently load and inspect the results. \n", - "\n", - "[Note: the code `agg.model.galaxies.lens.mass` corresponds to the fact that in the `Model` we named the model components \n", - "`galaxies`, `lens` and `mass`. If the `Model` had used a different name the code below would change correspondingly. \n", - "Models with multiple galaxies are therefore easily accessed via the database.]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = agg.model.galaxies.lens\n", - "agg_query = agg.query(lens.mass == al.mp.Isothermal)\n", - "print(\"Total Samples Objects via `Isothermal` model query = \", len(agg_query), \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Queries using the results of model-fitting are also supported. Below, we query the database to find all fits where the\n", - "inferred value of `einstein_radius` for the `Isothermal` mass of the lens is less than 1.5 (which returns only\n", - "the first of the three model-fits)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = agg.model.galaxies.lens.mass\n", - "agg_query = agg.query(mass.einstein_radius < 1.5)\n", - "print(\n", - " \"Total Samples Objects In Query `lens.mass.einstein_radius < 1.5` = \",\n", - " len(agg_query),\n", - " \"\\n\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Logic__\n", - "\n", - "Advanced queries can be constructed using logic, for example we below we combine the two queries above to find all\n", - "results which fitted an `Isothermal` mass model AND (using the & symbol) inferred a value of einstein radius above\n", - "1.0 for the lens's mass \n", - "\n", - "The OR logical clause is also supported via the symbol |." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = agg.model.galaxies.lens.mass\n", - "agg_query = agg.query((mass == al.mp.Isothermal) & (mass.einstein_radius > 1.0))\n", - "print(\n", - " \"Total Samples Objects In Query `Isothermal and einstein_radius > 3.0` = \",\n", - " len(agg_query),\n", - " \"\\n\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Queries\n", + "================\n", + "\n", + "Suppose we have the results of many fits in the database and we only wanted to load and inspect a specific set\n", + "of model-fits (e.g. the results of `tutorial_1_introduction`). We can use the database's querying tools to only load\n", + "the results we are interested in.\n", + "\n", + "The database also supports advanced querying, so that specific model-fits (e.g., which fit a certain model or dataset)\n", + "can be loaded.\n", + "\n", + "__Contents__\n", + "\n", + "- **Aggregator:** First, set up the aggregator as shown in `start_here.py`.\n", + "- **Unique Tag:** We can use the `Aggregator` to query the database and return only specific fits that we are.\n", + "- **Search Name:** We can also use the `name` of the search used to fit to the model as a query.\n", + "- **Model Queries:** We can also query based on the model fitted.\n", + "- **Logic:** Advanced queries can be constructed using logic, for example we below we combine the two queries." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Aggregator__\n", + "\n", + "First, set up the aggregator as shown in `start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "results_path = Path(\"output\") / \"results_folder\"\n", + "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "agg = Aggregator.from_directory(\n", + " directory=results_path,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Unique Tag__\n", + "\n", + "We can use the `Aggregator` to query the database and return only specific fits that we are interested in. We first \n", + "do this using the `unique_tag` which we can query to load the results of a specific `dataset_name` string we \n", + "input into the model-fit's search. \n", + "\n", + "By querying using the string `lens_sersic` the model-fit to only the second dataset is returned:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "unique_tag = agg.search.unique_tag\n", + "agg_query = agg.query(unique_tag == \"simple__no_lens_light\")\n", + "samples_gen = agg_query.values(\"samples\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As expected, this list now has only 1 `SamplesNest` corresponding to the second dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Directory Filtered DynestySampler Samples: \\n\")\n", + "print(\"Total Samples Objects via unique tag = \", len(list(samples_gen)), \"\\n\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we query using an incorrect dataset name we get no results:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "unique_tag = agg.search.unique_tag\n", + "agg_query = agg.query(unique_tag == \"incorrect_name\")\n", + "samples_gen = agg_query.values(\"samples\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Name__\n", + "\n", + "We can also use the `name` of the search used to fit to the model as a query. \n", + "\n", + "In this example, all three fits used the same search, which had the `name` `database_example`. Thus, using it as a \n", + "query in this example is somewhat pointless. However, querying based on the search name is very useful for model-fits\n", + "which use search chaining (see chapter 3 **HowToLens**), where the results of a particular fit in the chain can be\n", + "instantly loaded.\n", + "\n", + "As expected, this query contains all 3 results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "name = agg.search.name\n", + "agg_query = agg.query(name == \"database_example\")\n", + "print(\"Total Queried Results via search name = \", len(agg_query), \"\\n\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Queries__\n", + "\n", + "We can also query based on the model fitted. \n", + "\n", + "For example, we can load all results which fitted an `Isothermal` model-component, which in this simple \n", + "example is all 3 model-fits.\n", + "\n", + "The ability to query via the model is extremely powerful. It enables a user to fit many models to large samples \n", + "of lenses efficiently load and inspect the results. \n", + "\n", + "[Note: the code `agg.model.galaxies.lens.mass` corresponds to the fact that in the `Model` we named the model components \n", + "`galaxies`, `lens` and `mass`. If the `Model` had used a different name the code below would change correspondingly. \n", + "Models with multiple galaxies are therefore easily accessed via the database.]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = agg.model.galaxies.lens\n", + "agg_query = agg.query(lens.mass == al.mp.Isothermal)\n", + "print(\"Total Samples Objects via `Isothermal` model query = \", len(agg_query), \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Queries using the results of model-fitting are also supported. Below, we query the database to find all fits where the\n", + "inferred value of `einstein_radius` for the `Isothermal` mass of the lens is less than 1.5 (which returns only\n", + "the first of the three model-fits)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = agg.model.galaxies.lens.mass\n", + "agg_query = agg.query(mass.einstein_radius < 1.5)\n", + "print(\n", + " \"Total Samples Objects In Query `lens.mass.einstein_radius < 1.5` = \",\n", + " len(agg_query),\n", + " \"\\n\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Logic__\n", + "\n", + "Advanced queries can be constructed using logic, for example we below we combine the two queries above to find all\n", + "results which fitted an `Isothermal` mass model AND (using the & symbol) inferred a value of einstein radius above\n", + "1.0 for the lens's mass \n", + "\n", + "The OR logical clause is also supported via the symbol |." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = agg.model.galaxies.lens.mass\n", + "agg_query = agg.query((mass == al.mp.Isothermal) & (mass.einstein_radius > 1.0))\n", + "print(\n", + " \"Total Samples Objects In Query `Isothermal and einstein_radius > 3.0` = \",\n", + " len(agg_query),\n", + " \"\\n\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/aggregator/samples.ipynb b/notebooks/guides/results/aggregator/samples.ipynb index 8b1b1adc8..0c3e3bec5 100644 --- a/notebooks/guides/results/aggregator/samples.ipynb +++ b/notebooks/guides/results/aggregator/samples.ipynb @@ -1,973 +1,1010 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Samples\n", - "================\n", - "\n", - "After a non-linear search has completed, it returns a `Result` object that contains information on samples of\n", - "the non-linear search, such as the maximum likelihood model instance, the errors on each parameter and the\n", - "Bayesian evidence.\n", - "\n", - "This script illustrates how to use the result to inspect the non-linear search samples.\n", - "\n", - "__Contents__\n", - "\n", - "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", - "- **Model Fit:** Perform the model-fit using the search and analysis.\n", - "- **Info:** As seen throughout the workspace, the `info` attribute shows the result in a readable format.\n", - "- **Plot:** We now have the `Result` object we will cover in this script.\n", - "- **Samples:** The result contains a `Samples` object, which contains all samples of the non-linear search.\n", - "- **Parameters:** The parameters are stored as a list of lists, where.\n", - "- **Figures of Merit:** The `Samples` class contains the log likelihood, log prior, log posterior and weight_list of every.\n", - "- **Instances:** Many results can be returned as an instance of the model, using the Python class structure of the.\n", - "- **Errors:** Methods for computing error estimates on all parameters are provided.\n", - "- **Sample Instance:** A non-linear search retains every model that is accepted during the model-fit.\n", - "- **Search Plots:** The Probability Density Functions (PDF's) of the results can be plotted using the non-linear search.\n", - "- **Maximum Likelihood:** The maximum log likelihood value of the model-fit can be estimated by simple taking the maximum of.\n", - "- **Bayesian Evidence:** Nested sampling algorithms like Nautilus also estimate the Bayesian evidence (estimated via the.\n", - "- **Lists:** All results can alternatively be returned as a 1D list of values, by passing `as_instance=False`.\n", - "- **Latex:** If you are writing modeling results up in a paper, you can use inbuilt latex tools to create latex.\n", - "\n", - "__Units__\n", - "\n", - "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", - "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", - "\n", - "The guide `guides/units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", - "kiloparsecs, magnitudes and solar masses.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `results/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Quick Fit Auto-Trigger__\n", - "\n", - "If a previous fit has not been run yet, the shared helper ``_quick_fit.py`` is invoked to produce one.\n", - "The helper writes a capped Nautilus fit to ``output/results_folder/`` so this tutorial (and its siblings\n", - "in this folder) have results to work with. When that folder already exists the helper exits immediately,\n", - "so re-running this tutorial is cheap." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "results_path = Path(\"output\") / \"results_folder\"\n", - "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "To illustrate results, we need to perform a model-fit in order to create a `Result` object.\n", - "\n", - "The code below performs a model-fit using Nautilus. The helper above already wrote a completed fit to\n", - "``output/results_folder/``, so the ``search.fit(...)`` call below resumes from that checkpoint and\n", - "returns the in-memory ``Result`` object without redoing the search.\n", - "\n", - "You should be familiar with modeling already, if not read the `modeling/start_here.py` script before reading this one!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", - " source=af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.Sersic),\n", - " ),\n", - ")\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"results_folder\"),\n", - " name=\"results\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU batching and VRAM use explained in `modeling` examples.\n", - " n_like_max=300,\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Info__\n", - "\n", - "As seen throughout the workspace, the `info` attribute shows the result in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Plot__\n", - "\n", - "We now have the `Result` object we will cover in this script. \n", - "\n", - "As a reminder, in the `modeling` scripts we use the `max_log_likelihood_tracer` and `max_log_likelihood_fit` to plot \n", - "the results of the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_tracer(\n", - " tracer=result.max_log_likelihood_tracer, grid=mask.derive_grid.all_false\n", - ")\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results tutorials `tracer.py` and `fit.py` expand on the `max_log_likelihood_tracer` and `max_log_likelihood_fit`, \n", - "showing how they can be used to inspect many aspects of a model.\n", - "\n", - "__Samples__\n", - "\n", - "The result contains a `Samples` object, which contains all samples of the non-linear search.\n", - "\n", - "Each sample corresponds to a set of model parameters that were evaluated and accepted by the non linear search, \n", - "in this example `Nautilus`. \n", - "\n", - "This includes their log likelihoods, which are used for computing additional information about the model-fit,\n", - "for example the error on every parameter. \n", - "\n", - "Our model-fit used the nested sampling algorithm Nautilus, so the `Samples` object returned is a `SamplesNest` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "print(\"Nest Samples: \\n\")\n", - "print(samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parameters__\n", - "\n", - "The parameters are stored as a list of lists, where:\n", - "\n", - " - The outer list is the size of the total number of samples.\n", - " - The inner list is the size of the number of free parameters in the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"All parameters of the very first sample\")\n", - "print(samples.parameter_lists[0])\n", - "print(\"The fourth parameter of the tenth sample\")\n", - "print(samples.parameter_lists[9][3])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Figures of Merit__\n", - "\n", - "The `Samples` class contains the log likelihood, log prior, log posterior and weight_list of every accepted sample, where:\n", - "\n", - "- The `log_likelihood` is the value evaluated in the `log_likelihood_function`.\n", - "\n", - "- The `log_prior` encodes information on how parameter priors map log likelihood values to log posterior values.\n", - "\n", - "- The `log_posterior` is `log_likelihood + log_prior`.\n", - "\n", - "- The `weight` gives information on how samples are combined to estimate the posterior, which depends on type of search\n", - " used (for `Nautilus` they are all non-zero values which sum to 1).\n", - "\n", - "Lets inspect these values for the tenth sample." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"log(likelihood), log(prior), log(posterior) and weight of the tenth sample.\")\n", - "print(samples.log_likelihood_list[9])\n", - "print(samples.log_prior_list[9])\n", - "print(samples.log_posterior_list[9])\n", - "print(samples.weight_list[9])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Instances__\n", - "\n", - "Many results can be returned as an instance of the model, using the Python class structure of the model composition.\n", - "\n", - "For example, we can return the model parameters corresponding to the maximum log likelihood sample." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "instance = samples.max_log_likelihood()\n", - "print(\"Maximum Log Likelihood Model Instance: \\n\")\n", - "print(instance, \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The attributes of the `instance` (e.g. `galaxies`, `lens`) have these names due to how we composed the `Galaxy` and\n", - "its light and mass profiles via the `Collection` and `Model` above. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(instance.galaxies)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "These galaxies will be named according to the model fitted by the search (in this case, `lens` and `source`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(instance.galaxies.lens)\n", - "print(instance.galaxies.source)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Their light profiles are also named according to model composition allowing individual parameters to be printed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(instance.galaxies.lens.mass.einstein_radius)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can use this list of galaxies to create the maximum log likelihood `Tracer`, which is the property of the result \n", - "we've used up to now!\n", - "\n", - "Using this tracer is expanded upon in the `tracer.py` results tutorial.\n", - "\n", - "(If we had the `Imaging` available we could easily use this to create the maximum log likelihood `FitImaging`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "max_lh_tracer = al.Tracer(galaxies=instance.galaxies)\n", - "\n", - "print(max_lh_tracer)\n", - "print(mask.derive_grid.all_false)\n", - "\n", - "# Input to FitImaging to solve for linear light profile intensities, see `start_here.py` for details.\n", - "fit = al.FitImaging(dataset=dataset, tracer=max_lh_tracer)\n", - "max_lh_tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", - "\n", - "aplt.subplot_tracer(tracer=max_lh_tracer, grid=mask.derive_grid.all_false)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Posterior / PDF__\n", - "\n", - "The result contains the full posterior information of our non-linear search, which can be used for parameter \n", - "estimation. \n", - "\n", - "PDF stands for \"Probability Density Function\" and it quantifies probability of each model parameter having values\n", - "that are sampled. It therefore enables error estimation via a process called marginalization.\n", - "\n", - "The median pdf vector is available, which estimates every parameter via 1D marginalization of their PDFs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "instance = samples.median_pdf()\n", - "\n", - "print(\"Median PDF Model Instances: \\n\")\n", - "print(instance, \"\\n\")\n", - "print(instance.galaxies.source.bulge)\n", - "print()\n", - "\n", - "vector = samples.median_pdf(as_instance=False)\n", - "\n", - "print(\"Median PDF Model Parameter Lists: \\n\")\n", - "print(vector, \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors__\n", - "\n", - "Methods for computing error estimates on all parameters are provided. \n", - "\n", - "This again uses 1D marginalization, now at an input sigma confidence limit. \n", - "\n", - "By inputting `sigma=3.0` margnialization find the values spanning 99.7% of 1D PDF. Changing this to `sigma=1.0`\n", - "would give the errors at the 68.3% confidence limit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "instance_upper_sigma = samples.values_at_upper_sigma(sigma=3.0)\n", - "instance_lower_sigma = samples.values_at_lower_sigma(sigma=3.0)\n", - "\n", - "print(\"Errors Instances: \\n\")\n", - "print(instance_upper_sigma.galaxies.source.bulge, \"\\n\")\n", - "print(instance_lower_sigma.galaxies.source.bulge, \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "They can also be returned at the values of the parameters at their error values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "instance_upper_values = samples.errors_at_upper_sigma(sigma=3.0)\n", - "instance_lower_values = samples.errors_at_lower_sigma(sigma=3.0)\n", - "\n", - "print(\"Errors Instances: \\n\")\n", - "print(instance_upper_values.galaxies.source.bulge, \"\\n\")\n", - "print(instance_lower_values.galaxies.source.bulge, \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sample Instance__\n", - "\n", - "A non-linear search retains every model that is accepted during the model-fit.\n", - "\n", - "We can create an instance of any model -- below we create an instance of the last accepted model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "instance = samples.from_sample_index(sample_index=-1)\n", - "\n", - "print(instance.galaxies.lens.mass)\n", - "print(instance.galaxies.lens.mass)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Plots__\n", - "\n", - "The Probability Density Functions (PDF's) of the results can be plotted using the non-linear search in-built \n", - "visualization tools.\n", - "\n", - "This fit used `Nautilus` therefore we use the nested sampling visualization functions, which wrap in-built\n", - "visualization tools.\n", - "\n", - "The `autofit_workspace/*/plots` folder illustrates other packages that can be used to make these plots using\n", - "the standard output results formats (e.g. `GetDist.py`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Maximum Likelihood__\n", - "\n", - "The maximum log likelihood value of the model-fit can be estimated by simple taking the maximum of all log\n", - "likelihoods of the samples.\n", - "\n", - "If different models are fitted to the same dataset, this value can be compared to determine which model provides\n", - "the best fit (e.g. which model has the highest maximum likelihood)?" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Maximum Log Likelihood: \\n\")\n", - "print(max(samples.log_likelihood_list))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Bayesian Evidence__\n", - "\n", - "Nested sampling algorithms like Nautilus also estimate the Bayesian evidence (estimated via the nested sampling \n", - "algorithm).\n", - "\n", - "The Bayesian evidence accounts for \"Occam's Razor\", whereby it penalizes models for being more complex (e.g. if a model\n", - "has more parameters it needs to fit the da\n", - "\n", - "The Bayesian evidence is a better quantity to use to compare models, because it penalizes models with more parameters\n", - "for being more complex (\"Occam's Razor\"). Comparisons using the maximum likelihood value do not account for this and\n", - "therefore may unjustly favour more complex models.\n", - "\n", - "Using the Bayesian evidence for model comparison is well documented on the internet, for example the following\n", - "wikipedia page: https://en.wikipedia.org/wiki/Bayes_factor" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Maximum Log Likelihood and Log Evidence: \\n\")\n", - "print(samples.log_evidence)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lists__\n", - "\n", - "All results can alternatively be returned as a 1D list of values, by passing `as_instance=False`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "max_lh_list = samples.max_log_likelihood(as_instance=False)\n", - "print(\"Max Log Likelihood Model Parameters: \\n\")\n", - "print(max_lh_list, \"\\n\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The list above does not tell us which values correspond to which parameters.\n", - "\n", - "The following quantities are available in the `Model`, where the order of their entries correspond to the parameters \n", - "in the `ml_vector` above:\n", - "\n", - " - `paths`: a list of tuples which give the path of every parameter in the `Model`.\n", - " - `parameter_names`: a list of shorthand parameter names derived from the `paths`.\n", - " - `parameter_labels`: a list of parameter labels used when visualizing non-linear search results (see below).\n", - "\n", - "For simple models like the one fitted in this tutorial, the quantities below are somewhat redundant. For the\n", - "more complex models they are important for tracking the parameters of the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model = samples.model\n", - "\n", - "print(model.paths)\n", - "print(model.parameter_names)\n", - "print(model.parameter_labels)\n", - "print(model.model_component_and_parameter_names)\n", - "print(\"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "All the methods above are available as lists." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "instance = samples.median_pdf(as_instance=False)\n", - "values_at_upper_sigma = samples.values_at_upper_sigma(sigma=3.0, as_instance=False)\n", - "values_at_lower_sigma = samples.values_at_lower_sigma(sigma=3.0, as_instance=False)\n", - "errors_at_upper_sigma = samples.errors_at_upper_sigma(sigma=3.0, as_instance=False)\n", - "errors_at_lower_sigma = samples.errors_at_lower_sigma(sigma=3.0, as_instance=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Latex__\n", - "\n", - "If you are writing modeling results up in a paper, you can use inbuilt latex tools to create latex table code which \n", - "you can copy to your .tex document.\n", - "\n", - "By combining this with the filtering tools below, specific parameters can be included or removed from the latex.\n", - "\n", - "Remember that the superscripts of a parameter are loaded from the config file `notation/label.yaml`, providing high\n", - "levels of customization for how the parameter names appear in the latex table. This is especially useful if your model\n", - "uses the same model components with the same parameter, which therefore need to be distinguished via superscripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "latex = af.text.Samples.latex(\n", - " samples=result.samples,\n", - " median_pdf_model=True,\n", - " sigma=3.0,\n", - " name_to_label=True,\n", - " include_name=True,\n", - " include_quickmath=True,\n", - " prefix=\"Example Prefix \",\n", - " suffix=r\"\\\\[-2pt]\",\n", - ")\n", - "\n", - "print(latex)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Derived Errors (Advanced)__\n", - "\n", - "Computing the errors of a quantity like the `einstein_radius` is simple, because it is sampled by the non-linear \n", - "search. Errors are accessible using the `Samples` object's `errors_from` methods, which marginalize over the \n", - "parameters via the 1D Probability Density Function (PDF).\n", - "\n", - "Computing errors on derived quantities is more tricky, because they are not sampled directly by the non-linear search. \n", - "For example, what if we want the error on the axis-ratio of the mass model? In order to do this we need to create the \n", - "PDF of that derived quantity, which we can then marginalize over using the same function we use to marginalize model \n", - "parameters.\n", - "\n", - "Below, we compute the axis-ratio of every accepted model sampled by the non-linear search and use this determine the PDF \n", - "of the axis-ratio. When combining the axis-ratio's we weight each value by its `weight`. For Nautilus, a nested sampling \n", - "algorithm, the weight of every sample is different and thus must be included.\n", - "\n", - "In order to pass these samples to the function `marginalize`, which marginalizes over the PDF of the axis-ratio to \n", - "compute its error, we also pass the weight list of the samples.\n", - "\n", - "Note again how because when creating the model above using the input names `lens` and `mass` we access the instance\n", - "below using these." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "axis_ratio_list = []\n", - "\n", - "for sample in samples.sample_list:\n", - " instance = sample.instance_for_model(model=samples.model, ignore_assertions=True)\n", - "\n", - " ell_comps = instance.galaxies.lens.mass.ell_comps\n", - "\n", - " axis_ratio = al.convert.axis_ratio_from(ell_comps=ell_comps)\n", - "\n", - " axis_ratio_list.append(axis_ratio)\n", - "\n", - "median_axis_ratio, lower_axis_ratio, upper_axis_ratio = af.marginalize(\n", - " parameter_list=axis_ratio_list, sigma=3.0, weight_list=samples.weight_list\n", - ")\n", - "\n", - "print(f\"axis_ratio = {median_axis_ratio} ({upper_axis_ratio} {lower_axis_ratio}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The calculation above could be computationally expensive, if there are many samples and the derived quantity is\n", - "slow to compute.\n", - "\n", - "An alternative approach, which will provide comparable accuracy provided enough draws are used, is to sample \n", - "points randomy from the PDF of the model and use these to compute the derived quantity.\n", - "\n", - "Draws are from the PDF of the model, so the weights of the samples are accounted for and we therefore do not\n", - "pass them to the `marginalize` function (it essentially treats all samples as having equal weight)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "random_draws = 50\n", - "\n", - "axis_ratio_list = []\n", - "\n", - "for i in range(random_draws):\n", - " instance = samples.draw_randomly_via_pdf()\n", - "\n", - " ell_comps = instance.galaxies.lens.mass.ell_comps\n", - "\n", - " axis_ratio = al.convert.axis_ratio_from(ell_comps=ell_comps)\n", - "\n", - " axis_ratio_list.append(axis_ratio)\n", - "\n", - "median_axis_ratio, lower_axis_ratio, upper_axis_ratio = af.marginalize(\n", - " parameter_list=axis_ratio_list,\n", - " sigma=3.0,\n", - ")\n", - "\n", - "print(f\"axis_ratio = {median_axis_ratio} ({upper_axis_ratio} {lower_axis_ratio}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Samples Filtering (Advanced)__\n", - "\n", - "Our samples object has the results for all three parameters in our model. However, we might only be interested in the\n", - "results of a specific parameter.\n", - "\n", - "The basic form of filtering specifies parameters via their path, which was printed above via the model and is printed \n", - "again below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "print(\"Parameter paths in the model which are used for filtering:\")\n", - "print(samples.model.paths)\n", - "\n", - "print(\"All parameters of the very first sample\")\n", - "print(samples.parameter_lists[0])\n", - "\n", - "samples = samples.with_paths(\n", - " [\n", - " (\"galaxies\", \"lens\", \"mass\", \"einstein_radius\"),\n", - " (\"galaxies\", \"source\", \"bulge\", \"sersic_index\"),\n", - " ]\n", - ")\n", - "\n", - "print(\n", - " \"All parameters of the very first sample (containing only the lens mass's einstein radius and \"\n", - " \"source bulge's sersic index).\"\n", - ")\n", - "print(samples.parameter_lists[0])\n", - "\n", - "print(\n", - " \"Maximum Log Likelihood Model Instances (containing only the lens mass's einstein radius and \"\n", - " \"source bulge's sersic index):\\n\"\n", - ")\n", - "print(samples.max_log_likelihood(as_instance=False))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We specified each path as a list of tuples of strings. \n", - "\n", - "This is how the source code internally stores the path to different components of the model, but it is not \n", - "consistent with the API used to compose a model.\n", - "\n", - "We can alternatively use the following API:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "samples = samples.with_paths(\n", - " [\"galaxies.lens.mass.einstein_radius\", \"galaxies.source.bulge.sersic_index\"]\n", - ")\n", - "\n", - "print(\n", - " \"All parameters of the very first sample (containing only the lens mass's einstein radius and \"\n", - " \"source bulge's sersic index).\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can alternatively filter the `Samples` object by removing all parameters with a certain path. Below, we remove\n", - "the centres of the mass model to be left with 10 parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "print(\"Parameter paths in the model which are used for filtering:\")\n", - "print(samples.model.paths)\n", - "\n", - "print(\"Parameters of first sample\")\n", - "print(samples.parameter_lists[0])\n", - "\n", - "print(samples.model.total_free_parameters)\n", - "\n", - "samples = samples.without_paths(\n", - " [\n", - " # \"galaxies.lens.mass.centre\"),\n", - " \"galaxies.lens.mass.centre.centre_0\",\n", - " # \"galaxies.lens.mass.centre.centre_1),\n", - " ]\n", - ")\n", - "\n", - "print(\"Parameters of first sample without the lens mass centre.\")\n", - "print(samples.parameter_lists[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can keep and remove entire paths of the samples, for example keeping only the parameters of the lens or \n", - "removing all parameters of the source's bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "samples = samples.with_paths([\"galaxies.lens\"])\n", - "print(\"Parameters of the first sample of the lens galaxy\")\n", - "print(samples.parameter_lists[0])\n", - "\n", - "samples = result.samples\n", - "samples = samples.without_paths([\"galaxies.source.bulge\"])\n", - "print(\"Parameters of the first sample without the source's bulge\")\n", - "print(samples.parameter_lists[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Samples\n", + "================\n", + "\n", + "After a non-linear search has completed, it returns a `Result` object that contains information on samples of\n", + "the non-linear search, such as the maximum likelihood model instance, the errors on each parameter and the\n", + "Bayesian evidence.\n", + "\n", + "This script illustrates how to use the result to inspect the non-linear search samples.\n", + "\n", + "__Contents__\n", + "\n", + "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", + "- **Model Fit:** Perform the model-fit using the search and analysis.\n", + "- **Info:** As seen throughout the workspace, the `info` attribute shows the result in a readable format.\n", + "- **Plot:** We now have the `Result` object we will cover in this script.\n", + "- **Samples:** The result contains a `Samples` object, which contains all samples of the non-linear search.\n", + "- **Parameters:** The parameters are stored as a list of lists, where.\n", + "- **Figures of Merit:** The `Samples` class contains the log likelihood, log prior, log posterior and weight_list of every.\n", + "- **Instances:** Many results can be returned as an instance of the model, using the Python class structure of the.\n", + "- **Errors:** Methods for computing error estimates on all parameters are provided.\n", + "- **Sample Instance:** A non-linear search retains every model that is accepted during the model-fit.\n", + "- **Search Plots:** The Probability Density Functions (PDF's) of the results can be plotted using the non-linear search.\n", + "- **Maximum Likelihood:** The maximum log likelihood value of the model-fit can be estimated by simple taking the maximum of.\n", + "- **Bayesian Evidence:** Nested sampling algorithms like Nautilus also estimate the Bayesian evidence (estimated via the.\n", + "- **Lists:** All results can alternatively be returned as a 1D list of values, by passing `as_instance=False`.\n", + "- **Latex:** If you are writing modeling results up in a paper, you can use inbuilt latex tools to create latex.\n", + "\n", + "__Units__\n", + "\n", + "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", + "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", + "\n", + "The guide `guides/units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", + "kiloparsecs, magnitudes and solar masses.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `results/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Quick Fit Auto-Trigger__\n", + "\n", + "If a previous fit has not been run yet, the shared helper ``_quick_fit.py`` is invoked to produce one.\n", + "The helper writes a capped Nautilus fit to ``output/results_folder/`` so this tutorial (and its siblings\n", + "in this folder) have results to work with. When that folder already exists the helper exits immediately,\n", + "so re-running this tutorial is cheap." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "results_path = Path(\"output\") / \"results_folder\"\n", + "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "To illustrate results, we need to perform a model-fit in order to create a `Result` object.\n", + "\n", + "The code below performs a model-fit using Nautilus. The helper above already wrote a completed fit to\n", + "``output/results_folder/``, so the ``search.fit(...)`` call below resumes from that checkpoint and\n", + "returns the in-memory ``Result`` object without redoing the search.\n", + "\n", + "You should be familiar with modeling already, if not read the `modeling/start_here.py` script before reading this one!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", + " source=af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.Sersic),\n", + " ),\n", + ")\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"results_folder\"),\n", + " name=\"results\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU batching and VRAM use explained in `modeling` examples.\n", + " n_like_max=300,\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Info__\n", + "\n", + "As seen throughout the workspace, the `info` attribute shows the result in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Plot__\n", + "\n", + "We now have the `Result` object we will cover in this script. \n", + "\n", + "As a reminder, in the `modeling` scripts we use the `max_log_likelihood_tracer` and `max_log_likelihood_fit` to plot \n", + "the results of the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_tracer(\n", + " tracer=result.max_log_likelihood_tracer, grid=mask.derive_grid.all_false\n", + ")\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results tutorials `tracer.py` and `fit.py` expand on the `max_log_likelihood_tracer` and `max_log_likelihood_fit`, \n", + "showing how they can be used to inspect many aspects of a model.\n", + "\n", + "__Samples__\n", + "\n", + "The result contains a `Samples` object, which contains all samples of the non-linear search.\n", + "\n", + "Each sample corresponds to a set of model parameters that were evaluated and accepted by the non linear search, \n", + "in this example `Nautilus`. \n", + "\n", + "This includes their log likelihoods, which are used for computing additional information about the model-fit,\n", + "for example the error on every parameter. \n", + "\n", + "Our model-fit used the nested sampling algorithm Nautilus, so the `Samples` object returned is a `SamplesNest` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "print(\"Nest Samples: \\n\")\n", + "print(samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parameters__\n", + "\n", + "The parameters are stored as a list of lists, where:\n", + "\n", + " - The outer list is the size of the total number of samples.\n", + " - The inner list is the size of the number of free parameters in the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"All parameters of the very first sample\")\n", + "print(samples.parameter_lists[0])\n", + "print(\"The fourth parameter of the tenth sample\")\n", + "print(samples.parameter_lists[9][3])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Figures of Merit__\n", + "\n", + "The `Samples` class contains the log likelihood, log prior, log posterior and weight_list of every accepted sample, where:\n", + "\n", + "- The `log_likelihood` is the value evaluated in the `log_likelihood_function`.\n", + "\n", + "- The `log_prior` encodes information on how parameter priors map log likelihood values to log posterior values.\n", + "\n", + "- The `log_posterior` is `log_likelihood + log_prior`.\n", + "\n", + "- The `weight` gives information on how samples are combined to estimate the posterior, which depends on type of search\n", + " used (for `Nautilus` they are all non-zero values which sum to 1).\n", + "\n", + "Lets inspect these values for the tenth sample." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"log(likelihood), log(prior), log(posterior) and weight of the tenth sample.\")\n", + "print(samples.log_likelihood_list[9])\n", + "print(samples.log_prior_list[9])\n", + "print(samples.log_posterior_list[9])\n", + "print(samples.weight_list[9])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Instances__\n", + "\n", + "Many results can be returned as an instance of the model, using the Python class structure of the model composition.\n", + "\n", + "For example, we can return the model parameters corresponding to the maximum log likelihood sample." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "instance = samples.max_log_likelihood()\n", + "print(\"Maximum Log Likelihood Model Instance: \\n\")\n", + "print(instance, \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The attributes of the `instance` (e.g. `galaxies`, `lens`) have these names due to how we composed the `Galaxy` and\n", + "its light and mass profiles via the `Collection` and `Model` above. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(instance.galaxies)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "These galaxies will be named according to the model fitted by the search (in this case, `lens` and `source`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(instance.galaxies.lens)\n", + "print(instance.galaxies.source)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Their light profiles are also named according to model composition allowing individual parameters to be printed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(instance.galaxies.lens.mass.einstein_radius)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can use this list of galaxies to create the maximum log likelihood `Tracer`, which is the property of the result \n", + "we've used up to now!\n", + "\n", + "Using this tracer is expanded upon in the `tracer.py` results tutorial.\n", + "\n", + "(If we had the `Imaging` available we could easily use this to create the maximum log likelihood `FitImaging`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "max_lh_tracer = al.Tracer(galaxies=instance.galaxies)\n", + "\n", + "print(max_lh_tracer)\n", + "print(mask.derive_grid.all_false)\n", + "\n", + "# Input to FitImaging to solve for linear light profile intensities, see `start_here.py` for details.\n", + "fit = al.FitImaging(dataset=dataset, tracer=max_lh_tracer)\n", + "max_lh_tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", + "\n", + "aplt.subplot_tracer(tracer=max_lh_tracer, grid=mask.derive_grid.all_false)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Posterior / PDF__\n", + "\n", + "The result contains the full posterior information of our non-linear search, which can be used for parameter \n", + "estimation. \n", + "\n", + "PDF stands for \"Probability Density Function\" and it quantifies probability of each model parameter having values\n", + "that are sampled. It therefore enables error estimation via a process called marginalization.\n", + "\n", + "The median pdf vector is available, which estimates every parameter via 1D marginalization of their PDFs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "instance = samples.median_pdf()\n", + "\n", + "print(\"Median PDF Model Instances: \\n\")\n", + "print(instance, \"\\n\")\n", + "print(instance.galaxies.source.bulge)\n", + "print()\n", + "\n", + "vector = samples.median_pdf(as_instance=False)\n", + "\n", + "print(\"Median PDF Model Parameter Lists: \\n\")\n", + "print(vector, \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors__\n", + "\n", + "Methods for computing error estimates on all parameters are provided. \n", + "\n", + "This again uses 1D marginalization, now at an input sigma confidence limit. \n", + "\n", + "By inputting `sigma=3.0` margnialization find the values spanning 99.7% of 1D PDF. Changing this to `sigma=1.0`\n", + "would give the errors at the 68.3% confidence limit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "instance_upper_sigma = samples.values_at_upper_sigma(sigma=3.0)\n", + "instance_lower_sigma = samples.values_at_lower_sigma(sigma=3.0)\n", + "\n", + "print(\"Errors Instances: \\n\")\n", + "print(instance_upper_sigma.galaxies.source.bulge, \"\\n\")\n", + "print(instance_lower_sigma.galaxies.source.bulge, \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "They can also be returned at the values of the parameters at their error values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "instance_upper_values = samples.errors_at_upper_sigma(sigma=3.0)\n", + "instance_lower_values = samples.errors_at_lower_sigma(sigma=3.0)\n", + "\n", + "print(\"Errors Instances: \\n\")\n", + "print(instance_upper_values.galaxies.source.bulge, \"\\n\")\n", + "print(instance_lower_values.galaxies.source.bulge, \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sample Instance__\n", + "\n", + "A non-linear search retains every model that is accepted during the model-fit.\n", + "\n", + "We can create an instance of any model -- below we create an instance of the last accepted model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "instance = samples.from_sample_index(sample_index=-1)\n", + "\n", + "print(instance.galaxies.lens.mass)\n", + "print(instance.galaxies.lens.mass)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Plots__\n", + "\n", + "The Probability Density Functions (PDF's) of the results can be plotted using the non-linear search in-built \n", + "visualization tools.\n", + "\n", + "This fit used `Nautilus` therefore we use the nested sampling visualization functions, which wrap in-built\n", + "visualization tools.\n", + "\n", + "The `autofit_workspace/*/plots` folder illustrates other packages that can be used to make these plots using\n", + "the standard output results formats (e.g. `GetDist.py`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Maximum Likelihood__\n", + "\n", + "The maximum log likelihood value of the model-fit can be estimated by simple taking the maximum of all log\n", + "likelihoods of the samples.\n", + "\n", + "If different models are fitted to the same dataset, this value can be compared to determine which model provides\n", + "the best fit (e.g. which model has the highest maximum likelihood)?" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Maximum Log Likelihood: \\n\")\n", + "print(max(samples.log_likelihood_list))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Bayesian Evidence__\n", + "\n", + "Nested sampling algorithms like Nautilus also estimate the Bayesian evidence (estimated via the nested sampling \n", + "algorithm).\n", + "\n", + "The Bayesian evidence accounts for \"Occam's Razor\", whereby it penalizes models for being more complex (e.g. if a model\n", + "has more parameters it needs to fit the da\n", + "\n", + "The Bayesian evidence is a better quantity to use to compare models, because it penalizes models with more parameters\n", + "for being more complex (\"Occam's Razor\"). Comparisons using the maximum likelihood value do not account for this and\n", + "therefore may unjustly favour more complex models.\n", + "\n", + "Using the Bayesian evidence for model comparison is well documented on the internet, for example the following\n", + "wikipedia page: https://en.wikipedia.org/wiki/Bayes_factor" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Maximum Log Likelihood and Log Evidence: \\n\")\n", + "print(samples.log_evidence)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lists__\n", + "\n", + "All results can alternatively be returned as a 1D list of values, by passing `as_instance=False`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "max_lh_list = samples.max_log_likelihood(as_instance=False)\n", + "print(\"Max Log Likelihood Model Parameters: \\n\")\n", + "print(max_lh_list, \"\\n\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The list above does not tell us which values correspond to which parameters.\n", + "\n", + "The following quantities are available in the `Model`, where the order of their entries correspond to the parameters \n", + "in the `ml_vector` above:\n", + "\n", + " - `paths`: a list of tuples which give the path of every parameter in the `Model`.\n", + " - `parameter_names`: a list of shorthand parameter names derived from the `paths`.\n", + " - `parameter_labels`: a list of parameter labels used when visualizing non-linear search results (see below).\n", + "\n", + "For simple models like the one fitted in this tutorial, the quantities below are somewhat redundant. For the\n", + "more complex models they are important for tracking the parameters of the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model = samples.model\n", + "\n", + "print(model.paths)\n", + "print(model.parameter_names)\n", + "print(model.parameter_labels)\n", + "print(model.model_component_and_parameter_names)\n", + "print(\"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "All the methods above are available as lists." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "instance = samples.median_pdf(as_instance=False)\n", + "values_at_upper_sigma = samples.values_at_upper_sigma(sigma=3.0, as_instance=False)\n", + "values_at_lower_sigma = samples.values_at_lower_sigma(sigma=3.0, as_instance=False)\n", + "errors_at_upper_sigma = samples.errors_at_upper_sigma(sigma=3.0, as_instance=False)\n", + "errors_at_lower_sigma = samples.errors_at_lower_sigma(sigma=3.0, as_instance=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Latex__\n", + "\n", + "If you are writing modeling results up in a paper, you can use inbuilt latex tools to create latex table code which \n", + "you can copy to your .tex document.\n", + "\n", + "By combining this with the filtering tools below, specific parameters can be included or removed from the latex.\n", + "\n", + "Remember that the superscripts of a parameter are loaded from the config file `notation/label.yaml`, providing high\n", + "levels of customization for how the parameter names appear in the latex table. This is especially useful if your model\n", + "uses the same model components with the same parameter, which therefore need to be distinguished via superscripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "latex = af.text.Samples.latex(\n", + " samples=result.samples,\n", + " median_pdf_model=True,\n", + " sigma=3.0,\n", + " name_to_label=True,\n", + " include_name=True,\n", + " include_quickmath=True,\n", + " prefix=\"Example Prefix \",\n", + " suffix=r\"\\\\[-2pt]\",\n", + ")\n", + "\n", + "print(latex)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Derived Errors (Advanced)__\n", + "\n", + "Computing the errors of a quantity like the `einstein_radius` is simple, because it is sampled by the non-linear \n", + "search. Errors are accessible using the `Samples` object's `errors_from` methods, which marginalize over the \n", + "parameters via the 1D Probability Density Function (PDF).\n", + "\n", + "Computing errors on derived quantities is more tricky, because they are not sampled directly by the non-linear search. \n", + "For example, what if we want the error on the axis-ratio of the mass model? In order to do this we need to create the \n", + "PDF of that derived quantity, which we can then marginalize over using the same function we use to marginalize model \n", + "parameters.\n", + "\n", + "Below, we compute the axis-ratio of every accepted model sampled by the non-linear search and use this determine the PDF \n", + "of the axis-ratio. When combining the axis-ratio's we weight each value by its `weight`. For Nautilus, a nested sampling \n", + "algorithm, the weight of every sample is different and thus must be included.\n", + "\n", + "In order to pass these samples to the function `marginalize`, which marginalizes over the PDF of the axis-ratio to \n", + "compute its error, we also pass the weight list of the samples.\n", + "\n", + "Note again how because when creating the model above using the input names `lens` and `mass` we access the instance\n", + "below using these." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "axis_ratio_list = []\n", + "\n", + "for sample in samples.sample_list:\n", + " instance = sample.instance_for_model(model=samples.model, ignore_assertions=True)\n", + "\n", + " ell_comps = instance.galaxies.lens.mass.ell_comps\n", + "\n", + " axis_ratio = al.convert.axis_ratio_from(ell_comps=ell_comps)\n", + "\n", + " axis_ratio_list.append(axis_ratio)\n", + "\n", + "median_axis_ratio, lower_axis_ratio, upper_axis_ratio = af.marginalize(\n", + " parameter_list=axis_ratio_list, sigma=3.0, weight_list=samples.weight_list\n", + ")\n", + "\n", + "print(f\"axis_ratio = {median_axis_ratio} ({upper_axis_ratio} {lower_axis_ratio}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The calculation above could be computationally expensive, if there are many samples and the derived quantity is\n", + "slow to compute.\n", + "\n", + "An alternative approach, which will provide comparable accuracy provided enough draws are used, is to sample \n", + "points randomy from the PDF of the model and use these to compute the derived quantity.\n", + "\n", + "Draws are from the PDF of the model, so the weights of the samples are accounted for and we therefore do not\n", + "pass them to the `marginalize` function (it essentially treats all samples as having equal weight)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "random_draws = 50\n", + "\n", + "axis_ratio_list = []\n", + "\n", + "for i in range(random_draws):\n", + " instance = samples.draw_randomly_via_pdf()\n", + "\n", + " ell_comps = instance.galaxies.lens.mass.ell_comps\n", + "\n", + " axis_ratio = al.convert.axis_ratio_from(ell_comps=ell_comps)\n", + "\n", + " axis_ratio_list.append(axis_ratio)\n", + "\n", + "median_axis_ratio, lower_axis_ratio, upper_axis_ratio = af.marginalize(\n", + " parameter_list=axis_ratio_list,\n", + " sigma=3.0,\n", + ")\n", + "\n", + "print(f\"axis_ratio = {median_axis_ratio} ({upper_axis_ratio} {lower_axis_ratio}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Samples Filtering (Advanced)__\n", + "\n", + "Our samples object has the results for all three parameters in our model. However, we might only be interested in the\n", + "results of a specific parameter.\n", + "\n", + "The basic form of filtering specifies parameters via their path, which was printed above via the model and is printed \n", + "again below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "print(\"Parameter paths in the model which are used for filtering:\")\n", + "print(samples.model.paths)\n", + "\n", + "print(\"All parameters of the very first sample\")\n", + "print(samples.parameter_lists[0])\n", + "\n", + "samples = samples.with_paths(\n", + " [\n", + " (\"galaxies\", \"lens\", \"mass\", \"einstein_radius\"),\n", + " (\"galaxies\", \"source\", \"bulge\", \"sersic_index\"),\n", + " ]\n", + ")\n", + "\n", + "print(\n", + " \"All parameters of the very first sample (containing only the lens mass's einstein radius and \"\n", + " \"source bulge's sersic index).\"\n", + ")\n", + "print(samples.parameter_lists[0])\n", + "\n", + "print(\n", + " \"Maximum Log Likelihood Model Instances (containing only the lens mass's einstein radius and \"\n", + " \"source bulge's sersic index):\\n\"\n", + ")\n", + "print(samples.max_log_likelihood(as_instance=False))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We specified each path as a list of tuples of strings. \n", + "\n", + "This is how the source code internally stores the path to different components of the model, but it is not \n", + "consistent with the API used to compose a model.\n", + "\n", + "We can alternatively use the following API:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "samples = samples.with_paths(\n", + " [\"galaxies.lens.mass.einstein_radius\", \"galaxies.source.bulge.sersic_index\"]\n", + ")\n", + "\n", + "print(\n", + " \"All parameters of the very first sample (containing only the lens mass's einstein radius and \"\n", + " \"source bulge's sersic index).\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can alternatively filter the `Samples` object by removing all parameters with a certain path. Below, we remove\n", + "the centres of the mass model to be left with 10 parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "print(\"Parameter paths in the model which are used for filtering:\")\n", + "print(samples.model.paths)\n", + "\n", + "print(\"Parameters of first sample\")\n", + "print(samples.parameter_lists[0])\n", + "\n", + "print(samples.model.total_free_parameters)\n", + "\n", + "samples = samples.without_paths(\n", + " [\n", + " # \"galaxies.lens.mass.centre\"),\n", + " \"galaxies.lens.mass.centre.centre_0\",\n", + " # \"galaxies.lens.mass.centre.centre_1),\n", + " ]\n", + ")\n", + "\n", + "print(\"Parameters of first sample without the lens mass centre.\")\n", + "print(samples.parameter_lists[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can keep and remove entire paths of the samples, for example keeping only the parameters of the lens or \n", + "removing all parameters of the source's bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "samples = samples.with_paths([\"galaxies.lens\"])\n", + "print(\"Parameters of the first sample of the lens galaxy\")\n", + "print(samples.parameter_lists[0])\n", + "\n", + "samples = result.samples\n", + "samples = samples.without_paths([\"galaxies.source.bulge\"])\n", + "print(\"Parameters of the first sample without the source's bulge\")\n", + "print(samples.parameter_lists[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/aggregator/samples_via_aggregator.ipynb b/notebooks/guides/results/aggregator/samples_via_aggregator.ipynb index acd4963d3..fcde0d960 100644 --- a/notebooks/guides/results/aggregator/samples_via_aggregator.ipynb +++ b/notebooks/guides/results/aggregator/samples_via_aggregator.ipynb @@ -1,904 +1,941 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Samples via Aggregator\n", - "===============================\n", - "\n", - "In the script `autogalaxy_workspace/*/guides/results/aggregator/samples.py` we show how to inspect the non-linear\n", - "search samples from a result.\n", - "\n", - "We have also shown how to use the `Aggregator` to load the samples of a non-linear search from hard-disk or a\n", - ".sqllite database file.\n", - "\n", - "In this example, we'll load results via the aggregator and inspect the samples of the non-linear search. The\n", - "attributes we inspect are the same as those shown in the `samples.py` script.\n", - "\n", - "This script is simply an API cheat sheet for accessing the results of a non-linear search via the `Aggregator`, so you\n", - "can copy and paste the code to use in your own scripts!\n", - "\n", - "__Contents__\n", - "\n", - "- **Samples via Result:** A fraction of this example repeats the API for manipulating samples given in the.\n", - "- **Files:** In the `start_here.py` script, we discussed the `files` that are output by the non-linear search.\n", - "- **Aggregator:** First, set up the aggregator as shown in `start_here.py`.\n", - "- **Generators:** The `start_here.py` database example gives an explanation of what Python generators are and why and.\n", - "- **Samples:** The result contains a `Samples` object, which contains all samples of the non-linear search, which.\n", - "- **Parameters:** The parameters are stored as a list of lists, where.\n", - "- **Samples Info:** The samples info contains additional information on the samples, which depends on the non-linear.\n", - "- **Figures of Merit:** The `Samples` class contains the log likelihood, log prior, log posterior and weight_list of every.\n", - "- **Samples Summary:** The samples summary contains a subset of results access via the `Samples`, for example the maximum.\n", - "- **Maximum Likelihood Model:** We can use the outputs to create a list of the maximum log likelihood model of each fit to our.\n", - "- **Parameter Names:** Vectors return a lists of all model parameters, but do not tell us which values correspond to which.\n", - "- **Instances:** We can use the `Aggregator` to create a list of instances of the model, using the Python class.\n", - "- **Errors:** Methods for computing error estimates on all parameters are provided.\n", - "- **Sample Instance:** A non-linear search retains every model that is accepted during the model-fit.\n", - "- **Search Plots:** The Probability Density Functions (PDF's) of the results can be plotted using the non-linear search.\n", - "- **Maximum Likelihood:** The maximum log likelihood value of the model-fit can be estimated by simple taking the maximum of.\n", - "- **Bayesian Evidence:** Nested sampling algorithms like Nautilus also estimate the Bayesian evidence (estimated via the.\n", - "- **Lists:** All results can alternatively be returned as a 1D list of values, by passing `as_instance=False`.\n", - "- **Latex:** If you are writing modeling results up in a paper, you can use inbuilt latex tools to create latex.\n", - "- **Ordering:** The default ordering of the results can be a bit random, as it depends on how the sqlite database.\n", - "- **Samples Filtering:** The samples object has the results for all model parameter.\n", - "\n", - "__Samples via Result__\n", - "\n", - "A fraction of this example repeats the API for manipulating samples given in the\n", - "`autogalaxy_workspace/*/guides/results/aggregator/samples.py` example.\n", - "\n", - "This is done so users can directly copy and paste Python code which loads results from the database and manipulates\n", - "the samples." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Files__\n", - "\n", - "In the `start_here.py` script, we discussed the `files` that are output by the non-linear search. The \n", - "following files correspond to the information loaded when loading the non-linear search samples from the database:\n", - "\n", - " - `model`: The `model` defined above and used in the model-fit (`model.json`).\n", - " - `search`: The non-linear search settings (`search.json`).\n", - " - `samples`: The non-linear search samples (`samples.csv`).\n", - " - `samples_info`: Additional information about the samples (`samples_info.json`).\n", - " - `samples_summary`: A summary of key results of the samples (`samples_summary.json`).\n", - " - `info`: The info dictionary passed to the search (`info.json`).\n", - " - `covariance`: The inferred covariance matrix (`covariance.csv`).\n", - " - `cosmology`: The cosmology used by the fit (`cosmology.json`).\n", - " - `settings`: The settings associated with a inversion if used (`settings.json`).\n", - " - `dataset/data`: The data that is fitted (`data.fits`).\n", - " - `dataset/noise_map`: The noise-map (`noise_map.fits`).\n", - " - `dataset/psf`: The Point Spread Function (`psf.fits`).\n", - " - `dataset/mask`: The mask applied to the data (`mask.fits`).\n", - " - `dataset/settings`: The settings associated with the dataset (`settings.json`).\n", - "\n", - "The `samples` and `samples_summary` results contain a lot of repeated information. The `samples` result contains\n", - "the full non-linear search samples, for example every parameter sample and its log likelihood. The `samples_summary`\n", - "contains a summary of the results, for example the maximum log likelihood model and error estimates on parameters\n", - "at 1 and 3 sigma confidence.\n", - "\n", - "Accessing results via the `samples_summary` is much faster, because as it does reperform calculations using the full \n", - "list of samples. Therefore, if the result you want is accessible via the `samples_summary` you should use it\n", - "but if not you can revert to the `samples.\n", - "\n", - "__Aggregator__\n", - "\n", - "First, set up the aggregator as shown in `start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "results_path = Path(\"output\") / \"results_folder\"\n", - "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "agg = Aggregator.from_directory(\n", - " directory=results_path,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Generators__\n", - "\n", - "The `start_here.py` database example gives an explanation of what Python generators are and why and how they are used.\n", - "Refer back to that example if you are unsure." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples_gen = agg.values(\"samples\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Samples__\n", - "\n", - "The result contains a `Samples` object, which contains all samples of the non-linear search, which is accessible\n", - "via the database and aggregator.\n", - "\n", - "Each sample corresponds to a set of model parameters that were evaluated and accepted by the non linear search, \n", - "in this example `Nautilus`. \n", - "\n", - "This includes their log likelihoods, which are used for computing additional information about the model-fit,\n", - "for example the error on every parameter. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Samples: \\n\")\n", - "print(agg.values(\"samples\"))\n", - "print()\n", - "print(\"Total Samples Objects = \", len(agg), \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parameters__\n", - "\n", - "The parameters are stored as a list of lists, where:\n", - "\n", - " - The outer list is the size of the total number of samples.\n", - " - The inner list is the size of the number of free parameters in the fit" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " print(\"All parameters of the very first sample\")\n", - " print(samples.parameter_lists[0])\n", - " print(\"The third parameter of the tenth sample\")\n", - " print(samples.parameter_lists[9][2])\n", - " print()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Samples Info__\n", - "\n", - "The samples info contains additional information on the samples, which depends on the non-linear search used. \n", - "\n", - "For example, for a nested sampling algorithm it contains information on the number of live points, for a MCMC\n", - "algorithm it contains information on the number of steps, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples_info in agg.values(\"samples_info\"):\n", - " print(samples_info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Figures of Merit__\n", - "\n", - "The `Samples` class contains the log likelihood, log prior, log posterior and weight_list of every accepted sample, where:\n", - "\n", - "- The `log_likelihood` is the value evaluated in the `log_likelihood_function`.\n", - "\n", - "- The `log_prior` encodes information on how parameter priors map log likelihood values to log posterior values.\n", - "\n", - "- The `log_posterior` is `log_likelihood + log_prior`.\n", - "\n", - "- The `weight` gives information on how samples are combined to estimate the posterior, which depends on type of search\n", - " used (for `Nautilus` they are all non-zero values which sum to 1).\n", - "\n", - "Lets inspect these values for the tenth sample of each of the 3 model-fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " print(\"log(likelihood), log(prior), log(posterior) and weight of the tenth sample.\")\n", - " print(samples.log_likelihood_list[9])\n", - " print(samples.log_prior_list[9])\n", - " print(samples.log_posterior_list[9])\n", - " print(samples.weight_list[9])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Samples Summary__\n", - "\n", - "The samples summary contains a subset of results access via the `Samples`, for example the maximum likelihood model\n", - "and parameter error estimates.\n", - "\n", - "Using the samples method above can be slow, as the quantities have to be computed from all non-linear search samples\n", - "(e.g. computing errors requires that all samples are marginalized over). This information is stored directly in the\n", - "samples summary and can therefore be accessed instantly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples_summary in agg.values(\"samples_summary\"):\n", - " instance = samples_summary.max_log_likelihood()\n", - "\n", - " print(\"Max Log Likelihood Instance:\")\n", - " print(\"Centre = \", instance.galaxies.lens.mass.centre)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Maximum Likelihood Model__\n", - "\n", - "We can use the outputs to create a list of the maximum log likelihood model of each fit to our three images." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "ml_list = [\n", - " samps.max_log_likelihood(as_instance=False) for samps in agg.values(\"samples\")\n", - "]\n", - "\n", - "print(\"Max Log Likelihood Model Parameter Lists: \\n\")\n", - "print(ml_list, \"\\n\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parameter Names__\n", - "\n", - "Vectors return a lists of all model parameters, but do not tell us which values correspond to which parameters.\n", - "\n", - "The following quantities are available in the `Model`, where the order of their entries correspond to the parameters \n", - "in the `ml_list` above:\n", - "\n", - " - `paths`: a list of tuples which give the path of every parameter in the `Model`.\n", - " - `parameter_names`: a list of shorthand parameter names derived from the `paths`.\n", - " - `parameter_labels`: a list of parameter labels used when visualizing non-linear search results (see below)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " model = samples.model\n", - " print(model)\n", - " print(model.paths)\n", - " print(model.parameter_names)\n", - " print(model.parameter_labels)\n", - " print()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "These lists will be used later for visualization, how it is often more useful to create the model instance of every fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "ml_instances = [samps.max_log_likelihood() for samps in agg.values(\"samples\")]\n", - "print(\"Maximum Log Likelihood Model Instances: \\n\")\n", - "print(ml_instances, \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Instances__\n", - "\n", - "We can use the `Aggregator` to create a list of instances of the model, using the Python class structure of the \n", - "model composition.\n", - "\n", - "For example, we can return a list of the model instances corresponding to the maximum log likelihood sample." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(ml_instances[0].galaxies)\n", - "# print(ml_instances[1].galaxies)\n", - "# print(ml_instances[2].galaxies)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "These galaxies will be named according to the model composed and fitted by the search (in this case `lens` and `source`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(ml_instances[0].galaxies.lens)\n", - "print()\n", - "# print(ml_instances[1].galaxies.source)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Their light and mass profiles are also named according to model composition allowing individual parameters to be \n", - "printed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(ml_instances[0].galaxies.lens.mass.einstein_radius)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Posterior / PDF__\n", - "\n", - "The result contains the full posterior information of our non-linear search, which can be used for parameter \n", - "estimation. \n", - "\n", - "PDF stands for \"Probability Density Function\" and it quantifies probability of each model parameter having values\n", - "that are sampled. It therefore enables error estimation via a process called marginalization.\n", - "\n", - "The median pdf vector is available, which estimates every parameter via 1D marginalization of their PDFs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mp_instances = [samps.median_pdf() for samps in agg.values(\"samples\")]\n", - "\n", - "print(\"Median PDF Model Instances: \\n\")\n", - "print(mp_instances, \"\\n\")\n", - "print(mp_instances[0].galaxies.lens.mass)\n", - "print()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors__\n", - "\n", - "Methods for computing error estimates on all parameters are provided. \n", - "\n", - "This again uses 1D marginalization, now at an input sigma confidence limit. \n", - "\n", - "By inputting `sigma=3.0` margnialization find the values spanning 99.7% of 1D PDF. Changing this to `sigma=1.0`\n", - "would give the errors at the 68.3% confidence limit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "uv3_lists = [samps.values_at_upper_sigma(sigma=3.0) for samps in agg.values(\"samples\")]\n", - "\n", - "uv3_instances = [\n", - " samps.values_at_upper_sigma(sigma=3.0) for samps in agg.values(\"samples\")\n", - "]\n", - "\n", - "lv3_lists = [samps.values_at_lower_sigma(sigma=3.0) for samps in agg.values(\"samples\")]\n", - "\n", - "lv3_instances = [\n", - " samps.values_at_lower_sigma(sigma=3.0) for samps in agg.values(\"samples\")\n", - "]\n", - "\n", - "print(\"Errors Lists: \\n\")\n", - "print(uv3_lists, \"\\n\")\n", - "print(lv3_lists, \"\\n\")\n", - "print(\"Errors Instances: \\n\")\n", - "print(uv3_instances, \"\\n\")\n", - "print(lv3_instances, \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can compute the upper and lower errors on each parameter at a given sigma limit.\n", - "\n", - "The `ue3` below signifies the upper error at 3 sigma. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "ue3_lists = [samps.errors_at_upper_sigma(sigma=3.0) for samps in agg.values(\"samples\")]\n", - "\n", - "# ue3_instances = [\n", - "# samps.errors_at_upper_sigma(sigma=3.0) for samps in agg.values(\"samples\")\n", - "# ]\n", - "\n", - "le3_lists = [samps.errors_at_lower_sigma(sigma=3.0) for samps in agg.values(\"samples\")]\n", - "# le3_instances = [\n", - "# samps.errors_at_lower_sigma(sigma=3.0) for samps in agg.values(\"samples\")\n", - "# ]\n", - "\n", - "print(\"Errors Lists: \\n\")\n", - "print(ue3_lists, \"\\n\")\n", - "print(le3_lists, \"\\n\")\n", - "print(\"Errors Instances: \\n\")\n", - "# print(ue3_instances, \"\\n\")\n", - "# print(le3_instances, \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sample Instance__\n", - "\n", - "A non-linear search retains every model that is accepted during the model-fit.\n", - "\n", - "We can create an instance of any model -- below we create an instance of the last accepted model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " instance = samples.from_sample_index(sample_index=-1)\n", - "\n", - " print(instance.galaxies.source.bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Plots__\n", - "\n", - "The Probability Density Functions (PDF's) of the results can be plotted using the non-linear search in-built \n", - "visualization tools.\n", - "\n", - "This fit used `Nautilus` therefore we use the nested sampling visualization functions, which wrap `Nautilus`'s\n", - "in-built visualization tools.\n", - "\n", - "The `autofit_workspace/*/plots` folder illustrates other packages that can be used to make these plots using\n", - "the standard output results formats (e.g. `GetDist.py`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " aplt.corner_anesthetic(samples=samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Maximum Likelihood__\n", - "\n", - "The maximum log likelihood value of the model-fit can be estimated by simple taking the maximum of all log\n", - "likelihoods of the samples.\n", - "\n", - "If different models are fitted to the same dataset, this value can be compared to determine which model provides\n", - "the best fit (e.g. which model has the highest maximum likelihood)?" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print([max(samps.log_likelihood_list) for samps in agg.values(\"samples\")])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Bayesian Evidence__\n", - "\n", - "Nested sampling algorithms like Nautilus also estimate the Bayesian evidence (estimated via the nested sampling \n", - "algorithm).\n", - "\n", - "The Bayesian evidence accounts for \"Occam's Razor\", whereby it penalizes models for being more complex (e.g. if a model\n", - "has more parameters it needs to fit the da\n", - "\n", - "The Bayesian evidence is a better quantity to use to compare models, because it penalizes models with more parameters\n", - "for being more complex (\"Occam's Razor\"). Comparisons using the maximum likelihood value do not account for this and\n", - "therefore may unjustly favour more complex models.\n", - "\n", - "Using the Bayesian evidence for model comparison is well documented on the internet, for example the following\n", - "wikipedia page: https://en.wikipedia.org/wiki/Bayes_factor" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Log Evidences: \\n\")\n", - "print([samps.log_evidence for samps in agg.values(\"samples\")])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lists__\n", - "\n", - "All results can alternatively be returned as a 1D list of values, by passing `as_instance=False`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " max_lh_list = samples.max_log_likelihood(as_instance=False)\n", - " print(\"Max Log Likelihood Model Parameters: \\n\")\n", - " print(max_lh_list, \"\\n\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The list above does not tell us which values correspond to which parameters.\n", - "\n", - "The following quantities are available in the `Model`, where the order of their entries correspond to the parameters \n", - "in the `ml_vector` above:\n", - "\n", - " - `paths`: a list of tuples which give the path of every parameter in the `Model`.\n", - " - `parameter_names`: a list of shorthand parameter names derived from the `paths`.\n", - " - `parameter_labels`: a list of parameter labels used when visualizing non-linear search results (see below).\n", - "\n", - "For simple models like the one fitted in this tutorial, the quantities below are somewhat redundant. For the\n", - "more complex models they are important for tracking the parameters of the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for model in agg.values(\"model\"):\n", - " print(model.paths)\n", - " print(model.parameter_names)\n", - " print(model.parameter_labels)\n", - " print(model.model_component_and_parameter_names)\n", - " print(\"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Latex__\n", - "\n", - "If you are writing modeling results up in a paper, you can use inbuilt latex tools to create latex table code which \n", - "you can copy to your .tex document.\n", - "\n", - "By combining this with the filtering tools below, specific parameters can be included or removed from the latex.\n", - "\n", - "Remember that the superscripts of a parameter are loaded from the config file `notation/label.yaml`, providing high\n", - "levels of customization for how the parameter names appear in the latex table. This is especially useful if your model\n", - "uses the same model components with the same parameter, which therefore need to be distinguished via superscripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " latex = af.text.Samples.latex(\n", - " samples=samples,\n", - " median_pdf_model=True,\n", - " sigma=3.0,\n", - " name_to_label=True,\n", - " include_name=True,\n", - " include_quickmath=True,\n", - " prefix=\"Example Prefix \",\n", - " suffix=r\"\\\\[-2pt]\",\n", - " )\n", - "\n", - " print(latex)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ordering__\n", - "\n", - "The default ordering of the results can be a bit random, as it depends on how the sqlite database is built. \n", - "\n", - "The `order_by` method can be used to order by a property of the database that is a string, for example by ordering \n", - "using the `unique_tag` (which we set up in the search as the `dataset_name`) the database orders results alphabetically\n", - "according to dataset name.\n", - "\n", - "# agg = agg.order_by(agg.search.unique_tag)\n", - "\n", - "We can also order by a bool, for example making it so all completed results are at the front of the aggregator.\n", - "\n", - "# agg = agg.order_by(agg.search.is_complete)\n", - "\n", - "__Samples Filtering__\n", - "\n", - "The samples object has the results for all model parameter. It can be filtered to contain the results of specific \n", - "parameters of interest.\n", - "\n", - "The basic form of filtering specifies parameters via their path, which was printed above via the model and is printed \n", - "again below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = list(agg.values(\"samples\"))[0]\n", - "\n", - "print(\"Parameter paths in the model which are used for filtering:\")\n", - "print(samples.model.paths)\n", - "\n", - "print(\"All parameters of the very first sample\")\n", - "print(samples.parameter_lists[0])\n", - "\n", - "samples = samples.with_paths(\n", - " [\n", - " (\"galaxies\", \"lens\", \"mass\", \"einstein_radius\"),\n", - " (\"galaxies\", \"lens\", \"mass\", \"centre\", \"centre_0\"),\n", - " ]\n", - ")\n", - "\n", - "print(\n", - " \"All parameters of the very first sample (containing only the lens mass's einstein radius and \"\n", - " \"centre y-coordinate).\"\n", - ")\n", - "print(samples.parameter_lists[0])\n", - "\n", - "print(\n", - " \"Maximum Log Likelihood Model Instances (containing only the lens mass's einstein radius and \"\n", - " \"centre y-coordinate):\\n\"\n", - ")\n", - "print(samples.max_log_likelihood(as_instance=False))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Above, we specified each path as a list of tuples of strings. \n", - "\n", - "This is how the source code internally stores the path to different components of the model, but it is not in-line \n", - "with the PyAutoLens API used to compose a model.\n", - "\n", - "We can alternatively use the following API:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = list(agg.values(\"samples\"))[0]\n", - "\n", - "samples = samples.with_paths(\n", - " [\"galaxies.lens.mass.einstein_radius\", \"galaxies.lens.mass.centre.centre_0\"]\n", - ")\n", - "\n", - "print(\n", - " \"All parameters of the very first sample (containing only the lens mass's einstein radius and \"\n", - " \"centre y-coordinate).\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Above, we filtered the `Samples` but asking for all parameters which included the\n", - "path (\"galaxies\", \"lens\", \"mass\", \"einstein_radius\").\n", - "\n", - "We can alternatively filter the `Samples` object by removing all parameters with a certain path. Below, we remove\n", - "the centres of the mass model to be left with 10 parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = list(agg.values(\"samples\"))[0]\n", - "\n", - "print(\"Parameter paths in the model which are used for filtering:\")\n", - "print(samples.model.paths)\n", - "\n", - "print(\"Parameters of first sample\")\n", - "print(samples.parameter_lists[0])\n", - "\n", - "print(samples.model.total_free_parameters)\n", - "\n", - "samples = samples.without_paths(\n", - " [\n", - " # \"galaxies.lens.mass.centre\"),\n", - " \"galaxies.lens.mass.centre.centre_0\",\n", - " # \"galaxies.lens.mass.centre.centre_1),\n", - " ]\n", - ")\n", - "\n", - "print(\"Parameters of first sample without the lens mass centre.\")\n", - "print(samples.parameter_lists[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can keep and remove entire paths of the samples, for example keeping only the parameters of the lens or \n", - "removing all parameters of the source's bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = list(agg.values(\"samples\"))[0]\n", - "samples = samples.with_paths([\"galaxies.lens\"])\n", - "print(\"Parameters of the first sample of the lens galaxy\")\n", - "print(samples.parameter_lists[0])\n", - "\n", - "samples = list(agg.values(\"samples\"))[0]\n", - "samples = samples.with_paths([\"galaxies.source.bulge\"])\n", - "print(\"Parameters of the first sample without the source's bulge\")\n", - "print(samples.parameter_lists[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Samples via Aggregator\n", + "===============================\n", + "\n", + "In the script `autogalaxy_workspace/*/guides/results/aggregator/samples.py` we show how to inspect the non-linear\n", + "search samples from a result.\n", + "\n", + "We have also shown how to use the `Aggregator` to load the samples of a non-linear search from hard-disk or a\n", + ".sqllite database file.\n", + "\n", + "In this example, we'll load results via the aggregator and inspect the samples of the non-linear search. The\n", + "attributes we inspect are the same as those shown in the `samples.py` script.\n", + "\n", + "This script is simply an API cheat sheet for accessing the results of a non-linear search via the `Aggregator`, so you\n", + "can copy and paste the code to use in your own scripts!\n", + "\n", + "__Contents__\n", + "\n", + "- **Samples via Result:** A fraction of this example repeats the API for manipulating samples given in the.\n", + "- **Files:** In the `start_here.py` script, we discussed the `files` that are output by the non-linear search.\n", + "- **Aggregator:** First, set up the aggregator as shown in `start_here.py`.\n", + "- **Generators:** The `start_here.py` database example gives an explanation of what Python generators are and why and.\n", + "- **Samples:** The result contains a `Samples` object, which contains all samples of the non-linear search, which.\n", + "- **Parameters:** The parameters are stored as a list of lists, where.\n", + "- **Samples Info:** The samples info contains additional information on the samples, which depends on the non-linear.\n", + "- **Figures of Merit:** The `Samples` class contains the log likelihood, log prior, log posterior and weight_list of every.\n", + "- **Samples Summary:** The samples summary contains a subset of results access via the `Samples`, for example the maximum.\n", + "- **Maximum Likelihood Model:** We can use the outputs to create a list of the maximum log likelihood model of each fit to our.\n", + "- **Parameter Names:** Vectors return a lists of all model parameters, but do not tell us which values correspond to which.\n", + "- **Instances:** We can use the `Aggregator` to create a list of instances of the model, using the Python class.\n", + "- **Errors:** Methods for computing error estimates on all parameters are provided.\n", + "- **Sample Instance:** A non-linear search retains every model that is accepted during the model-fit.\n", + "- **Search Plots:** The Probability Density Functions (PDF's) of the results can be plotted using the non-linear search.\n", + "- **Maximum Likelihood:** The maximum log likelihood value of the model-fit can be estimated by simple taking the maximum of.\n", + "- **Bayesian Evidence:** Nested sampling algorithms like Nautilus also estimate the Bayesian evidence (estimated via the.\n", + "- **Lists:** All results can alternatively be returned as a 1D list of values, by passing `as_instance=False`.\n", + "- **Latex:** If you are writing modeling results up in a paper, you can use inbuilt latex tools to create latex.\n", + "- **Ordering:** The default ordering of the results can be a bit random, as it depends on how the sqlite database.\n", + "- **Samples Filtering:** The samples object has the results for all model parameter.\n", + "\n", + "__Samples via Result__\n", + "\n", + "A fraction of this example repeats the API for manipulating samples given in the\n", + "`autogalaxy_workspace/*/guides/results/aggregator/samples.py` example.\n", + "\n", + "This is done so users can directly copy and paste Python code which loads results from the database and manipulates\n", + "the samples." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Files__\n", + "\n", + "In the `start_here.py` script, we discussed the `files` that are output by the non-linear search. The \n", + "following files correspond to the information loaded when loading the non-linear search samples from the database:\n", + "\n", + " - `model`: The `model` defined above and used in the model-fit (`model.json`).\n", + " - `search`: The non-linear search settings (`search.json`).\n", + " - `samples`: The non-linear search samples (`samples.csv`).\n", + " - `samples_info`: Additional information about the samples (`samples_info.json`).\n", + " - `samples_summary`: A summary of key results of the samples (`samples_summary.json`).\n", + " - `info`: The info dictionary passed to the search (`info.json`).\n", + " - `covariance`: The inferred covariance matrix (`covariance.csv`).\n", + " - `cosmology`: The cosmology used by the fit (`cosmology.json`).\n", + " - `settings`: The settings associated with a inversion if used (`settings.json`).\n", + " - `dataset/data`: The data that is fitted (`data.fits`).\n", + " - `dataset/noise_map`: The noise-map (`noise_map.fits`).\n", + " - `dataset/psf`: The Point Spread Function (`psf.fits`).\n", + " - `dataset/mask`: The mask applied to the data (`mask.fits`).\n", + " - `dataset/settings`: The settings associated with the dataset (`settings.json`).\n", + "\n", + "The `samples` and `samples_summary` results contain a lot of repeated information. The `samples` result contains\n", + "the full non-linear search samples, for example every parameter sample and its log likelihood. The `samples_summary`\n", + "contains a summary of the results, for example the maximum log likelihood model and error estimates on parameters\n", + "at 1 and 3 sigma confidence.\n", + "\n", + "Accessing results via the `samples_summary` is much faster, because as it does reperform calculations using the full \n", + "list of samples. Therefore, if the result you want is accessible via the `samples_summary` you should use it\n", + "but if not you can revert to the `samples.\n", + "\n", + "__Aggregator__\n", + "\n", + "First, set up the aggregator as shown in `start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "results_path = Path(\"output\") / \"results_folder\"\n", + "if not any(results_path.glob(\"**/image/dataset.fits\")):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "agg = Aggregator.from_directory(\n", + " directory=results_path,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Generators__\n", + "\n", + "The `start_here.py` database example gives an explanation of what Python generators are and why and how they are used.\n", + "Refer back to that example if you are unsure." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples_gen = agg.values(\"samples\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Samples__\n", + "\n", + "The result contains a `Samples` object, which contains all samples of the non-linear search, which is accessible\n", + "via the database and aggregator.\n", + "\n", + "Each sample corresponds to a set of model parameters that were evaluated and accepted by the non linear search, \n", + "in this example `Nautilus`. \n", + "\n", + "This includes their log likelihoods, which are used for computing additional information about the model-fit,\n", + "for example the error on every parameter. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Samples: \\n\")\n", + "print(agg.values(\"samples\"))\n", + "print()\n", + "print(\"Total Samples Objects = \", len(agg), \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parameters__\n", + "\n", + "The parameters are stored as a list of lists, where:\n", + "\n", + " - The outer list is the size of the total number of samples.\n", + " - The inner list is the size of the number of free parameters in the fit" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " print(\"All parameters of the very first sample\")\n", + " print(samples.parameter_lists[0])\n", + " print(\"The third parameter of the tenth sample\")\n", + " print(samples.parameter_lists[9][2])\n", + " print()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Samples Info__\n", + "\n", + "The samples info contains additional information on the samples, which depends on the non-linear search used. \n", + "\n", + "For example, for a nested sampling algorithm it contains information on the number of live points, for a MCMC\n", + "algorithm it contains information on the number of steps, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples_info in agg.values(\"samples_info\"):\n", + " print(samples_info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Figures of Merit__\n", + "\n", + "The `Samples` class contains the log likelihood, log prior, log posterior and weight_list of every accepted sample, where:\n", + "\n", + "- The `log_likelihood` is the value evaluated in the `log_likelihood_function`.\n", + "\n", + "- The `log_prior` encodes information on how parameter priors map log likelihood values to log posterior values.\n", + "\n", + "- The `log_posterior` is `log_likelihood + log_prior`.\n", + "\n", + "- The `weight` gives information on how samples are combined to estimate the posterior, which depends on type of search\n", + " used (for `Nautilus` they are all non-zero values which sum to 1).\n", + "\n", + "Lets inspect these values for the tenth sample of each of the 3 model-fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " print(\"log(likelihood), log(prior), log(posterior) and weight of the tenth sample.\")\n", + " print(samples.log_likelihood_list[9])\n", + " print(samples.log_prior_list[9])\n", + " print(samples.log_posterior_list[9])\n", + " print(samples.weight_list[9])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Samples Summary__\n", + "\n", + "The samples summary contains a subset of results access via the `Samples`, for example the maximum likelihood model\n", + "and parameter error estimates.\n", + "\n", + "Using the samples method above can be slow, as the quantities have to be computed from all non-linear search samples\n", + "(e.g. computing errors requires that all samples are marginalized over). This information is stored directly in the\n", + "samples summary and can therefore be accessed instantly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples_summary in agg.values(\"samples_summary\"):\n", + " instance = samples_summary.max_log_likelihood()\n", + "\n", + " print(\"Max Log Likelihood Instance:\")\n", + " print(\"Centre = \", instance.galaxies.lens.mass.centre)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Maximum Likelihood Model__\n", + "\n", + "We can use the outputs to create a list of the maximum log likelihood model of each fit to our three images." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "ml_list = [\n", + " samps.max_log_likelihood(as_instance=False) for samps in agg.values(\"samples\")\n", + "]\n", + "\n", + "print(\"Max Log Likelihood Model Parameter Lists: \\n\")\n", + "print(ml_list, \"\\n\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parameter Names__\n", + "\n", + "Vectors return a lists of all model parameters, but do not tell us which values correspond to which parameters.\n", + "\n", + "The following quantities are available in the `Model`, where the order of their entries correspond to the parameters \n", + "in the `ml_list` above:\n", + "\n", + " - `paths`: a list of tuples which give the path of every parameter in the `Model`.\n", + " - `parameter_names`: a list of shorthand parameter names derived from the `paths`.\n", + " - `parameter_labels`: a list of parameter labels used when visualizing non-linear search results (see below)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " model = samples.model\n", + " print(model)\n", + " print(model.paths)\n", + " print(model.parameter_names)\n", + " print(model.parameter_labels)\n", + " print()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "These lists will be used later for visualization, how it is often more useful to create the model instance of every fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "ml_instances = [samps.max_log_likelihood() for samps in agg.values(\"samples\")]\n", + "print(\"Maximum Log Likelihood Model Instances: \\n\")\n", + "print(ml_instances, \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Instances__\n", + "\n", + "We can use the `Aggregator` to create a list of instances of the model, using the Python class structure of the \n", + "model composition.\n", + "\n", + "For example, we can return a list of the model instances corresponding to the maximum log likelihood sample." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(ml_instances[0].galaxies)\n", + "# print(ml_instances[1].galaxies)\n", + "# print(ml_instances[2].galaxies)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "These galaxies will be named according to the model composed and fitted by the search (in this case `lens` and `source`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(ml_instances[0].galaxies.lens)\n", + "print()\n", + "# print(ml_instances[1].galaxies.source)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Their light and mass profiles are also named according to model composition allowing individual parameters to be \n", + "printed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(ml_instances[0].galaxies.lens.mass.einstein_radius)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Posterior / PDF__\n", + "\n", + "The result contains the full posterior information of our non-linear search, which can be used for parameter \n", + "estimation. \n", + "\n", + "PDF stands for \"Probability Density Function\" and it quantifies probability of each model parameter having values\n", + "that are sampled. It therefore enables error estimation via a process called marginalization.\n", + "\n", + "The median pdf vector is available, which estimates every parameter via 1D marginalization of their PDFs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mp_instances = [samps.median_pdf() for samps in agg.values(\"samples\")]\n", + "\n", + "print(\"Median PDF Model Instances: \\n\")\n", + "print(mp_instances, \"\\n\")\n", + "print(mp_instances[0].galaxies.lens.mass)\n", + "print()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors__\n", + "\n", + "Methods for computing error estimates on all parameters are provided. \n", + "\n", + "This again uses 1D marginalization, now at an input sigma confidence limit. \n", + "\n", + "By inputting `sigma=3.0` margnialization find the values spanning 99.7% of 1D PDF. Changing this to `sigma=1.0`\n", + "would give the errors at the 68.3% confidence limit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "uv3_lists = [samps.values_at_upper_sigma(sigma=3.0) for samps in agg.values(\"samples\")]\n", + "\n", + "uv3_instances = [\n", + " samps.values_at_upper_sigma(sigma=3.0) for samps in agg.values(\"samples\")\n", + "]\n", + "\n", + "lv3_lists = [samps.values_at_lower_sigma(sigma=3.0) for samps in agg.values(\"samples\")]\n", + "\n", + "lv3_instances = [\n", + " samps.values_at_lower_sigma(sigma=3.0) for samps in agg.values(\"samples\")\n", + "]\n", + "\n", + "print(\"Errors Lists: \\n\")\n", + "print(uv3_lists, \"\\n\")\n", + "print(lv3_lists, \"\\n\")\n", + "print(\"Errors Instances: \\n\")\n", + "print(uv3_instances, \"\\n\")\n", + "print(lv3_instances, \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can compute the upper and lower errors on each parameter at a given sigma limit.\n", + "\n", + "The `ue3` below signifies the upper error at 3 sigma. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "ue3_lists = [samps.errors_at_upper_sigma(sigma=3.0) for samps in agg.values(\"samples\")]\n", + "\n", + "# ue3_instances = [\n", + "# samps.errors_at_upper_sigma(sigma=3.0) for samps in agg.values(\"samples\")\n", + "# ]\n", + "\n", + "le3_lists = [samps.errors_at_lower_sigma(sigma=3.0) for samps in agg.values(\"samples\")]\n", + "# le3_instances = [\n", + "# samps.errors_at_lower_sigma(sigma=3.0) for samps in agg.values(\"samples\")\n", + "# ]\n", + "\n", + "print(\"Errors Lists: \\n\")\n", + "print(ue3_lists, \"\\n\")\n", + "print(le3_lists, \"\\n\")\n", + "print(\"Errors Instances: \\n\")\n", + "# print(ue3_instances, \"\\n\")\n", + "# print(le3_instances, \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sample Instance__\n", + "\n", + "A non-linear search retains every model that is accepted during the model-fit.\n", + "\n", + "We can create an instance of any model -- below we create an instance of the last accepted model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " instance = samples.from_sample_index(sample_index=-1)\n", + "\n", + " print(instance.galaxies.source.bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Plots__\n", + "\n", + "The Probability Density Functions (PDF's) of the results can be plotted using the non-linear search in-built \n", + "visualization tools.\n", + "\n", + "This fit used `Nautilus` therefore we use the nested sampling visualization functions, which wrap `Nautilus`'s\n", + "in-built visualization tools.\n", + "\n", + "The `autofit_workspace/*/plots` folder illustrates other packages that can be used to make these plots using\n", + "the standard output results formats (e.g. `GetDist.py`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " aplt.corner_anesthetic(samples=samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Maximum Likelihood__\n", + "\n", + "The maximum log likelihood value of the model-fit can be estimated by simple taking the maximum of all log\n", + "likelihoods of the samples.\n", + "\n", + "If different models are fitted to the same dataset, this value can be compared to determine which model provides\n", + "the best fit (e.g. which model has the highest maximum likelihood)?" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print([max(samps.log_likelihood_list) for samps in agg.values(\"samples\")])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Bayesian Evidence__\n", + "\n", + "Nested sampling algorithms like Nautilus also estimate the Bayesian evidence (estimated via the nested sampling \n", + "algorithm).\n", + "\n", + "The Bayesian evidence accounts for \"Occam's Razor\", whereby it penalizes models for being more complex (e.g. if a model\n", + "has more parameters it needs to fit the da\n", + "\n", + "The Bayesian evidence is a better quantity to use to compare models, because it penalizes models with more parameters\n", + "for being more complex (\"Occam's Razor\"). Comparisons using the maximum likelihood value do not account for this and\n", + "therefore may unjustly favour more complex models.\n", + "\n", + "Using the Bayesian evidence for model comparison is well documented on the internet, for example the following\n", + "wikipedia page: https://en.wikipedia.org/wiki/Bayes_factor" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Log Evidences: \\n\")\n", + "print([samps.log_evidence for samps in agg.values(\"samples\")])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lists__\n", + "\n", + "All results can alternatively be returned as a 1D list of values, by passing `as_instance=False`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " max_lh_list = samples.max_log_likelihood(as_instance=False)\n", + " print(\"Max Log Likelihood Model Parameters: \\n\")\n", + " print(max_lh_list, \"\\n\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The list above does not tell us which values correspond to which parameters.\n", + "\n", + "The following quantities are available in the `Model`, where the order of their entries correspond to the parameters \n", + "in the `ml_vector` above:\n", + "\n", + " - `paths`: a list of tuples which give the path of every parameter in the `Model`.\n", + " - `parameter_names`: a list of shorthand parameter names derived from the `paths`.\n", + " - `parameter_labels`: a list of parameter labels used when visualizing non-linear search results (see below).\n", + "\n", + "For simple models like the one fitted in this tutorial, the quantities below are somewhat redundant. For the\n", + "more complex models they are important for tracking the parameters of the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for model in agg.values(\"model\"):\n", + " print(model.paths)\n", + " print(model.parameter_names)\n", + " print(model.parameter_labels)\n", + " print(model.model_component_and_parameter_names)\n", + " print(\"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Latex__\n", + "\n", + "If you are writing modeling results up in a paper, you can use inbuilt latex tools to create latex table code which \n", + "you can copy to your .tex document.\n", + "\n", + "By combining this with the filtering tools below, specific parameters can be included or removed from the latex.\n", + "\n", + "Remember that the superscripts of a parameter are loaded from the config file `notation/label.yaml`, providing high\n", + "levels of customization for how the parameter names appear in the latex table. This is especially useful if your model\n", + "uses the same model components with the same parameter, which therefore need to be distinguished via superscripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " latex = af.text.Samples.latex(\n", + " samples=samples,\n", + " median_pdf_model=True,\n", + " sigma=3.0,\n", + " name_to_label=True,\n", + " include_name=True,\n", + " include_quickmath=True,\n", + " prefix=\"Example Prefix \",\n", + " suffix=r\"\\\\[-2pt]\",\n", + " )\n", + "\n", + " print(latex)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ordering__\n", + "\n", + "The default ordering of the results can be a bit random, as it depends on how the sqlite database is built. \n", + "\n", + "The `order_by` method can be used to order by a property of the database that is a string, for example by ordering \n", + "using the `unique_tag` (which we set up in the search as the `dataset_name`) the database orders results alphabetically\n", + "according to dataset name.\n", + "\n", + "# agg = agg.order_by(agg.search.unique_tag)\n", + "\n", + "We can also order by a bool, for example making it so all completed results are at the front of the aggregator.\n", + "\n", + "# agg = agg.order_by(agg.search.is_complete)\n", + "\n", + "__Samples Filtering__\n", + "\n", + "The samples object has the results for all model parameter. It can be filtered to contain the results of specific \n", + "parameters of interest.\n", + "\n", + "The basic form of filtering specifies parameters via their path, which was printed above via the model and is printed \n", + "again below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = list(agg.values(\"samples\"))[0]\n", + "\n", + "print(\"Parameter paths in the model which are used for filtering:\")\n", + "print(samples.model.paths)\n", + "\n", + "print(\"All parameters of the very first sample\")\n", + "print(samples.parameter_lists[0])\n", + "\n", + "samples = samples.with_paths(\n", + " [\n", + " (\"galaxies\", \"lens\", \"mass\", \"einstein_radius\"),\n", + " (\"galaxies\", \"lens\", \"mass\", \"centre\", \"centre_0\"),\n", + " ]\n", + ")\n", + "\n", + "print(\n", + " \"All parameters of the very first sample (containing only the lens mass's einstein radius and \"\n", + " \"centre y-coordinate).\"\n", + ")\n", + "print(samples.parameter_lists[0])\n", + "\n", + "print(\n", + " \"Maximum Log Likelihood Model Instances (containing only the lens mass's einstein radius and \"\n", + " \"centre y-coordinate):\\n\"\n", + ")\n", + "print(samples.max_log_likelihood(as_instance=False))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Above, we specified each path as a list of tuples of strings. \n", + "\n", + "This is how the source code internally stores the path to different components of the model, but it is not in-line \n", + "with the PyAutoLens API used to compose a model.\n", + "\n", + "We can alternatively use the following API:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = list(agg.values(\"samples\"))[0]\n", + "\n", + "samples = samples.with_paths(\n", + " [\"galaxies.lens.mass.einstein_radius\", \"galaxies.lens.mass.centre.centre_0\"]\n", + ")\n", + "\n", + "print(\n", + " \"All parameters of the very first sample (containing only the lens mass's einstein radius and \"\n", + " \"centre y-coordinate).\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Above, we filtered the `Samples` but asking for all parameters which included the\n", + "path (\"galaxies\", \"lens\", \"mass\", \"einstein_radius\").\n", + "\n", + "We can alternatively filter the `Samples` object by removing all parameters with a certain path. Below, we remove\n", + "the centres of the mass model to be left with 10 parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = list(agg.values(\"samples\"))[0]\n", + "\n", + "print(\"Parameter paths in the model which are used for filtering:\")\n", + "print(samples.model.paths)\n", + "\n", + "print(\"Parameters of first sample\")\n", + "print(samples.parameter_lists[0])\n", + "\n", + "print(samples.model.total_free_parameters)\n", + "\n", + "samples = samples.without_paths(\n", + " [\n", + " # \"galaxies.lens.mass.centre\"),\n", + " \"galaxies.lens.mass.centre.centre_0\",\n", + " # \"galaxies.lens.mass.centre.centre_1),\n", + " ]\n", + ")\n", + "\n", + "print(\"Parameters of first sample without the lens mass centre.\")\n", + "print(samples.parameter_lists[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can keep and remove entire paths of the samples, for example keeping only the parameters of the lens or \n", + "removing all parameters of the source's bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = list(agg.values(\"samples\"))[0]\n", + "samples = samples.with_paths([\"galaxies.lens\"])\n", + "print(\"Parameters of the first sample of the lens galaxy\")\n", + "print(samples.parameter_lists[0])\n", + "\n", + "samples = list(agg.values(\"samples\"))[0]\n", + "samples = samples.with_paths([\"galaxies.source.bulge\"])\n", + "print(\"Parameters of the first sample without the source's bulge\")\n", + "print(samples.parameter_lists[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/database/start_here.ipynb b/notebooks/guides/results/database/start_here.ipynb index 35ccf29a2..7bbcdbdae 100644 --- a/notebooks/guides/results/database/start_here.ipynb +++ b/notebooks/guides/results/database/start_here.ipynb @@ -1,508 +1,545 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Database: Introduction\n", - "======================\n", - "\n", - "The default behaviour of model-fitting results output is to be written to hard-disc in folders. These are simple to\n", - "navigate and manually check.\n", - "\n", - "For small model-fitting tasks this is sufficient, however it does not scale well when performing many model fits to\n", - "large datasets, because manual inspection of results becomes time consuming.\n", - "\n", - "All results can therefore be output to an sqlite3 (https://docs.python.org/3/library/sqlite3.html) relational database,\n", - "meaning that results can be loaded into a Jupyter notebook or Python script for inspection, analysis and interpretation.\n", - "This database supports advanced querying, so that specific model-fits (e.g., which fit a certain model or dataset) can\n", - "be loaded.\n", - "\n", - "This script fits a sample of three simulated strong lenses using the same non-linear search. The results will be used\n", - "to illustrate the database in the database tutorials that follow.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Unique Identifiers:** Results output to hard-disk are contained in a folder named via a unique identifier (a random.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Results From Hard Disk:** In this example, results will be first be written to hard-disk using the standard output directory.\n", - "- **Building a Database File From an Output Folder:** The fits above wrote the results to hard-disk in folders, not as an .sqlite database file.\n", - "- **Writing Directly To Database:** Results can be written directly to the .sqlite database file, skipping output to hard-disk.\n", - "- **Files:** When performing fits which output results to hard-disc, a `files` folder is created containing.\n", - "- **Generators:** Before using the aggregator to inspect results, lets discuss Python generators.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Samples:** The `Samples` class contains all information on the non-linear search samples, for example the.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Model__\n", - "\n", - "The search fits each lens with:\n", - "\n", - " - An `Isothermal` `MassProfile` for the lens galaxy's mass.\n", - " - An `Sersic` `LightProfile` for the source galaxy's light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import json\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Unique Identifiers__\n", - "\n", - "Results output to hard-disk are contained in a folder named via a unique identifier (a \n", - "random collection of characters, e.g. `8hds89fhndlsiuhnfiusdh`). The unique identifier changes if the model or \n", - "search change, to ensure different fits to not overwrite one another on hard-disk.\n", - "\n", - "Each unique identifier is used to define every entry of the database as it is built. Unique identifiers therefore play \n", - "the same vital role for the database of ensuring that every set of results written to it are unique.\n", - "\n", - "In this example, we fit 3 different datasets with the same search and model. Each `dataset_name` is therefore passed\n", - "in as the search's `unique_tag` to ensure 3 separate sets of results for each model-fit are written to the .sqlite\n", - "database.\n", - "\n", - "__Dataset__\n", - "\n", - "For each dataset we load it from hard-disc, set up its `Analysis` class and fit it with a non-linear search. \n", - "\n", - "We want each results to be stored in the database with an entry specific to the dataset. We'll use the `Dataset`'s name \n", - "string to do this, so lets create a list of the 3 dataset names." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_names = [\n", - " \"simple\",\n", - " \"lens_sersic\",\n", - " \"simple__no_lens_light\",\n", - "]\n", - "\n", - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Results From Hard Disk__\n", - "\n", - "In this example, results will be first be written to hard-disk using the standard output directory structure and we\n", - "will then build the database from these results. This behaviour is governed by us inputting `session=None`.\n", - "\n", - "If you have existing results you wish to build a database for, you can therefore adapt this example you to do this.\n", - "\n", - "Later in this example we show how results can also also be output directly to an .sqlite database, saving on hard-disk \n", - "space. This will be acheived by setting `session` to something that is not `None`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "session = None\n", - "\n", - "for dataset_name in dataset_names:\n", - " \"\"\"\n", - " __Paths__\n", - "\n", - " Set up the config and output paths.\n", - " \"\"\"\n", - " dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - " \"\"\"\n", - " __Dataset__\n", - "\n", - " Using the dataset path, load the data (image, noise-map, PSF) as an `Imaging` object from .fits files.\n", - "\n", - " This `Imaging` object will be available via the aggregator. Note also that we give the dataset a `name` via the\n", - " command `name=dataset_name`. we'll use this name in the aggregator tutorials.\n", - " \"\"\"\n", - " dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=pixel_scales,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Mask__\n", - "\n", - " The `Mask2D` we fit this data-set with, which will be available via the aggregator.\n", - " \"\"\"\n", - " mask_radius = 3.0\n", - "\n", - " mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=mask)\n", - "\n", - " \"\"\"\n", - " __Info__\n", - "\n", - " Information about the model-fit that is not part included in the model-fit itself can be made accessible via the \n", - " database by passing an `info` dictionary. \n", - "\n", - " Below we write info on the dataset`s (hypothetical) data of observation and exposure time, which we will later show\n", - " the database can access. \n", - "\n", - " For fits to large datasets this ensures that all relevant information for interpreting results is accessible.\n", - " \"\"\"\n", - " with open(Path(dataset_path, \"info.json\")) as json_file:\n", - " info = json.load(json_file)\n", - "\n", - " \"\"\"\n", - " __Model__\n", - "\n", - " Set up the model as per usual, and will see in tutorial 3 why we have included `disk=None`.\n", - " \"\"\"\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", - " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=None),\n", - " )\n", - " )\n", - "\n", - " \"\"\"\n", - " The `unique_tag` below uses the `dataset_name` to alter the unique identifier, which as we have seen is also \n", - " generated depending on the search settings and model. In this example, all three model fits use an identical \n", - " search and model, so this `unique_tag` is key for ensuring 3 separate sets of results for each model-fit are \n", - " stored in the output folder and written to the .sqlite database. \n", - " \"\"\"\n", - " search = af.Nautilus(\n", - " path_prefix=Path(\"database\"),\n", - " name=\"database_example\",\n", - " unique_tag=dataset_name, # This makes the unique identifier use the dataset name\n", - " session=session, # This can instruct the search to write to the .sqlite database.\n", - " n_live=100,\n", - " n_batch=50, # GPU batching and VRAM use explained in `modeling` examples.\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " search.fit(analysis=analysis, model=model, info=info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Building a Database File From an Output Folder__\n", - "\n", - "The fits above wrote the results to hard-disk in folders, not as an .sqlite database file. \n", - "\n", - "We build the database below, where the `database_name` corresponds to the name of your output folder and is also the \n", - "name of the `.sqlite` database file that is created.\n", - "\n", - "If you are fitting a relatively small number of datasets (e.g. 10-100) having all results written to hard-disk (e.g. \n", - "for quick visual inspection) and using the database for sample wide analysis is beneficial.\n", - "\n", - "We can optionally only include completed model-fits but setting `completed_only=True`.\n", - "\n", - "If you inspect the `output` folder, you will see a `database.sqlite` file which contains the results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "database_name = \"database\"\n", - "\n", - "agg = af.Aggregator.from_database(\n", - " filename=f\"{database_name}.sqlite\", completed_only=False\n", - ")\n", - "\n", - "agg.add_directory(directory=Path(\"output\") / database_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Writing Directly To Database__\n", - "\n", - "Results can be written directly to the .sqlite database file, skipping output to hard-disk entirely, by creating\n", - "a session and passing this to the non-linear search.\n", - "\n", - "The code below shows how to do this, but it is commented out to avoid rerunning the non-linear searches.\n", - "\n", - "This is ideal for tasks where model-fits to hundreds or thousands of datasets are performed, as it becomes unfeasible\n", - "to inspect the results of all fits on the hard-disk. \n", - "\n", - "Our recommended workflow is to set up database analysis scripts using ~10 model-fits, and then scaling these up\n", - "to large samples by writing directly to the database." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# session = af.db.open_database(\"database.sqlite\")\n", - "#\n", - "# search = af.Nautilus(\n", - "# path_prefix=Path(\"database\"),\n", - "# name=\"database_example\",\n", - "# unique_tag=dataset_name, # This makes the unique identifier use the dataset name\n", - "# session=session, # This can instruct the search to write to the .sqlite database.\n", - "# n_live=100,\n", - "# n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - "# )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Files__\n", - "\n", - "When performing fits which output results to hard-disc, a `files` folder is created containing .json / .csv files of \n", - "the model, samples, search, etc.\n", - "\n", - "These are the files that are written to the database, which the aggregator loads via the database in order to make \n", - "them accessible in a Python script or Jupyter notebook.\n", - "\n", - "You can checkout the output folder created by this fit to see these files.\n", - "\n", - "Below, we will access these results using the aggregator's `values` method. A full list of what can be loaded is\n", - "as follows:\n", - "\n", - " - `model`: The `model` defined above and used in the model-fit (`model.json`).\n", - " - `search`: The non-linear search settings (`search.json`).\n", - " - `samples`: The non-linear search samples (`samples.csv`).\n", - " - `samples_info`: Additional information about the samples (`samples_info.json`).\n", - " - `samples_summary`: A summary of key results of the samples (`samples_summary.json`).\n", - " - `info`: The info dictionary passed to the search (`info.json`).\n", - " - `covariance`: The inferred covariance matrix (`covariance.csv`).\n", - " - `cosmology`: The cosmology used by the fit (`cosmology.json`).\n", - " - `settings`: The settings associated with a inversion if used (`settings.json`).\n", - " - `dataset/data`: The data that is fitted (`data.fits`).\n", - " - `dataset/noise_map`: The noise-map (`noise_map.fits`).\n", - " - `dataset/psf`: The Point Spread Function (`psf.fits`).\n", - " - `dataset/mask`: The mask applied to the data (`mask.fits`).\n", - " - `dataset/settings`: The settings associated with the dataset (`settings.json`).\n", - "\n", - "The `samples` and `samples_summary` results contain a lot of repeated information. The `samples` result contains\n", - "the full non-linear search samples, for example every parameter sample and its log likelihood. The `samples_summary`\n", - "contains a summary of the results, for example the maximum log likelihood model and error estimates on parameters\n", - "at 1 and 3 sigma confidence.\n", - "\n", - "Accessing results via the `samples_summary` is much faster, because as it does reperform calculations using the full \n", - "list of samples. Therefore, if the result you want is accessible via the `samples_summary` you should use it\n", - "but if not you can revert to the `samples.\n", - "\n", - "__Generators__\n", - "\n", - "Before using the aggregator to inspect results, lets discuss Python generators. \n", - "\n", - "A generator is an object that iterates over a function when it is called. The aggregator creates all of the objects \n", - "that it loads from the database as generators (as opposed to a list, or dictionary, or another Python type).\n", - "\n", - "This is because generators are memory efficient, as they do not store the entries of the database in memory \n", - "simultaneously. This contrasts objects like lists and dictionaries, which store all entries in memory all at once. \n", - "If you fit a large number of datasets, lists and dictionaries will use a lot of memory and could crash your computer!\n", - "\n", - "Once we use a generator in the Python code, it cannot be used again. To perform the same task twice, the \n", - "generator must be remade it. This cookbook therefore rarely stores generators as variables and instead uses the \n", - "aggregator to create each generator at the point of use.\n", - "\n", - "To create a generator of a specific set of results, we use the `values` method. This takes the `name` of the\n", - "object we want to create a generator of, for example inputting `name=samples` will return the results `Samples`\n", - "object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples_gen = agg.values(\"samples\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By converting this generator to a list and printing it, it is a list of 3 `SamplesNest` objects, corresponding to \n", - "the 3 model-fits performed above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Samples:\\n\")\n", - "print(samples_gen)\n", - "print(\"Total Samples Objects = \", len(agg), \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "The model used to perform the model fit for each of the 3 datasets can be loaded via the aggregator and printed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_gen = agg.values(\"model\")\n", - "\n", - "for model in model_gen:\n", - " print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The non-linear search used to perform the model fit can be loaded via the aggregator and printed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_gen = agg.values(\"search\")\n", - "\n", - "for search in search_gen:\n", - " print(search)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Samples__\n", - "\n", - "The `Samples` class contains all information on the non-linear search samples, for example the value of every parameter\n", - "sampled using the fit or an instance of the maximum likelihood model.\n", - "\n", - "The `Samples` class is described fully in the results cookbook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " print(\"The tenth sample`s third parameter\")\n", - " print(samples.parameter_lists[9][2], \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Therefore, by loading the `Samples` via the database we can now access the results of the fit to each dataset.\n", - "\n", - "For example, we can plot the maximum likelihood model for each of the 3 model-fits performed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "ml_vector = [\n", - " samps.max_log_likelihood(as_instance=False) for samps in agg.values(\"samples\")\n", - "]\n", - "\n", - "print(\"Max Log Likelihood Model Parameter Lists: \\n\")\n", - "print(ml_vector, \"\\n\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "All remaining methods accessible by `agg.values` are described in the other database examples.\n", - "\n", - "__Wrap Up__\n", - "\n", - "This example illustrates how to use the database.\n", - "\n", - "The API above can be combined with the `results/aggregator` scripts in order to use the database to load results and\n", - "perform analysis on them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Database: Introduction\n", + "======================\n", + "\n", + "The default behaviour of model-fitting results output is to be written to hard-disc in folders. These are simple to\n", + "navigate and manually check.\n", + "\n", + "For small model-fitting tasks this is sufficient, however it does not scale well when performing many model fits to\n", + "large datasets, because manual inspection of results becomes time consuming.\n", + "\n", + "All results can therefore be output to an sqlite3 (https://docs.python.org/3/library/sqlite3.html) relational database,\n", + "meaning that results can be loaded into a Jupyter notebook or Python script for inspection, analysis and interpretation.\n", + "This database supports advanced querying, so that specific model-fits (e.g., which fit a certain model or dataset) can\n", + "be loaded.\n", + "\n", + "This script fits a sample of three simulated strong lenses using the same non-linear search. The results will be used\n", + "to illustrate the database in the database tutorials that follow.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Unique Identifiers:** Results output to hard-disk are contained in a folder named via a unique identifier (a random.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Results From Hard Disk:** In this example, results will be first be written to hard-disk using the standard output directory.\n", + "- **Building a Database File From an Output Folder:** The fits above wrote the results to hard-disk in folders, not as an .sqlite database file.\n", + "- **Writing Directly To Database:** Results can be written directly to the .sqlite database file, skipping output to hard-disk.\n", + "- **Files:** When performing fits which output results to hard-disc, a `files` folder is created containing.\n", + "- **Generators:** Before using the aggregator to inspect results, lets discuss Python generators.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Samples:** The `Samples` class contains all information on the non-linear search samples, for example the.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Model__\n", + "\n", + "The search fits each lens with:\n", + "\n", + " - An `Isothermal` `MassProfile` for the lens galaxy's mass.\n", + " - An `Sersic` `LightProfile` for the source galaxy's light." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import json\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Unique Identifiers__\n", + "\n", + "Results output to hard-disk are contained in a folder named via a unique identifier (a \n", + "random collection of characters, e.g. `8hds89fhndlsiuhnfiusdh`). The unique identifier changes if the model or \n", + "search change, to ensure different fits to not overwrite one another on hard-disk.\n", + "\n", + "Each unique identifier is used to define every entry of the database as it is built. Unique identifiers therefore play \n", + "the same vital role for the database of ensuring that every set of results written to it are unique.\n", + "\n", + "In this example, we fit 3 different datasets with the same search and model. Each `dataset_name` is therefore passed\n", + "in as the search's `unique_tag` to ensure 3 separate sets of results for each model-fit are written to the .sqlite\n", + "database.\n", + "\n", + "__Dataset__\n", + "\n", + "For each dataset we load it from hard-disc, set up its `Analysis` class and fit it with a non-linear search. \n", + "\n", + "We want each results to be stored in the database with an entry specific to the dataset. We'll use the `Dataset`'s name \n", + "string to do this, so lets create a list of the 3 dataset names." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_names = [\n", + " \"simple\",\n", + " \"lens_sersic\",\n", + " \"simple__no_lens_light\",\n", + "]\n", + "\n", + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Results From Hard Disk__\n", + "\n", + "In this example, results will be first be written to hard-disk using the standard output directory structure and we\n", + "will then build the database from these results. This behaviour is governed by us inputting `session=None`.\n", + "\n", + "If you have existing results you wish to build a database for, you can therefore adapt this example you to do this.\n", + "\n", + "Later in this example we show how results can also also be output directly to an .sqlite database, saving on hard-disk \n", + "space. This will be acheived by setting `session` to something that is not `None`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "session = None\n", + "\n", + "for dataset_name in dataset_names:\n", + " \"\"\"\n", + " __Paths__\n", + "\n", + " Set up the config and output paths.\n", + " \"\"\"\n", + " dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + " \"\"\"\n", + " __Dataset__\n", + "\n", + " Using the dataset path, load the data (image, noise-map, PSF) as an `Imaging` object from .fits files.\n", + "\n", + " This `Imaging` object will be available via the aggregator. Note also that we give the dataset a `name` via the\n", + " command `name=dataset_name`. we'll use this name in the aggregator tutorials.\n", + " \"\"\"\n", + " dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=pixel_scales,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Mask__\n", + "\n", + " The `Mask2D` we fit this data-set with, which will be available via the aggregator.\n", + " \"\"\"\n", + " mask_radius = 3.0\n", + "\n", + " mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=mask)\n", + "\n", + " \"\"\"\n", + " __Info__\n", + "\n", + " Information about the model-fit that is not part included in the model-fit itself can be made accessible via the \n", + " database by passing an `info` dictionary. \n", + "\n", + " Below we write info on the dataset`s (hypothetical) data of observation and exposure time, which we will later show\n", + " the database can access. \n", + "\n", + " For fits to large datasets this ensures that all relevant information for interpreting results is accessible.\n", + " \"\"\"\n", + " with open(Path(dataset_path, \"info.json\")) as json_file:\n", + " info = json.load(json_file)\n", + "\n", + " \"\"\"\n", + " __Model__\n", + "\n", + " Set up the model as per usual, and will see in tutorial 3 why we have included `disk=None`.\n", + " \"\"\"\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", + " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=None),\n", + " )\n", + " )\n", + "\n", + " \"\"\"\n", + " The `unique_tag` below uses the `dataset_name` to alter the unique identifier, which as we have seen is also \n", + " generated depending on the search settings and model. In this example, all three model fits use an identical \n", + " search and model, so this `unique_tag` is key for ensuring 3 separate sets of results for each model-fit are \n", + " stored in the output folder and written to the .sqlite database. \n", + " \"\"\"\n", + " search = af.Nautilus(\n", + " path_prefix=Path(\"database\"),\n", + " name=\"database_example\",\n", + " unique_tag=dataset_name, # This makes the unique identifier use the dataset name\n", + " session=session, # This can instruct the search to write to the .sqlite database.\n", + " n_live=100,\n", + " n_batch=50, # GPU batching and VRAM use explained in `modeling` examples.\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " search.fit(analysis=analysis, model=model, info=info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Building a Database File From an Output Folder__\n", + "\n", + "The fits above wrote the results to hard-disk in folders, not as an .sqlite database file. \n", + "\n", + "We build the database below, where the `database_name` corresponds to the name of your output folder and is also the \n", + "name of the `.sqlite` database file that is created.\n", + "\n", + "If you are fitting a relatively small number of datasets (e.g. 10-100) having all results written to hard-disk (e.g. \n", + "for quick visual inspection) and using the database for sample wide analysis is beneficial.\n", + "\n", + "We can optionally only include completed model-fits but setting `completed_only=True`.\n", + "\n", + "If you inspect the `output` folder, you will see a `database.sqlite` file which contains the results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "database_name = \"database\"\n", + "\n", + "agg = af.Aggregator.from_database(\n", + " filename=f\"{database_name}.sqlite\", completed_only=False\n", + ")\n", + "\n", + "agg.add_directory(directory=Path(\"output\") / database_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Writing Directly To Database__\n", + "\n", + "Results can be written directly to the .sqlite database file, skipping output to hard-disk entirely, by creating\n", + "a session and passing this to the non-linear search.\n", + "\n", + "The code below shows how to do this, but it is commented out to avoid rerunning the non-linear searches.\n", + "\n", + "This is ideal for tasks where model-fits to hundreds or thousands of datasets are performed, as it becomes unfeasible\n", + "to inspect the results of all fits on the hard-disk. \n", + "\n", + "Our recommended workflow is to set up database analysis scripts using ~10 model-fits, and then scaling these up\n", + "to large samples by writing directly to the database." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# session = af.db.open_database(\"database.sqlite\")\n", + "#\n", + "# search = af.Nautilus(\n", + "# path_prefix=Path(\"database\"),\n", + "# name=\"database_example\",\n", + "# unique_tag=dataset_name, # This makes the unique identifier use the dataset name\n", + "# session=session, # This can instruct the search to write to the .sqlite database.\n", + "# n_live=100,\n", + "# n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + "# )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Files__\n", + "\n", + "When performing fits which output results to hard-disc, a `files` folder is created containing .json / .csv files of \n", + "the model, samples, search, etc.\n", + "\n", + "These are the files that are written to the database, which the aggregator loads via the database in order to make \n", + "them accessible in a Python script or Jupyter notebook.\n", + "\n", + "You can checkout the output folder created by this fit to see these files.\n", + "\n", + "Below, we will access these results using the aggregator's `values` method. A full list of what can be loaded is\n", + "as follows:\n", + "\n", + " - `model`: The `model` defined above and used in the model-fit (`model.json`).\n", + " - `search`: The non-linear search settings (`search.json`).\n", + " - `samples`: The non-linear search samples (`samples.csv`).\n", + " - `samples_info`: Additional information about the samples (`samples_info.json`).\n", + " - `samples_summary`: A summary of key results of the samples (`samples_summary.json`).\n", + " - `info`: The info dictionary passed to the search (`info.json`).\n", + " - `covariance`: The inferred covariance matrix (`covariance.csv`).\n", + " - `cosmology`: The cosmology used by the fit (`cosmology.json`).\n", + " - `settings`: The settings associated with a inversion if used (`settings.json`).\n", + " - `dataset/data`: The data that is fitted (`data.fits`).\n", + " - `dataset/noise_map`: The noise-map (`noise_map.fits`).\n", + " - `dataset/psf`: The Point Spread Function (`psf.fits`).\n", + " - `dataset/mask`: The mask applied to the data (`mask.fits`).\n", + " - `dataset/settings`: The settings associated with the dataset (`settings.json`).\n", + "\n", + "The `samples` and `samples_summary` results contain a lot of repeated information. The `samples` result contains\n", + "the full non-linear search samples, for example every parameter sample and its log likelihood. The `samples_summary`\n", + "contains a summary of the results, for example the maximum log likelihood model and error estimates on parameters\n", + "at 1 and 3 sigma confidence.\n", + "\n", + "Accessing results via the `samples_summary` is much faster, because as it does reperform calculations using the full \n", + "list of samples. Therefore, if the result you want is accessible via the `samples_summary` you should use it\n", + "but if not you can revert to the `samples.\n", + "\n", + "__Generators__\n", + "\n", + "Before using the aggregator to inspect results, lets discuss Python generators. \n", + "\n", + "A generator is an object that iterates over a function when it is called. The aggregator creates all of the objects \n", + "that it loads from the database as generators (as opposed to a list, or dictionary, or another Python type).\n", + "\n", + "This is because generators are memory efficient, as they do not store the entries of the database in memory \n", + "simultaneously. This contrasts objects like lists and dictionaries, which store all entries in memory all at once. \n", + "If you fit a large number of datasets, lists and dictionaries will use a lot of memory and could crash your computer!\n", + "\n", + "Once we use a generator in the Python code, it cannot be used again. To perform the same task twice, the \n", + "generator must be remade it. This cookbook therefore rarely stores generators as variables and instead uses the \n", + "aggregator to create each generator at the point of use.\n", + "\n", + "To create a generator of a specific set of results, we use the `values` method. This takes the `name` of the\n", + "object we want to create a generator of, for example inputting `name=samples` will return the results `Samples`\n", + "object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples_gen = agg.values(\"samples\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By converting this generator to a list and printing it, it is a list of 3 `SamplesNest` objects, corresponding to \n", + "the 3 model-fits performed above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Samples:\\n\")\n", + "print(samples_gen)\n", + "print(\"Total Samples Objects = \", len(agg), \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "The model used to perform the model fit for each of the 3 datasets can be loaded via the aggregator and printed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_gen = agg.values(\"model\")\n", + "\n", + "for model in model_gen:\n", + " print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The non-linear search used to perform the model fit can be loaded via the aggregator and printed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_gen = agg.values(\"search\")\n", + "\n", + "for search in search_gen:\n", + " print(search)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Samples__\n", + "\n", + "The `Samples` class contains all information on the non-linear search samples, for example the value of every parameter\n", + "sampled using the fit or an instance of the maximum likelihood model.\n", + "\n", + "The `Samples` class is described fully in the results cookbook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " print(\"The tenth sample`s third parameter\")\n", + " print(samples.parameter_lists[9][2], \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Therefore, by loading the `Samples` via the database we can now access the results of the fit to each dataset.\n", + "\n", + "For example, we can plot the maximum likelihood model for each of the 3 model-fits performed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "ml_vector = [\n", + " samps.max_log_likelihood(as_instance=False) for samps in agg.values(\"samples\")\n", + "]\n", + "\n", + "print(\"Max Log Likelihood Model Parameter Lists: \\n\")\n", + "print(ml_vector, \"\\n\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "All remaining methods accessible by `agg.values` are described in the other database examples.\n", + "\n", + "__Wrap Up__\n", + "\n", + "This example illustrates how to use the database.\n", + "\n", + "The API above can be combined with the `results/aggregator` scripts in order to use the database to load results and\n", + "perform analysis on them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/latent_variables.ipynb b/notebooks/guides/results/latent_variables.ipynb index 9d0ffc6d6..3c8149247 100644 --- a/notebooks/guides/results/latent_variables.ipynb +++ b/notebooks/guides/results/latent_variables.ipynb @@ -1,419 +1,456 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Latent Variables\n", - "=========================\n", - "\n", - "PyAutoLens ships a curated catalogue of lensing-specific latent variables \u2014 quantities derived from the lens\n", - "model that aren't sampled directly but are computed from each posterior draw to give you Bayesian uncertainties\n", - "on the science quantities you actually care about. This tutorial shows the catalogue, how to toggle individual\n", - "latents via the workspace config, how to load latent results from a completed fit, and how to extend\n", - "``al.AnalysisImaging`` with your own derived quantities.\n", - "\n", - "The conceptual underpinning \u2014 what a latent variable IS in the Bayesian sense, what its 1\u03c3/3\u03c3 errors actually\n", - "mean (empirical posterior quantiles, NOT analytic Gaussian propagation), the trade-off between every-sample\n", - "and N-draws-from-PDF output modes \u2014 lives in the autofit_workspace foundational tutorial at\n", - "``../../../autofit_workspace/scripts/cookbooks/latent_variables.py``. Read that first if any of those terms\n", - "look unfamiliar; the tutorial here focuses on the lensing-specific catalogue and the workspace ergonomics.\n", - "\n", - "__Contents__\n", - "\n", - " - Lensing Latents in PyAutoLens: The eight library-shipped latents and what each one means physically.\n", - " - Toggling Latents: The workspace ``config/latent.yaml`` override.\n", - " - Model Fit: Reuse the shared quick fit (``_quick_fit.py``) that produces real latent output.\n", - " - Loading Latent Results: ``analysis.compute_latent_samples`` over a subset of PDF draws, and the two\n", - " config surfaces (``latent.yaml`` / ``output.yaml``) that control which latents and how many draws.\n", - " - Extending with a Custom Latent: Subclass ``al.AnalysisImaging`` to add lens-mass derived quantities.\n", - " - Contributing Upstream: When your custom latent is general enough, promote it to the library." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lensing Latents in PyAutoLens__\n", - "\n", - "The library ships a flat registry of named latent functions at ``autolens.analysis.latent.LATENT_FUNCTIONS``,\n", - "backed by the toggle file ``autolens/config/latent.yaml``. Each entry maps a snake-case latent name to a Python\n", - "function that takes a fit, magzero, and ``xp`` and returns a scalar value. The catalogue splits into three\n", - "groups: raw-flux latents (no instrument inputs, default-on), microjansky variants (require ``magzero``,\n", - "default-off), and the dimensionless lensing latents (``magnification``, ``effective_einstein_radius``).\n", - "\n", - "Raw-flux latents \u2014 sum the relevant model image in the fit's raw image units. Same units as\n", - "``dataset.data.array`` (typically e- s^-1 for HST, MJy/sr for JWST). See\n", - "``scripts/guides/units/flux.py`` for how to convert to microjanskies or AB magnitudes in post.\n", - "\n", - " - ``total_lens_flux`` \u2014 total integrated flux of the lens galaxy. Sum of\n", - " ``fit.galaxy_image_dict[fit.tracer.galaxies[0]].array``. Returns NaN when the lens has no light profile.\n", - "\n", - " - ``total_lensed_source_flux`` \u2014 image-plane integrated flux of the source galaxy after lensing (with\n", - " magnification baked in). Sum of ``fit.galaxy_image_dict[fit.tracer.galaxies[-1]].array``.\n", - "\n", - " - ``total_source_flux`` \u2014 the source's intrinsic flux in the source plane (before lensing). Computed from\n", - " the source's light profile evaluated on the workspace's light-profile grid. Critically uses\n", - " ``fit.tracer_linear_light_profiles_to_light_profiles`` rather than ``fit.tracer`` so MGE / linear light\n", - " profiles report the inversion-solved intensities rather than zero.\n", - "\n", - "Microjansky variants \u2014 same image sources as the raw-flux trio, but with the AB-mag \u2192 \u00b5Jy conversion baked\n", - "in. Each one requires ``magzero`` on the analysis; if it's missing, the latent returns NaN and the library\n", - "emits a single warning per process per latent name (your search still completes \u2014 the cost is just an empty\n", - "column).\n", - "\n", - " - ``total_lens_flux_mujy`` \u2014 ``total_lens_flux`` in microjanskies. Useful for stellar-mass-light scaling and\n", - " photometric inference.\n", - "\n", - " - ``total_lensed_source_flux_mujy`` \u2014 ``total_lensed_source_flux`` in microjanskies.\n", - "\n", - " - ``total_source_flux_mujy`` \u2014 ``total_source_flux`` in microjanskies.\n", - "\n", - "Dimensionless lensing latents \u2014 no instrument inputs, no \u00b5Jy variant.\n", - "\n", - " - ``magnification`` \u2014 dimensionless ratio of ``total_lensed_source_flux_mujy / total_source_flux_mujy``.\n", - " This is the empirical flux-amplification factor implied by the lens model and source light profile. Source-\n", - " plane errors propagate non-linearly here because the source brightness and the lens magnification both\n", - " contribute multiplicatively to the lensed flux. The empirical posterior on ``magnification`` (samples\n", - " transformed through the ratio) captures that non-linearity faithfully; analytic error propagation through\n", - " the ratio would not.\n", - "\n", - " - ``effective_einstein_radius`` \u2014 the Einstein radius in arcseconds, computed via\n", - " ``LensCalc.einstein_radius_jit_from`` which traces the zero-contour of the tangential eigenvalue of the\n", - " deflection field. \"Effective\" here means the radius of the circle with the same enclosed area as the\n", - " tangential critical curve \u2014 the critical curve isn't circular for non-spherical mass models, but its\n", - " enclosed area is the right physical quantity for mass-within-Einstein-radius estimates.\n", - "\n", - "The raw-flux latents default to ``true`` in the library yaml \u2014 they cost essentially nothing and produce a\n", - "universally useful column. The \u00b5Jy variants and the two dimensionless latents default to ``false`` so existing\n", - "fits and instrument-naive workflows stay unchanged on upgrade." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Toggling Latents__\n", - "\n", - "The library defaults the three raw-flux latents to ``true`` and everything else to ``false`` for the\n", - "regression-safety reasons described above. To opt in to the \u00b5Jy variants, dimensionless lensing latents, or to\n", - "disable a default-on raw-flux latent, edit your workspace's ``config/latent.yaml`` and set the keys you want.\n", - "This workspace ships such a file at ``autolens_workspace/config/latent.yaml`` with all eight lensing latents\n", - "enabled.\n", - "\n", - "Workspace ``config/`` values shadow the library defaults \u2014 PyAutoFit's ``conf.instance`` searches the workspace\n", - "``config/`` directory first, so toggling a latent in your workspace yaml is enough to enable it without modifying\n", - "the library install. To disable a specific latent for a particular fit (e.g. you're profiling and don't want\n", - "to incur the latent computation cost on every search update), flip it to ``false`` in\n", - "``autolens_workspace/config/latent.yaml`` or override locally with ``conf.instance.push(...)``.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Model Fit__\n", - "\n", - "The loading and extending sections need a completed fit to read latents from. Rather than run a bespoke\n", - "fit here, we reuse the shared quick fit that the other results guides use: ``_quick_fit.py`` writes an\n", - "Isothermal-lens + MGE-source fit of the ``simple__no_lens_light`` dataset to ``output/results_folder/``.\n", - "It is idempotent \u2014 it returns immediately if those results already exist \u2014 so across the whole guide\n", - "suite the expensive non-linear search is paid once rather than repeated in every example.\n", - "\n", - "We then load that fit's samples via the aggregator, exactly as ``start_here.py`` and\n", - "``aggregator/models.py`` do.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import subprocess\n", - "import sys\n", - "\n", - "subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - ")\n", - "\n", - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "agg = Aggregator.from_directory(directory=Path(\"output\") / \"results_folder\")\n", - "samples = list(agg)[0].samples" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The samples carry the parameter posterior; the ``analysis`` carries the machinery that turns each\n", - "posterior draw into latent values. We rebuild the dataset and an ``al.AnalysisImaging`` with\n", - "``magzero=25.0`` so the three \u00b5Jy latents populate with real values (without it they'd be NaN and the\n", - "library would log a single warning per latent name per process). The raw-flux trio, Einstein-radius and\n", - "magnification latents don't need ``magzero``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=3.0,\n", - ")\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=False, magzero=25.0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Latent Results__\n", - "\n", - "``analysis.compute_latent_samples(samples)`` returns a ``Samples`` object whose API matches the parameter\n", - "``Samples`` \u2014 ``median_pdf``, ``max_log_likelihood``, ``values_at_sigma_1``, etc. \u2014 but reports on the induced\n", - "latent posterior.\n", - "\n", - "__Controlling the Cost via Config__\n", - "\n", - "Latents are computed by reconstructing a fit for every posterior sample, so the cost scales with the number\n", - "of samples \u2014 and the two dimensionless lensing latents (``magnification``, ``effective_einstein_radius``) add\n", - "a zero-contour critical-curve solve on top. Two workspace config files control this, and both are worth\n", - "knowing:\n", - "\n", - " - ``config/latent.yaml`` \u2014 controls *which* latents are computed (the eight toggles; this workspace enables\n", - " all of them). Disable the ones you don't need to save their per-sample cost.\n", - "\n", - " - ``config/output.yaml`` \u2014 ``latent_draw_via_pdf`` / ``latent_draw_via_pdf_size`` control *how many* posterior\n", - " draws the latents are computed over when a live search updates. Drawing a representative subset from the PDF\n", - " gives faithful latent errors at a fraction of the every-sample cost.\n", - "\n", - "Here we mirror that draw-from-PDF behaviour explicitly with ``samples.samples_drawn_randomly_via_pdf_from``,\n", - "computing the latents over 20 PDF draws so this guide runs quickly while still producing a real, representative\n", - "latent posterior. For a publication-quality result, compute over all samples (or a larger number of draws)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "latent_draws = samples.samples_drawn_randomly_via_pdf_from(total_draws=20)\n", - "latent_samples = analysis.compute_latent_samples(latent_draws)\n", - "\n", - "median_instance = latent_samples.median_pdf()\n", - "print(f\"Median PDF magnification: {median_instance.magnification}\")\n", - "print(\n", - " f\"Median PDF effective_einstein_radius: {median_instance.effective_einstein_radius}\"\n", - ")\n", - "print(f\"Median PDF total_source_flux_mujy: {median_instance.total_source_flux_mujy}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The 1\u03c3 / 3\u03c3 intervals on these latents are *empirical quantiles of the induced posterior*. For magnification\n", - "specifically, this is the right thing \u2014 magnification is a strongly non-linear function of the mass parameters\n", - "and the source position, so the symmetric Gaussian propagation of parameter errors would underreport the tail\n", - "risk. The empirical posterior captures the full non-linearity faithfully. See the autofit foundational tutorial\n", - "for the full treatment.\n", - "\n", - "The ``effective_einstein_radius`` latent can produce noisy 1D posteriors when the zero-contour solver converges\n", - "to slightly different critical-curve estimates across samples. If you see step-like artefacts in a corner plot,\n", - "that's the discrete contour-finding cadence interacting with the continuous parameter posterior \u2014 the median\n", - "and 1\u03c3 intervals remain reliable, but use caution when reading the tails." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Extending with a Custom Latent__\n", - "\n", - "The library catalogue is intentionally narrow. If you want a different derived quantity \u2014 the lens mass's\n", - "axis ratio, the source-plane S\u00e9rsic effective radius, the time-delay between multiply-imaged source pixels \u2014\n", - "subclass ``al.LatentLens`` (override ``keys`` and ``variables``) and declare it on your analysis through the\n", - "``Latent`` class attribute. This mirrors exactly how you customise visualisation with a ``Visualizer`` subclass:\n", - "the latent catalogue is a first-class, swappable component of the analysis rather than a pair of methods you\n", - "monkey-patch.\n", - "\n", - "The example below adds ``mass_axis_ratio`` (the lens-mass axis ratio derived from the ``Isothermal`` mass's\n", - "``ell_comps``). The same pattern works for any function of the lens model instance. Calling\n", - "``al.LatentLens.keys(analysis)`` / ``al.LatentLens.variables(...)`` from your overrides keeps the library\n", - "latents alongside your custom ones \u2014 the Euclid pipeline\n", - "(``euclid_strong_lens_modeling_pipeline/util.py``, ``LatentEuclid``) uses this exact composition pattern in\n", - "production to append its FWHM aperture-flux latents.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "import numpy as np\n", - "\n", - "\n", - "class LatentMassAxisRatio(al.LatentLens):\n", - " \"\"\"\n", - " The library lensing latents plus a custom ``mass_axis_ratio`` \u2014 the axis\n", - " ratio of the lens galaxy's Isothermal mass profile, derived from its\n", - " ``ell_comps``. Demonstrates adding a user-defined latent without modifying\n", - " the library: subclass ``al.LatentLens`` and compose its ``keys`` /\n", - " ``variables`` static methods.\n", - " \"\"\"\n", - "\n", - " @staticmethod\n", - " def keys(analysis):\n", - " return list(al.LatentLens.keys(analysis)) + [\"mass_axis_ratio\"]\n", - "\n", - " @staticmethod\n", - " def variables(analysis, parameters, model):\n", - " library_values = al.LatentLens.variables(analysis, parameters, model)\n", - "\n", - " instance = model.instance_from_vector(vector=parameters)\n", - " xp = analysis._xp\n", - " try:\n", - " ell_y, ell_x = instance.galaxies.lens.mass.ell_comps\n", - " axis_ratio = (1.0 - np.sqrt(ell_y**2 + ell_x**2)) / (\n", - " 1.0 + np.sqrt(ell_y**2 + ell_x**2)\n", - " )\n", - " except AttributeError:\n", - " axis_ratio = xp.nan\n", - "\n", - " return library_values + (axis_ratio,)\n", - "\n", - "\n", - "class AnalysisImagingWithMassAxisRatio(al.AnalysisImaging):\n", - " \"\"\"\n", - " ``AnalysisImaging`` that swaps in the custom ``LatentMassAxisRatio`` catalogue\n", - " via the ``Latent`` class attribute \u2014 the same one-line mechanism used to\n", - " declare a custom ``Visualizer`` or ``Result``. No library code is modified.\n", - " \"\"\"\n", - "\n", - " Latent = LatentMassAxisRatio\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A fit that uses ``AnalysisImagingWithMassAxisRatio`` produces a ``latent.csv`` with one extra column\n", - "(``mass_axis_ratio``) on top of the eight library defaults (raw-flux + \u00b5Jy + dimensionless lensing). We don't\n", - "run a second fit here \u2014 the pattern above is the full recipe." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Contributing Upstream__\n", - "\n", - "If your custom latent is general enough that other PyAutoLens users would benefit (a SLACS-style external-\n", - "convergence proxy, a critical-curve perimeter, an enclosed-mass-at-fixed-radius), please consider promoting it\n", - "to the library:\n", - "\n", - " 1. Add the function to ``autolens/analysis/latent.py``, following the signature\n", - " ``(fit, magzero, xp=np) -> scalar``. Use NaN as the fallback when the function can't apply (no lens / no\n", - " source / singular mass model).\n", - " 2. Register it in the module-level ``LATENT_FUNCTIONS`` dict.\n", - " 3. Add an entry to ``autolens/config/latent.yaml`` defaulting it to ``false`` (the workspace yaml opts users in).\n", - " 4. Add a unit test under ``test_autolens/analysis/test_latent.py``.\n", - " 5. Open a PR.\n", - "\n", - "Pipeline-specific latents that need non-standard kwargs (PSF-relative aperture fluxes, dataset-specific\n", - "photometric quantities, multi-band colour terms) belong in your pipeline's local Analysis subclass. The Euclid\n", - "pipeline (``euclid_strong_lens_modeling_pipeline/util.py``) demonstrates this \u2014 its FWHM aperture-flux latents\n", - "stay pipeline-local and the rest inherit from the library.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Latent Variables\n", + "=========================\n", + "\n", + "PyAutoLens ships a curated catalogue of lensing-specific latent variables \u2014 quantities derived from the lens\n", + "model that aren't sampled directly but are computed from each posterior draw to give you Bayesian uncertainties\n", + "on the science quantities you actually care about. This tutorial shows the catalogue, how to toggle individual\n", + "latents via the workspace config, how to load latent results from a completed fit, and how to extend\n", + "``al.AnalysisImaging`` with your own derived quantities.\n", + "\n", + "The conceptual underpinning \u2014 what a latent variable IS in the Bayesian sense, what its 1\u03c3/3\u03c3 errors actually\n", + "mean (empirical posterior quantiles, NOT analytic Gaussian propagation), the trade-off between every-sample\n", + "and N-draws-from-PDF output modes \u2014 lives in the autofit_workspace foundational tutorial at\n", + "``../../../autofit_workspace/scripts/cookbooks/latent_variables.py``. Read that first if any of those terms\n", + "look unfamiliar; the tutorial here focuses on the lensing-specific catalogue and the workspace ergonomics.\n", + "\n", + "__Contents__\n", + "\n", + " - Lensing Latents in PyAutoLens: The eight library-shipped latents and what each one means physically.\n", + " - Toggling Latents: The workspace ``config/latent.yaml`` override.\n", + " - Model Fit: Reuse the shared quick fit (``_quick_fit.py``) that produces real latent output.\n", + " - Loading Latent Results: ``analysis.compute_latent_samples`` over a subset of PDF draws, and the two\n", + " config surfaces (``latent.yaml`` / ``output.yaml``) that control which latents and how many draws.\n", + " - Extending with a Custom Latent: Subclass ``al.AnalysisImaging`` to add lens-mass derived quantities.\n", + " - Contributing Upstream: When your custom latent is general enough, promote it to the library." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lensing Latents in PyAutoLens__\n", + "\n", + "The library ships a flat registry of named latent functions at ``autolens.analysis.latent.LATENT_FUNCTIONS``,\n", + "backed by the toggle file ``autolens/config/latent.yaml``. Each entry maps a snake-case latent name to a Python\n", + "function that takes a fit, magzero, and ``xp`` and returns a scalar value. The catalogue splits into three\n", + "groups: raw-flux latents (no instrument inputs, default-on), microjansky variants (require ``magzero``,\n", + "default-off), and the dimensionless lensing latents (``magnification``, ``effective_einstein_radius``).\n", + "\n", + "Raw-flux latents \u2014 sum the relevant model image in the fit's raw image units. Same units as\n", + "``dataset.data.array`` (typically e- s^-1 for HST, MJy/sr for JWST). See\n", + "``scripts/guides/units/flux.py`` for how to convert to microjanskies or AB magnitudes in post.\n", + "\n", + " - ``total_lens_flux`` \u2014 total integrated flux of the lens galaxy. Sum of\n", + " ``fit.galaxy_image_dict[fit.tracer.galaxies[0]].array``. Returns NaN when the lens has no light profile.\n", + "\n", + " - ``total_lensed_source_flux`` \u2014 image-plane integrated flux of the source galaxy after lensing (with\n", + " magnification baked in). Sum of ``fit.galaxy_image_dict[fit.tracer.galaxies[-1]].array``.\n", + "\n", + " - ``total_source_flux`` \u2014 the source's intrinsic flux in the source plane (before lensing). Computed from\n", + " the source's light profile evaluated on the workspace's light-profile grid. Critically uses\n", + " ``fit.tracer_linear_light_profiles_to_light_profiles`` rather than ``fit.tracer`` so MGE / linear light\n", + " profiles report the inversion-solved intensities rather than zero.\n", + "\n", + "Microjansky variants \u2014 same image sources as the raw-flux trio, but with the AB-mag \u2192 \u00b5Jy conversion baked\n", + "in. Each one requires ``magzero`` on the analysis; if it's missing, the latent returns NaN and the library\n", + "emits a single warning per process per latent name (your search still completes \u2014 the cost is just an empty\n", + "column).\n", + "\n", + " - ``total_lens_flux_mujy`` \u2014 ``total_lens_flux`` in microjanskies. Useful for stellar-mass-light scaling and\n", + " photometric inference.\n", + "\n", + " - ``total_lensed_source_flux_mujy`` \u2014 ``total_lensed_source_flux`` in microjanskies.\n", + "\n", + " - ``total_source_flux_mujy`` \u2014 ``total_source_flux`` in microjanskies.\n", + "\n", + "Dimensionless lensing latents \u2014 no instrument inputs, no \u00b5Jy variant.\n", + "\n", + " - ``magnification`` \u2014 dimensionless ratio of ``total_lensed_source_flux_mujy / total_source_flux_mujy``.\n", + " This is the empirical flux-amplification factor implied by the lens model and source light profile. Source-\n", + " plane errors propagate non-linearly here because the source brightness and the lens magnification both\n", + " contribute multiplicatively to the lensed flux. The empirical posterior on ``magnification`` (samples\n", + " transformed through the ratio) captures that non-linearity faithfully; analytic error propagation through\n", + " the ratio would not.\n", + "\n", + " - ``effective_einstein_radius`` \u2014 the Einstein radius in arcseconds, computed via\n", + " ``LensCalc.einstein_radius_jit_from`` which traces the zero-contour of the tangential eigenvalue of the\n", + " deflection field. \"Effective\" here means the radius of the circle with the same enclosed area as the\n", + " tangential critical curve \u2014 the critical curve isn't circular for non-spherical mass models, but its\n", + " enclosed area is the right physical quantity for mass-within-Einstein-radius estimates.\n", + "\n", + "The raw-flux latents default to ``true`` in the library yaml \u2014 they cost essentially nothing and produce a\n", + "universally useful column. The \u00b5Jy variants and the two dimensionless latents default to ``false`` so existing\n", + "fits and instrument-naive workflows stay unchanged on upgrade." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Toggling Latents__\n", + "\n", + "The library defaults the three raw-flux latents to ``true`` and everything else to ``false`` for the\n", + "regression-safety reasons described above. To opt in to the \u00b5Jy variants, dimensionless lensing latents, or to\n", + "disable a default-on raw-flux latent, edit your workspace's ``config/latent.yaml`` and set the keys you want.\n", + "This workspace ships such a file at ``autolens_workspace/config/latent.yaml`` with all eight lensing latents\n", + "enabled.\n", + "\n", + "Workspace ``config/`` values shadow the library defaults \u2014 PyAutoFit's ``conf.instance`` searches the workspace\n", + "``config/`` directory first, so toggling a latent in your workspace yaml is enough to enable it without modifying\n", + "the library install. To disable a specific latent for a particular fit (e.g. you're profiling and don't want\n", + "to incur the latent computation cost on every search update), flip it to ``false`` in\n", + "``autolens_workspace/config/latent.yaml`` or override locally with ``conf.instance.push(...)``.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Model Fit__\n", + "\n", + "The loading and extending sections need a completed fit to read latents from. Rather than run a bespoke\n", + "fit here, we reuse the shared quick fit that the other results guides use: ``_quick_fit.py`` writes an\n", + "Isothermal-lens + MGE-source fit of the ``simple__no_lens_light`` dataset to ``output/results_folder/``.\n", + "It is idempotent \u2014 it returns immediately if those results already exist \u2014 so across the whole guide\n", + "suite the expensive non-linear search is paid once rather than repeated in every example.\n", + "\n", + "We then load that fit's samples via the aggregator, exactly as ``start_here.py`` and\n", + "``aggregator/models.py`` do.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import subprocess\n", + "import sys\n", + "\n", + "subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + ")\n", + "\n", + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "agg = Aggregator.from_directory(directory=Path(\"output\") / \"results_folder\")\n", + "samples = list(agg)[0].samples" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The samples carry the parameter posterior; the ``analysis`` carries the machinery that turns each\n", + "posterior draw into latent values. We rebuild the dataset and an ``al.AnalysisImaging`` with\n", + "``magzero=25.0`` so the three \u00b5Jy latents populate with real values (without it they'd be NaN and the\n", + "library would log a single warning per latent name per process). The raw-flux trio, Einstein-radius and\n", + "magnification latents don't need ``magzero``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=3.0,\n", + ")\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=False, magzero=25.0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Latent Results__\n", + "\n", + "``analysis.compute_latent_samples(samples)`` returns a ``Samples`` object whose API matches the parameter\n", + "``Samples`` \u2014 ``median_pdf``, ``max_log_likelihood``, ``values_at_sigma_1``, etc. \u2014 but reports on the induced\n", + "latent posterior.\n", + "\n", + "__Controlling the Cost via Config__\n", + "\n", + "Latents are computed by reconstructing a fit for every posterior sample, so the cost scales with the number\n", + "of samples \u2014 and the two dimensionless lensing latents (``magnification``, ``effective_einstein_radius``) add\n", + "a zero-contour critical-curve solve on top. Two workspace config files control this, and both are worth\n", + "knowing:\n", + "\n", + " - ``config/latent.yaml`` \u2014 controls *which* latents are computed (the eight toggles; this workspace enables\n", + " all of them). Disable the ones you don't need to save their per-sample cost.\n", + "\n", + " - ``config/output.yaml`` \u2014 ``latent_draw_via_pdf`` / ``latent_draw_via_pdf_size`` control *how many* posterior\n", + " draws the latents are computed over when a live search updates. Drawing a representative subset from the PDF\n", + " gives faithful latent errors at a fraction of the every-sample cost.\n", + "\n", + "Here we mirror that draw-from-PDF behaviour explicitly with ``samples.samples_drawn_randomly_via_pdf_from``,\n", + "computing the latents over 20 PDF draws so this guide runs quickly while still producing a real, representative\n", + "latent posterior. For a publication-quality result, compute over all samples (or a larger number of draws)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "latent_draws = samples.samples_drawn_randomly_via_pdf_from(total_draws=20)\n", + "latent_samples = analysis.compute_latent_samples(latent_draws)\n", + "\n", + "median_instance = latent_samples.median_pdf()\n", + "print(f\"Median PDF magnification: {median_instance.magnification}\")\n", + "print(\n", + " f\"Median PDF effective_einstein_radius: {median_instance.effective_einstein_radius}\"\n", + ")\n", + "print(f\"Median PDF total_source_flux_mujy: {median_instance.total_source_flux_mujy}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The 1\u03c3 / 3\u03c3 intervals on these latents are *empirical quantiles of the induced posterior*. For magnification\n", + "specifically, this is the right thing \u2014 magnification is a strongly non-linear function of the mass parameters\n", + "and the source position, so the symmetric Gaussian propagation of parameter errors would underreport the tail\n", + "risk. The empirical posterior captures the full non-linearity faithfully. See the autofit foundational tutorial\n", + "for the full treatment.\n", + "\n", + "The ``effective_einstein_radius`` latent can produce noisy 1D posteriors when the zero-contour solver converges\n", + "to slightly different critical-curve estimates across samples. If you see step-like artefacts in a corner plot,\n", + "that's the discrete contour-finding cadence interacting with the continuous parameter posterior \u2014 the median\n", + "and 1\u03c3 intervals remain reliable, but use caution when reading the tails." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Extending with a Custom Latent__\n", + "\n", + "The library catalogue is intentionally narrow. If you want a different derived quantity \u2014 the lens mass's\n", + "axis ratio, the source-plane S\u00e9rsic effective radius, the time-delay between multiply-imaged source pixels \u2014\n", + "subclass ``al.LatentLens`` (override ``keys`` and ``variables``) and declare it on your analysis through the\n", + "``Latent`` class attribute. This mirrors exactly how you customise visualisation with a ``Visualizer`` subclass:\n", + "the latent catalogue is a first-class, swappable component of the analysis rather than a pair of methods you\n", + "monkey-patch.\n", + "\n", + "The example below adds ``mass_axis_ratio`` (the lens-mass axis ratio derived from the ``Isothermal`` mass's\n", + "``ell_comps``). The same pattern works for any function of the lens model instance. Calling\n", + "``al.LatentLens.keys(analysis)`` / ``al.LatentLens.variables(...)`` from your overrides keeps the library\n", + "latents alongside your custom ones \u2014 the Euclid pipeline\n", + "(``euclid_strong_lens_modeling_pipeline/util.py``, ``LatentEuclid``) uses this exact composition pattern in\n", + "production to append its FWHM aperture-flux latents.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "import numpy as np\n", + "\n", + "\n", + "class LatentMassAxisRatio(al.LatentLens):\n", + " \"\"\"\n", + " The library lensing latents plus a custom ``mass_axis_ratio`` \u2014 the axis\n", + " ratio of the lens galaxy's Isothermal mass profile, derived from its\n", + " ``ell_comps``. Demonstrates adding a user-defined latent without modifying\n", + " the library: subclass ``al.LatentLens`` and compose its ``keys`` /\n", + " ``variables`` static methods.\n", + " \"\"\"\n", + "\n", + " @staticmethod\n", + " def keys(analysis):\n", + " return list(al.LatentLens.keys(analysis)) + [\"mass_axis_ratio\"]\n", + "\n", + " @staticmethod\n", + " def variables(analysis, parameters, model):\n", + " library_values = al.LatentLens.variables(analysis, parameters, model)\n", + "\n", + " instance = model.instance_from_vector(vector=parameters)\n", + " xp = analysis._xp\n", + " try:\n", + " ell_y, ell_x = instance.galaxies.lens.mass.ell_comps\n", + " axis_ratio = (1.0 - np.sqrt(ell_y**2 + ell_x**2)) / (\n", + " 1.0 + np.sqrt(ell_y**2 + ell_x**2)\n", + " )\n", + " except AttributeError:\n", + " axis_ratio = xp.nan\n", + "\n", + " return library_values + (axis_ratio,)\n", + "\n", + "\n", + "class AnalysisImagingWithMassAxisRatio(al.AnalysisImaging):\n", + " \"\"\"\n", + " ``AnalysisImaging`` that swaps in the custom ``LatentMassAxisRatio`` catalogue\n", + " via the ``Latent`` class attribute \u2014 the same one-line mechanism used to\n", + " declare a custom ``Visualizer`` or ``Result``. No library code is modified.\n", + " \"\"\"\n", + "\n", + " Latent = LatentMassAxisRatio\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A fit that uses ``AnalysisImagingWithMassAxisRatio`` produces a ``latent.csv`` with one extra column\n", + "(``mass_axis_ratio``) on top of the eight library defaults (raw-flux + \u00b5Jy + dimensionless lensing). We don't\n", + "run a second fit here \u2014 the pattern above is the full recipe." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Contributing Upstream__\n", + "\n", + "If your custom latent is general enough that other PyAutoLens users would benefit (a SLACS-style external-\n", + "convergence proxy, a critical-curve perimeter, an enclosed-mass-at-fixed-radius), please consider promoting it\n", + "to the library:\n", + "\n", + " 1. Add the function to ``autolens/analysis/latent.py``, following the signature\n", + " ``(fit, magzero, xp=np) -> scalar``. Use NaN as the fallback when the function can't apply (no lens / no\n", + " source / singular mass model).\n", + " 2. Register it in the module-level ``LATENT_FUNCTIONS`` dict.\n", + " 3. Add an entry to ``autolens/config/latent.yaml`` defaulting it to ``false`` (the workspace yaml opts users in).\n", + " 4. Add a unit test under ``test_autolens/analysis/test_latent.py``.\n", + " 5. Open a PR.\n", + "\n", + "Pipeline-specific latents that need non-standard kwargs (PSF-relative aperture fluxes, dataset-specific\n", + "photometric quantities, multi-band colour terms) belong in your pipeline's local Analysis subclass. The Euclid\n", + "pipeline (``euclid_strong_lens_modeling_pipeline/util.py``) demonstrates this \u2014 its FWHM aperture-flux latents\n", + "stay pipeline-local and the rest inherit from the library.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/start_here.ipynb b/notebooks/guides/results/start_here.ipynb index df6bd6fa6..f91b6baef 100644 --- a/notebooks/guides/results/start_here.ipynb +++ b/notebooks/guides/results/start_here.ipynb @@ -1,864 +1,901 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Start Here\n", - "===================\n", - "\n", - "After a lens model-fit completes, nearly everything a user could need is written to the `output/` folder. Most of it\n", - "can be loaded back into full Python objects with a single line of code, via either `.json` files (for model objects\n", - "like the `Tracer`, `Model` and samples) or `.fits` files (for imaging products like the model image, residuals and\n", - "source-plane images).\n", - "\n", - "This guide shows two complementary ways to get those results back into Python:\n", - "\n", - " - **Simple loading** \u2014 point at a single fit's `output/...` folder and load `.json` / `.fits` files directly with\n", - " `from_json(...)` and `al.Imaging.from_fits(...)`. The objects that come back behave exactly like the in-memory\n", - " `Result` returned by `search.fit()`. This is the fastest way to inspect one fit, but everything you load sits\n", - " in memory.\n", - " - **Aggregator** \u2014 scrape a directory of completed fits and yield the same objects (`Tracer`, `Samples`, `Model`,\n", - " ...) via Python generators, so you can iterate over hundreds of fits without holding them all in memory at\n", - " once. This is the right tool when you want to analyse a sample of lenses together.\n", - "\n", - "Both sections appear below in that order. To keep them runnable from a fresh checkout, this script first\n", - "performs a quick model-fit so a results directory exists for both halves to read from. The aggregator\n", - "section then walks through the deeper API (samples, fits, tracer, units, pixelization). Where each\n", - "aggregator section reaches a result that the simple-loading API also exposes, this is noted \u2014 both routes\n", - "return the same PyAutoFit / PyAutoLens objects, just sourced from disk in different ways.\n", - "\n", - "__Output Folder Layout__\n", - "\n", - "Each completed fit lives at a path like::\n", - "\n", - " output/imaging//modeling//\n", - " files/ <- JSON + CSV: loadable Python objects\n", - " tracer.json <- max log likelihood Tracer\n", - " model.json <- fitted af.Collection model\n", - " samples.csv <- full Nautilus samples\n", - " samples_summary.json <- max log likelihood parameter values + errors\n", - " samples_info.json <- metadata about the samples\n", - " search.json <- non-linear search configuration\n", - " settings.json <- search settings\n", - " cosmology.json <- cosmology used for the fit\n", - " covariance.csv <- parameter covariance matrix\n", - " image/ <- FITS: imaging products\n", - " dataset.fits <- data, noise-map and PSF\n", - " fit.fits <- model image, residuals, chi-squared map\n", - " tracer.fits <- tracer image-plane images per galaxy\n", - " source_plane_images.fits <- source plane reconstructions\n", - " model_galaxy_images.fits <- per-galaxy model images\n", - " galaxy_images.fits <- per-galaxy images\n", - " dataset.png, fit.png, tracer.png <- visualisations\n", - " model.info <- human-readable model summary\n", - " model.results <- human-readable fit summary\n", - " search.summary <- search run summary\n", - " metadata <- run metadata\n", - "\n", - "__Contents__\n", - "\n", - "- **Model Fit:** Run a quick fit once so both halves of this guide have a real result on disk to read from.\n", - "- **Info:** Print the result in a readable format.\n", - "\n", - "**Simple Loading (one fit at a time):**\n", - "\n", - "**Tracer:** Load the maximum log likelihood `Tracer` from `tracer.json`.\n", - "**Model:** Load the fitted `af.Collection` model from `model.json`.\n", - "**Samples:** Load the non-linear search samples from `samples.csv` / `samples_summary.json`.\n", - "**FITS Files:** Load imaging products (data, fit, tracer images) from the `image/` folder.\n", - "\n", - "**Aggregator (many fits, generator-based):**\n", - "\n", - "**Loading From Hard-disk:** Use `Aggregator.from_directory(...)` to scrape `output/`.\n", - "**Generators:** How Python generators give the aggregator its memory efficiency.\n", - "**Database File:** Loading from a `.sqlite` database for very large samples.\n", - "**Workflow Examples:** Building scientific workflows (CSV / PNG / FITS makers).\n", - "**Result:** Working with the in-memory `Result` returned by `search.fit()`.\n", - "**Samples:** Median-PDF model and parameter errors from the `Samples` object.\n", - "**Linear Light Profiles:** Reading `intensity` values solved by linear algebra.\n", - "**Tracer:** Producing images and lensing quantities from the maximum likelihood `Tracer`.\n", - "**Fits:** Inspecting the `FitImaging` object (chi-squared, log likelihood).\n", - "**Galaxies:** Accessing individual lens / source galaxies inside the `Tracer`.\n", - "**Units and Cosmological Quantities:** Converting parameters to physical units.\n", - "**Linear Light Profiles / Basis Objects:** Specific functionality for linear light profiles and basis functions.\n", - "**Pixelization:** Pixelized source reconstructions on a Voronoi mesh." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "from autoconf.dictable import from_json\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import os\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "==============================================================================\n", - " MODEL FIT\n", - "==============================================================================\n", - "\n", - "Both halves of this guide \u2014 simple loading and the aggregator \u2014 need a real fit on disk to read from. We\n", - "perform that fit once here. The rest of the file then uses the resulting `search` object (via\n", - "`search.paths.output_path`) and the in-memory `result`.\n", - "\n", - "The model and search match `aggregator/samples.py` and `_quick_fit.py`, so re-running this tutorial resumes\n", - "from the cached fit instead of redoing the search.\n", - "\n", - "__Quick Fit Auto-Trigger__\n", - "\n", - "If a previous fit has not been run yet, the shared helper ``_quick_fit.py`` is invoked to produce one.\n", - "The helper writes a capped Nautilus fit to ``output/results_folder/`` so this tutorial has results to\n", - "work with. When that folder already exists the helper exits immediately, so re-running this tutorial is\n", - "cheap." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "results_path = Path(\"output\") / \"results_folder\"\n", - "if not results_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset and Model__\n", - "\n", - "Set up the same dataset and model as `_quick_fit.py`, then call `search.fit(...)`. Because the search has\n", - "already run, this resumes from the checkpoint and returns the in-memory `Result` object instantly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", - " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=None),\n", - " ),\n", - ")\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"results_folder\"),\n", - " name=\"results\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU batching and VRAM use explained in `modeling` examples.\n", - " n_like_max=300,\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Info__\n", - "\n", - "As seen throughout the workspace, the `info` attribute shows the result in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "==============================================================================\n", - " SIMPLE LOADING\n", - "==============================================================================\n", - "\n", - "The first half of this guide loads a single fit directly from `output/`. This is the fastest way to inspect\n", - "one fit \u2014 every file under `files/` and `image/` is a Python object away.\n", - "\n", - "__Result Path__\n", - "\n", - "Point at the fit's output folder. Because the fit ran above, ``search.paths.output_path`` already points at\n", - "the right location \u2014 there is no need to construct the path manually or know the unique hash." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result_path = search.paths.output_path\n", - "\n", - "files_path = result_path / \"files\"\n", - "image_path = result_path / \"image\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "The maximum log likelihood `Tracer` is saved to `files/tracer.json` and can be loaded in one line.\n", - "\n", - "The `Tracer` contains every `Galaxy`, light profile and mass profile at their max log likelihood values, so it can\n", - "be used directly to compute convergence maps, deflection angles, source-plane images and more \u2014 exactly as if it\n", - "had been returned by `search.fit()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if (files_path / \"tracer.json\").exists():\n", - " tracer = from_json(file_path=files_path / \"tracer.json\")\n", - "\n", - " print(tracer)\n", - " print(tracer.galaxies)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "The fitted `af.Collection` model is saved to `files/model.json`. This is the *prior* model (with free parameters),\n", - "not the max log likelihood instance \u2014 useful for inspecting the structure of what was fitted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if (files_path / \"model.json\").exists():\n", - " model = from_json(file_path=files_path / \"model.json\")\n", - " print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Samples__\n", - "\n", - "The full set of non-linear search samples is saved to `files/samples.csv` and its summary to\n", - "`files/samples_summary.json`. Both can be loaded without re-running the search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if (files_path / \"samples.csv\").exists() and (files_path / \"model.json\").exists():\n", - " model = from_json(file_path=files_path / \"model.json\")\n", - " samples = af.SamplesNest.from_table(\n", - " filename=files_path / \"samples.csv\", model=model\n", - " )\n", - " print(samples.max_log_likelihood())" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__FITS Files__\n", - "\n", - "The `image/` folder contains the imaging products of the fit as `.fits` files. These load with the standard\n", - "`al.Imaging` / `al.Array2D` APIs and can be plotted with `aplt`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if (image_path / \"dataset.fits\").exists():\n", - " dataset = al.Imaging.from_fits(\n", - " data_path=image_path / \"dataset.fits\",\n", - " noise_map_path=image_path / \"dataset.fits\",\n", - " psf_path=image_path / \"dataset.fits\",\n", - " data_hdu=1,\n", - " noise_map_hdu=2,\n", - " psf_hdu=3,\n", - " pixel_scales=0.1,\n", - " check_noise_map=False,\n", - " )\n", - "\n", - "if (image_path / \"tracer.fits\").exists():\n", - " tracer_images = al.Array2D.from_fits(\n", - " file_path=image_path / \"tracer.fits\", hdu=0, pixel_scales=0.1\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "==============================================================================\n", - " AGGREGATOR\n", - "==============================================================================\n", - "\n", - "Simple loading is the right tool when you have a single fit and you want to pull a specific object off\n", - "disk. The **aggregator** is a different tool \u2014 it scrapes a directory of fits and yields the same objects\n", - "(`Tracer`, `Samples`, `Model`, ...) via **Python generators**, so memory use stays bounded no matter how\n", - "many fits the directory contains. That generator-based design is its core feature, and it's what lets the\n", - "aggregator scale from one fit to hundreds.\n", - "\n", - "The two routes are complementary, not a hierarchy:\n", - "\n", - " - **Simple loading** is the most direct way to inspect one fit you already have a path to \u2014 one Python\n", - " object per call, all in memory.\n", - " - **Aggregator** is the right tool when you want to *iterate* over fits \u2014 even a single fit \u2014 and rely\n", - " on lazy evaluation, query filtering across many runs, or a `.sqlite` database back-end for very large\n", - " samples. It is also the API used by the workflow tools (`csv_maker.py`, `png_maker.py`, `fits_maker.py`)\n", - " that build scientific summaries of large fit samples.\n", - "\n", - "Anything reached via `from_json(...)` in the simple-loading section above can also be reached through the\n", - "aggregator below \u2014 both APIs return the same PyAutoFit / PyAutoLens objects. After reading this section,\n", - "the sibling files in `aggregator/` provide deeper examples for samples, fits, queries and database use.\n", - "\n", - "If you are not familiar with the lens modeling API, see `autolens_workspace/*/examples/modeling/` first.\n", - "\n", - "__Loading From Hard-disk__\n", - "\n", - "When performing fits which output results to hard-disk, a `files` folder is created containing .json / .csv files of\n", - "the model, samples, search, etc. You should check it out now for a completed fit on your hard-disk if you have\n", - "not already!\n", - "\n", - "These files can be loaded from hard-disk to Python variables via the aggregator, making them accessible in a\n", - "Python script or Jupyter notebook. They are loaded as the internal **PyAutoFit** objects we are familiar with,\n", - "for example the `model` is loaded as the same `Model` object the simple-loading section above reached via\n", - "`from_json(file_path=\"files/model.json\")`.\n", - "\n", - "Below, we will access these results using the aggregator's `values` method. A full list of what can be loaded is\n", - "as follows:\n", - "\n", - " - `model`: The `model` defined above and used in the model-fit (`model.json`).\n", - " - `search`: The non-linear search settings (`search.json`).\n", - " - `samples`: The non-linear search samples (`samples.csv`).\n", - " - `samples_info`: Additional information about the samples (`samples_info.json`).\n", - " - `samples_summary`: A summary of key results of the samples (`samples_summary.json`).\n", - " - `info`: The info dictionary passed to the search (`info.json`).\n", - " - `covariance`: The inferred covariance matrix (`covariance.csv`).\n", - " - `cosmology`: The cosmology used by the fit (`cosmology.json`).\n", - " - `settings`: The settings associated with a inversion if used (`settings.json`).\n", - " - `dataset/data`: The data that is fitted (`data.fits`).\n", - " - `dataset/noise_map`: The noise-map (`noise_map.fits`).\n", - " - `dataset/psf`: The Point Spread Function (`psf.fits`).\n", - " - `dataset/mask`: The mask applied to the data (`mask.fits`).\n", - " - `dataset/settings`: The settings associated with the dataset (`settings.json`).\n", - "\n", - "The `samples` and `samples_summary` results contain a lot of repeated information. The `samples` result contains\n", - "the full non-linear search samples, for example every parameter sample and its log likelihood. The `samples_summary`\n", - "contains a summary of the results, for example the maximum log likelihood model and error estimates on parameters\n", - "at 1 and 3 sigma confidence.\n", - "\n", - "Accessing results via the `samples_summary` is much faster, because as it does not reperform calculations using the full\n", - "list of samples. Therefore, if the result you want is accessible via the `samples_summary` you should use it\n", - "but if not you can revert to the `samples." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "agg = Aggregator.from_directory(\n", - " directory=Path(\"output\") / \"results_folder\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Generators__\n", - "\n", - "Before using the aggregator to inspect results, lets discuss Python generators.\n", - "\n", - "A generator is an object that iterates over a function when it is called. The aggregator creates all of the objects\n", - "that it loads from the database as generators (as opposed to a list, or dictionary, or another Python type).\n", - "\n", - "This is because generators are memory efficient, as they do not store the entries of the database in memory\n", - "simultaneously. This contrasts objects like lists and dictionaries, which store all entries in memory all at once.\n", - "If you fit a large number of datasets, lists and dictionaries will use a lot of memory and could crash your computer!\n", - "\n", - "Once we use a generator in the Python code, it cannot be used again. To perform the same task twice, the\n", - "generator must be remade it. This cookbook therefore rarely stores generators as variables and instead uses the\n", - "aggregator to create each generator at the point of use.\n", - "\n", - "To create a generator of a specific set of results, we use the `values` method. This takes the `name` of the\n", - "object we want to create a generator of, for example inputting `name=samples` will return the results `Samples`\n", - "object (which is illustrated in detail below)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for samples in agg.values(\"samples\"):\n", - " print(samples.parameter_lists[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Database File__\n", - "\n", - "The aggregator can also load results from a `.sqlite` database file.\n", - "\n", - "This is beneficial when loading results for large numbers of model-fits (e.g. more than hundreds)\n", - "because it is optimized for fast querying of results.\n", - "\n", - "It is recommended you use hard-disk loading to begin, as it is simpler and easier to use.\n", - "\n", - "See the package `results/database` for a full description of how to set up the database and the benefits it provides,\n", - "especially if loading results from hard-disk is slow.\n", - "\n", - "__Workflow Examples__\n", - "\n", - "The `results/workflow` folder contains examples describing how to build a scientific workflow using the results\n", - "of model-fits, in order to quickly and easily inspect and interpret results.\n", - "\n", - "These examples use functionality designed for modeling large dataset samples, with the following examples:\n", - "\n", - "- `csv_maker.py`: Make .csv files from the modeling results which summarize the results of a large sample of fits.\n", - "- `png_maker.py`: Make .png files of every fit, to quickly check the quality of the fit and interpret the results.\n", - "- `fits_maker.py`: Make .fits files of every fit, to quickly check the quality of the fit and interpret the results.\n", - "\n", - "The above examples work on the raw outputs of the model-fits that are stored in the `output` folder, for example\n", - "the visualization .png files, the .fits files containing results and parameter inferences which make the .csv files.\n", - "\n", - "They are therefore often quick to run and allow you to make a large number of checks on the results of your model-fits\n", - "in a short period of time.\n", - "\n", - "Below is a quick example, where we use code from the `csv_maker.py` scripts to create a .csv file from the fit above,\n", - "containing the inferred Einstein radius, in a folder you can inspect quickly.\n", - "\n", - "The `workflow_path` specifies where these files are output, in this case the .csv files which summarise the results,\n", - "and the code below can easily be adapted to output the .png and .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "workflow_path = Path(\"output\") / \"results_folder\" / \"workflow_make_example\"\n", - "\n", - "agg_csv = af.AggregateCSV(aggregator=agg)\n", - "agg_csv.add_variable(\n", - " argument=\"galaxies.lens.mass.einstein_radius\"\n", - ") # Example of adding a column\n", - "agg_csv.save(path=workflow_path / \"csv_very_simple.csv\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "From here on we will use attributes contained in the `result` passed from the `search.fit` method above, as opposed\n", - "to using the aggregator. This is because things will run faster, but all of the results we use can be loaded using\n", - "the aggregator as shown above (or via the simple-loading API in the first half of this file).\n", - "\n", - "__Samples__\n", - "\n", - "The result's `Samples` object contains the complete set of non-linear search Nautilus samples, where each sample\n", - "corresponds to a set of a model parameters that were evaluated and accepted. This is the same `Samples` object that\n", - "`af.SamplesNest.from_csv(file_path=\"files/samples.csv\", model=model)` returned in the simple-loading section.\n", - "\n", - "The examples script `autolens_workspace/*/guides/results/aggregator/samples.py` provides a detailed description of\n", - "this object, including:\n", - "\n", - " - Extracting the maximum likelihood lens model.\n", - " - Using marginalized PDFs to estimate errors on the lens model parameters.\n", - " - Deriving errors on derived quantities, such as the Einstein radius.\n", - "\n", - "Below, is an example of how to use the `Samples` object to estimate the lens mass model parameters which are\n", - "the median of the probability distribution function and its errors at 3 sigma confidence intervals." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "\n", - "median_pdf_instance = samples.median_pdf()\n", - "\n", - "print(\"Median PDF Model Instances: \\n\")\n", - "print(median_pdf_instance.galaxies.lens.mass)\n", - "print()\n", - "\n", - "ue3_instance = samples.values_at_upper_sigma(sigma=3.0)\n", - "le3_instance = samples.values_at_lower_sigma(sigma=3.0)\n", - "\n", - "print(\"Errors Instances: \\n\")\n", - "print(ue3_instance.galaxies.lens.mass, \"\\n\")\n", - "print(le3_instance.galaxies.lens.mass, \"\\n\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Light Profiles__\n", - "\n", - "In the model fit, linear light profiles are used, solving for the `intensity` of each profile through linear algebra.\n", - "\n", - "The `intensity` value is not a free parameter of the linear light profiles in the model, meaning that in the `Samples`\n", - "object the `intensity` are always defaulted to values of 1.0 in the `Samples` object.\n", - "\n", - "You can observe this by comparing the `intensity` values in the `Samples` object to those in\n", - "the `result.max_log_likelihood_galaxies` instance and `result.max_log_likelihood_fit` instance." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples = result.samples\n", - "ml_instance = samples.max_log_likelihood()\n", - "\n", - "print(\n", - " \"Intensity of source galaxy's bulge in the Samples object (before solving linear algebra):\"\n", - ")\n", - "print(ml_instance.galaxies.source.bulge.intensity)\n", - "\n", - "print(\n", - " \"Intensity of source galaxy's bulge in the max log likelihood galaxy (after solving linear algebra):\"\n", - ")\n", - "print(result.max_log_likelihood_tracer.planes[1][0].bulge.intensity)\n", - "print(\n", - " result.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles.planes[\n", - " 1\n", - " ][0].bulge.intensity\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To interpret results associated with the linear light profiles, you must input the `Samples` object into a `FitImaging`,\n", - "which converts the linear light profiles to standard light profiles with `intensity` values solved for using the linear\n", - "algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "ml_instance = samples.max_log_likelihood()\n", - "\n", - "tracer = al.Tracer(galaxies=ml_instance.galaxies)\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", - "\n", - "print(\"Intensity of source galaxy's bulge after conversion using FitImaging:\")\n", - "print(tracer.planes[1][0].bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Whenever possible, the result already containing the solved `intensity` values is used, for example\n", - "the `Result` object returned by a search.\n", - "\n", - "However, when manually loading results from the `Samples` object, you must use the `FitImaging` object to convert\n", - "the linear light profiles to their correct `intensity` values.\n", - "\n", - "__Tracer__\n", - "\n", - "The result's maximum likelihood `Tracer` object contains everything necessary to perform ray-tracing and other\n", - "calculations with the lens model. It is the same `Tracer` that the simple-loading section above reached via\n", - "`from_json(file_path=\"files/tracer.json\")`.\n", - "\n", - "The guide `autolens_workspace/*/guides/tracer.py` provides a detailed description of this object, including:\n", - "\n", - " - Producing individual images of the strong lens from a tracer.\n", - " - Inspecting mass model components like the convergence, potential and deflection angles.\n", - " - Other lensing quantities like the critical curve and caustics.\n", - "\n", - "The examples script `autolens_workspace/*/guides/results/aggregator/galaxies_fit.py` show how to use\n", - "model-fitting results specific functionality of galaxies, including:\n", - "\n", - " - Drawing tracers from the samples and plotting their images.\n", - " - Producing 1D plots of the galaxy's light and mass profiles with error bars.\n", - "\n", - "Below, is an example of how to use the `Tracer` object to calculate the image of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "image = tracer.image_2d_from(grid=dataset.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fits__\n", - "\n", - "The result's maximum likelihood `FitImaging` object contains everything necessary to inspect the lens model fit to the\n", - "data.\n", - "\n", - "The guide `autolens_workspace/*/guides/fits.py` provides a detailed description of this object, including:\n", - "\n", - " - Performing a fit to data with galaxies.\n", - " - Inspecting the model data, residual-map, chi-squared, noise-map of the fit.\n", - " - Other properties of the fit that inspect how good it is.\n", - "\n", - "The examples script `autolens_workspace/*/guides/results/aggregator/galaxies_fits.py` provides a detailed description of this\n", - "object, including:\n", - "\n", - " - Repeating fits using the results contained in the samples.\n", - "\n", - "Below, is an example of how to use the `FitImaging` object to print the maximum likelihood chi-squared and\n", - "log likelihood values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = result.max_log_likelihood_fit\n", - "\n", - "print(fit.chi_squared)\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "The result's maximum likelihood `Galaxy` objects contained within the `Tracer` contain everything necessary to\n", - "inspect the individual properties of the lens and source galaxies.\n", - "\n", - "The guide `autolens_workspace/*/guides/fits.py` provides a detailed description of this, including:\n", - "\n", - " - Extracting the lens and source galaixes from a tracer.\n", - " - Extracting the individual light and mass profiles of the galaxies.\n", - "\n", - "The examples script `autolens_workspace/*/guides/results/aggregator/galaxies_fits.py` shows how to use\n", - "model-fitting results specific functionality of galaxies, including:\n", - "\n", - " - Repeating fits using the results contained in the samples.\n", - "\n", - "Below, is an example of how to use the `Galaxy` objects to plot the source galaxy's source-plane image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "source = tracer.planes[1][0]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Units and Cosmological Quantities__\n", - "\n", - "The maximum likelihood model includes cosmological quantities, which can be computed via the result.\n", - "\n", - "The examples script `autolens_workspace/*/guides/units_and_cosmology.py` provides a detailed\n", - "description of this object, including:\n", - "\n", - " - Calculating the Einstein radius of the lens galaxy.\n", - " - Converting quantities like the Einstein radius or effective radius from arcseconds to kiloparsecs.\n", - " - Computing the Einstein mass of the lens galaxy in solar masses.\n", - "\n", - "This guide is not in the `results` package but the `guides` package, as it is a general guide to the\n", - "**PyAutoLens** API. However, it may be useful when inspecting results.\n", - "\n", - "Below, is an example of how to convert the y centre of the source galaxy from arcseconds to kiloparsecs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "cosmology = al.cosmo.Planck15()\n", - "\n", - "source = tracer.planes[1][0]\n", - "source_kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=source.redshift)\n", - "source_centre_0_kpc = source.bulge.centre[0] * source_kpc_per_arcsec" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Light Profiles / Basis Objects__\n", - "\n", - "A lens model can be fitted using a linear light profile, which is a light profile whose `intensity` parameter is\n", - "sovled for via linear algebra.\n", - "\n", - "This includes Basis objects such as a Multi-Gaussian expansion of Shapelets.\n", - "\n", - "These objects mostly behave identically to ordinary light profiles, but due to the linear algebra have their own\n", - "specific functionality.\n", - "\n", - "The example script `autolens_workspace/*/features/linear_light_profiles.py` provides a detailed description of\n", - "using linear light profile results including:\n", - "\n", - " - Extracting individual quantities from the linear light profile, such as the coefficients of the basis functions.\n", - " - Extracting the intensity of the linear light profiles after they have been computed via linear algebra.\n", - " - Plotting the linear light profiles.\n", - "\n", - "Therefore if your results contain a linear light profile, checkout the example script above for a detailed description\n", - "of how to use their results.\n", - "\n", - "__Pixelization__\n", - "\n", - "The lens model can reconstruct the source galaxy using a pixelization, for example on a Voronoi mesh.\n", - "\n", - "The example script `autolens_workspace/*/features/pixelization.py` describes using pixelization results including:\n", - "\n", - " - Producing source reconstructions using the Voronoi mesh, RectangularAdaptDensity triangulation or whichever mesh is used.\n", - " - Inspecting the evidence terms of the fit, which quantify how well the pixelization reconstructs fits the data whilst\n", - " accounting for the complexity of the pixelization.\n", - " - Estimating the magnification of the source galaxy's image using the pixelization.\n", - "\n", - "Therefore if your results contain a pixelization, checkout the example script above for a detailed description\n", - "of how to use their results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Start Here\n", + "===================\n", + "\n", + "After a lens model-fit completes, nearly everything a user could need is written to the `output/` folder. Most of it\n", + "can be loaded back into full Python objects with a single line of code, via either `.json` files (for model objects\n", + "like the `Tracer`, `Model` and samples) or `.fits` files (for imaging products like the model image, residuals and\n", + "source-plane images).\n", + "\n", + "This guide shows two complementary ways to get those results back into Python:\n", + "\n", + " - **Simple loading** \u2014 point at a single fit's `output/...` folder and load `.json` / `.fits` files directly with\n", + " `from_json(...)` and `al.Imaging.from_fits(...)`. The objects that come back behave exactly like the in-memory\n", + " `Result` returned by `search.fit()`. This is the fastest way to inspect one fit, but everything you load sits\n", + " in memory.\n", + " - **Aggregator** \u2014 scrape a directory of completed fits and yield the same objects (`Tracer`, `Samples`, `Model`,\n", + " ...) via Python generators, so you can iterate over hundreds of fits without holding them all in memory at\n", + " once. This is the right tool when you want to analyse a sample of lenses together.\n", + "\n", + "Both sections appear below in that order. To keep them runnable from a fresh checkout, this script first\n", + "performs a quick model-fit so a results directory exists for both halves to read from. The aggregator\n", + "section then walks through the deeper API (samples, fits, tracer, units, pixelization). Where each\n", + "aggregator section reaches a result that the simple-loading API also exposes, this is noted \u2014 both routes\n", + "return the same PyAutoFit / PyAutoLens objects, just sourced from disk in different ways.\n", + "\n", + "__Output Folder Layout__\n", + "\n", + "Each completed fit lives at a path like::\n", + "\n", + " output/imaging//modeling//\n", + " files/ <- JSON + CSV: loadable Python objects\n", + " tracer.json <- max log likelihood Tracer\n", + " model.json <- fitted af.Collection model\n", + " samples.csv <- full Nautilus samples\n", + " samples_summary.json <- max log likelihood parameter values + errors\n", + " samples_info.json <- metadata about the samples\n", + " search.json <- non-linear search configuration\n", + " settings.json <- search settings\n", + " cosmology.json <- cosmology used for the fit\n", + " covariance.csv <- parameter covariance matrix\n", + " image/ <- FITS: imaging products\n", + " dataset.fits <- data, noise-map and PSF\n", + " fit.fits <- model image, residuals, chi-squared map\n", + " tracer.fits <- tracer image-plane images per galaxy\n", + " source_plane_images.fits <- source plane reconstructions\n", + " model_galaxy_images.fits <- per-galaxy model images\n", + " galaxy_images.fits <- per-galaxy images\n", + " dataset.png, fit.png, tracer.png <- visualisations\n", + " model.info <- human-readable model summary\n", + " model.results <- human-readable fit summary\n", + " search.summary <- search run summary\n", + " metadata <- run metadata\n", + "\n", + "__Contents__\n", + "\n", + "- **Model Fit:** Run a quick fit once so both halves of this guide have a real result on disk to read from.\n", + "- **Info:** Print the result in a readable format.\n", + "\n", + "**Simple Loading (one fit at a time):**\n", + "\n", + "**Tracer:** Load the maximum log likelihood `Tracer` from `tracer.json`.\n", + "**Model:** Load the fitted `af.Collection` model from `model.json`.\n", + "**Samples:** Load the non-linear search samples from `samples.csv` / `samples_summary.json`.\n", + "**FITS Files:** Load imaging products (data, fit, tracer images) from the `image/` folder.\n", + "\n", + "**Aggregator (many fits, generator-based):**\n", + "\n", + "**Loading From Hard-disk:** Use `Aggregator.from_directory(...)` to scrape `output/`.\n", + "**Generators:** How Python generators give the aggregator its memory efficiency.\n", + "**Database File:** Loading from a `.sqlite` database for very large samples.\n", + "**Workflow Examples:** Building scientific workflows (CSV / PNG / FITS makers).\n", + "**Result:** Working with the in-memory `Result` returned by `search.fit()`.\n", + "**Samples:** Median-PDF model and parameter errors from the `Samples` object.\n", + "**Linear Light Profiles:** Reading `intensity` values solved by linear algebra.\n", + "**Tracer:** Producing images and lensing quantities from the maximum likelihood `Tracer`.\n", + "**Fits:** Inspecting the `FitImaging` object (chi-squared, log likelihood).\n", + "**Galaxies:** Accessing individual lens / source galaxies inside the `Tracer`.\n", + "**Units and Cosmological Quantities:** Converting parameters to physical units.\n", + "**Linear Light Profiles / Basis Objects:** Specific functionality for linear light profiles and basis functions.\n", + "**Pixelization:** Pixelized source reconstructions on a Voronoi mesh." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "from autoconf.dictable import from_json\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import os\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "==============================================================================\n", + " MODEL FIT\n", + "==============================================================================\n", + "\n", + "Both halves of this guide \u2014 simple loading and the aggregator \u2014 need a real fit on disk to read from. We\n", + "perform that fit once here. The rest of the file then uses the resulting `search` object (via\n", + "`search.paths.output_path`) and the in-memory `result`.\n", + "\n", + "The model and search match `aggregator/samples.py` and `_quick_fit.py`, so re-running this tutorial resumes\n", + "from the cached fit instead of redoing the search.\n", + "\n", + "__Quick Fit Auto-Trigger__\n", + "\n", + "If a previous fit has not been run yet, the shared helper ``_quick_fit.py`` is invoked to produce one.\n", + "The helper writes a capped Nautilus fit to ``output/results_folder/`` so this tutorial has results to\n", + "work with. When that folder already exists the helper exits immediately, so re-running this tutorial is\n", + "cheap." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "results_path = Path(\"output\") / \"results_folder\"\n", + "if not results_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset and Model__\n", + "\n", + "Set up the same dataset and model as `_quick_fit.py`, then call `search.fit(...)`. Because the search has\n", + "already run, this resumes from the checkpoint and returns the in-memory `Result` object instantly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", + " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge, disk=None),\n", + " ),\n", + ")\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"results_folder\"),\n", + " name=\"results\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU batching and VRAM use explained in `modeling` examples.\n", + " n_like_max=300,\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Info__\n", + "\n", + "As seen throughout the workspace, the `info` attribute shows the result in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "==============================================================================\n", + " SIMPLE LOADING\n", + "==============================================================================\n", + "\n", + "The first half of this guide loads a single fit directly from `output/`. This is the fastest way to inspect\n", + "one fit \u2014 every file under `files/` and `image/` is a Python object away.\n", + "\n", + "__Result Path__\n", + "\n", + "Point at the fit's output folder. Because the fit ran above, ``search.paths.output_path`` already points at\n", + "the right location \u2014 there is no need to construct the path manually or know the unique hash." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result_path = search.paths.output_path\n", + "\n", + "files_path = result_path / \"files\"\n", + "image_path = result_path / \"image\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "The maximum log likelihood `Tracer` is saved to `files/tracer.json` and can be loaded in one line.\n", + "\n", + "The `Tracer` contains every `Galaxy`, light profile and mass profile at their max log likelihood values, so it can\n", + "be used directly to compute convergence maps, deflection angles, source-plane images and more \u2014 exactly as if it\n", + "had been returned by `search.fit()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if (files_path / \"tracer.json\").exists():\n", + " tracer = from_json(file_path=files_path / \"tracer.json\")\n", + "\n", + " print(tracer)\n", + " print(tracer.galaxies)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "The fitted `af.Collection` model is saved to `files/model.json`. This is the *prior* model (with free parameters),\n", + "not the max log likelihood instance \u2014 useful for inspecting the structure of what was fitted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if (files_path / \"model.json\").exists():\n", + " model = from_json(file_path=files_path / \"model.json\")\n", + " print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Samples__\n", + "\n", + "The full set of non-linear search samples is saved to `files/samples.csv` and its summary to\n", + "`files/samples_summary.json`. Both can be loaded without re-running the search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if (files_path / \"samples.csv\").exists() and (files_path / \"model.json\").exists():\n", + " model = from_json(file_path=files_path / \"model.json\")\n", + " samples = af.SamplesNest.from_table(\n", + " filename=files_path / \"samples.csv\", model=model\n", + " )\n", + " print(samples.max_log_likelihood())" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__FITS Files__\n", + "\n", + "The `image/` folder contains the imaging products of the fit as `.fits` files. These load with the standard\n", + "`al.Imaging` / `al.Array2D` APIs and can be plotted with `aplt`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if (image_path / \"dataset.fits\").exists():\n", + " dataset = al.Imaging.from_fits(\n", + " data_path=image_path / \"dataset.fits\",\n", + " noise_map_path=image_path / \"dataset.fits\",\n", + " psf_path=image_path / \"dataset.fits\",\n", + " data_hdu=1,\n", + " noise_map_hdu=2,\n", + " psf_hdu=3,\n", + " pixel_scales=0.1,\n", + " check_noise_map=False,\n", + " )\n", + "\n", + "if (image_path / \"tracer.fits\").exists():\n", + " tracer_images = al.Array2D.from_fits(\n", + " file_path=image_path / \"tracer.fits\", hdu=0, pixel_scales=0.1\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "==============================================================================\n", + " AGGREGATOR\n", + "==============================================================================\n", + "\n", + "Simple loading is the right tool when you have a single fit and you want to pull a specific object off\n", + "disk. The **aggregator** is a different tool \u2014 it scrapes a directory of fits and yields the same objects\n", + "(`Tracer`, `Samples`, `Model`, ...) via **Python generators**, so memory use stays bounded no matter how\n", + "many fits the directory contains. That generator-based design is its core feature, and it's what lets the\n", + "aggregator scale from one fit to hundreds.\n", + "\n", + "The two routes are complementary, not a hierarchy:\n", + "\n", + " - **Simple loading** is the most direct way to inspect one fit you already have a path to \u2014 one Python\n", + " object per call, all in memory.\n", + " - **Aggregator** is the right tool when you want to *iterate* over fits \u2014 even a single fit \u2014 and rely\n", + " on lazy evaluation, query filtering across many runs, or a `.sqlite` database back-end for very large\n", + " samples. It is also the API used by the workflow tools (`csv_maker.py`, `png_maker.py`, `fits_maker.py`)\n", + " that build scientific summaries of large fit samples.\n", + "\n", + "Anything reached via `from_json(...)` in the simple-loading section above can also be reached through the\n", + "aggregator below \u2014 both APIs return the same PyAutoFit / PyAutoLens objects. After reading this section,\n", + "the sibling files in `aggregator/` provide deeper examples for samples, fits, queries and database use.\n", + "\n", + "If you are not familiar with the lens modeling API, see `autolens_workspace/*/examples/modeling/` first.\n", + "\n", + "__Loading From Hard-disk__\n", + "\n", + "When performing fits which output results to hard-disk, a `files` folder is created containing .json / .csv files of\n", + "the model, samples, search, etc. You should check it out now for a completed fit on your hard-disk if you have\n", + "not already!\n", + "\n", + "These files can be loaded from hard-disk to Python variables via the aggregator, making them accessible in a\n", + "Python script or Jupyter notebook. They are loaded as the internal **PyAutoFit** objects we are familiar with,\n", + "for example the `model` is loaded as the same `Model` object the simple-loading section above reached via\n", + "`from_json(file_path=\"files/model.json\")`.\n", + "\n", + "Below, we will access these results using the aggregator's `values` method. A full list of what can be loaded is\n", + "as follows:\n", + "\n", + " - `model`: The `model` defined above and used in the model-fit (`model.json`).\n", + " - `search`: The non-linear search settings (`search.json`).\n", + " - `samples`: The non-linear search samples (`samples.csv`).\n", + " - `samples_info`: Additional information about the samples (`samples_info.json`).\n", + " - `samples_summary`: A summary of key results of the samples (`samples_summary.json`).\n", + " - `info`: The info dictionary passed to the search (`info.json`).\n", + " - `covariance`: The inferred covariance matrix (`covariance.csv`).\n", + " - `cosmology`: The cosmology used by the fit (`cosmology.json`).\n", + " - `settings`: The settings associated with a inversion if used (`settings.json`).\n", + " - `dataset/data`: The data that is fitted (`data.fits`).\n", + " - `dataset/noise_map`: The noise-map (`noise_map.fits`).\n", + " - `dataset/psf`: The Point Spread Function (`psf.fits`).\n", + " - `dataset/mask`: The mask applied to the data (`mask.fits`).\n", + " - `dataset/settings`: The settings associated with the dataset (`settings.json`).\n", + "\n", + "The `samples` and `samples_summary` results contain a lot of repeated information. The `samples` result contains\n", + "the full non-linear search samples, for example every parameter sample and its log likelihood. The `samples_summary`\n", + "contains a summary of the results, for example the maximum log likelihood model and error estimates on parameters\n", + "at 1 and 3 sigma confidence.\n", + "\n", + "Accessing results via the `samples_summary` is much faster, because as it does not reperform calculations using the full\n", + "list of samples. Therefore, if the result you want is accessible via the `samples_summary` you should use it\n", + "but if not you can revert to the `samples." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "agg = Aggregator.from_directory(\n", + " directory=Path(\"output\") / \"results_folder\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Generators__\n", + "\n", + "Before using the aggregator to inspect results, lets discuss Python generators.\n", + "\n", + "A generator is an object that iterates over a function when it is called. The aggregator creates all of the objects\n", + "that it loads from the database as generators (as opposed to a list, or dictionary, or another Python type).\n", + "\n", + "This is because generators are memory efficient, as they do not store the entries of the database in memory\n", + "simultaneously. This contrasts objects like lists and dictionaries, which store all entries in memory all at once.\n", + "If you fit a large number of datasets, lists and dictionaries will use a lot of memory and could crash your computer!\n", + "\n", + "Once we use a generator in the Python code, it cannot be used again. To perform the same task twice, the\n", + "generator must be remade it. This cookbook therefore rarely stores generators as variables and instead uses the\n", + "aggregator to create each generator at the point of use.\n", + "\n", + "To create a generator of a specific set of results, we use the `values` method. This takes the `name` of the\n", + "object we want to create a generator of, for example inputting `name=samples` will return the results `Samples`\n", + "object (which is illustrated in detail below)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for samples in agg.values(\"samples\"):\n", + " print(samples.parameter_lists[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Database File__\n", + "\n", + "The aggregator can also load results from a `.sqlite` database file.\n", + "\n", + "This is beneficial when loading results for large numbers of model-fits (e.g. more than hundreds)\n", + "because it is optimized for fast querying of results.\n", + "\n", + "It is recommended you use hard-disk loading to begin, as it is simpler and easier to use.\n", + "\n", + "See the package `results/database` for a full description of how to set up the database and the benefits it provides,\n", + "especially if loading results from hard-disk is slow.\n", + "\n", + "__Workflow Examples__\n", + "\n", + "The `results/workflow` folder contains examples describing how to build a scientific workflow using the results\n", + "of model-fits, in order to quickly and easily inspect and interpret results.\n", + "\n", + "These examples use functionality designed for modeling large dataset samples, with the following examples:\n", + "\n", + "- `csv_maker.py`: Make .csv files from the modeling results which summarize the results of a large sample of fits.\n", + "- `png_maker.py`: Make .png files of every fit, to quickly check the quality of the fit and interpret the results.\n", + "- `fits_maker.py`: Make .fits files of every fit, to quickly check the quality of the fit and interpret the results.\n", + "\n", + "The above examples work on the raw outputs of the model-fits that are stored in the `output` folder, for example\n", + "the visualization .png files, the .fits files containing results and parameter inferences which make the .csv files.\n", + "\n", + "They are therefore often quick to run and allow you to make a large number of checks on the results of your model-fits\n", + "in a short period of time.\n", + "\n", + "Below is a quick example, where we use code from the `csv_maker.py` scripts to create a .csv file from the fit above,\n", + "containing the inferred Einstein radius, in a folder you can inspect quickly.\n", + "\n", + "The `workflow_path` specifies where these files are output, in this case the .csv files which summarise the results,\n", + "and the code below can easily be adapted to output the .png and .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "workflow_path = Path(\"output\") / \"results_folder\" / \"workflow_make_example\"\n", + "\n", + "agg_csv = af.AggregateCSV(aggregator=agg)\n", + "agg_csv.add_variable(\n", + " argument=\"galaxies.lens.mass.einstein_radius\"\n", + ") # Example of adding a column\n", + "agg_csv.save(path=workflow_path / \"csv_very_simple.csv\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "From here on we will use attributes contained in the `result` passed from the `search.fit` method above, as opposed\n", + "to using the aggregator. This is because things will run faster, but all of the results we use can be loaded using\n", + "the aggregator as shown above (or via the simple-loading API in the first half of this file).\n", + "\n", + "__Samples__\n", + "\n", + "The result's `Samples` object contains the complete set of non-linear search Nautilus samples, where each sample\n", + "corresponds to a set of a model parameters that were evaluated and accepted. This is the same `Samples` object that\n", + "`af.SamplesNest.from_csv(file_path=\"files/samples.csv\", model=model)` returned in the simple-loading section.\n", + "\n", + "The examples script `autolens_workspace/*/guides/results/aggregator/samples.py` provides a detailed description of\n", + "this object, including:\n", + "\n", + " - Extracting the maximum likelihood lens model.\n", + " - Using marginalized PDFs to estimate errors on the lens model parameters.\n", + " - Deriving errors on derived quantities, such as the Einstein radius.\n", + "\n", + "Below, is an example of how to use the `Samples` object to estimate the lens mass model parameters which are\n", + "the median of the probability distribution function and its errors at 3 sigma confidence intervals." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "\n", + "median_pdf_instance = samples.median_pdf()\n", + "\n", + "print(\"Median PDF Model Instances: \\n\")\n", + "print(median_pdf_instance.galaxies.lens.mass)\n", + "print()\n", + "\n", + "ue3_instance = samples.values_at_upper_sigma(sigma=3.0)\n", + "le3_instance = samples.values_at_lower_sigma(sigma=3.0)\n", + "\n", + "print(\"Errors Instances: \\n\")\n", + "print(ue3_instance.galaxies.lens.mass, \"\\n\")\n", + "print(le3_instance.galaxies.lens.mass, \"\\n\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Light Profiles__\n", + "\n", + "In the model fit, linear light profiles are used, solving for the `intensity` of each profile through linear algebra.\n", + "\n", + "The `intensity` value is not a free parameter of the linear light profiles in the model, meaning that in the `Samples`\n", + "object the `intensity` are always defaulted to values of 1.0 in the `Samples` object.\n", + "\n", + "You can observe this by comparing the `intensity` values in the `Samples` object to those in\n", + "the `result.max_log_likelihood_galaxies` instance and `result.max_log_likelihood_fit` instance." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples = result.samples\n", + "ml_instance = samples.max_log_likelihood()\n", + "\n", + "print(\n", + " \"Intensity of source galaxy's bulge in the Samples object (before solving linear algebra):\"\n", + ")\n", + "print(ml_instance.galaxies.source.bulge.intensity)\n", + "\n", + "print(\n", + " \"Intensity of source galaxy's bulge in the max log likelihood galaxy (after solving linear algebra):\"\n", + ")\n", + "print(result.max_log_likelihood_tracer.planes[1][0].bulge.intensity)\n", + "print(\n", + " result.max_log_likelihood_fit.tracer_linear_light_profiles_to_light_profiles.planes[\n", + " 1\n", + " ][0].bulge.intensity\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To interpret results associated with the linear light profiles, you must input the `Samples` object into a `FitImaging`,\n", + "which converts the linear light profiles to standard light profiles with `intensity` values solved for using the linear\n", + "algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "ml_instance = samples.max_log_likelihood()\n", + "\n", + "tracer = al.Tracer(galaxies=ml_instance.galaxies)\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "tracer = fit.tracer_linear_light_profiles_to_light_profiles\n", + "\n", + "print(\"Intensity of source galaxy's bulge after conversion using FitImaging:\")\n", + "print(tracer.planes[1][0].bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Whenever possible, the result already containing the solved `intensity` values is used, for example\n", + "the `Result` object returned by a search.\n", + "\n", + "However, when manually loading results from the `Samples` object, you must use the `FitImaging` object to convert\n", + "the linear light profiles to their correct `intensity` values.\n", + "\n", + "__Tracer__\n", + "\n", + "The result's maximum likelihood `Tracer` object contains everything necessary to perform ray-tracing and other\n", + "calculations with the lens model. It is the same `Tracer` that the simple-loading section above reached via\n", + "`from_json(file_path=\"files/tracer.json\")`.\n", + "\n", + "The guide `autolens_workspace/*/guides/tracer.py` provides a detailed description of this object, including:\n", + "\n", + " - Producing individual images of the strong lens from a tracer.\n", + " - Inspecting mass model components like the convergence, potential and deflection angles.\n", + " - Other lensing quantities like the critical curve and caustics.\n", + "\n", + "The examples script `autolens_workspace/*/guides/results/aggregator/galaxies_fit.py` show how to use\n", + "model-fitting results specific functionality of galaxies, including:\n", + "\n", + " - Drawing tracers from the samples and plotting their images.\n", + " - Producing 1D plots of the galaxy's light and mass profiles with error bars.\n", + "\n", + "Below, is an example of how to use the `Tracer` object to calculate the image of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "image = tracer.image_2d_from(grid=dataset.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fits__\n", + "\n", + "The result's maximum likelihood `FitImaging` object contains everything necessary to inspect the lens model fit to the\n", + "data.\n", + "\n", + "The guide `autolens_workspace/*/guides/fits.py` provides a detailed description of this object, including:\n", + "\n", + " - Performing a fit to data with galaxies.\n", + " - Inspecting the model data, residual-map, chi-squared, noise-map of the fit.\n", + " - Other properties of the fit that inspect how good it is.\n", + "\n", + "The examples script `autolens_workspace/*/guides/results/aggregator/galaxies_fits.py` provides a detailed description of this\n", + "object, including:\n", + "\n", + " - Repeating fits using the results contained in the samples.\n", + "\n", + "Below, is an example of how to use the `FitImaging` object to print the maximum likelihood chi-squared and\n", + "log likelihood values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = result.max_log_likelihood_fit\n", + "\n", + "print(fit.chi_squared)\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "The result's maximum likelihood `Galaxy` objects contained within the `Tracer` contain everything necessary to\n", + "inspect the individual properties of the lens and source galaxies.\n", + "\n", + "The guide `autolens_workspace/*/guides/fits.py` provides a detailed description of this, including:\n", + "\n", + " - Extracting the lens and source galaixes from a tracer.\n", + " - Extracting the individual light and mass profiles of the galaxies.\n", + "\n", + "The examples script `autolens_workspace/*/guides/results/aggregator/galaxies_fits.py` shows how to use\n", + "model-fitting results specific functionality of galaxies, including:\n", + "\n", + " - Repeating fits using the results contained in the samples.\n", + "\n", + "Below, is an example of how to use the `Galaxy` objects to plot the source galaxy's source-plane image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "source = tracer.planes[1][0]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Units and Cosmological Quantities__\n", + "\n", + "The maximum likelihood model includes cosmological quantities, which can be computed via the result.\n", + "\n", + "The examples script `autolens_workspace/*/guides/units_and_cosmology.py` provides a detailed\n", + "description of this object, including:\n", + "\n", + " - Calculating the Einstein radius of the lens galaxy.\n", + " - Converting quantities like the Einstein radius or effective radius from arcseconds to kiloparsecs.\n", + " - Computing the Einstein mass of the lens galaxy in solar masses.\n", + "\n", + "This guide is not in the `results` package but the `guides` package, as it is a general guide to the\n", + "**PyAutoLens** API. However, it may be useful when inspecting results.\n", + "\n", + "Below, is an example of how to convert the y centre of the source galaxy from arcseconds to kiloparsecs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "cosmology = al.cosmo.Planck15()\n", + "\n", + "source = tracer.planes[1][0]\n", + "source_kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=source.redshift)\n", + "source_centre_0_kpc = source.bulge.centre[0] * source_kpc_per_arcsec" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Light Profiles / Basis Objects__\n", + "\n", + "A lens model can be fitted using a linear light profile, which is a light profile whose `intensity` parameter is\n", + "sovled for via linear algebra.\n", + "\n", + "This includes Basis objects such as a Multi-Gaussian expansion of Shapelets.\n", + "\n", + "These objects mostly behave identically to ordinary light profiles, but due to the linear algebra have their own\n", + "specific functionality.\n", + "\n", + "The example script `autolens_workspace/*/features/linear_light_profiles.py` provides a detailed description of\n", + "using linear light profile results including:\n", + "\n", + " - Extracting individual quantities from the linear light profile, such as the coefficients of the basis functions.\n", + " - Extracting the intensity of the linear light profiles after they have been computed via linear algebra.\n", + " - Plotting the linear light profiles.\n", + "\n", + "Therefore if your results contain a linear light profile, checkout the example script above for a detailed description\n", + "of how to use their results.\n", + "\n", + "__Pixelization__\n", + "\n", + "The lens model can reconstruct the source galaxy using a pixelization, for example on a Voronoi mesh.\n", + "\n", + "The example script `autolens_workspace/*/features/pixelization.py` describes using pixelization results including:\n", + "\n", + " - Producing source reconstructions using the Voronoi mesh, RectangularAdaptDensity triangulation or whichever mesh is used.\n", + " - Inspecting the evidence terms of the fit, which quantify how well the pixelization reconstructs fits the data whilst\n", + " accounting for the complexity of the pixelization.\n", + " - Estimating the magnification of the source galaxy's image using the pixelization.\n", + "\n", + "Therefore if your results contain a pixelization, checkout the example script above for a detailed description\n", + "of how to use their results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/workflow/csv_make.ipynb b/notebooks/guides/results/workflow/csv_make.ipynb index 12901ffdc..46c0335b9 100644 --- a/notebooks/guides/results/workflow/csv_make.ipynb +++ b/notebooks/guides/results/workflow/csv_make.ipynb @@ -1,474 +1,511 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: CSV\n", - "============\n", - "\n", - "This example is a results workflow example, which means it provides tool to set up an effective workflow inspecting\n", - "and interpreting the large libraries of modeling results.\n", - "\n", - "In this tutorial, we use the aggregator to load the results of model-fits and output them in a single .csv file.\n", - "\n", - "This enables the results of many model-fits to be concisely summarised and inspected in a single table, which\n", - "can also be easily passed on to other collaborators.\n", - "\n", - "__CSV, Png and Fits__\n", - "\n", - "Workflow functionality closely mirrors the `png_make.py` and `fits_make.py` examples, which load results of\n", - "model-fits and output th em as .png files and .fits files to quickly summarise results.\n", - "\n", - "The shared `_quick_fit.py` helper creates these results in `results_folder`. If you have older outputs under `results_folder_csv_png_fits`, use `results_folder` for these examples instead.\n", - "\n", - "__Contents__\n", - "\n", - "- **Interferometer:** This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", - "- **Database File:** The aggregator can also load results from a `.sqlite` database file.\n", - "- **Model Fit:** Run the shared quick-fit helper if results have not already been created.\n", - "- **Workflow Paths:** The workflow examples are designed to take large libraries of results and distill them down to the.\n", - "- **Aggregator:** Set up the aggregator as shown in `start_here.py`.\n", - "- **Model Paths:** The paths are the tuples which define how model parameters are accessed from the model.\n", - "- **Adding CSV Columns:** We first make a simple .csv which contains two columns, corresponding to the inferred median PDF.\n", - "- **Saving the CSV:** We can now output the results of all our model-fits to the .csv file, using the `save` method.\n", - "- **Customizing CSV Headers:** The headers of the .csv file are by default the argument input above based on the model.\n", - "- **Maximum Likelihood Values:** We can also output the maximum likelihood values of each parameter to the .csv file, using the.\n", - "- **Errors:** We can also output PDF values at a given sigma confidence of each parameter to the .csv file, using.\n", - "- **Column Label List:** We can add a list of values to the .csv file, provided the list is the same length as the number of.\n", - "- **Latent Variables:** Latent variables are not free model parameters but can be derived from the model, and they are.\n", - "- **Computed Columns:** We can also add columns to the .csv file that are computed from the non-linear search samples (e.g.\n", - "\n", - "__Interferometer__\n", - "\n", - "This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", - "\n", - "The only entries that needs changing are:\n", - "\n", - " - `ImagingAgg` -> `InterferometerAgg`.\n", - " - `FitImagingAgg` -> `FitInterferometerAgg`.\n", - " - `aplt.subplot_imaging_dataset` -> `aplt.subplot_interferometer_dirty_images`.\n", - " - `aplt.subplot_fit_imaging` -> `aplt.subplot_fit_interferometer`.\n", - "\n", - "Quantities specific to an interfometer, for example its uv-wavelengths real space mask, are accessed using the same API\n", - "(e.g. `values(\"dataset.uv_wavelengths\")` and `.values{\"dataset.real_space_mask\")).\n", - "\n", - "__Database File__\n", - "\n", - "The aggregator can also load results from a `.sqlite` database file.\n", - "\n", - "This is beneficial when loading results for large numbers of model-fits (e.g. more than hundreds)\n", - "because it is optimized for fast querying of results.\n", - "\n", - "See the package `results/database` for a full description of how to set up the database and the benefits it provides,\n", - "especially if loading results from hard-disk is slow." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "These workflow examples reuse the shared ``_quick_fit.py`` helper instead of\n", - "performing model-fits in every script. The helper creates two capped Nautilus\n", - "fits, including the latent quantities used below, in ``output/results_folder/``.\n", - "\n", - "Older versions of these workflow examples used ``output/results_folder_csv_png_fits/``;\n", - "use ``output/results_folder/`` for the centralized setup here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import subprocess\n", - "import sys\n", - "\n", - "results_path = Path(\"output\") / \"results_folder\"\n", - "if (\n", - " len(list(results_path.glob(\"**/image/dataset.fits\"))) < 2\n", - " or len(list(results_path.glob(\"**/files/latent/latent_summary.json\"))) < 2\n", - " or len(list(results_path.glob(\"**/image/fit.png\"))) < 2\n", - " or len(list(results_path.glob(\"**/image/fit.fits\"))) < 2\n", - "):\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Workflow Paths__\n", - "\n", - "The workflow examples are designed to take large libraries of results and distill them down to the key information\n", - "required for your science, which are therefore placed in a single path for easy access.\n", - "\n", - "The `workflow_path` specifies where these files are output, in this case the .csv files which summarise the results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "workflow_path = Path(\"output\") / \"results_folder\" / \"workflow_make_example\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Aggregator__\n", - "\n", - "Set up the aggregator as shown in `start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "agg = Aggregator.from_directory(\n", - " directory=Path(\"output\") / \"results_folder\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Extract the `AggregateCSV` object, which has specific functions for outputting results in a CSV format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv = af.AggregateCSV(aggregator=agg)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Paths__\n", - "\n", - "The paths are the tuples which define how model parameters are accessed from the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model = [model for model in agg.values(\"model\")][0]\n", - "print(model.paths)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Adding CSV Columns__\n", - "\n", - "We first make a simple .csv which contains two columns, corresponding to the inferred median PDF values for\n", - "the y centre of the mass of the lens galaxy and its einstein radius.\n", - "\n", - "To do this, we use the `add_variable` method, which adds a column to the .csv file we write at the end. Every time\n", - "we call `add_variable` we add a new column to the .csv file.\n", - "\n", - "Note the API for the `centre`, which is a tuple parameter and therefore needs for `centre_0` to be specified.\n", - "\n", - "The `results_folder` contains two model-fits, meaning that each `add_variable` call adds two rows.\n", - "\n", - "This adds the median PDF value of the parameter to the .csv file, we show how to add other values later in this script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv.add_variable(argument=\"galaxies.lens.mass.centre.centre_0\")\n", - "agg_csv.add_variable(argument=\"galaxies.lens.mass.einstein_radius\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Saving the CSV__\n", - "\n", - "We can now output the results of all our model-fits to the .csv file, using the `save` method.\n", - "\n", - "This will output in your current working directory (e.g. the `autolens_workspace/output/results_folder`)\n", - "as a .csv file containing the median PDF values of the parameters, have a quick look now to see the format of \n", - "the .csv file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv.save(path=workflow_path / \"csv_simple.csv\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Customizing CSV Headers__\n", - "\n", - "The headers of the .csv file are by default the argument input above based on the model. \n", - "\n", - "However, we can customize these headers using the `name` input of the `add_variable` method, for example making them\n", - "shorter or more readable.\n", - "\n", - "We recreate the `agg_csv` first, so that we begin adding columns to a new .csv file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv = af.AggregateCSV(aggregator=agg)\n", - "\n", - "agg_csv.add_variable(\n", - " argument=\"galaxies.lens.mass.centre.centre_0\",\n", - " name=\"mass_centre_0\",\n", - ")\n", - "agg_csv.add_variable(\n", - " argument=\"galaxies.lens.mass.einstein_radius\",\n", - " name=\"mass_einstein_radius\",\n", - ")\n", - "\n", - "agg_csv.save(path=workflow_path / \"csv_simple_custom_headers.csv\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Maximum Likelihood Values__\n", - "\n", - "We can also output the maximum likelihood values of each parameter to the .csv file, using the `use_max_log_likelihood`\n", - "input." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv = af.AggregateCSV(aggregator=agg)\n", - "\n", - "agg_csv.add_variable(\n", - " argument=\"galaxies.lens.mass.einstein_radius\",\n", - " name=\"mass_einstein_radius_max_lh\",\n", - " value_types=[af.ValueType.MaxLogLikelihood],\n", - ")\n", - "\n", - "agg_csv.save(path=workflow_path / \"csv_simple_max_likelihood.csv\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors__\n", - "\n", - "We can also output PDF values at a given sigma confidence of each parameter to the .csv file, using \n", - "the `af.ValueType.ValuesAt3Sigma` input and specifying the sigma confidence.\n", - "\n", - "Below, we add the values at 3.0 sigma confidence to the .csv file, in order to compute the errors you would \n", - "subtract the median value from these values. We add this after the median value, so that the overall inferred\n", - "uncertainty of the parameter is clear.\n", - "\n", - "The method below adds three columns to the .csv file, corresponding to the values at the median, lower and upper sigma \n", - "values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv = af.AggregateCSV(aggregator=agg)\n", - "\n", - "agg_csv.add_variable(\n", - " argument=\"galaxies.lens.mass.einstein_radius\",\n", - " name=\"mass_einstein_radius\",\n", - " value_types=[af.ValueType.Median, af.ValueType.ValuesAt3Sigma],\n", - ")\n", - "\n", - "agg_csv.save(path=workflow_path / \"csv_simple_errors.csv\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Column Label List__\n", - "\n", - "We can add a list of values to the .csv file, provided the list is the same length as the number of model-fits\n", - "in the aggregator.\n", - "\n", - "A useful example is adding the name of every dataset to the .csv file in a column on the left, indicating \n", - "which dataset each row corresponds to.\n", - "\n", - "To make this list, we use the `Aggregator` to loop over the `search` objects and extract their `unique_tag`'s, which \n", - "which the helper set from the dataset names. This API can also be used to extract the `name` or `path_prefix`\n", - "of the search and build an informative list for the names of the subplots.\n", - "\n", - "We then pass the column `name` and this list to the `add_label_column` method, which will add a column to the .csv file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv = af.AggregateCSV(aggregator=agg)\n", - "\n", - "unique_tag_list = [search.unique_tag for search in agg.values(\"search\")]\n", - "\n", - "agg_csv.add_label_column(\n", - " name=\"lens_name\",\n", - " values=unique_tag_list,\n", - ")\n", - "\n", - "agg_csv.save(path=workflow_path / \"csv_simple_dataset_name.csv\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Latent Variables__\n", - "\n", - "Latent variables are not free model parameters but can be derived from the model, and they are described fully in\n", - "?.\n", - "\n", - "This example was run with a latent variable for the shear magnitude, and below we show that this latent variable\n", - "can be added to the .csv file using the same API as above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv = af.AggregateCSV(aggregator=agg)\n", - "\n", - "agg_csv.add_variable(\n", - " argument=\"galaxies.lens.shear.magnitude\",\n", - ")\n", - "\n", - "agg_csv.save(path=workflow_path / \"csv_example_latent.csv\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Computed Columns__\n", - "\n", - "We can also add columns to the .csv file that are computed from the non-linear search samples (e.g. the nested sampling\n", - "samples), for example a value derived from the median PDF instance values of the model.\n", - "\n", - "To do this, we write a function which is input into the `add_computed_column` method, where this function takes the\n", - "median PDF instance as input and returns the computed value.\n", - "\n", - "Below, we add a trivial example of a computed column, where the median PDF value that is twice lens Einstein radius\n", - "is computed and added to the .csv file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_csv = af.AggregateCSV(aggregator=agg)\n", - "\n", - "\n", - "def einstein_radius_x2_from(result):\n", - "\n", - " samples = result.samples\n", - "\n", - " instance = samples.median_pdf()\n", - "\n", - " return 2.0 * instance.galaxies.lens.mass.einstein_radius\n", - "\n", - "\n", - "agg_csv.add_computed_column(\n", - " name=\"bulge_einstein_radius_x2_computed\",\n", - " compute=einstein_radius_x2_from,\n", - ")\n", - "\n", - "agg_csv.save(path=workflow_path / \"csv_computed_columns.csv\")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: CSV\n", + "============\n", + "\n", + "This example is a results workflow example, which means it provides tool to set up an effective workflow inspecting\n", + "and interpreting the large libraries of modeling results.\n", + "\n", + "In this tutorial, we use the aggregator to load the results of model-fits and output them in a single .csv file.\n", + "\n", + "This enables the results of many model-fits to be concisely summarised and inspected in a single table, which\n", + "can also be easily passed on to other collaborators.\n", + "\n", + "__CSV, Png and Fits__\n", + "\n", + "Workflow functionality closely mirrors the `png_make.py` and `fits_make.py` examples, which load results of\n", + "model-fits and output th em as .png files and .fits files to quickly summarise results.\n", + "\n", + "The shared `_quick_fit.py` helper creates these results in `results_folder`. If you have older outputs under `results_folder_csv_png_fits`, use `results_folder` for these examples instead.\n", + "\n", + "__Contents__\n", + "\n", + "- **Interferometer:** This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", + "- **Database File:** The aggregator can also load results from a `.sqlite` database file.\n", + "- **Model Fit:** Run the shared quick-fit helper if results have not already been created.\n", + "- **Workflow Paths:** The workflow examples are designed to take large libraries of results and distill them down to the.\n", + "- **Aggregator:** Set up the aggregator as shown in `start_here.py`.\n", + "- **Model Paths:** The paths are the tuples which define how model parameters are accessed from the model.\n", + "- **Adding CSV Columns:** We first make a simple .csv which contains two columns, corresponding to the inferred median PDF.\n", + "- **Saving the CSV:** We can now output the results of all our model-fits to the .csv file, using the `save` method.\n", + "- **Customizing CSV Headers:** The headers of the .csv file are by default the argument input above based on the model.\n", + "- **Maximum Likelihood Values:** We can also output the maximum likelihood values of each parameter to the .csv file, using the.\n", + "- **Errors:** We can also output PDF values at a given sigma confidence of each parameter to the .csv file, using.\n", + "- **Column Label List:** We can add a list of values to the .csv file, provided the list is the same length as the number of.\n", + "- **Latent Variables:** Latent variables are not free model parameters but can be derived from the model, and they are.\n", + "- **Computed Columns:** We can also add columns to the .csv file that are computed from the non-linear search samples (e.g.\n", + "\n", + "__Interferometer__\n", + "\n", + "This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", + "\n", + "The only entries that needs changing are:\n", + "\n", + " - `ImagingAgg` -> `InterferometerAgg`.\n", + " - `FitImagingAgg` -> `FitInterferometerAgg`.\n", + " - `aplt.subplot_imaging_dataset` -> `aplt.subplot_interferometer_dirty_images`.\n", + " - `aplt.subplot_fit_imaging` -> `aplt.subplot_fit_interferometer`.\n", + "\n", + "Quantities specific to an interfometer, for example its uv-wavelengths real space mask, are accessed using the same API\n", + "(e.g. `values(\"dataset.uv_wavelengths\")` and `.values{\"dataset.real_space_mask\")).\n", + "\n", + "__Database File__\n", + "\n", + "The aggregator can also load results from a `.sqlite` database file.\n", + "\n", + "This is beneficial when loading results for large numbers of model-fits (e.g. more than hundreds)\n", + "because it is optimized for fast querying of results.\n", + "\n", + "See the package `results/database` for a full description of how to set up the database and the benefits it provides,\n", + "especially if loading results from hard-disk is slow." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "These workflow examples reuse the shared ``_quick_fit.py`` helper instead of\n", + "performing model-fits in every script. The helper creates two capped Nautilus\n", + "fits, including the latent quantities used below, in ``output/results_folder/``.\n", + "\n", + "Older versions of these workflow examples used ``output/results_folder_csv_png_fits/``;\n", + "use ``output/results_folder/`` for the centralized setup here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import subprocess\n", + "import sys\n", + "\n", + "results_path = Path(\"output\") / \"results_folder\"\n", + "if (\n", + " len(list(results_path.glob(\"**/image/dataset.fits\"))) < 2\n", + " or len(list(results_path.glob(\"**/files/latent/latent_summary.json\"))) < 2\n", + " or len(list(results_path.glob(\"**/image/fit.png\"))) < 2\n", + " or len(list(results_path.glob(\"**/image/fit.fits\"))) < 2\n", + "):\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Workflow Paths__\n", + "\n", + "The workflow examples are designed to take large libraries of results and distill them down to the key information\n", + "required for your science, which are therefore placed in a single path for easy access.\n", + "\n", + "The `workflow_path` specifies where these files are output, in this case the .csv files which summarise the results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "workflow_path = Path(\"output\") / \"results_folder\" / \"workflow_make_example\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Aggregator__\n", + "\n", + "Set up the aggregator as shown in `start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "agg = Aggregator.from_directory(\n", + " directory=Path(\"output\") / \"results_folder\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Extract the `AggregateCSV` object, which has specific functions for outputting results in a CSV format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv = af.AggregateCSV(aggregator=agg)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Paths__\n", + "\n", + "The paths are the tuples which define how model parameters are accessed from the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model = [model for model in agg.values(\"model\")][0]\n", + "print(model.paths)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Adding CSV Columns__\n", + "\n", + "We first make a simple .csv which contains two columns, corresponding to the inferred median PDF values for\n", + "the y centre of the mass of the lens galaxy and its einstein radius.\n", + "\n", + "To do this, we use the `add_variable` method, which adds a column to the .csv file we write at the end. Every time\n", + "we call `add_variable` we add a new column to the .csv file.\n", + "\n", + "Note the API for the `centre`, which is a tuple parameter and therefore needs for `centre_0` to be specified.\n", + "\n", + "The `results_folder` contains two model-fits, meaning that each `add_variable` call adds two rows.\n", + "\n", + "This adds the median PDF value of the parameter to the .csv file, we show how to add other values later in this script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv.add_variable(argument=\"galaxies.lens.mass.centre.centre_0\")\n", + "agg_csv.add_variable(argument=\"galaxies.lens.mass.einstein_radius\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Saving the CSV__\n", + "\n", + "We can now output the results of all our model-fits to the .csv file, using the `save` method.\n", + "\n", + "This will output in your current working directory (e.g. the `autolens_workspace/output/results_folder`)\n", + "as a .csv file containing the median PDF values of the parameters, have a quick look now to see the format of \n", + "the .csv file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv.save(path=workflow_path / \"csv_simple.csv\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Customizing CSV Headers__\n", + "\n", + "The headers of the .csv file are by default the argument input above based on the model. \n", + "\n", + "However, we can customize these headers using the `name` input of the `add_variable` method, for example making them\n", + "shorter or more readable.\n", + "\n", + "We recreate the `agg_csv` first, so that we begin adding columns to a new .csv file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv = af.AggregateCSV(aggregator=agg)\n", + "\n", + "agg_csv.add_variable(\n", + " argument=\"galaxies.lens.mass.centre.centre_0\",\n", + " name=\"mass_centre_0\",\n", + ")\n", + "agg_csv.add_variable(\n", + " argument=\"galaxies.lens.mass.einstein_radius\",\n", + " name=\"mass_einstein_radius\",\n", + ")\n", + "\n", + "agg_csv.save(path=workflow_path / \"csv_simple_custom_headers.csv\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Maximum Likelihood Values__\n", + "\n", + "We can also output the maximum likelihood values of each parameter to the .csv file, using the `use_max_log_likelihood`\n", + "input." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv = af.AggregateCSV(aggregator=agg)\n", + "\n", + "agg_csv.add_variable(\n", + " argument=\"galaxies.lens.mass.einstein_radius\",\n", + " name=\"mass_einstein_radius_max_lh\",\n", + " value_types=[af.ValueType.MaxLogLikelihood],\n", + ")\n", + "\n", + "agg_csv.save(path=workflow_path / \"csv_simple_max_likelihood.csv\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors__\n", + "\n", + "We can also output PDF values at a given sigma confidence of each parameter to the .csv file, using \n", + "the `af.ValueType.ValuesAt3Sigma` input and specifying the sigma confidence.\n", + "\n", + "Below, we add the values at 3.0 sigma confidence to the .csv file, in order to compute the errors you would \n", + "subtract the median value from these values. We add this after the median value, so that the overall inferred\n", + "uncertainty of the parameter is clear.\n", + "\n", + "The method below adds three columns to the .csv file, corresponding to the values at the median, lower and upper sigma \n", + "values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv = af.AggregateCSV(aggregator=agg)\n", + "\n", + "agg_csv.add_variable(\n", + " argument=\"galaxies.lens.mass.einstein_radius\",\n", + " name=\"mass_einstein_radius\",\n", + " value_types=[af.ValueType.Median, af.ValueType.ValuesAt3Sigma],\n", + ")\n", + "\n", + "agg_csv.save(path=workflow_path / \"csv_simple_errors.csv\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Column Label List__\n", + "\n", + "We can add a list of values to the .csv file, provided the list is the same length as the number of model-fits\n", + "in the aggregator.\n", + "\n", + "A useful example is adding the name of every dataset to the .csv file in a column on the left, indicating \n", + "which dataset each row corresponds to.\n", + "\n", + "To make this list, we use the `Aggregator` to loop over the `search` objects and extract their `unique_tag`'s, which \n", + "which the helper set from the dataset names. This API can also be used to extract the `name` or `path_prefix`\n", + "of the search and build an informative list for the names of the subplots.\n", + "\n", + "We then pass the column `name` and this list to the `add_label_column` method, which will add a column to the .csv file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv = af.AggregateCSV(aggregator=agg)\n", + "\n", + "unique_tag_list = [search.unique_tag for search in agg.values(\"search\")]\n", + "\n", + "agg_csv.add_label_column(\n", + " name=\"lens_name\",\n", + " values=unique_tag_list,\n", + ")\n", + "\n", + "agg_csv.save(path=workflow_path / \"csv_simple_dataset_name.csv\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Latent Variables__\n", + "\n", + "Latent variables are not free model parameters but can be derived from the model, and they are described fully in\n", + "?.\n", + "\n", + "This example was run with a latent variable for the shear magnitude, and below we show that this latent variable\n", + "can be added to the .csv file using the same API as above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv = af.AggregateCSV(aggregator=agg)\n", + "\n", + "agg_csv.add_variable(\n", + " argument=\"galaxies.lens.shear.magnitude\",\n", + ")\n", + "\n", + "agg_csv.save(path=workflow_path / \"csv_example_latent.csv\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Computed Columns__\n", + "\n", + "We can also add columns to the .csv file that are computed from the non-linear search samples (e.g. the nested sampling\n", + "samples), for example a value derived from the median PDF instance values of the model.\n", + "\n", + "To do this, we write a function which is input into the `add_computed_column` method, where this function takes the\n", + "median PDF instance as input and returns the computed value.\n", + "\n", + "Below, we add a trivial example of a computed column, where the median PDF value that is twice lens Einstein radius\n", + "is computed and added to the .csv file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_csv = af.AggregateCSV(aggregator=agg)\n", + "\n", + "\n", + "def einstein_radius_x2_from(result):\n", + "\n", + " samples = result.samples\n", + "\n", + " instance = samples.median_pdf()\n", + "\n", + " return 2.0 * instance.galaxies.lens.mass.einstein_radius\n", + "\n", + "\n", + "agg_csv.add_computed_column(\n", + " name=\"bulge_einstein_radius_x2_computed\",\n", + " compute=einstein_radius_x2_from,\n", + ")\n", + "\n", + "agg_csv.save(path=workflow_path / \"csv_computed_columns.csv\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/workflow/fits_make.ipynb b/notebooks/guides/results/workflow/fits_make.ipynb index b1771d059..331565e16 100644 --- a/notebooks/guides/results/workflow/fits_make.ipynb +++ b/notebooks/guides/results/workflow/fits_make.ipynb @@ -1,460 +1,497 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: Fits Make\n", - "==================\n", - "\n", - "This example is a results workflow example, which means it provides tool to set up an effective workflow inspecting\n", - "and interpreting the large libraries of modeling results.\n", - "\n", - "In this tutorial, we use the aggregator to load .fits files output by a model-fit, extract hdu images and create\n", - "new .fits files, for example all to a single folder on your hard-disk.\n", - "\n", - "For example, a common use case is extracting an image from `model_galaxy_images.fits` of many fits and putting them\n", - "into a single .fits file on your hard-disk. If you have modeled 100+ datasets, you can then inspect all model images\n", - "in DS9 in .fits format n a single folder, which is more efficient than clicking throughout the `output` open each\n", - ".fits file one-by-one.\n", - "\n", - "The most common use of .fits splciing is where multiple observations of the same galaxy are analysed, for example\n", - "at different wavelengths, where each fit outputs a different .fits files. The model images of each fit to each\n", - "wavelength can then be packaged up into a single .fits file.\n", - "\n", - "This enables the results of many model-fits to be concisely visualized and inspected, which can also be easily passed\n", - "on to other collaborators.\n", - "\n", - "Internally, splicing uses standard Astorpy functions to open, edit and save .fit files.\n", - "\n", - "__CSV, Png and Fits__\n", - "\n", - "Workflow functionality closely mirrors the `png_make.py` and `fits_make.py` examples, which load results of\n", - "model-fits and output th em as .png files and .fits files to quickly summarise results.\n", - "\n", - "The shared `_quick_fit.py` helper creates these results in `results_folder`. If you have older outputs under `results_folder_csv_png_fits`, use `results_folder` for these examples instead.\n", - "\n", - "__Contents__\n", - "\n", - "- **Interferometer:** This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", - "- **Database File:** The aggregator can also load results from a `.sqlite` database file.\n", - "- **Model Fit:** Run the shared quick-fit helper if results have not already been created.\n", - "- **Workflow Paths:** The workflow examples are designed to take large libraries of results and distill them down to the.\n", - "- **Aggregator:** Set up the aggregator as shown in `start_here.py`.\n", - "- **Extract Images:** We now extract 2 images from the `fit.fits` file and combine them together into a single .fits file.\n", - "- **Output Single Fits:** The `image` object which has been extracted is an `astropy` `Fits` object, which we use to save the.\n", - "- **Output to Folder:** An alternative way to output the .fits files is to output them as single .fits files for each.\n", - "- **Naming Convention:** We require a naming convention for the output files.\n", - "- **CSV Files:** In the results `image` folder .csv files containing the information to visualize aspects of a.\n", - "- **Add Extra Fits:** We can also add an extra .fits image to the extracted .fits file, for example an RGB image of the.\n", - "- **Custom Fits Files in Analysis:** Describe how a user can extend the `Analysis` class to compute custom images that are output to the.\n", - "- **Path Navigation:** Example combinig `fit.fits` from `source_lp[1]` and `mass_total[0]`.\n", - "\n", - "__Interferometer__\n", - "\n", - "This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", - "\n", - "The only entries that needs changing are:\n", - "\n", - " - `ImagingAgg` -> `InterferometerAgg`.\n", - " - `FitImagingAgg` -> `FitInterferometerAgg`.\n", - " - `aplt.subplot_imaging_dataset` -> `aplt.subplot_interferometer_dirty_images`.\n", - " - `aplt.subplot_fit_imaging` -> `aplt.subplot_fit_interferometer`.\n", - "\n", - "Quantities specific to an interfometer, for example its uv-wavelengths real space mask, are accessed using the same API\n", - "(e.g. `values(\"dataset.uv_wavelengths\")` and `.values{\"dataset.real_space_mask\")).\n", - "\n", - "__Database File__\n", - "\n", - "The aggregator can also load results from a `.sqlite` database file.\n", - "\n", - "This is beneficial when loading results for large numbers of model-fits (e.g. more than hundreds)\n", - "because it is optimized for fast querying of results.\n", - "\n", - "See the package `results/database` for a full description of how to set up the database and the benefits it provides,\n", - "especially if loading results from hard-disk is slow." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "These workflow examples reuse the shared ``_quick_fit.py`` helper instead of\n", - "performing model-fits in every script. The helper creates two capped Nautilus\n", - "fits, including the latent quantities used below, in ``output/results_folder/``.\n", - "\n", - "Older versions of these workflow examples used ``output/results_folder_csv_png_fits/``;\n", - "use ``output/results_folder/`` for the centralized setup here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import subprocess\n", - "import sys\n", - "\n", - "results_path = Path(\"output\") / \"results_folder\"\n", - "if (\n", - " len(list(results_path.glob(\"**/image/dataset.fits\"))) < 2\n", - " or len(list(results_path.glob(\"**/files/latent/latent_summary.json\"))) < 2\n", - " or len(list(results_path.glob(\"**/image/fit.png\"))) < 2\n", - " or len(list(results_path.glob(\"**/image/fit.fits\"))) < 2\n", - "):\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Workflow Paths__\n", - "\n", - "The workflow examples are designed to take large libraries of results and distill them down to the key information\n", - "required for your science, which are therefore placed in a single path for easy access.\n", - "\n", - "The `workflow_path` specifies where these files are output, in this case the .fits files containing the key \n", - "results we require." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "workflow_path = Path(\"output\") / \"results_folder\" / \"workflow_make_example\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Aggregator__\n", - "\n", - "Set up the aggregator as shown in `start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "agg = Aggregator.from_directory(\n", - " directory=Path(\"output\") / \"results_folder\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Extract the `AggregateFITS` object, which has specific functions for loading .fits files and outputting results in \n", - ".fits format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_fits = af.AggregateFITS(aggregator=agg)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extract Images__\n", - "\n", - "We now extract 2 images from the `fit.fits` file and combine them together into a single .fits file.\n", - "\n", - "We will extract the `model_image` and `residual_map` images, which are images you are used to\n", - "plotting and inspecting in the `output` folder of a model-fit and can load and inspect in DS9 from the file\n", - "`fit.fits`.\n", - "\n", - "By inspecting `fit.fits` you will see it contains four images which each have a an `ext_name`: `model_image`,\n", - "`residual_map`, `normalized_residual_map`, `chi_squared_map`.\n", - "\n", - "We do this by simply passing the `agg_fits.extract_fits` method the name of the fits file we load from `fits.fit`\n", - "and the `ext_name` of what we extract.\n", - "\n", - "This runs on all results the `Aggregator` object has loaded from the `output` folder, meaning that for this example\n", - "where two model-fits are loaded, the `image` object contains two images." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "hdu_list = agg_fits.extract_fits(\n", - " hdus=[al.agg.fits_fit.model_data, al.agg.fits_fit.residual_map],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output Single Fits__\n", - "\n", - "The `image` object which has been extracted is an `astropy` `Fits` object, which we use to save the .fits to the \n", - "hard-disk.\n", - "\n", - "The .fits has 4 hdus, the `model_image` and `residual_map` for the two datasets fitted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "hdu_list.writeto(\"fits_make_single.fits\", overwrite=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output to Folder__\n", - "\n", - "An alternative way to output the .fits files is to output them as single .fits files for each model-fit in a single \n", - "folder, which is done using the `output_to_folder` method.\n", - "\n", - "It can sometimes be easier and quicker to inspect the results of many model-fits when they are output to individual\n", - "files in a folder, as using an IDE you can click load and flick through the images. This contrasts a single .png\n", - "file you scroll through, which may be slower to load and inspect.\n", - "\n", - "__Naming Convention__\n", - "\n", - "We require a naming convention for the output files. In this example, we have two model-fits, therefore two .fits\n", - "files are going to be output.\n", - "\n", - "One way to name the .fits files is to use the `unique_tag` of the search, which is unique to every model-fit. For\n", - "the helper-generated `unique_tag` values are `simple__no_lens_light_0` and `simple__no_lens_light_1`, therefore this will informatively name the .fits\n", - "files the names of the datasets.\n", - "\n", - "We achieve this behaviour by inputting `name=\"unique_tag\"` to the `output_to_folder` method. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_fits.output_to_folder(\n", - " folder=workflow_path,\n", - " name=\"unique_tag\",\n", - " hdus=[al.agg.fits_fit.model_data, al.agg.fits_fit.residual_map],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `name` can be any search attribute, for example the `name` of the search, the `path_prefix` of the search, etc,\n", - "if they will give informative names to the .fits files.\n", - "\n", - "You can also manually input a list of names, one for each fit, if you want to name the .fits files something else.\n", - "However, the list must be the same length as the number of fits in the aggregator, and you may not be certain of the\n", - "order of fits in the aggregator and therefore will need to extract this information, for example by printing the\n", - "`unique_tag` of each search (or another attribute containing the dataset name)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print([search.unique_tag for search in agg.values(\"search\")])\n", - "\n", - "agg_fits.output_to_folder(\n", - " folder=workflow_path,\n", - " name=[\"hi_0.fits\", \"hi_1.fits\"],\n", - " hdus=[al.agg.fits_fit.model_data, al.agg.fits_fit.residual_map],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__CSV Files__\n", - "\n", - "In the results `image` folder .csv files containing the information to visualize aspects of a result may be present.\n", - "\n", - "A common example is the file `source_plane_reconstruction_0.csv`, which contains the y and x coordinates of the \n", - "pixelization mesh, the reconstruct values and the noise map of these values.\n", - "\n", - "The `AggregateFITS` object has a method `extract_csv` which extracts this table from each .csv file in the results,\n", - "returning the extracted data as a list of dictionaries. This can then be used to visualize the data, and output\n", - "it to a .fits file elsewhere.\n", - "\n", - "Below, we demonstrate a common use case for a pixelization. Each .csv file is loaded, benefitting from the fact\n", - "that because it stores the irregular mesh values it is the most accurate way to store the data whilst also using\n", - "much less hard-disk space than, for example. converting it to a 2D array and .fits file. We then use the\n", - "loaded values to interpolate the data onto a regular grid and output it to .fits files in a folder.\n", - "\n", - "The code below is commented out because the model does not use a pixelization, but it does work if a\n", - "pixelization is used." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# reconstruction_dict_list = agg_fits.extract_csv(\n", - "# filename=\"source_plane_reconstruction_0\",\n", - "# )\n", - "#\n", - "# from scipy.interpolate import griddata\n", - "#\n", - "# for i, reconstruction_dict in enumerate(reconstruction_dict_list):\n", - "#\n", - "# y = reconstruction_dict[\"y\"]\n", - "# x = reconstruction_dict[\"x\"]\n", - "# values = reconstruction_dict[\"reconstruction\"]\n", - "#\n", - "# points = np.stack(\n", - "# arrays=(reconstruction_dict[\"x\"], reconstruction_dict[\"y\"]), axis=-1\n", - "# )\n", - "#\n", - "# interpolation_grid = al.Grid2D.from_extent(\n", - "# extent=(-1.0, 1.0, -1.0, 1.0), shape_native=(201, 201)\n", - "# )\n", - "#\n", - "# interpolated_array = griddata(points=points, values=values, xi=interpolation_grid)\n", - "#\n", - "# al.output_to_fits(\n", - "# values=interpolated_array,\n", - "# file_path=workflow_path / f\"interpolated_reconstruction_{i}.fits\",\n", - "# )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Add Extra Fits__\n", - "\n", - "We can also add an extra .fits image to the extracted .fits file, for example an RGB image of the dataset.\n", - "\n", - "We create an image of shape (1, 2) and add the RGB image to the left of the subplot, so that the new subplot has\n", - "shape (1, 3).\n", - "\n", - "When we add a single .png, we cannot extract or make it, it simply gets added to the subplot." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# image_rgb = Image.open(Path(dataset_path, \"rgb.png\"))\n", - "#\n", - "# image = agg_fits.extract_fits(\n", - "# al.agg.subplot_dataset.data,\n", - "# al.agg.subplot_dataset.psf_log_10,\n", - "# subplot_shape=(1, 2),\n", - "# )\n", - "\n", - "# image = al.add_image_to_left(image, additional_img)\n", - "\n", - "# image.save(\"png_make_with_rgb.png\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Custom Fits Files in Analysis__\n", - "\n", - "Describe how a user can extend the `Analysis` class to compute custom images that are output to the .png files,\n", - "which they can then extract and make together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Path Navigation__\n", - "\n", - "Example combinig `fit.fits` from `source_lp[1]` and `mass_total[0]`.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: Fits Make\n", + "==================\n", + "\n", + "This example is a results workflow example, which means it provides tool to set up an effective workflow inspecting\n", + "and interpreting the large libraries of modeling results.\n", + "\n", + "In this tutorial, we use the aggregator to load .fits files output by a model-fit, extract hdu images and create\n", + "new .fits files, for example all to a single folder on your hard-disk.\n", + "\n", + "For example, a common use case is extracting an image from `model_galaxy_images.fits` of many fits and putting them\n", + "into a single .fits file on your hard-disk. If you have modeled 100+ datasets, you can then inspect all model images\n", + "in DS9 in .fits format n a single folder, which is more efficient than clicking throughout the `output` open each\n", + ".fits file one-by-one.\n", + "\n", + "The most common use of .fits splciing is where multiple observations of the same galaxy are analysed, for example\n", + "at different wavelengths, where each fit outputs a different .fits files. The model images of each fit to each\n", + "wavelength can then be packaged up into a single .fits file.\n", + "\n", + "This enables the results of many model-fits to be concisely visualized and inspected, which can also be easily passed\n", + "on to other collaborators.\n", + "\n", + "Internally, splicing uses standard Astorpy functions to open, edit and save .fit files.\n", + "\n", + "__CSV, Png and Fits__\n", + "\n", + "Workflow functionality closely mirrors the `png_make.py` and `fits_make.py` examples, which load results of\n", + "model-fits and output th em as .png files and .fits files to quickly summarise results.\n", + "\n", + "The shared `_quick_fit.py` helper creates these results in `results_folder`. If you have older outputs under `results_folder_csv_png_fits`, use `results_folder` for these examples instead.\n", + "\n", + "__Contents__\n", + "\n", + "- **Interferometer:** This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", + "- **Database File:** The aggregator can also load results from a `.sqlite` database file.\n", + "- **Model Fit:** Run the shared quick-fit helper if results have not already been created.\n", + "- **Workflow Paths:** The workflow examples are designed to take large libraries of results and distill them down to the.\n", + "- **Aggregator:** Set up the aggregator as shown in `start_here.py`.\n", + "- **Extract Images:** We now extract 2 images from the `fit.fits` file and combine them together into a single .fits file.\n", + "- **Output Single Fits:** The `image` object which has been extracted is an `astropy` `Fits` object, which we use to save the.\n", + "- **Output to Folder:** An alternative way to output the .fits files is to output them as single .fits files for each.\n", + "- **Naming Convention:** We require a naming convention for the output files.\n", + "- **CSV Files:** In the results `image` folder .csv files containing the information to visualize aspects of a.\n", + "- **Add Extra Fits:** We can also add an extra .fits image to the extracted .fits file, for example an RGB image of the.\n", + "- **Custom Fits Files in Analysis:** Describe how a user can extend the `Analysis` class to compute custom images that are output to the.\n", + "- **Path Navigation:** Example combinig `fit.fits` from `source_lp[1]` and `mass_total[0]`.\n", + "\n", + "__Interferometer__\n", + "\n", + "This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", + "\n", + "The only entries that needs changing are:\n", + "\n", + " - `ImagingAgg` -> `InterferometerAgg`.\n", + " - `FitImagingAgg` -> `FitInterferometerAgg`.\n", + " - `aplt.subplot_imaging_dataset` -> `aplt.subplot_interferometer_dirty_images`.\n", + " - `aplt.subplot_fit_imaging` -> `aplt.subplot_fit_interferometer`.\n", + "\n", + "Quantities specific to an interfometer, for example its uv-wavelengths real space mask, are accessed using the same API\n", + "(e.g. `values(\"dataset.uv_wavelengths\")` and `.values{\"dataset.real_space_mask\")).\n", + "\n", + "__Database File__\n", + "\n", + "The aggregator can also load results from a `.sqlite` database file.\n", + "\n", + "This is beneficial when loading results for large numbers of model-fits (e.g. more than hundreds)\n", + "because it is optimized for fast querying of results.\n", + "\n", + "See the package `results/database` for a full description of how to set up the database and the benefits it provides,\n", + "especially if loading results from hard-disk is slow." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "These workflow examples reuse the shared ``_quick_fit.py`` helper instead of\n", + "performing model-fits in every script. The helper creates two capped Nautilus\n", + "fits, including the latent quantities used below, in ``output/results_folder/``.\n", + "\n", + "Older versions of these workflow examples used ``output/results_folder_csv_png_fits/``;\n", + "use ``output/results_folder/`` for the centralized setup here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import subprocess\n", + "import sys\n", + "\n", + "results_path = Path(\"output\") / \"results_folder\"\n", + "if (\n", + " len(list(results_path.glob(\"**/image/dataset.fits\"))) < 2\n", + " or len(list(results_path.glob(\"**/files/latent/latent_summary.json\"))) < 2\n", + " or len(list(results_path.glob(\"**/image/fit.png\"))) < 2\n", + " or len(list(results_path.glob(\"**/image/fit.fits\"))) < 2\n", + "):\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Workflow Paths__\n", + "\n", + "The workflow examples are designed to take large libraries of results and distill them down to the key information\n", + "required for your science, which are therefore placed in a single path for easy access.\n", + "\n", + "The `workflow_path` specifies where these files are output, in this case the .fits files containing the key \n", + "results we require." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "workflow_path = Path(\"output\") / \"results_folder\" / \"workflow_make_example\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Aggregator__\n", + "\n", + "Set up the aggregator as shown in `start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "agg = Aggregator.from_directory(\n", + " directory=Path(\"output\") / \"results_folder\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Extract the `AggregateFITS` object, which has specific functions for loading .fits files and outputting results in \n", + ".fits format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_fits = af.AggregateFITS(aggregator=agg)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extract Images__\n", + "\n", + "We now extract 2 images from the `fit.fits` file and combine them together into a single .fits file.\n", + "\n", + "We will extract the `model_image` and `residual_map` images, which are images you are used to\n", + "plotting and inspecting in the `output` folder of a model-fit and can load and inspect in DS9 from the file\n", + "`fit.fits`.\n", + "\n", + "By inspecting `fit.fits` you will see it contains four images which each have a an `ext_name`: `model_image`,\n", + "`residual_map`, `normalized_residual_map`, `chi_squared_map`.\n", + "\n", + "We do this by simply passing the `agg_fits.extract_fits` method the name of the fits file we load from `fits.fit`\n", + "and the `ext_name` of what we extract.\n", + "\n", + "This runs on all results the `Aggregator` object has loaded from the `output` folder, meaning that for this example\n", + "where two model-fits are loaded, the `image` object contains two images." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "hdu_list = agg_fits.extract_fits(\n", + " hdus=[al.agg.fits_fit.model_data, al.agg.fits_fit.residual_map],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output Single Fits__\n", + "\n", + "The `image` object which has been extracted is an `astropy` `Fits` object, which we use to save the .fits to the \n", + "hard-disk.\n", + "\n", + "The .fits has 4 hdus, the `model_image` and `residual_map` for the two datasets fitted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "hdu_list.writeto(\"fits_make_single.fits\", overwrite=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output to Folder__\n", + "\n", + "An alternative way to output the .fits files is to output them as single .fits files for each model-fit in a single \n", + "folder, which is done using the `output_to_folder` method.\n", + "\n", + "It can sometimes be easier and quicker to inspect the results of many model-fits when they are output to individual\n", + "files in a folder, as using an IDE you can click load and flick through the images. This contrasts a single .png\n", + "file you scroll through, which may be slower to load and inspect.\n", + "\n", + "__Naming Convention__\n", + "\n", + "We require a naming convention for the output files. In this example, we have two model-fits, therefore two .fits\n", + "files are going to be output.\n", + "\n", + "One way to name the .fits files is to use the `unique_tag` of the search, which is unique to every model-fit. For\n", + "the helper-generated `unique_tag` values are `simple__no_lens_light_0` and `simple__no_lens_light_1`, therefore this will informatively name the .fits\n", + "files the names of the datasets.\n", + "\n", + "We achieve this behaviour by inputting `name=\"unique_tag\"` to the `output_to_folder` method. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_fits.output_to_folder(\n", + " folder=workflow_path,\n", + " name=\"unique_tag\",\n", + " hdus=[al.agg.fits_fit.model_data, al.agg.fits_fit.residual_map],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `name` can be any search attribute, for example the `name` of the search, the `path_prefix` of the search, etc,\n", + "if they will give informative names to the .fits files.\n", + "\n", + "You can also manually input a list of names, one for each fit, if you want to name the .fits files something else.\n", + "However, the list must be the same length as the number of fits in the aggregator, and you may not be certain of the\n", + "order of fits in the aggregator and therefore will need to extract this information, for example by printing the\n", + "`unique_tag` of each search (or another attribute containing the dataset name)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print([search.unique_tag for search in agg.values(\"search\")])\n", + "\n", + "agg_fits.output_to_folder(\n", + " folder=workflow_path,\n", + " name=[\"hi_0.fits\", \"hi_1.fits\"],\n", + " hdus=[al.agg.fits_fit.model_data, al.agg.fits_fit.residual_map],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__CSV Files__\n", + "\n", + "In the results `image` folder .csv files containing the information to visualize aspects of a result may be present.\n", + "\n", + "A common example is the file `source_plane_reconstruction_0.csv`, which contains the y and x coordinates of the \n", + "pixelization mesh, the reconstruct values and the noise map of these values.\n", + "\n", + "The `AggregateFITS` object has a method `extract_csv` which extracts this table from each .csv file in the results,\n", + "returning the extracted data as a list of dictionaries. This can then be used to visualize the data, and output\n", + "it to a .fits file elsewhere.\n", + "\n", + "Below, we demonstrate a common use case for a pixelization. Each .csv file is loaded, benefitting from the fact\n", + "that because it stores the irregular mesh values it is the most accurate way to store the data whilst also using\n", + "much less hard-disk space than, for example. converting it to a 2D array and .fits file. We then use the\n", + "loaded values to interpolate the data onto a regular grid and output it to .fits files in a folder.\n", + "\n", + "The code below is commented out because the model does not use a pixelization, but it does work if a\n", + "pixelization is used." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# reconstruction_dict_list = agg_fits.extract_csv(\n", + "# filename=\"source_plane_reconstruction_0\",\n", + "# )\n", + "#\n", + "# from scipy.interpolate import griddata\n", + "#\n", + "# for i, reconstruction_dict in enumerate(reconstruction_dict_list):\n", + "#\n", + "# y = reconstruction_dict[\"y\"]\n", + "# x = reconstruction_dict[\"x\"]\n", + "# values = reconstruction_dict[\"reconstruction\"]\n", + "#\n", + "# points = np.stack(\n", + "# arrays=(reconstruction_dict[\"x\"], reconstruction_dict[\"y\"]), axis=-1\n", + "# )\n", + "#\n", + "# interpolation_grid = al.Grid2D.from_extent(\n", + "# extent=(-1.0, 1.0, -1.0, 1.0), shape_native=(201, 201)\n", + "# )\n", + "#\n", + "# interpolated_array = griddata(points=points, values=values, xi=interpolation_grid)\n", + "#\n", + "# al.output_to_fits(\n", + "# values=interpolated_array,\n", + "# file_path=workflow_path / f\"interpolated_reconstruction_{i}.fits\",\n", + "# )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Add Extra Fits__\n", + "\n", + "We can also add an extra .fits image to the extracted .fits file, for example an RGB image of the dataset.\n", + "\n", + "We create an image of shape (1, 2) and add the RGB image to the left of the subplot, so that the new subplot has\n", + "shape (1, 3).\n", + "\n", + "When we add a single .png, we cannot extract or make it, it simply gets added to the subplot." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# image_rgb = Image.open(Path(dataset_path, \"rgb.png\"))\n", + "#\n", + "# image = agg_fits.extract_fits(\n", + "# al.agg.subplot_dataset.data,\n", + "# al.agg.subplot_dataset.psf_log_10,\n", + "# subplot_shape=(1, 2),\n", + "# )\n", + "\n", + "# image = al.add_image_to_left(image, additional_img)\n", + "\n", + "# image.save(\"png_make_with_rgb.png\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Custom Fits Files in Analysis__\n", + "\n", + "Describe how a user can extend the `Analysis` class to compute custom images that are output to the .png files,\n", + "which they can then extract and make together." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Path Navigation__\n", + "\n", + "Example combinig `fit.fits` from `source_lp[1]` and `mass_total[0]`.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/results/workflow/png_make.ipynb b/notebooks/guides/results/workflow/png_make.ipynb index 85aa37fbb..9deed795e 100644 --- a/notebooks/guides/results/workflow/png_make.ipynb +++ b/notebooks/guides/results/workflow/png_make.ipynb @@ -1,406 +1,443 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Results: PNG Make\n", - "=================\n", - "\n", - "This example is a results workflow example, which means it provides tool to set up an effective workflow inspecting\n", - "and interpreting the large libraries of modeling results.\n", - "\n", - "In this tutorial, we use the aggregator to load .png files output by a model-fit, make them together to create\n", - "new .png images and then output them all to a single folder on your hard-disk.\n", - "\n", - "For example, a common use case is extracting a subset of 3 or 4 images from `subplot_fit.png` which show the model-fit\n", - "quality, put them on a single line .png subplot and output them all to a single folder on your hard-disk. If you have\n", - "modeled 100+ datasets, you can then inspect all fits as .pngs in a single folder (or make a single. png file of all of\n", - "them which you scroll down), which is more efficient than clicking throughout the `output` folder to inspect\n", - "each lens result one-by-one.\n", - "\n", - "Different .png images can be combined together, for example the goodness-of-fit images from `subplot.png`,\n", - "RGB images of each galaxy in the `dataset` folder and other images.\n", - "\n", - "This enables the results of many model-fits to be concisely visualized and inspected, which can also be easily passed\n", - "on to other collaborators.\n", - "\n", - "Internally, splicing uses the Python Imaging Library (PIL) to open, edit and save .png files. This is a Python library\n", - "that provides extensive file format support, an efficient internal representation and powerful image-processing\n", - "capabilities.\n", - "\n", - "__CSV, Png and Fits__\n", - "\n", - "Workflow functionality closely mirrors the `png_make.py` and `fits_make.py` examples, which load results of\n", - "model-fits and output th em as .png files and .fits files to quickly summarise results.\n", - "\n", - "The shared `_quick_fit.py` helper creates these results in `results_folder`. If you have older outputs under `results_folder_csv_png_fits`, use `results_folder` for these examples instead.\n", - "\n", - "__Contents__\n", - "\n", - "- **Interferometer:** This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", - "- **Database File:** The aggregator can also load results from a `.sqlite` database file.\n", - "- **Model Fit:** Run the shared quick-fit helper if results have not already been created.\n", - "- **Workflow Paths:** The workflow examples are designed to take large libraries of results and distill them down to the.\n", - "- **Aggregator:** Set up the aggregator as shown in `start_here.py`.\n", - "- **Extract Images:** We now extract 3 images from the `subplot_fit.png` file and make them together into a single image.\n", - "- **Output Single Png:** The `image` object which has been extracted is a `Image` object from the Python package `PIL`.\n", - "- **Output to Folder:** An alternative way to output the image is to output them as single .png files for each model-fit in.\n", - "- **Naming Convention:** We require a naming convention for the output files.\n", - "- **Combine Images From Subplots:** We now combine images from two different subplots into a single image, which we will save to the.\n", - "- **Custom Subplots in Analysis:** Describe how a user can extend the `Analysis` class to compute custom images that are output to the.\n", - "- **Path Navigation:** Example combinng `subplot_fit.png` from `source_lp[1]` and `mass_total[0]`.\n", - "\n", - "__Interferometer__\n", - "\n", - "This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", - "\n", - "The only entries that needs changing are:\n", - "\n", - " - `ImagingAgg` -> `InterferometerAgg`.\n", - " - `FitImagingAgg` -> `FitInterferometerAgg`.\n", - " - `aplt.subplot_imaging_dataset` -> `aplt.subplot_interferometer_dirty_images`.\n", - " - `aplt.subplot_fit_imaging` -> `aplt.subplot_fit_interferometer`.\n", - "\n", - "Quantities specific to an interfometer, for example its uv-wavelengths real space mask, are accessed using the same API\n", - "(e.g. `values(\"dataset.uv_wavelengths\")` and `.values{\"dataset.real_space_mask\")).\n", - "\n", - "__Database File__\n", - "\n", - "The aggregator can also load results from a `.sqlite` database file.\n", - "\n", - "This is beneficial when loading results for large numbers of model-fits (e.g. more than hundreds)\n", - "because it is optimized for fast querying of results.\n", - "\n", - "See the package `results/database` for a full description of how to set up the database and the benefits it provides,\n", - "especially if loading results from hard-disk is slow." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "These workflow examples reuse the shared ``_quick_fit.py`` helper instead of\n", - "performing model-fits in every script. The helper creates two capped Nautilus\n", - "fits, including the latent quantities used below, in ``output/results_folder/``.\n", - "\n", - "Older versions of these workflow examples used ``output/results_folder_csv_png_fits/``;\n", - "use ``output/results_folder/`` for the centralized setup here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import subprocess\n", - "import sys\n", - "\n", - "results_path = Path(\"output\") / \"results_folder\"\n", - "if (\n", - " len(list(results_path.glob(\"**/image/dataset.fits\"))) < 2\n", - " or len(list(results_path.glob(\"**/files/latent/latent_summary.json\"))) < 2\n", - " or len(list(results_path.glob(\"**/image/fit.png\"))) < 2\n", - " or len(list(results_path.glob(\"**/image/fit.fits\"))) < 2\n", - "):\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Workflow Paths__\n", - "\n", - "The workflow examples are designed to take large libraries of results and distill them down to the key information\n", - "required for your science, which are therefore placed in a single path for easy access.\n", - "\n", - "The `workflow_path` specifies where these files are output, in this case the .png files containing the key \n", - "results we require." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "workflow_path = Path(\"output\") / \"results_folder\" / \"workflow_make_example\"\n", - "folder_path = workflow_path.parent if workflow_path.suffix else workflow_path\n", - "folder_path.mkdir(parents=True, exist_ok=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Aggregator__\n", - "\n", - "Set up the aggregator as shown in `start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autofit.aggregator.aggregator import Aggregator\n", - "\n", - "agg = Aggregator.from_directory(\n", - " directory=Path(\"output\") / \"results_folder\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Extract the `AggregateImages` object, which has specific functions for loading image files (e.g. .png, .pdf) and\n", - "outputting results in an image format (e.g. .png, .pdf)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_image = af.AggregateImages(aggregator=agg)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extract Images__\n", - "\n", - "We now extract 3 images from the `subplot_fit.png` file and make them together into a single image.\n", - "\n", - "We will extract the `data`, `model_data` and `normalized_residual_map` images, which are images you are used to\n", - "plotting and inspecting in the `output` folder of a model-fit.\n", - "\n", - "We do this by simply passing the `agg_image.extract_image` method the `al.agg` attribute for each image we want to\n", - "extract.\n", - "\n", - "This runs on all results the `Aggregator` object has loaded from the `output` folder, meaning that for this example\n", - "where two model-fits are loaded, the `image` object contains two images.\n", - "\n", - "The `subplot_shape` input above determines the layout of the subplots in the final image, which for the example below\n", - "is a single row of 3 subplots." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = agg_image.extract_image(\n", - " subplots=[\n", - " al.agg.subplot_fit.data,\n", - " al.agg.subplot_fit.model_data,\n", - " al.agg.subplot_fit.normalized_residual_map,\n", - " ],\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output Single Png__\n", - "\n", - "The `image` object which has been extracted is a `Image` object from the Python package `PIL`, which we use\n", - "to save the image to the hard-disk as a .png file.\n", - "\n", - "The .png is a single subplot of two rows, where each subplot is the data, model data and residual-map of a model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image.save(workflow_path / \"png_make_single_subplot.png\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output to Folder__\n", - "\n", - "An alternative way to output the image is to output them as single .png files for each model-fit in a single folder,\n", - "which is done using the `output_to_folder` method.\n", - "\n", - "It can sometimes be easier and quicker to inspect the results of many model-fits when they are output to individual\n", - "files in a folder, as using an IDE you can click load and flick through the images. This contrasts a single .png\n", - "file you scroll through, which may be slower to load and inspect.\n", - "\n", - "__Naming Convention__\n", - "\n", - "We require a naming convention for the output files. In this example, we have two model-fits, therefore two .png\n", - "files are going to be output.\n", - "\n", - "One way to name the .png files is to use the `unique_tag` of the search, which is unique to every model-fit. For\n", - "the helper-generated `unique_tag` values are `simple__no_lens_light_0` and `simple__no_lens_light_1`, therefore this will informatively name the .png\n", - "files the names of the datasets.\n", - "\n", - "We achieve this behaviour by inputting `name=\"unique_tag\"` to the `output_to_folder` method. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_image.output_to_folder(\n", - " folder=workflow_path,\n", - " name=\"unique_tag\",\n", - " subplots=[\n", - " al.agg.subplot_fit.data,\n", - " al.agg.subplot_fit.model_data,\n", - " al.agg.subplot_fit.normalized_residual_map,\n", - " ],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `name` can be any search attribute, for example the `name` of the search, the `path_prefix` of the search, etc,\n", - "if they will give informative names to the .png files.\n", - "\n", - "You can also manually input a list of names, one for each fit, if you want to name the .png files something else.\n", - "However, the list must be the same length as the number of fits in the aggregator, and you may not be certain of the\n", - "order of fits in the aggregator and therefore will need to extract this information, for example by printing the\n", - "`unique_tag` of each search (or another attribute containing the dataset name)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print([search.unique_tag for search in agg.values(\"search\")])\n", - "\n", - "agg_image.output_to_folder(\n", - " folder=workflow_path,\n", - " name=\"unique_tag\",\n", - " subplots=[\n", - " al.agg.subplot_fit.data,\n", - " al.agg.subplot_fit.model_data,\n", - " al.agg.subplot_fit.normalized_residual_map,\n", - " ],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Combine Images From Subplots__\n", - "\n", - "We now combine images from two different subplots into a single image, which we will save to the hard-disk as a .png\n", - "file.\n", - "\n", - "We will extract images from the `subplot_dataset.png` and `subplot_fit.png` images, which are images you are used to \n", - "plotting and inspecting in the `output` folder of a model-fit.\n", - "\n", - "We extract the `data` and `psf_log10` from the dataset and the `model_data` and `chi_squared_map` from the fit,\n", - "and combine them into a subplot with an overall shape of (2, 2)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = agg_image.extract_image(\n", - " subplots=[\n", - " al.agg.subplot_dataset.data,\n", - " al.agg.subplot_dataset.psf_log_10,\n", - " al.agg.subplot_fit.model_data,\n", - " al.agg.subplot_fit.chi_squared_map,\n", - " ]\n", - " # subplot_shape=(2, 2),\n", - ")\n", - "\n", - "image.save(workflow_path / \"png_make_multi_subplot.png\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Custom Subplots in Analysis__\n", - "\n", - "Describe how a user can extend the `Analysis` class to compute custom images that are output to the .png files,\n", - "which they can then extract and make together.\n", - "\n", - "__Path Navigation__\n", - "\n", - "Example combinng `subplot_fit.png` from `source_lp[1]` and `mass_total[0]`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Results: PNG Make\n", + "=================\n", + "\n", + "This example is a results workflow example, which means it provides tool to set up an effective workflow inspecting\n", + "and interpreting the large libraries of modeling results.\n", + "\n", + "In this tutorial, we use the aggregator to load .png files output by a model-fit, make them together to create\n", + "new .png images and then output them all to a single folder on your hard-disk.\n", + "\n", + "For example, a common use case is extracting a subset of 3 or 4 images from `subplot_fit.png` which show the model-fit\n", + "quality, put them on a single line .png subplot and output them all to a single folder on your hard-disk. If you have\n", + "modeled 100+ datasets, you can then inspect all fits as .pngs in a single folder (or make a single. png file of all of\n", + "them which you scroll down), which is more efficient than clicking throughout the `output` folder to inspect\n", + "each lens result one-by-one.\n", + "\n", + "Different .png images can be combined together, for example the goodness-of-fit images from `subplot.png`,\n", + "RGB images of each galaxy in the `dataset` folder and other images.\n", + "\n", + "This enables the results of many model-fits to be concisely visualized and inspected, which can also be easily passed\n", + "on to other collaborators.\n", + "\n", + "Internally, splicing uses the Python Imaging Library (PIL) to open, edit and save .png files. This is a Python library\n", + "that provides extensive file format support, an efficient internal representation and powerful image-processing\n", + "capabilities.\n", + "\n", + "__CSV, Png and Fits__\n", + "\n", + "Workflow functionality closely mirrors the `png_make.py` and `fits_make.py` examples, which load results of\n", + "model-fits and output th em as .png files and .fits files to quickly summarise results.\n", + "\n", + "The shared `_quick_fit.py` helper creates these results in `results_folder`. If you have older outputs under `results_folder_csv_png_fits`, use `results_folder` for these examples instead.\n", + "\n", + "__Contents__\n", + "\n", + "- **Interferometer:** This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", + "- **Database File:** The aggregator can also load results from a `.sqlite` database file.\n", + "- **Model Fit:** Run the shared quick-fit helper if results have not already been created.\n", + "- **Workflow Paths:** The workflow examples are designed to take large libraries of results and distill them down to the.\n", + "- **Aggregator:** Set up the aggregator as shown in `start_here.py`.\n", + "- **Extract Images:** We now extract 3 images from the `subplot_fit.png` file and make them together into a single image.\n", + "- **Output Single Png:** The `image` object which has been extracted is a `Image` object from the Python package `PIL`.\n", + "- **Output to Folder:** An alternative way to output the image is to output them as single .png files for each model-fit in.\n", + "- **Naming Convention:** We require a naming convention for the output files.\n", + "- **Combine Images From Subplots:** We now combine images from two different subplots into a single image, which we will save to the.\n", + "- **Custom Subplots in Analysis:** Describe how a user can extend the `Analysis` class to compute custom images that are output to the.\n", + "- **Path Navigation:** Example combinng `subplot_fit.png` from `source_lp[1]` and `mass_total[0]`.\n", + "\n", + "__Interferometer__\n", + "\n", + "This script can easily be adapted to analyse the results of charge injection imaging model-fits.\n", + "\n", + "The only entries that needs changing are:\n", + "\n", + " - `ImagingAgg` -> `InterferometerAgg`.\n", + " - `FitImagingAgg` -> `FitInterferometerAgg`.\n", + " - `aplt.subplot_imaging_dataset` -> `aplt.subplot_interferometer_dirty_images`.\n", + " - `aplt.subplot_fit_imaging` -> `aplt.subplot_fit_interferometer`.\n", + "\n", + "Quantities specific to an interfometer, for example its uv-wavelengths real space mask, are accessed using the same API\n", + "(e.g. `values(\"dataset.uv_wavelengths\")` and `.values{\"dataset.real_space_mask\")).\n", + "\n", + "__Database File__\n", + "\n", + "The aggregator can also load results from a `.sqlite` database file.\n", + "\n", + "This is beneficial when loading results for large numbers of model-fits (e.g. more than hundreds)\n", + "because it is optimized for fast querying of results.\n", + "\n", + "See the package `results/database` for a full description of how to set up the database and the benefits it provides,\n", + "especially if loading results from hard-disk is slow." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "These workflow examples reuse the shared ``_quick_fit.py`` helper instead of\n", + "performing model-fits in every script. The helper creates two capped Nautilus\n", + "fits, including the latent quantities used below, in ``output/results_folder/``.\n", + "\n", + "Older versions of these workflow examples used ``output/results_folder_csv_png_fits/``;\n", + "use ``output/results_folder/`` for the centralized setup here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import subprocess\n", + "import sys\n", + "\n", + "results_path = Path(\"output\") / \"results_folder\"\n", + "if (\n", + " len(list(results_path.glob(\"**/image/dataset.fits\"))) < 2\n", + " or len(list(results_path.glob(\"**/files/latent/latent_summary.json\"))) < 2\n", + " or len(list(results_path.glob(\"**/image/fit.png\"))) < 2\n", + " or len(list(results_path.glob(\"**/image/fit.fits\"))) < 2\n", + "):\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/guides/results/_quick_fit.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Workflow Paths__\n", + "\n", + "The workflow examples are designed to take large libraries of results and distill them down to the key information\n", + "required for your science, which are therefore placed in a single path for easy access.\n", + "\n", + "The `workflow_path` specifies where these files are output, in this case the .png files containing the key \n", + "results we require." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "workflow_path = Path(\"output\") / \"results_folder\" / \"workflow_make_example\"\n", + "folder_path = workflow_path.parent if workflow_path.suffix else workflow_path\n", + "folder_path.mkdir(parents=True, exist_ok=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Aggregator__\n", + "\n", + "Set up the aggregator as shown in `start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autofit.aggregator.aggregator import Aggregator\n", + "\n", + "agg = Aggregator.from_directory(\n", + " directory=Path(\"output\") / \"results_folder\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Extract the `AggregateImages` object, which has specific functions for loading image files (e.g. .png, .pdf) and\n", + "outputting results in an image format (e.g. .png, .pdf)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_image = af.AggregateImages(aggregator=agg)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extract Images__\n", + "\n", + "We now extract 3 images from the `subplot_fit.png` file and make them together into a single image.\n", + "\n", + "We will extract the `data`, `model_data` and `normalized_residual_map` images, which are images you are used to\n", + "plotting and inspecting in the `output` folder of a model-fit.\n", + "\n", + "We do this by simply passing the `agg_image.extract_image` method the `al.agg` attribute for each image we want to\n", + "extract.\n", + "\n", + "This runs on all results the `Aggregator` object has loaded from the `output` folder, meaning that for this example\n", + "where two model-fits are loaded, the `image` object contains two images.\n", + "\n", + "The `subplot_shape` input above determines the layout of the subplots in the final image, which for the example below\n", + "is a single row of 3 subplots." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = agg_image.extract_image(\n", + " subplots=[\n", + " al.agg.subplot_fit.data,\n", + " al.agg.subplot_fit.model_data,\n", + " al.agg.subplot_fit.normalized_residual_map,\n", + " ],\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output Single Png__\n", + "\n", + "The `image` object which has been extracted is a `Image` object from the Python package `PIL`, which we use\n", + "to save the image to the hard-disk as a .png file.\n", + "\n", + "The .png is a single subplot of two rows, where each subplot is the data, model data and residual-map of a model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image.save(workflow_path / \"png_make_single_subplot.png\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output to Folder__\n", + "\n", + "An alternative way to output the image is to output them as single .png files for each model-fit in a single folder,\n", + "which is done using the `output_to_folder` method.\n", + "\n", + "It can sometimes be easier and quicker to inspect the results of many model-fits when they are output to individual\n", + "files in a folder, as using an IDE you can click load and flick through the images. This contrasts a single .png\n", + "file you scroll through, which may be slower to load and inspect.\n", + "\n", + "__Naming Convention__\n", + "\n", + "We require a naming convention for the output files. In this example, we have two model-fits, therefore two .png\n", + "files are going to be output.\n", + "\n", + "One way to name the .png files is to use the `unique_tag` of the search, which is unique to every model-fit. For\n", + "the helper-generated `unique_tag` values are `simple__no_lens_light_0` and `simple__no_lens_light_1`, therefore this will informatively name the .png\n", + "files the names of the datasets.\n", + "\n", + "We achieve this behaviour by inputting `name=\"unique_tag\"` to the `output_to_folder` method. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_image.output_to_folder(\n", + " folder=workflow_path,\n", + " name=\"unique_tag\",\n", + " subplots=[\n", + " al.agg.subplot_fit.data,\n", + " al.agg.subplot_fit.model_data,\n", + " al.agg.subplot_fit.normalized_residual_map,\n", + " ],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `name` can be any search attribute, for example the `name` of the search, the `path_prefix` of the search, etc,\n", + "if they will give informative names to the .png files.\n", + "\n", + "You can also manually input a list of names, one for each fit, if you want to name the .png files something else.\n", + "However, the list must be the same length as the number of fits in the aggregator, and you may not be certain of the\n", + "order of fits in the aggregator and therefore will need to extract this information, for example by printing the\n", + "`unique_tag` of each search (or another attribute containing the dataset name)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print([search.unique_tag for search in agg.values(\"search\")])\n", + "\n", + "agg_image.output_to_folder(\n", + " folder=workflow_path,\n", + " name=\"unique_tag\",\n", + " subplots=[\n", + " al.agg.subplot_fit.data,\n", + " al.agg.subplot_fit.model_data,\n", + " al.agg.subplot_fit.normalized_residual_map,\n", + " ],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Combine Images From Subplots__\n", + "\n", + "We now combine images from two different subplots into a single image, which we will save to the hard-disk as a .png\n", + "file.\n", + "\n", + "We will extract images from the `subplot_dataset.png` and `subplot_fit.png` images, which are images you are used to \n", + "plotting and inspecting in the `output` folder of a model-fit.\n", + "\n", + "We extract the `data` and `psf_log10` from the dataset and the `model_data` and `chi_squared_map` from the fit,\n", + "and combine them into a subplot with an overall shape of (2, 2)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = agg_image.extract_image(\n", + " subplots=[\n", + " al.agg.subplot_dataset.data,\n", + " al.agg.subplot_dataset.psf_log_10,\n", + " al.agg.subplot_fit.model_data,\n", + " al.agg.subplot_fit.chi_squared_map,\n", + " ]\n", + " # subplot_shape=(2, 2),\n", + ")\n", + "\n", + "image.save(workflow_path / \"png_make_multi_subplot.png\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Custom Subplots in Analysis__\n", + "\n", + "Describe how a user can extend the `Analysis` class to compute custom images that are output to the .png files,\n", + "which they can then extract and make together.\n", + "\n", + "__Path Navigation__\n", + "\n", + "Example combinng `subplot_fit.png` from `source_lp[1]` and `mass_total[0]`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/tracer.ipynb b/notebooks/guides/tracer.ipynb index a55310c04..c7dc4e942 100644 --- a/notebooks/guides/tracer.ipynb +++ b/notebooks/guides/tracer.ipynb @@ -1,972 +1,1009 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fits\n", - "====\n", - "\n", - "This tutorial inspects an inferred model using the `Tracer` object inferred by the non-linear search.\n", - "This allows us to visualize and interpret its results.\n", - "\n", - "The first half of this tutorial repeats the over example `overview/overview_1_lensing.py` and contains the\n", - "following:\n", - "\n", - "This tutorial focuses on explaining how to use the inferred tracer to compute results as numpy arrays and only\n", - "briefly discusses visualization.\n", - "\n", - "__Contents__\n", - "\n", - "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", - "- **Data Structures:** Arrays inspected in this example use bespoke data structures for storing arrays, grids, vectors and.\n", - "- **Other Models:** This tutorial does not use a pixelized source reconstruction or linear light profiles, which have.\n", - "- **Grids:** To describe the deflection of light, **PyAutoLens** uses `Grid2D` data structures, which are.\n", - "- **Light Profiles:** We will ray-trace this `Grid2D`'s coordinates to calculate how the lens galaxy's mass deflects the.\n", - "- **Mass Profiles:** **PyAutoLens** uses `MassProfile` objects to represent a galaxy's mass distribution and perform.\n", - "- **Galaxies:** A `Galaxy` object is a collection of `LightProfile` and `MassProfile` objects at a given redshift.\n", - "- **Ray Tracing:** We can now create the image of the strong lens system!\n", - "- **Log10:** The light and masss distributions of galaxies are closer to a log10 distribution than a linear one.\n", - "- **Extending Objects:** The API has been designed such that all of the objects introduced above are extensible.\n", - "- **Attributes:** Printing individual attributes of the max log likelihood tracer gives us access to the inferred.\n", - "- **Lensing Quantities:** The maximum log likelihood tracer contains a lot of information about the inferred model.\n", - "- **Grid Choices:** We can input a different grid, which is not masked, to evaluate the image everywhere of interest.\n", - "- **Sub Gridding:** A grid can also have a sub-grid, defined via its `sub_size`, which defines how each pixel on the 2D.\n", - "- **Positions Grid:** We may want the image at specific (y,x) coordinates.\n", - "- **Scalar Lensing Quantities:** The tracer has many scalar lensing quantities, which are all returned using an `Array2D` and.\n", - "- **Vector Quantities:** Many lensing quantities are vectors.\n", - "- **Other Vector Lensing Quantities:** The tracer has other vector lensing quantities, which use the same interface described above.\n", - "- **Other Quantities:** Many more quantities are shown below.\n", - "\n", - "__Units__\n", - "\n", - "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", - "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", - "\n", - "The guide `guides/units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", - "kiloparsecs, magnitudes and solar masses.\n", - "\n", - "__Data Structures__\n", - "\n", - "Arrays inspected in this example use bespoke data structures for storing arrays, grids,\n", - "vectors and other 1D and 2D quantities. These use the `slim` and `native` API to toggle between representing the\n", - "data in 1D numpy arrays or high dimension numpy arrays.\n", - "\n", - "This tutorial will only use the `slim` properties which show results in 1D numpy arrays of\n", - "shape [total_unmasked_pixels]. This is a slimmed-down representation of the data in 1D that contains only the\n", - "unmasked data points\n", - "\n", - "These are documented fully in the `autolens_workspace/*/guides/data_structures.ipynb` guide.\n", - "\n", - "__Other Models__\n", - "\n", - "This tutorial does not use a pixelized source reconstruction or linear light profiles, which have their own dediciated\n", - "functionality that interfacts with the `FitImaging` object.\n", - "\n", - "These are described in the dedicated example scripts `results/aggregator/linear.py` and `results/aggregator/pixelizaiton.py`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `results/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import autolens as al\n", - "import autoarray as aa\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grids__\n", - "\n", - "To describe the deflection of light, **PyAutoLens** uses `Grid2D` data structures, which are two-dimensional\n", - "Cartesian grids of (y,x) coordinates. \n", - "\n", - "Below, we make and plot a uniform Cartesian grid in units of arcseconds. \n", - "\n", - "All quantities which are distance units (e.g. coordinate centre's radii) are in units of arc-seconds, as this is the\n", - "most convenient unit to represent lensing quantities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.05, # The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "aplt.plot_grid(grid=grid, title=\"Cartesian (y,x) Grid of Coordinates\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Light Profiles__\n", - "\n", - "We will ray-trace this `Grid2D`'s coordinates to calculate how the lens galaxy's mass deflects the source \n", - "galaxy's light. We therefore need analytic functions representing a galaxy's light and mass distributions. \n", - "\n", - "This requires analytic functions representing the light and mass distributions of galaxies, for example the \n", - "elliptical `Sersic` `LightProfile`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sersic_light_profile = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.2, 0.1),\n", - " intensity=0.005,\n", - " effective_radius=2.0,\n", - " sersic_index=2.5,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By passing this profile a `Grid2D`, we can evaluate the light at every (y,x) coordinate on the `Grid2D` and create an \n", - "image of the Sersic.\n", - "\n", - "All images in **PyAutoLens** are in units of electrons per second." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = sersic_light_profile.image_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The **PyAutoLens** plot module provides methods for plotting objects and their properties, like light profile's image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=sersic_light_profile.image_2d_from(grid=grid),\n", - " title=\"Image of Sersic Light Profile\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mass Profiles__\n", - "\n", - "**PyAutoLens** uses `MassProfile` objects to represent a galaxy's mass distribution and perform ray-tracing\n", - "calculations. \n", - "\n", - "Below we create an `Isothermal` mass profile and compute its deflection angles on our Cartesian grid, which describe\n", - "how the source galaxy's light rays are deflected as they pass this mass distribution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "isothermal_mass_profile = al.mp.Isothermal(\n", - " centre=(0.0, 0.0), ell_comps=(0.1, 0.0), einstein_radius=1.6\n", - ")\n", - "deflections = isothermal_mass_profile.deflections_yx_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets plot the isothermal mass profile's deflection angle map.\n", - "\n", - "The black curve on the figure is the tangential critical curve of the mass profile, if you do not know what this is\n", - "don't worry about it for now!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections = isothermal_mass_profile.deflections_yx_2d_from(grid=grid)\n", - "deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=grid.mask)\n", - "aplt.plot_array(array=deflections_y, title=\"Deflections Y\")\n", - "deflections = isothermal_mass_profile.deflections_yx_2d_from(grid=grid)\n", - "deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=grid.mask)\n", - "aplt.plot_array(array=deflections_x, title=\"Deflections X\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There are many other lensing quantities which can be plotted, for example the convergence and gravitational\n", - "potential.\n", - "\n", - "If you are not familiar with gravitational lensing and therefore are unclear on what the convergence and potential \n", - "are, don't worry for now!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=isothermal_mass_profile.convergence_2d_from(grid=grid),\n", - " title=\"Isothermal Mass Convergence\",\n", - ")\n", - "aplt.plot_array(\n", - " array=isothermal_mass_profile.potential_2d_from(grid=grid),\n", - " title=\"Isothermal Mass Potential\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "A `Galaxy` object is a collection of `LightProfile` and `MassProfile` objects at a given redshift. \n", - "\n", - "The code below creates two galaxies representing the lens and source galaxies shown in the strong lensing diagram above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5, bulge=sersic_light_profile, mass=isothermal_mass_profile\n", - ")\n", - "\n", - "source_light_profile = al.lp.ExponentialCore(\n", - " centre=(0.3, 0.2), ell_comps=(0.1, 0.0), intensity=0.1, effective_radius=0.5\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, bulge=source_light_profile)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The geometry of the strong lens system depends on the cosmological distances between the Earth, the lens galaxy and \n", - "the source galaxy. It there depends on the redshifts of the `Galaxy` objects. \n", - "\n", - "By passing these `Galaxy` objects to a `Tracer` with a `Cosmology` object, **PyAutoLens** uses these galaxy redshifts \n", - "and a cosmological model to create the appropriate strong lens system." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy], cosmology=al.cosmo.Planck15())" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "We can now create the image of the strong lens system! \n", - "\n", - "When calculating this image, the `Tracer` performs all ray-tracing for the strong lens system. This includes using the \n", - "lens galaxy's total mass distribution to deflect the light-rays that are traced to the source galaxy. As a result, \n", - "the source's appears as a multiply imaged and strongly lensed Einstein ring." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = tracer.image_2d_from(grid=grid)\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid), title=\"Image of Strong Lens System\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log10__\n", - "\n", - "The light and masss distributions of galaxies are closer to a log10 distribution than a linear one. \n", - "\n", - "This means that when we plot an image of a light profile, its appearance is better highlighted when we take the\n", - "logarithm of its values and plot it in log10 space.\n", - "\n", - "The `plot_array`/`subplot_\\*` object has an input `use_log10`, which will do this automatically when we call the `plot_array` method.\n", - "Below, we can see that the image plotted now appears more clearly, with the outskirts of the light profile more visible." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid.mask.derive_grid.all_false), title=\"Image\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `aplt.subplot_tracer` includes the mass quantities we plotted previously, which can be plotted as a subplot \n", - "that plots all these quantities simultaneously.\n", - "\n", - "The black and white lines in the source-plane image are the tangential and radial caustics of the mass, which again\n", - "you do not need to worry about for now if you don't know what that is!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_tracer(tracer=tracer, grid=grid.mask.derive_grid.all_false)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The tracer is composed of planes. The system above has two planes, an image-plane (at redshift=0.5) and a \n", - "source-plane (at redshift=1.0). \n", - "\n", - "When creating an image via a Tracer, the mass profiles are used to ray-trace the image-plane grid (plotted above) \n", - "to a source-plane grid, via the mass profile's deflection angles.\n", - "\n", - "We can use the tracer`s `traced_grid_2d_list_from` method to calculate and plot the image-plane and source-plane grids." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "aplt.plot_grid(grid=traced_grid_list[0], title=\"Image Plane Grid\")\n", - "\n", - "aplt.plot_grid(grid=traced_grid_list[1], title=\"Source Plane Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extending Objects__\n", - "\n", - "The API has been designed such that all of the objects introduced above are extensible. `Galaxy` \n", - "objects can take many `LightProfile`'s and `MassProfile`'s. `Tracer`' objects can take many `Galaxy`'s. \n", - "\n", - "If the galaxies are at different redshifts a strong lensing system with multiple lens planes will be created, \n", - "performing complex multi-plane ray-tracing calculations.\n", - "\n", - "To finish, lets create a `Tracer` with 3 galaxies at 3 different redshifts, forming a system with two distinct Einstein\n", - "rings! The mass distribution of the first galaxy also has separate components for its stellar mass and dark matter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lmp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.05),\n", - " intensity=0.5,\n", - " effective_radius=0.3,\n", - " sersic_index=3.5,\n", - " mass_to_light_ratio=0.6,\n", - " ),\n", - " disk=al.lmp.Exponential(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.1),\n", - " intensity=1.0,\n", - " effective_radius=2.0,\n", - " mass_to_light_ratio=0.2,\n", - " ),\n", - " dark=al.mp.NFWSph(centre=(0.0, 0.0), kappa_s=0.08, scale_radius=30.0),\n", - ")\n", - "\n", - "lens_galaxy_1 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Exponential(\n", - " centre=(0.00, 0.00),\n", - " ell_comps=(0.05, 0.05),\n", - " intensity=1.2,\n", - " effective_radius=0.1,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0), ell_comps=(0.05, 0.05), einstein_radius=0.3\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=2.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.111111),\n", - " intensity=1.4,\n", - " effective_radius=0.1,\n", - " sersic_index=1.5,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy_0, lens_galaxy_1, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This is what the lens looks like. \n", - "\n", - "Note how crazy the critical curves are!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid), title=\"Image of Complex Strong Lens\"\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Attributes__\n", - "\n", - "Printing individual attributes of the max log likelihood tracer gives us access to the inferred parameters of the\n", - "lens and source galaxies.\n", - "\n", - "The tracer contains the galaxies as both a list and an instance of the model used to fit it. This means we can\n", - "access the same values in two ways, either indexing the galaxies list index or by the name used in model composition.\n", - "\n", - "It can be difficult to track which galaxy is which index in the list, so it is recommended to use the model\n", - "composition to access the galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Einstein Radius via list index = {tracer.galaxies[1].mass.einstein_radius}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lensing Quantities__\n", - "\n", - "The maximum log likelihood tracer contains a lot of information about the inferred model.\n", - "\n", - "For example, by passing it a 2D grid of (y,x) coordinates we can return a numpy array containing its 2D image. This\n", - "includes the lens light and lensed source images.\n", - "\n", - "Below, we use the grid of the `imaging` to computed the image on, which is the grid used to fit to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = tracer.image_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Data Structures Slim / Native__\n", - "\n", - "The image above is returned as a 1D numpy array. \n", - "\n", - "**PyAutoLens** includes dedicated functionality for manipulating this array, for example mapping it to 2D or\n", - "performing the calculation on a high resolution sub-grid which is then binned up. \n", - "\n", - "This uses the data structure API, which is described in the `results/aggregator/data_structures.py` example. This \n", - "tutorial will avoid using this API, but if you need to manipulate results in more detail you should check it out." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(image.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid Choices__\n", - "\n", - "We can input a different grid, which is not masked, to evaluate the image everywhere of interest. We can also change\n", - "the grid's resolution from that used in the model-fit.\n", - "\n", - "The examples uses a grid with `shape_native=(3,3)`. This is much lower resolution than one would typically use to \n", - "perform ray tracing, but is chosen here so that the `print()` statements display in a concise and readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(5, 5), pixel_scales=0.1)\n", - "\n", - "image = tracer.image_2d_from(grid=grid)\n", - "\n", - "print(image.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sub Gridding__\n", - "\n", - "A grid can also have a sub-grid, defined via its `sub_size`, which defines how each pixel on the 2D grid is split \n", - "into sub-pixels of size (`sub_size` x `sub_size`). \n", - "\n", - "The calculation below shows how to use a sub-grid and return an image which has already been binned up. \n", - "\n", - "Full details of the API for this calculation are given in the `guides/over_sampling.py` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=grid.shape_native,\n", - " pixel_scales=grid.pixel_scales,\n", - " over_sample_size=2,\n", - ")\n", - "\n", - "grid_sub = al.Grid2D.uniform(shape_native=(3, 3), pixel_scales=0.1)\n", - "\n", - "image = tracer.image_2d_from(grid=grid_sub)\n", - "\n", - "print(image)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions Grid__\n", - "\n", - "We may want the image at specific (y,x) coordinates.\n", - "\n", - "We can use an irregular 2D (y,x) grid of coordinates for this. The grid below evaluates the image at:\n", - "\n", - "- y = 1.0, x = 1.0.\n", - "- y = 1.0, x = 2.0.\n", - "- y = 2.0, x = 2.0." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_irregular = al.Grid2DIrregular(values=[[1.0, 1.0], [1.0, 2.0], [2.0, 2.0]])\n", - "\n", - "image = tracer.image_2d_from(grid=grid_irregular)\n", - "\n", - "print(image)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scalar Lensing Quantities__\n", - "\n", - "The tracer has many scalar lensing quantities, which are all returned using an `Array2D` and therefore use the same \n", - "interface as images, described above.\n", - "\n", - "For example, we can compute the `Tracer`'s convergence using all of the grids above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "convergence_2d = tracer.convergence_2d_from(grid=grid)\n", - "print(convergence_2d)\n", - "\n", - "convergence_2d = tracer.convergence_2d_from(grid=grid_sub)\n", - "print(convergence_2d)\n", - "\n", - "convergence_2d = tracer.convergence_2d_from(grid=grid_irregular)\n", - "print(convergence_2d)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This is the convergence of every galaxy in the tracer summed together. It may not be appropriate if your lens model \n", - "performs multi-plane ray-tracing (e.g. there are more than 2 redshifts containing galaxies). Later results tutorials\n", - "provide tools that are more appropriate for multi-plane tracers.\n", - "\n", - "There are other scalar quantities accessible via the tracer (those not familiar with strong lensing mathematical \n", - "formalism may not recognise what these quantities are -- don't worry about it for now!):" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "potential_2d = tracer.potential_2d_from(grid=grid)\n", - "\n", - "lens_calc = al.LensCalc.from_tracer(tracer)\n", - "\n", - "tangential_eigen_value = lens_calc.tangential_eigen_value_from(grid=grid)\n", - "radial_eigen_value = lens_calc.radial_eigen_value_from(grid=grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D magnification map is available, which using only the ray-tracing and therefore mass model quantities how much\n", - "light rays are focus at a given point in the image-plane.\n", - "\n", - "If you are studying a strongly lensed source galaxy and want to know how much the galaxy itself is magnified, the\n", - "magnification below is not of too much use too you. In the result tutorial `galaxies.py` we explain how the \n", - "magnification of the source can be quantified." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnification_2d = lens_calc.magnification_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Vector Quantities__\n", - "\n", - "Many lensing quantities are vectors. That is, they are (y,x) coordinates that have 2 values representing their\n", - "magnitudes in both the y and x directions.\n", - "\n", - "These quantities also have a dedicated data structure which is described fully in \n", - "the `results/aggregator/data_structure.py` example.\n", - "\n", - "The most obvious of these is the deflection angles, which are used throughout lens modeling to ray-trace grids\n", - "from the image-plane to the source-plane via a lens galaxy mass model.\n", - "\n", - "To indicate that a quantities is a vector, **PyAutoLens** uses the label `_yx`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_yx_2d = tracer.deflections_yx_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For vector quantities the has shape `2`, corresponding to the y and x vectors respectively." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(deflections_yx_2d[0, :])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `VectorYX2D` object has a built in method to return the magnitude of each vector, which is a scalar quantity\n", - "and therefore returned using a 1D Numpy array." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflection_magnitudes_2d = deflections_yx_2d.magnitudes\n", - "print(deflection_magnitudes_2d)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Other Vector Lensing Quantities__\n", - "\n", - "The tracer has other vector lensing quantities, which use the same interface described above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "shear_yx_2d = lens_calc.shear_yx_2d_via_hessian_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Other Quantities__\n", - "\n", - "Many more quantities are shown below.\n", - "\n", - "A full description of each can be found in the docstring of the source code of each function:\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/operate/deflections.py" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tangential_critical_curve = lens_calc.tangential_critical_curve_list_from(grid=grid)\n", - "\n", - "radial_critical_curve = lens_calc.radial_critical_curve_list_from(grid=grid)\n", - "\n", - "tangential_caustic = lens_calc.tangential_caustic_list_from(grid=grid)\n", - "\n", - "radial_caustic = lens_calc.radial_caustic_list_from(grid=grid)\n", - "\n", - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy_0, lens_galaxy_1]\n", - ") # No support for multi plane tracer time delays yet\n", - "\n", - "time_delay = tracer.time_delays_from(grid=grid)\n", - "\n", - "### You should be able to comment this out and it work fine ###\n", - "\n", - "# area_within_tangential_critical_curve = (\n", - "# lens_calc.tangential_critical_curve_area_list_from(grid=grid)\n", - "# )\n", - "#\n", - "# einstein_radius = tracer.einstein_radius_from(grid=grid)\n", - "#\n", - "# einstein_mass_angular = tracer.einstein_mass_angular_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__JAX__\n", - "\n", - "`Tracer` ray-tracing is the most JAX-friendly part of PyAutoLens \u2014 pure\n", - "numerical kernels with no data-dependent control flow. Typical speedups\n", - "for `tracer.image_2d_from(grid)` and related ops on a large grid:\n", - "10-30\u00d7 for galaxy-scale models, 30-100\u00d7 for cluster-scale on GPU.\n", - "\n", - "You access this in two ways.\n", - "\n", - "__1. The implicit path: `Analysis` and `Simulator`__\n", - "\n", - "`AnalysisImaging(use_jax=True)` (the default) and\n", - "`SimulatorImaging(use_jax=True)` both JAX-accelerate the tracer\n", - "internally. Pytree registration runs as a side effect of the first\n", - "`fit_from` / `via_tracer_from` call; you write nothing JAX-specific.\n", - "\n", - "__2. The explicit path: your own `@jax.jit`__\n", - "\n", - "For parameter sweeps, custom forward models, or batch figure generation\n", - "where you want fine-grained control:\n", - "\n", - "```python\n", - "import jax\n", - "import jax.numpy as jnp\n", - "from autolens.jax import register_tracer_classes\n", - "\n", - "register_tracer_classes(tracer) # one-time pytree registration\n", - "\n", - "@jax.jit\n", - "def image_fn(tracer, grid):\n", - " return tracer.image_2d_from(grid=grid, xp=jnp).array\n", - "\n", - "image = image_fn(tracer, grid)\n", - "```\n", - "\n", - "Two rules:\n", - "\n", - "- **`@jax.jit` + `xp=jnp` pair up.** Forgetting `xp=jnp` either\n", - " silently host-transfers (slow) or fails at the boundary; the library\n", - " now raises a clear `ValueError` on the easy mismatch (see\n", - " `lens_calc.py` for the rationale and `AbstractMaker.__init__`'s\n", - " guard).\n", - "- **`.array` unwrap inside the jit, rewrap outside.** Wrapper types\n", - " (`aa.Array2D`, `aa.Grid2DIrregular`) aren't reliably pytree for\n", - " return-from-JIT \u2014 return raw arrays, rewrap on the host.\n", - "\n", - "__Multi-plane traces under JIT__\n", - "\n", - "The recursive multi-plane lens equation in\n", - "`tracer.traced_grid_2d_list_from(grid)` is pure numerical with no\n", - "data-dependent control flow, so it JITs cleanly. For multi-plane point-\n", - "source solving (forward-solving multiple-image positions through several\n", - "planes), use the higher-level `al.PointSolver(use_jax=True)` \u2014 see\n", - "`scripts/point_source/simulator.py` `__JAX Variant__`.\n", - "\n", - "__Performance expectations__\n", - "\n", - "Tracer image generation on JAX-GPU typically beats NumPy-CPU by:\n", - "\n", - "- 10-30\u00d7 for galaxy-scale (single lens galaxy, single source).\n", - "- 30-100\u00d7 for cluster-scale (many galaxies, multi-plane).\n", - "\n", - "Actual speedup depends on grid size, profile complexity, and GPU\n", - "hardware. `autolens_workspace_developer/jax_profiling/` carries measured\n", - "numbers for representative configurations.\n", - "\n", - "For the full \"JIT-it-yourself\" deep-dive (bound-method form, cache-\n", - "identity considerations, closure-captured `self` vs traced-argument),\n", - "see `scripts/guides/lens_calc.py`. `scripts/guides/galaxies.py` covers\n", - "the pytree registration mechanics. `scripts/guides/data_structures.py`\n", - "covers the `.array` host-transfer story." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "Fin.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fits\n", + "====\n", + "\n", + "This tutorial inspects an inferred model using the `Tracer` object inferred by the non-linear search.\n", + "This allows us to visualize and interpret its results.\n", + "\n", + "The first half of this tutorial repeats the over example `overview/overview_1_lensing.py` and contains the\n", + "following:\n", + "\n", + "This tutorial focuses on explaining how to use the inferred tracer to compute results as numpy arrays and only\n", + "briefly discusses visualization.\n", + "\n", + "__Contents__\n", + "\n", + "- **Units:** In this example, all quantities use the source code's internal unit coordinates, with spatial.\n", + "- **Data Structures:** Arrays inspected in this example use bespoke data structures for storing arrays, grids, vectors and.\n", + "- **Other Models:** This tutorial does not use a pixelized source reconstruction or linear light profiles, which have.\n", + "- **Grids:** To describe the deflection of light, **PyAutoLens** uses `Grid2D` data structures, which are.\n", + "- **Light Profiles:** We will ray-trace this `Grid2D`'s coordinates to calculate how the lens galaxy's mass deflects the.\n", + "- **Mass Profiles:** **PyAutoLens** uses `MassProfile` objects to represent a galaxy's mass distribution and perform.\n", + "- **Galaxies:** A `Galaxy` object is a collection of `LightProfile` and `MassProfile` objects at a given redshift.\n", + "- **Ray Tracing:** We can now create the image of the strong lens system!\n", + "- **Log10:** The light and masss distributions of galaxies are closer to a log10 distribution than a linear one.\n", + "- **Extending Objects:** The API has been designed such that all of the objects introduced above are extensible.\n", + "- **Attributes:** Printing individual attributes of the max log likelihood tracer gives us access to the inferred.\n", + "- **Lensing Quantities:** The maximum log likelihood tracer contains a lot of information about the inferred model.\n", + "- **Grid Choices:** We can input a different grid, which is not masked, to evaluate the image everywhere of interest.\n", + "- **Sub Gridding:** A grid can also have a sub-grid, defined via its `sub_size`, which defines how each pixel on the 2D.\n", + "- **Positions Grid:** We may want the image at specific (y,x) coordinates.\n", + "- **Scalar Lensing Quantities:** The tracer has many scalar lensing quantities, which are all returned using an `Array2D` and.\n", + "- **Vector Quantities:** Many lensing quantities are vectors.\n", + "- **Other Vector Lensing Quantities:** The tracer has other vector lensing quantities, which use the same interface described above.\n", + "- **Other Quantities:** Many more quantities are shown below.\n", + "\n", + "__Units__\n", + "\n", + "In this example, all quantities use the source code's internal unit coordinates, with spatial coordinates in\n", + "arc seconds, luminosities in electrons per second and mass quantities (e.g. convergence) are dimensionless.\n", + "\n", + "The guide `guides/units_and_cosmology.ipynb` illustrates how to convert these quantities to physical units like\n", + "kiloparsecs, magnitudes and solar masses.\n", + "\n", + "__Data Structures__\n", + "\n", + "Arrays inspected in this example use bespoke data structures for storing arrays, grids,\n", + "vectors and other 1D and 2D quantities. These use the `slim` and `native` API to toggle between representing the\n", + "data in 1D numpy arrays or high dimension numpy arrays.\n", + "\n", + "This tutorial will only use the `slim` properties which show results in 1D numpy arrays of\n", + "shape [total_unmasked_pixels]. This is a slimmed-down representation of the data in 1D that contains only the\n", + "unmasked data points\n", + "\n", + "These are documented fully in the `autolens_workspace/*/guides/data_structures.ipynb` guide.\n", + "\n", + "__Other Models__\n", + "\n", + "This tutorial does not use a pixelized source reconstruction or linear light profiles, which have their own dediciated\n", + "functionality that interfacts with the `FitImaging` object.\n", + "\n", + "These are described in the dedicated example scripts `results/aggregator/linear.py` and `results/aggregator/pixelizaiton.py`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `results/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import autolens as al\n", + "import autoarray as aa\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grids__\n", + "\n", + "To describe the deflection of light, **PyAutoLens** uses `Grid2D` data structures, which are two-dimensional\n", + "Cartesian grids of (y,x) coordinates. \n", + "\n", + "Below, we make and plot a uniform Cartesian grid in units of arcseconds. \n", + "\n", + "All quantities which are distance units (e.g. coordinate centre's radii) are in units of arc-seconds, as this is the\n", + "most convenient unit to represent lensing quantities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.05, # The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "aplt.plot_grid(grid=grid, title=\"Cartesian (y,x) Grid of Coordinates\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Light Profiles__\n", + "\n", + "We will ray-trace this `Grid2D`'s coordinates to calculate how the lens galaxy's mass deflects the source \n", + "galaxy's light. We therefore need analytic functions representing a galaxy's light and mass distributions. \n", + "\n", + "This requires analytic functions representing the light and mass distributions of galaxies, for example the \n", + "elliptical `Sersic` `LightProfile`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sersic_light_profile = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.2, 0.1),\n", + " intensity=0.005,\n", + " effective_radius=2.0,\n", + " sersic_index=2.5,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By passing this profile a `Grid2D`, we can evaluate the light at every (y,x) coordinate on the `Grid2D` and create an \n", + "image of the Sersic.\n", + "\n", + "All images in **PyAutoLens** are in units of electrons per second." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = sersic_light_profile.image_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The **PyAutoLens** plot module provides methods for plotting objects and their properties, like light profile's image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=sersic_light_profile.image_2d_from(grid=grid),\n", + " title=\"Image of Sersic Light Profile\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mass Profiles__\n", + "\n", + "**PyAutoLens** uses `MassProfile` objects to represent a galaxy's mass distribution and perform ray-tracing\n", + "calculations. \n", + "\n", + "Below we create an `Isothermal` mass profile and compute its deflection angles on our Cartesian grid, which describe\n", + "how the source galaxy's light rays are deflected as they pass this mass distribution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "isothermal_mass_profile = al.mp.Isothermal(\n", + " centre=(0.0, 0.0), ell_comps=(0.1, 0.0), einstein_radius=1.6\n", + ")\n", + "deflections = isothermal_mass_profile.deflections_yx_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets plot the isothermal mass profile's deflection angle map.\n", + "\n", + "The black curve on the figure is the tangential critical curve of the mass profile, if you do not know what this is\n", + "don't worry about it for now!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections = isothermal_mass_profile.deflections_yx_2d_from(grid=grid)\n", + "deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=grid.mask)\n", + "aplt.plot_array(array=deflections_y, title=\"Deflections Y\")\n", + "deflections = isothermal_mass_profile.deflections_yx_2d_from(grid=grid)\n", + "deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=grid.mask)\n", + "aplt.plot_array(array=deflections_x, title=\"Deflections X\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There are many other lensing quantities which can be plotted, for example the convergence and gravitational\n", + "potential.\n", + "\n", + "If you are not familiar with gravitational lensing and therefore are unclear on what the convergence and potential \n", + "are, don't worry for now!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=isothermal_mass_profile.convergence_2d_from(grid=grid),\n", + " title=\"Isothermal Mass Convergence\",\n", + ")\n", + "aplt.plot_array(\n", + " array=isothermal_mass_profile.potential_2d_from(grid=grid),\n", + " title=\"Isothermal Mass Potential\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "A `Galaxy` object is a collection of `LightProfile` and `MassProfile` objects at a given redshift. \n", + "\n", + "The code below creates two galaxies representing the lens and source galaxies shown in the strong lensing diagram above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5, bulge=sersic_light_profile, mass=isothermal_mass_profile\n", + ")\n", + "\n", + "source_light_profile = al.lp.ExponentialCore(\n", + " centre=(0.3, 0.2), ell_comps=(0.1, 0.0), intensity=0.1, effective_radius=0.5\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, bulge=source_light_profile)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The geometry of the strong lens system depends on the cosmological distances between the Earth, the lens galaxy and \n", + "the source galaxy. It there depends on the redshifts of the `Galaxy` objects. \n", + "\n", + "By passing these `Galaxy` objects to a `Tracer` with a `Cosmology` object, **PyAutoLens** uses these galaxy redshifts \n", + "and a cosmological model to create the appropriate strong lens system." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy], cosmology=al.cosmo.Planck15())" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "We can now create the image of the strong lens system! \n", + "\n", + "When calculating this image, the `Tracer` performs all ray-tracing for the strong lens system. This includes using the \n", + "lens galaxy's total mass distribution to deflect the light-rays that are traced to the source galaxy. As a result, \n", + "the source's appears as a multiply imaged and strongly lensed Einstein ring." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = tracer.image_2d_from(grid=grid)\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid), title=\"Image of Strong Lens System\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log10__\n", + "\n", + "The light and masss distributions of galaxies are closer to a log10 distribution than a linear one. \n", + "\n", + "This means that when we plot an image of a light profile, its appearance is better highlighted when we take the\n", + "logarithm of its values and plot it in log10 space.\n", + "\n", + "The `plot_array`/`subplot_\\*` object has an input `use_log10`, which will do this automatically when we call the `plot_array` method.\n", + "Below, we can see that the image plotted now appears more clearly, with the outskirts of the light profile more visible." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid.mask.derive_grid.all_false), title=\"Image\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `aplt.subplot_tracer` includes the mass quantities we plotted previously, which can be plotted as a subplot \n", + "that plots all these quantities simultaneously.\n", + "\n", + "The black and white lines in the source-plane image are the tangential and radial caustics of the mass, which again\n", + "you do not need to worry about for now if you don't know what that is!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_tracer(tracer=tracer, grid=grid.mask.derive_grid.all_false)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The tracer is composed of planes. The system above has two planes, an image-plane (at redshift=0.5) and a \n", + "source-plane (at redshift=1.0). \n", + "\n", + "When creating an image via a Tracer, the mass profiles are used to ray-trace the image-plane grid (plotted above) \n", + "to a source-plane grid, via the mass profile's deflection angles.\n", + "\n", + "We can use the tracer`s `traced_grid_2d_list_from` method to calculate and plot the image-plane and source-plane grids." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "aplt.plot_grid(grid=traced_grid_list[0], title=\"Image Plane Grid\")\n", + "\n", + "aplt.plot_grid(grid=traced_grid_list[1], title=\"Source Plane Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extending Objects__\n", + "\n", + "The API has been designed such that all of the objects introduced above are extensible. `Galaxy` \n", + "objects can take many `LightProfile`'s and `MassProfile`'s. `Tracer`' objects can take many `Galaxy`'s. \n", + "\n", + "If the galaxies are at different redshifts a strong lensing system with multiple lens planes will be created, \n", + "performing complex multi-plane ray-tracing calculations.\n", + "\n", + "To finish, lets create a `Tracer` with 3 galaxies at 3 different redshifts, forming a system with two distinct Einstein\n", + "rings! The mass distribution of the first galaxy also has separate components for its stellar mass and dark matter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lmp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.05),\n", + " intensity=0.5,\n", + " effective_radius=0.3,\n", + " sersic_index=3.5,\n", + " mass_to_light_ratio=0.6,\n", + " ),\n", + " disk=al.lmp.Exponential(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.1),\n", + " intensity=1.0,\n", + " effective_radius=2.0,\n", + " mass_to_light_ratio=0.2,\n", + " ),\n", + " dark=al.mp.NFWSph(centre=(0.0, 0.0), kappa_s=0.08, scale_radius=30.0),\n", + ")\n", + "\n", + "lens_galaxy_1 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Exponential(\n", + " centre=(0.00, 0.00),\n", + " ell_comps=(0.05, 0.05),\n", + " intensity=1.2,\n", + " effective_radius=0.1,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0), ell_comps=(0.05, 0.05), einstein_radius=0.3\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=2.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.111111),\n", + " intensity=1.4,\n", + " effective_radius=0.1,\n", + " sersic_index=1.5,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy_0, lens_galaxy_1, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This is what the lens looks like. \n", + "\n", + "Note how crazy the critical curves are!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid), title=\"Image of Complex Strong Lens\"\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Attributes__\n", + "\n", + "Printing individual attributes of the max log likelihood tracer gives us access to the inferred parameters of the\n", + "lens and source galaxies.\n", + "\n", + "The tracer contains the galaxies as both a list and an instance of the model used to fit it. This means we can\n", + "access the same values in two ways, either indexing the galaxies list index or by the name used in model composition.\n", + "\n", + "It can be difficult to track which galaxy is which index in the list, so it is recommended to use the model\n", + "composition to access the galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Einstein Radius via list index = {tracer.galaxies[1].mass.einstein_radius}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lensing Quantities__\n", + "\n", + "The maximum log likelihood tracer contains a lot of information about the inferred model.\n", + "\n", + "For example, by passing it a 2D grid of (y,x) coordinates we can return a numpy array containing its 2D image. This\n", + "includes the lens light and lensed source images.\n", + "\n", + "Below, we use the grid of the `imaging` to computed the image on, which is the grid used to fit to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = tracer.image_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Data Structures Slim / Native__\n", + "\n", + "The image above is returned as a 1D numpy array. \n", + "\n", + "**PyAutoLens** includes dedicated functionality for manipulating this array, for example mapping it to 2D or\n", + "performing the calculation on a high resolution sub-grid which is then binned up. \n", + "\n", + "This uses the data structure API, which is described in the `results/aggregator/data_structures.py` example. This \n", + "tutorial will avoid using this API, but if you need to manipulate results in more detail you should check it out." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(image.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid Choices__\n", + "\n", + "We can input a different grid, which is not masked, to evaluate the image everywhere of interest. We can also change\n", + "the grid's resolution from that used in the model-fit.\n", + "\n", + "The examples uses a grid with `shape_native=(3,3)`. This is much lower resolution than one would typically use to \n", + "perform ray tracing, but is chosen here so that the `print()` statements display in a concise and readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(5, 5), pixel_scales=0.1)\n", + "\n", + "image = tracer.image_2d_from(grid=grid)\n", + "\n", + "print(image.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sub Gridding__\n", + "\n", + "A grid can also have a sub-grid, defined via its `sub_size`, which defines how each pixel on the 2D grid is split \n", + "into sub-pixels of size (`sub_size` x `sub_size`). \n", + "\n", + "The calculation below shows how to use a sub-grid and return an image which has already been binned up. \n", + "\n", + "Full details of the API for this calculation are given in the `guides/over_sampling.py` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=grid.shape_native,\n", + " pixel_scales=grid.pixel_scales,\n", + " over_sample_size=2,\n", + ")\n", + "\n", + "grid_sub = al.Grid2D.uniform(shape_native=(3, 3), pixel_scales=0.1)\n", + "\n", + "image = tracer.image_2d_from(grid=grid_sub)\n", + "\n", + "print(image)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions Grid__\n", + "\n", + "We may want the image at specific (y,x) coordinates.\n", + "\n", + "We can use an irregular 2D (y,x) grid of coordinates for this. The grid below evaluates the image at:\n", + "\n", + "- y = 1.0, x = 1.0.\n", + "- y = 1.0, x = 2.0.\n", + "- y = 2.0, x = 2.0." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_irregular = al.Grid2DIrregular(values=[[1.0, 1.0], [1.0, 2.0], [2.0, 2.0]])\n", + "\n", + "image = tracer.image_2d_from(grid=grid_irregular)\n", + "\n", + "print(image)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scalar Lensing Quantities__\n", + "\n", + "The tracer has many scalar lensing quantities, which are all returned using an `Array2D` and therefore use the same \n", + "interface as images, described above.\n", + "\n", + "For example, we can compute the `Tracer`'s convergence using all of the grids above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "convergence_2d = tracer.convergence_2d_from(grid=grid)\n", + "print(convergence_2d)\n", + "\n", + "convergence_2d = tracer.convergence_2d_from(grid=grid_sub)\n", + "print(convergence_2d)\n", + "\n", + "convergence_2d = tracer.convergence_2d_from(grid=grid_irregular)\n", + "print(convergence_2d)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This is the convergence of every galaxy in the tracer summed together. It may not be appropriate if your lens model \n", + "performs multi-plane ray-tracing (e.g. there are more than 2 redshifts containing galaxies). Later results tutorials\n", + "provide tools that are more appropriate for multi-plane tracers.\n", + "\n", + "There are other scalar quantities accessible via the tracer (those not familiar with strong lensing mathematical \n", + "formalism may not recognise what these quantities are -- don't worry about it for now!):" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "potential_2d = tracer.potential_2d_from(grid=grid)\n", + "\n", + "lens_calc = al.LensCalc.from_tracer(tracer)\n", + "\n", + "tangential_eigen_value = lens_calc.tangential_eigen_value_from(grid=grid)\n", + "radial_eigen_value = lens_calc.radial_eigen_value_from(grid=grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D magnification map is available, which using only the ray-tracing and therefore mass model quantities how much\n", + "light rays are focus at a given point in the image-plane.\n", + "\n", + "If you are studying a strongly lensed source galaxy and want to know how much the galaxy itself is magnified, the\n", + "magnification below is not of too much use too you. In the result tutorial `galaxies.py` we explain how the \n", + "magnification of the source can be quantified." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnification_2d = lens_calc.magnification_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Vector Quantities__\n", + "\n", + "Many lensing quantities are vectors. That is, they are (y,x) coordinates that have 2 values representing their\n", + "magnitudes in both the y and x directions.\n", + "\n", + "These quantities also have a dedicated data structure which is described fully in \n", + "the `results/aggregator/data_structure.py` example.\n", + "\n", + "The most obvious of these is the deflection angles, which are used throughout lens modeling to ray-trace grids\n", + "from the image-plane to the source-plane via a lens galaxy mass model.\n", + "\n", + "To indicate that a quantities is a vector, **PyAutoLens** uses the label `_yx`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_yx_2d = tracer.deflections_yx_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For vector quantities the has shape `2`, corresponding to the y and x vectors respectively." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(deflections_yx_2d[0, :])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `VectorYX2D` object has a built in method to return the magnitude of each vector, which is a scalar quantity\n", + "and therefore returned using a 1D Numpy array." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflection_magnitudes_2d = deflections_yx_2d.magnitudes\n", + "print(deflection_magnitudes_2d)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Other Vector Lensing Quantities__\n", + "\n", + "The tracer has other vector lensing quantities, which use the same interface described above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "shear_yx_2d = lens_calc.shear_yx_2d_via_hessian_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Other Quantities__\n", + "\n", + "Many more quantities are shown below.\n", + "\n", + "A full description of each can be found in the docstring of the source code of each function:\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoGalaxy/blob/main/autogalaxy/operate/deflections.py" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tangential_critical_curve = lens_calc.tangential_critical_curve_list_from(grid=grid)\n", + "\n", + "radial_critical_curve = lens_calc.radial_critical_curve_list_from(grid=grid)\n", + "\n", + "tangential_caustic = lens_calc.tangential_caustic_list_from(grid=grid)\n", + "\n", + "radial_caustic = lens_calc.radial_caustic_list_from(grid=grid)\n", + "\n", + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy_0, lens_galaxy_1]\n", + ") # No support for multi plane tracer time delays yet\n", + "\n", + "time_delay = tracer.time_delays_from(grid=grid)\n", + "\n", + "### You should be able to comment this out and it work fine ###\n", + "\n", + "# area_within_tangential_critical_curve = (\n", + "# lens_calc.tangential_critical_curve_area_list_from(grid=grid)\n", + "# )\n", + "#\n", + "# einstein_radius = tracer.einstein_radius_from(grid=grid)\n", + "#\n", + "# einstein_mass_angular = tracer.einstein_mass_angular_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__JAX__\n", + "\n", + "`Tracer` ray-tracing is the most JAX-friendly part of PyAutoLens \u2014 pure\n", + "numerical kernels with no data-dependent control flow. Typical speedups\n", + "for `tracer.image_2d_from(grid)` and related ops on a large grid:\n", + "10-30\u00d7 for galaxy-scale models, 30-100\u00d7 for cluster-scale on GPU.\n", + "\n", + "You access this in two ways.\n", + "\n", + "__1. The implicit path: `Analysis` and `Simulator`__\n", + "\n", + "`AnalysisImaging(use_jax=True)` (the default) and\n", + "`SimulatorImaging(use_jax=True)` both JAX-accelerate the tracer\n", + "internally. Pytree registration runs as a side effect of the first\n", + "`fit_from` / `via_tracer_from` call; you write nothing JAX-specific.\n", + "\n", + "__2. The explicit path: your own `@jax.jit`__\n", + "\n", + "For parameter sweeps, custom forward models, or batch figure generation\n", + "where you want fine-grained control:\n", + "\n", + "```python\n", + "import jax\n", + "import jax.numpy as jnp\n", + "from autolens.jax import register_tracer_classes\n", + "\n", + "register_tracer_classes(tracer) # one-time pytree registration\n", + "\n", + "@jax.jit\n", + "def image_fn(tracer, grid):\n", + " return tracer.image_2d_from(grid=grid, xp=jnp).array\n", + "\n", + "image = image_fn(tracer, grid)\n", + "```\n", + "\n", + "Two rules:\n", + "\n", + "- **`@jax.jit` + `xp=jnp` pair up.** Forgetting `xp=jnp` either\n", + " silently host-transfers (slow) or fails at the boundary; the library\n", + " now raises a clear `ValueError` on the easy mismatch (see\n", + " `lens_calc.py` for the rationale and `AbstractMaker.__init__`'s\n", + " guard).\n", + "- **`.array` unwrap inside the jit, rewrap outside.** Wrapper types\n", + " (`aa.Array2D`, `aa.Grid2DIrregular`) aren't reliably pytree for\n", + " return-from-JIT \u2014 return raw arrays, rewrap on the host.\n", + "\n", + "__Multi-plane traces under JIT__\n", + "\n", + "The recursive multi-plane lens equation in\n", + "`tracer.traced_grid_2d_list_from(grid)` is pure numerical with no\n", + "data-dependent control flow, so it JITs cleanly. For multi-plane point-\n", + "source solving (forward-solving multiple-image positions through several\n", + "planes), use the higher-level `al.PointSolver(use_jax=True)` \u2014 see\n", + "`scripts/point_source/simulator.py` `__JAX Variant__`.\n", + "\n", + "__Performance expectations__\n", + "\n", + "Tracer image generation on JAX-GPU typically beats NumPy-CPU by:\n", + "\n", + "- 10-30\u00d7 for galaxy-scale (single lens galaxy, single source).\n", + "- 30-100\u00d7 for cluster-scale (many galaxies, multi-plane).\n", + "\n", + "Actual speedup depends on grid size, profile complexity, and GPU\n", + "hardware. `autolens_workspace_developer/jax_profiling/` carries measured\n", + "numbers for representative configurations.\n", + "\n", + "For the full \"JIT-it-yourself\" deep-dive (bound-method form, cache-\n", + "identity considerations, closure-captured `self` vs traced-argument),\n", + "see `scripts/guides/lens_calc.py`. `scripts/guides/galaxies.py` covers\n", + "the pytree registration mechanics. `scripts/guides/data_structures.py`\n", + "covers the `.array` host-transfer story." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "Fin.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/units/cosmology.ipynb b/notebooks/guides/units/cosmology.ipynb index c44671d79..6443e568c 100644 --- a/notebooks/guides/units/cosmology.ipynb +++ b/notebooks/guides/units/cosmology.ipynb @@ -1,468 +1,505 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Units and Cosmology\n", - "===================\n", - "\n", - "This tutorial illustrates how to perform unit conversions from **PyAutoLens**'s internal units (e.g. arc-seconds,\n", - "electrons per second, dimensionless mass units) to physical units (e.g. kiloparsecs, magnitudes, solar masses).\n", - "\n", - "This is used on a variety of important lens model cosmological quantities for example the lens's Einstein radius and\n", - "Mass or the effective radii of the galaxies in the lens model.\n", - "\n", - "__Contents__\n", - "\n", - "- **Errors:** To produce errors on unit converted quantities, you`ll may need to perform marginalization over.\n", - "- **Tracer:** We set up a simple strong lens tracer and grid which will illustrate the unit conversion.\n", - "- **Arcsec to Kiloparsec:** The majority of distance quantities in **PyAutoLens** are in arcseconds, because this means that.\n", - "- **Einstein Radius:** Given a tracer, galaxy or mass profile we can compute its Einstein Radius, which is defined as the.\n", - "- **Einstein Mass:** The Einstein mass can also be computed from a tracer, galaxy or mass profile.\n", - "- **Convergence:** The `colorbar_convert_factor` and `colorbar_label` inputs above can also be used to convert the.\n", - "\n", - "__Errors__\n", - "\n", - "To produce errors on unit converted quantities, you`ll may need to perform marginalization over samples of these\n", - "converted quantities (see `results/aggregator/samples.ipynb`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "We set up a simple strong lens tracer and grid which will illustrate the unit conversion functionality. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.1, 0.0),\n", - " einstein_radius=1.6,\n", - " ),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.1, 0.0),\n", - " intensity=1.0,\n", - " effective_radius=1.0,\n", - " sersic_index=4.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Arcsec to Kiloparsec__\n", - "\n", - "The majority of distance quantities in **PyAutoLens** are in arcseconds, because this means that known redshifts are\n", - "not required in order to compose the lens model.\n", - "\n", - "By assuming redshifts for the lens and source galaxies we can convert their quantities from arcseconds to kiloparsecs.\n", - "\n", - "Below, we compute the effective radii of the source in kiloparsecs. To do this, we assume a cosmology which \n", - "allows us to compute the conversion factor `kpc_per_arcsec`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "cosmology = al.cosmo.Planck15()\n", - "\n", - "source = tracer.planes[1][0]\n", - "source_plane_kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=source.redshift)\n", - "source_effective_radius_kpc = (\n", - " source.bulge.effective_radius * source_plane_kpc_per_arcsec\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This `kpc_per_arcsec` can be used as a conversion factor between arcseconds and kiloparsecs when plotting images of\n", - "galaxies.\n", - "\n", - "Below, we compute this value in both the image-plane and source-plane, and plot the images in both planes in their\n", - "respectively converted units of kilo-parsec.\n", - "\n", - "This passes the plotting modules `Units` object a `ticks_convert_factor` and manually specified the new units of the\n", - "plot ticks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = tracer.planes[0][0]\n", - "image_plane_kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=lens.redshift)\n", - "\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", - "\n", - "\n", - "aplt.plot_grid(\n", - " grid=tracer.traced_grid_2d_list_from(grid=grid)[-1], title=\"Source Plane\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Einstein Radius__\n", - "\n", - "Given a tracer, galaxy or mass profile we can compute its Einstein Radius, which is defined as the area within the \n", - "tangential critical curve. \n", - "\n", - "These are calculated from the functions: \n", - "\n", - " - `einstein_radius_from`. \n", - " - `einstein_mass_via_tangential_critical_curve`.\n", - "\n", - "Although these quantities should not depend on the grid we input, they are calculated using the input grid. Thus,\n", - "we must specify a grid which matches the scale of the lens model, which would typically be the grid of image-pixels\n", - "that we use to model our data.\n", - "\n", - "Lets print the Einstein Radius, which is returned in the default internal **PyAutoLens** units of arc-seconds." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.1)\n", - "lens_calc = al.LensCalc.from_tracer(tracer=tracer)\n", - "einstein_radius = lens_calc.einstein_radius_from(grid=grid)\n", - "\n", - "print(\"Einstein Radius (arcsec) = \", einstein_radius)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we know the redshift of the lens galaxy and assume an cosmology we can convert this to kilo-parsecs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "cosmology = al.cosmo.Planck15()\n", - "\n", - "kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=tracer.planes[0].redshift)\n", - "einstein_radius_kpc = einstein_radius * kpc_per_arcsec\n", - "print(\"Einstein Radius (kpc) = \", einstein_radius_kpc)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also compute the Einstein radius of individual planes, galaxies and mass profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " al.LensCalc.from_mass_obj(mass_obj=tracer.planes[0][0]).einstein_radius_from(\n", - " grid=grid\n", - " )\n", - ")\n", - "print(\n", - " al.LensCalc.from_mass_obj(mass_obj=tracer.planes[0][0].mass).einstein_radius_from(\n", - " grid=grid\n", - " )\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Einstein Mass__\n", - "\n", - "The Einstein mass can also be computed from a tracer, galaxy or mass profile.\n", - "\n", - "The default units of an Einstein mass are angular units; this is because to convert it to physical units (e.g. solar\n", - "masses) one must assume redsfhits for the lens and source galaxies.\n", - "\n", - "The mass in angular units is given by: `pi * einstein_radius (arcsec) ** 2.0`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "einstein_mass = lens_calc.einstein_mass_angular_from(grid=grid)\n", - "print(\"Einstein Mass (angular) = \", einstein_mass)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To convert this mass to solar masses, we need the critical surface mass density of the strong lens, which relies on \n", - "it being a strong lens with not only a lens redshift (e.g. the redshift of the profile) but also a source redshift.\n", - "\n", - "If we use the `tracer`'s galaxies for the redshifts, where the lens is at redshift 0.5 and it is lensing a source at \n", - "redshift 1.0, we can compute its mass in solar masses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "cosmology = al.cosmo.Planck15()\n", - "\n", - "critical_surface_density = cosmology.critical_surface_density_between_redshifts_from(\n", - " redshift_0=tracer.planes[0].redshift, redshift_1=tracer.planes[1].redshift\n", - ")\n", - "einstein_mass_solar_mass = einstein_mass * critical_surface_density\n", - "print(\"Einstein Mass (solMass) = \", einstein_mass_solar_mass)\n", - "print(\"Einstein Mass (solMass) = \", \"{:.4e}\".format(einstein_mass_solar_mass))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can compute Einstein masses of individual components:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " al.LensCalc.from_mass_obj(mass_obj=tracer.planes[0][0]).einstein_mass_angular_from(\n", - " grid=grid\n", - " )\n", - ")\n", - "print(\n", - " al.LensCalc.from_mass_obj(\n", - " mass_obj=tracer.planes[0][0].mass\n", - " ).einstein_mass_angular_from(grid=grid)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In principle, the Einstein Mass of a `Tracer` should be readily accessible in a `Tracer` object, given this contains\n", - "all of the galaxies in a strong lens system (and thus has their redshifts) as well as an input Cosmology.\n", - "\n", - "However, we do not provide methods with this quantity and require that you, the user, compute the Einstein mass \n", - "(in angular or solar masses) using examples above. This is because for systems with multiple galaxies or planes, the \n", - "definition of an Einstein Radius / Mass become less clear. \n", - "\n", - "We feel it is better that a user explicitly computes these quantities from a `Tracer` so if it has multiple galaxies \n", - "or planes you are aware of this.\n", - "\n", - "__Brightness Units / Luminosity__\n", - "\n", - "When plotting the image of a galaxy, each pixel value is also plotted in electrons / second, which is the unit values\n", - "displayed in the colorbar. \n", - "\n", - "A conversion factor between electrons per second and another unit can be input when plotting images of galaxies.\n", - "\n", - "Below, we pass the exposure time of the image, which converts the units of the image from `electrons / second` to\n", - "electrons. \n", - "\n", - "Note that this input `ticks_convert_factor_values` is the same input parameter used above to convert mass plots like the \n", - "convergence to physical units." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "exposure_time_seconds = 2000.0\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The luminosity of a galaxy is the total amount of light it emits, which is computed by integrating the light profile.\n", - "This integral is performed over the entire light profile, or within a specified radius.\n", - "\n", - "Lets compute the luminosity of the source galaxy in the default internal **PyAutoLens** units of `electrons / second`.\n", - "Below, we compute the luminosity to infinite radius, which is the total luminosity of the galaxy, but one could\n", - "easily compute the luminosity within a specified radius instead." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source = tracer.planes[1][0]\n", - "\n", - "luminosity = source.luminosity_within_circle_from(radius=np.inf)\n", - "print(\"Luminosity (electrons / second) = \", luminosity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From a luminosity in `electrons / second`, we can convert it to other units, such as `Jansky` or `erg / second`. \n", - "This can also be used to compute the magnitude of the galaxy, which is the apparent brightness of the galaxy in a\n", - "given bandpass.\n", - "\n", - "This functionality is not currently implemented in **PyAutoLens**, but would be fairly simple for you to do\n", - "yourself (e.g. using the `astropy` package). If you want to contribute to **PyAutoLens**, this would be a great\n", - "first issue to tackle, so please get in touch on SLACK!\n", - "\n", - "__Convergence__\n", - "\n", - "The `colorbar_convert_factor` and `colorbar_label` inputs above can also be used to convert the units of mass\n", - "profiles images. \n", - "\n", - "For example, we can convert the convergence from its dimensionless lensing units to a physical surface density\n", - "in units of solar masses per kpc^2." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "critical_surface_density = cosmology.critical_surface_density_between_redshifts_from(\n", - " redshift_0=tracer.planes[0].redshift, redshift_1=tracer.planes[1].redshift\n", - ")\n", - "\n", - "convergence = tracer.convergence_2d_from(grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "With the convergence in units of MSun / kpc^2, we can easily compute the total mass associated with it in a specifc\n", - "area.\n", - "\n", - "For example, in a single pixel of convergence in these units, we can compute the mass by simply multiplying it by the\n", - "area of the pixel in kpc^2." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_area_kpc = (\n", - " grid.pixel_scales[0] * grid.pixel_scales[1] * image_plane_kpc_per_arcsec**2\n", - ")\n", - "\n", - "print(\n", - " f\"Total mass in central pixel: {convergence.native[convergence.shape_native[0] // 2, convergence.shape_native[1] // 2] * critical_surface_density * pixel_area_kpc} MSun\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The total mass of the convergence map is the sum of all these masses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"Total mass in convergence map: {np.sum(convergence * critical_surface_density * pixel_area_kpc)} MSun\"\n", - ")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Units and Cosmology\n", + "===================\n", + "\n", + "This tutorial illustrates how to perform unit conversions from **PyAutoLens**'s internal units (e.g. arc-seconds,\n", + "electrons per second, dimensionless mass units) to physical units (e.g. kiloparsecs, magnitudes, solar masses).\n", + "\n", + "This is used on a variety of important lens model cosmological quantities for example the lens's Einstein radius and\n", + "Mass or the effective radii of the galaxies in the lens model.\n", + "\n", + "__Contents__\n", + "\n", + "- **Errors:** To produce errors on unit converted quantities, you`ll may need to perform marginalization over.\n", + "- **Tracer:** We set up a simple strong lens tracer and grid which will illustrate the unit conversion.\n", + "- **Arcsec to Kiloparsec:** The majority of distance quantities in **PyAutoLens** are in arcseconds, because this means that.\n", + "- **Einstein Radius:** Given a tracer, galaxy or mass profile we can compute its Einstein Radius, which is defined as the.\n", + "- **Einstein Mass:** The Einstein mass can also be computed from a tracer, galaxy or mass profile.\n", + "- **Convergence:** The `colorbar_convert_factor` and `colorbar_label` inputs above can also be used to convert the.\n", + "\n", + "__Errors__\n", + "\n", + "To produce errors on unit converted quantities, you`ll may need to perform marginalization over samples of these\n", + "converted quantities (see `results/aggregator/samples.ipynb`)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "We set up a simple strong lens tracer and grid which will illustrate the unit conversion functionality. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.1, 0.0),\n", + " einstein_radius=1.6,\n", + " ),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.1, 0.0),\n", + " intensity=1.0,\n", + " effective_radius=1.0,\n", + " sersic_index=4.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Arcsec to Kiloparsec__\n", + "\n", + "The majority of distance quantities in **PyAutoLens** are in arcseconds, because this means that known redshifts are\n", + "not required in order to compose the lens model.\n", + "\n", + "By assuming redshifts for the lens and source galaxies we can convert their quantities from arcseconds to kiloparsecs.\n", + "\n", + "Below, we compute the effective radii of the source in kiloparsecs. To do this, we assume a cosmology which \n", + "allows us to compute the conversion factor `kpc_per_arcsec`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cosmology = al.cosmo.Planck15()\n", + "\n", + "source = tracer.planes[1][0]\n", + "source_plane_kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=source.redshift)\n", + "source_effective_radius_kpc = (\n", + " source.bulge.effective_radius * source_plane_kpc_per_arcsec\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `kpc_per_arcsec` can be used as a conversion factor between arcseconds and kiloparsecs when plotting images of\n", + "galaxies.\n", + "\n", + "Below, we compute this value in both the image-plane and source-plane, and plot the images in both planes in their\n", + "respectively converted units of kilo-parsec.\n", + "\n", + "This passes the plotting modules `Units` object a `ticks_convert_factor` and manually specified the new units of the\n", + "plot ticks." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = tracer.planes[0][0]\n", + "image_plane_kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=lens.redshift)\n", + "\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", + "\n", + "\n", + "aplt.plot_grid(\n", + " grid=tracer.traced_grid_2d_list_from(grid=grid)[-1], title=\"Source Plane\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Einstein Radius__\n", + "\n", + "Given a tracer, galaxy or mass profile we can compute its Einstein Radius, which is defined as the area within the \n", + "tangential critical curve. \n", + "\n", + "These are calculated from the functions: \n", + "\n", + " - `einstein_radius_from`. \n", + " - `einstein_mass_via_tangential_critical_curve`.\n", + "\n", + "Although these quantities should not depend on the grid we input, they are calculated using the input grid. Thus,\n", + "we must specify a grid which matches the scale of the lens model, which would typically be the grid of image-pixels\n", + "that we use to model our data.\n", + "\n", + "Lets print the Einstein Radius, which is returned in the default internal **PyAutoLens** units of arc-seconds." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.1)\n", + "lens_calc = al.LensCalc.from_tracer(tracer=tracer)\n", + "einstein_radius = lens_calc.einstein_radius_from(grid=grid)\n", + "\n", + "print(\"Einstein Radius (arcsec) = \", einstein_radius)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we know the redshift of the lens galaxy and assume an cosmology we can convert this to kilo-parsecs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cosmology = al.cosmo.Planck15()\n", + "\n", + "kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=tracer.planes[0].redshift)\n", + "einstein_radius_kpc = einstein_radius * kpc_per_arcsec\n", + "print(\"Einstein Radius (kpc) = \", einstein_radius_kpc)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also compute the Einstein radius of individual planes, galaxies and mass profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " al.LensCalc.from_mass_obj(mass_obj=tracer.planes[0][0]).einstein_radius_from(\n", + " grid=grid\n", + " )\n", + ")\n", + "print(\n", + " al.LensCalc.from_mass_obj(mass_obj=tracer.planes[0][0].mass).einstein_radius_from(\n", + " grid=grid\n", + " )\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Einstein Mass__\n", + "\n", + "The Einstein mass can also be computed from a tracer, galaxy or mass profile.\n", + "\n", + "The default units of an Einstein mass are angular units; this is because to convert it to physical units (e.g. solar\n", + "masses) one must assume redsfhits for the lens and source galaxies.\n", + "\n", + "The mass in angular units is given by: `pi * einstein_radius (arcsec) ** 2.0`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "einstein_mass = lens_calc.einstein_mass_angular_from(grid=grid)\n", + "print(\"Einstein Mass (angular) = \", einstein_mass)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To convert this mass to solar masses, we need the critical surface mass density of the strong lens, which relies on \n", + "it being a strong lens with not only a lens redshift (e.g. the redshift of the profile) but also a source redshift.\n", + "\n", + "If we use the `tracer`'s galaxies for the redshifts, where the lens is at redshift 0.5 and it is lensing a source at \n", + "redshift 1.0, we can compute its mass in solar masses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cosmology = al.cosmo.Planck15()\n", + "\n", + "critical_surface_density = cosmology.critical_surface_density_between_redshifts_from(\n", + " redshift_0=tracer.planes[0].redshift, redshift_1=tracer.planes[1].redshift\n", + ")\n", + "einstein_mass_solar_mass = einstein_mass * critical_surface_density\n", + "print(\"Einstein Mass (solMass) = \", einstein_mass_solar_mass)\n", + "print(\"Einstein Mass (solMass) = \", \"{:.4e}\".format(einstein_mass_solar_mass))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can compute Einstein masses of individual components:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " al.LensCalc.from_mass_obj(mass_obj=tracer.planes[0][0]).einstein_mass_angular_from(\n", + " grid=grid\n", + " )\n", + ")\n", + "print(\n", + " al.LensCalc.from_mass_obj(\n", + " mass_obj=tracer.planes[0][0].mass\n", + " ).einstein_mass_angular_from(grid=grid)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In principle, the Einstein Mass of a `Tracer` should be readily accessible in a `Tracer` object, given this contains\n", + "all of the galaxies in a strong lens system (and thus has their redshifts) as well as an input Cosmology.\n", + "\n", + "However, we do not provide methods with this quantity and require that you, the user, compute the Einstein mass \n", + "(in angular or solar masses) using examples above. This is because for systems with multiple galaxies or planes, the \n", + "definition of an Einstein Radius / Mass become less clear. \n", + "\n", + "We feel it is better that a user explicitly computes these quantities from a `Tracer` so if it has multiple galaxies \n", + "or planes you are aware of this.\n", + "\n", + "__Brightness Units / Luminosity__\n", + "\n", + "When plotting the image of a galaxy, each pixel value is also plotted in electrons / second, which is the unit values\n", + "displayed in the colorbar. \n", + "\n", + "A conversion factor between electrons per second and another unit can be input when plotting images of galaxies.\n", + "\n", + "Below, we pass the exposure time of the image, which converts the units of the image from `electrons / second` to\n", + "electrons. \n", + "\n", + "Note that this input `ticks_convert_factor_values` is the same input parameter used above to convert mass plots like the \n", + "convergence to physical units." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "exposure_time_seconds = 2000.0\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The luminosity of a galaxy is the total amount of light it emits, which is computed by integrating the light profile.\n", + "This integral is performed over the entire light profile, or within a specified radius.\n", + "\n", + "Lets compute the luminosity of the source galaxy in the default internal **PyAutoLens** units of `electrons / second`.\n", + "Below, we compute the luminosity to infinite radius, which is the total luminosity of the galaxy, but one could\n", + "easily compute the luminosity within a specified radius instead." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source = tracer.planes[1][0]\n", + "\n", + "luminosity = source.luminosity_within_circle_from(radius=np.inf)\n", + "print(\"Luminosity (electrons / second) = \", luminosity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From a luminosity in `electrons / second`, we can convert it to other units, such as `Jansky` or `erg / second`. \n", + "This can also be used to compute the magnitude of the galaxy, which is the apparent brightness of the galaxy in a\n", + "given bandpass.\n", + "\n", + "This functionality is not currently implemented in **PyAutoLens**, but would be fairly simple for you to do\n", + "yourself (e.g. using the `astropy` package). If you want to contribute to **PyAutoLens**, this would be a great\n", + "first issue to tackle, so please get in touch on SLACK!\n", + "\n", + "__Convergence__\n", + "\n", + "The `colorbar_convert_factor` and `colorbar_label` inputs above can also be used to convert the units of mass\n", + "profiles images. \n", + "\n", + "For example, we can convert the convergence from its dimensionless lensing units to a physical surface density\n", + "in units of solar masses per kpc^2." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "critical_surface_density = cosmology.critical_surface_density_between_redshifts_from(\n", + " redshift_0=tracer.planes[0].redshift, redshift_1=tracer.planes[1].redshift\n", + ")\n", + "\n", + "convergence = tracer.convergence_2d_from(grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "With the convergence in units of MSun / kpc^2, we can easily compute the total mass associated with it in a specifc\n", + "area.\n", + "\n", + "For example, in a single pixel of convergence in these units, we can compute the mass by simply multiplying it by the\n", + "area of the pixel in kpc^2." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_area_kpc = (\n", + " grid.pixel_scales[0] * grid.pixel_scales[1] * image_plane_kpc_per_arcsec**2\n", + ")\n", + "\n", + "print(\n", + " f\"Total mass in central pixel: {convergence.native[convergence.shape_native[0] // 2, convergence.shape_native[1] // 2] * critical_surface_density * pixel_area_kpc} MSun\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The total mass of the convergence map is the sum of all these masses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"Total mass in convergence map: {np.sum(convergence * critical_surface_density * pixel_area_kpc)} MSun\"\n", + ")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/units/flux.ipynb b/notebooks/guides/units/flux.ipynb index daf7e3327..14fec7b1e 100644 --- a/notebooks/guides/units/flux.ipynb +++ b/notebooks/guides/units/flux.ipynb @@ -1,468 +1,505 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Flux\n", - "====\n", - "\n", - "Absolute flux calibration in Astronomy is the process of converting the number of photons detected by a telescope into\n", - "a physical unit of luminosity or a magnitude. For example, a luminosity might be given in units of solar luminosities\n", - "or the brightness of a galaxy quoted as a magnitude in units of AB magnitudes.\n", - "\n", - "The conversion of a light profile, that has been fitted to data, to physical units can be non-trivial, as careful\n", - "consideration must be given to the units that are involved.\n", - "\n", - "The key quantity is the `intensity` of the light profile, the units of which match the units of the data that is fitted.\n", - "For example, if the data is in units of electrons per second, the intensity will also be in units of electrons per\n", - "second per pixel.\n", - "\n", - "The conversion of this intensity to a physical unit, like solar luminosities, therefore requires us to make a number\n", - "of conversion steps that go from electrons per second to the desired physical unit or magnitude.\n", - "\n", - "This guide gives example conversions for units commonly used in astronomy, such as converting the intensity of a\n", - "light profile from electrons per second to solar luminosities or AB magnitudes. Once we have values in a more standard\n", - "unit, like a solar luminosity or AB magnitude, it becomes a lot more straightforward to follow Astropy tutorials\n", - "(or other resources) to convert these values to other units or perform calculations with them.\n", - "\n", - "__Contents__\n", - "\n", - "- **Zero Point:** In astronomy, a zero point refers to a reference value used in photometry and spectroscopy to.\n", - "- **Total Flux:** A key quantity for computing the magnitudes of galaxies is the total flux of a light profile.\n", - "- **Latent Variables:** Reading the same total flux directly from the `latent.csv` of a completed fit.\n", - "\n", - "__Zero Point__\n", - "\n", - "In astronomy, a zero point refers to a reference value used in photometry and spectroscopy to calibrate the brightness\n", - "of celestial objects. It sets the baseline for a magnitude system, allowing astronomers to compare the brightness of\n", - "different stars, galaxies, or other objects.\n", - "\n", - "For example, the zero point in a photometric system corresponds to the magnitude that a standard star (or a theoretical\n", - "object) would have if it produced a specific amount of light at a particular wavelength. It provides a way to convert\n", - "the raw measurements of light received by a telescope into meaningful values of brightness (magnitudes).\n", - "\n", - "The conversions below all require a zero point, which is typically provided in the documentation of the telescope or\n", - "instrument that was used to observe the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "from scipy.special import gamma\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Total Flux__\n", - "\n", - "A key quantity for computing the magnitudes of galaxies is the total flux of a light profile.\n", - "\n", - "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", - "we compute the image of the light profile, and then sum the image. \n", - "\n", - "The units of the light profile `intensity` are the units of the data the light profile was fitted to. For example, \n", - "HST data is often electrons per second, so the intensity is in units of electrons per second per pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "light = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=2.0, # in units of e- pix^-1 s^-1, representative of HST data in units of electrons per second\n", - " effective_radius=0.1,\n", - " sersic_index=3.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The total flux, in units of electrons per second, is computed by summing the image of the light profile over all pixels.\n", - "\n", - "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", - "and high enough resolution that it captures all the light from the light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "image = light.image_2d_from(grid=grid)\n", - "\n", - "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For a spherical Sersic function, there is an analytic expression for the total flux, shown below.\n", - "\n", - "However, because the light profile is in units of pix^-1, the total flux computed via this expression is in slightly\n", - "strange units we need to account for afterwards of e- s^-1 arcsec^2 pix^-1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_flux_strange_units = (\n", - " light.intensity\n", - " * (light.effective_radius**2)\n", - " * 2\n", - " * np.pi\n", - " * light.sersic_index\n", - " * (\n", - " np.exp(light.sersic_constant)\n", - " / (light.sersic_constant ** (2 * light.sersic_index))\n", - " )\n", - " * gamma(2 * light.sersic_index)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To get the total flux in units of e- s^-1, we divide by the total grid area (in arcsec^2) and multiply by the total\n", - "number of pixels, which are provided by the grid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_flux = (total_flux_strange_units / grid.total_area) * grid.total_pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The two calculations come out very close to one another, and become closer the more pixels we use in the grid to \n", - "compute the total flux.\n", - "\n", - "If possible, you should use analytic expressions to compute the total flux of a light profile, as this is exact, \n", - "especially if computing magnitudes precisely is important for your science case.\n", - "\n", - "However, for many light profiles the total flux cannot easily be computed analytically, and the summed image approach\n", - "sufficient.\n", - "\n", - "__Mega Janskys / steradian (MJy/sr): James Webb Space Telescope__\n", - "\n", - "James Webb Space Telescope (JWST) NIRCam data is often provided in units of Mega Janskys per steradian (MJy/sr).\n", - "We therefore show how to convert the intensity of a light profile from MJy/sr to absolute AB magnitudes.\n", - "\n", - "This calculation is well documented in the JWST documentation, and we are following the steps in the following\n", - "webpage:\n", - "\n", - "https://jwst-docs.stsci.edu/jwst-near-infrared-camera/nircam-performance/nircam-absolute-flux-calibration-and-zeropoints#gsc.tab=0\n", - "\n", - "First, we need a light profile, which we'll assume is a Sersic profilee. If you're analyzing real JWST data, you'll\n", - "need to use the light profile that was fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "light = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=2.0, # in units of MJy sr^-1 pix^-1\n", - " effective_radius=0.1,\n", - " sersic_index=3.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "According to the document above, flux density in MJy/sr can be converted to AB magnitude using the following formula:\n", - "\n", - " mag_AB = -6.10 - 2.5 * log10(flux[MJy/sr]*PIXAR_SR[sr/pix] ) = ZP_AB \u2013 2.5 log10(flux[MJy/sr])\n", - "\n", - "Where ZP_AB is the zeropoint: \n", - "\n", - " ZP_AB = \u20136.10 \u2013 2.5 log10(PIXAR_SR[sr/pix]). \n", - "\n", - "For example, ZP_AB = 28.0 for PIXAR_SR = 2.29e-14 (corresponding to pixel size 0.0312\").\n", - "\n", - "For data in units of MJy/sr, computing the total flux that goes into the log10 term is straightforward, it is\n", - "simply the sum of the image of the light profile. \n", - "\n", - "We compute this using a grid, which must be large enough that all light from the light profile is included. Below,\n", - "we use a grid which extends to 10\" from the centre of the light profile, which is sufficient for this example,\n", - "but you may need to increase this size for your own data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "image = light.image_2d_from(grid=grid)\n", - "\n", - "total_flux = np.sum(image) # In units of MJy sr^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now convert this total flux to an AB magnitude using the zero point of the JWST NIRCam filter we are analyzing.\n", - "\n", - "As stated above, the zero point is given by:\n", - "\n", - " ZP_AB = \u20136.10 \u2013 2.5 log10(PIXAR_SR[sr/pix])\n", - " \n", - "Where the value of PIXAR_SR is provided in the JWST documentation for the filter you are analyzing. \n", - "\n", - "The Pixar_SR values for JWST (James Webb Space Telescope) NIRCam filters refer to the pixel scale in steradians (sr) \n", - "for each filter, which is a measure of the solid angle covered by each pixel. These values are important for \n", - "calibrating and understanding how light is captured by the instrument.\n", - "\n", - "For the F444W filter, which we are using in this example, the value is 2.29e-14 (corresponding to a pixel size o\n", - "f 0.0312\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixar_sr = 2.29e-14\n", - "\n", - "zero_point = -6.10 - 2.5 * np.log10(pixar_sr)\n", - "\n", - "magnitude_ab = zero_point - 2.5 * np.log10(total_flux)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "With an absolute magnitude and quantity of light in physical units, you should now be able to convert these values to\n", - "whatever units you need for your science case.\n", - "\n", - "__Electrons Per Second (e s^-1): Hubble Space Telescope__\n", - "\n", - "Hubble Space Telescope (HST) data is often provided in units of electrons per second (e- s^-1). \n", - "\n", - "We therefore show how to convert the intensity of a light profile from electrons per second to an absolute magnitude." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "light = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=2.0, # in units of e- pix^-1 s^-1\n", - " effective_radius=0.1,\n", - " sersic_index=3.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We first compute the total flux in electrons per second by summing the image of the light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "image = light.image_2d_from(grid=grid) # in units e- s^-1 as summed over pixels\n", - "\n", - "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now use the zero point of the HST filter we are analyzing to convert this total flux to an AB magnitude.\n", - "\n", - "The zero point for the F814W filter, which we are using in this example, is 25.943.\n", - "\n", - "Zero points of the HST ACS filter are provided here: https://acszeropoints.stsci.edu, for other filters you should\n", - "consult the HST documentation.\n", - "\n", - "The zero point below is defined in units such that it converts the total flux from input units of electrons per second,\n", - "you should make sure your HST data is in these units and that the zero point you are using follows the same convention." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "zero_point_f814w = 25.943\n", - "\n", - "magnitude_ab = zero_point_f814w - 2.5 * np.log10(total_flux)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "With an absolute magnitude and quantity of light in physical units, you should now be able to convert these values to\n", - "whatever units you need for your science case.\n", - "\n", - "For HST, a few quantitites that may be useful and worth looking into are:\n", - "\n", - "- The HST PHOTFLAM value, in units of erg cm^-2 s^-1 A^-1 e-^-1, which is used to convert to ergs, which radio \n", - " astronomers may be interested in.\n", - " \n", - "- The HST PHOTNU value, in units of Jy (e s^-1), which converts to Janskys, which is often used by SED fitting\n", - " software.\n", - "\n", - "__Latent Variables: Total Flux Directly from the Fit__\n", - "\n", - "The examples above all computed total flux by hand: build a light profile, sample it on a grid, sum the image, then\n", - "apply the zero point. PyAutoLens does exactly this automatically as part of every fit and records the result as a\n", - "latent variable in the `latent/samples.csv` file beside the search output. You can skip the manual recipe entirely\n", - "and just read the column.\n", - "\n", - "Three lensing flux latents ship default-on (they need no instrument inputs and run on every fit unless disabled in\n", - "`config/latent.yaml`):\n", - "\n", - "- `total_lens_flux` \u2014 total integrated flux of the lens galaxy (the sum of `fit.tracer.galaxies[0]`'s model image),\n", - " in the *raw* image units the fit was performed in. For HST data in e- s^-1, this is e- s^-1; for JWST data in\n", - " MJy/sr, this is MJy/sr.\n", - "\n", - "- `total_lensed_source_flux` \u2014 image-plane integrated flux of the source galaxy after lensing. This is the source\n", - " flux as observed: stretched, multiplied, and distorted by the lens. Same raw units.\n", - "\n", - "- `total_source_flux` \u2014 source-plane intrinsic flux of the source galaxy: the flux the source would have if the\n", - " lens weren't there. Computed via `fit.tracer_linear_light_profiles_to_light_profiles` so that linear light\n", - " profiles (MGEs, pixelizations) work correctly. The ratio\n", - " `total_lensed_source_flux / total_source_flux` is the integrated magnification (also recorded directly as the\n", - " `magnification` latent).\n", - "\n", - "To convert any of these to AB magnitudes or microjanskies, apply the same zero-point recipe used above. Suppose\n", - "you have an HST F814W fit with the F814W zero point of 25.943; reading the lens galaxy flux from your result and\n", - "converting goes:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autogalaxy.imaging.model.latent import (\n", - " ab_mag_via_flux_from,\n", - " flux_mujy_via_ab_mag_from,\n", - ")\n", - "\n", - "# Stand-in for what you'd read from `latent.csv` \u2014 in a real script this is one column of one row, e.g.\n", - "# total_lens_flux = pd.read_csv(search.paths.output_path / \"latent\" / \"samples.csv\")[\"total_lens_flux\"].iloc[-1]\n", - "total_lens_flux = 1234.5 # e- s^-1\n", - "\n", - "zero_point_f814w = 25.943\n", - "ab_mag_lens = ab_mag_via_flux_from(flux=total_lens_flux, magzero=zero_point_f814w)\n", - "flux_mujy_lens = flux_mujy_via_ab_mag_from(ab_mag=ab_mag_lens)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The two helpers used above are the same ones the library uses internally to populate the `_mujy` variants of the\n", - "latents (`total_lens_flux_mujy`, `total_lensed_source_flux_mujy`, `total_source_flux_mujy`). Those variants are\n", - "default-off because they need a `magzero` you supply per-instrument. If you have a single fixed zero-point you can\n", - "flip them on by:\n", - "\n", - "1. Setting `total_lens_flux_mujy: true` (and friends) in your project's `config/latent.yaml`.\n", - "2. Passing `magzero=` when constructing the analysis:\n", - " `analysis = al.AnalysisImaging(dataset=dataset, magzero=25.943)`.\n", - "\n", - "The latent dispatcher then writes the converted \u00b5Jy columns into `latent/samples.csv` directly, so you don't have\n", - "to run the conversion in post. If you enable the `_mujy` latents but forget the `magzero` keyword, the columns are\n", - "populated with NaN and a single warning per process notes that the conversion was skipped \u2014 the fit itself is\n", - "unaffected.\n", - "\n", - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Flux\n", + "====\n", + "\n", + "Absolute flux calibration in Astronomy is the process of converting the number of photons detected by a telescope into\n", + "a physical unit of luminosity or a magnitude. For example, a luminosity might be given in units of solar luminosities\n", + "or the brightness of a galaxy quoted as a magnitude in units of AB magnitudes.\n", + "\n", + "The conversion of a light profile, that has been fitted to data, to physical units can be non-trivial, as careful\n", + "consideration must be given to the units that are involved.\n", + "\n", + "The key quantity is the `intensity` of the light profile, the units of which match the units of the data that is fitted.\n", + "For example, if the data is in units of electrons per second, the intensity will also be in units of electrons per\n", + "second per pixel.\n", + "\n", + "The conversion of this intensity to a physical unit, like solar luminosities, therefore requires us to make a number\n", + "of conversion steps that go from electrons per second to the desired physical unit or magnitude.\n", + "\n", + "This guide gives example conversions for units commonly used in astronomy, such as converting the intensity of a\n", + "light profile from electrons per second to solar luminosities or AB magnitudes. Once we have values in a more standard\n", + "unit, like a solar luminosity or AB magnitude, it becomes a lot more straightforward to follow Astropy tutorials\n", + "(or other resources) to convert these values to other units or perform calculations with them.\n", + "\n", + "__Contents__\n", + "\n", + "- **Zero Point:** In astronomy, a zero point refers to a reference value used in photometry and spectroscopy to.\n", + "- **Total Flux:** A key quantity for computing the magnitudes of galaxies is the total flux of a light profile.\n", + "- **Latent Variables:** Reading the same total flux directly from the `latent.csv` of a completed fit.\n", + "\n", + "__Zero Point__\n", + "\n", + "In astronomy, a zero point refers to a reference value used in photometry and spectroscopy to calibrate the brightness\n", + "of celestial objects. It sets the baseline for a magnitude system, allowing astronomers to compare the brightness of\n", + "different stars, galaxies, or other objects.\n", + "\n", + "For example, the zero point in a photometric system corresponds to the magnitude that a standard star (or a theoretical\n", + "object) would have if it produced a specific amount of light at a particular wavelength. It provides a way to convert\n", + "the raw measurements of light received by a telescope into meaningful values of brightness (magnitudes).\n", + "\n", + "The conversions below all require a zero point, which is typically provided in the documentation of the telescope or\n", + "instrument that was used to observe the data." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "from scipy.special import gamma\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Total Flux__\n", + "\n", + "A key quantity for computing the magnitudes of galaxies is the total flux of a light profile.\n", + "\n", + "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", + "we compute the image of the light profile, and then sum the image. \n", + "\n", + "The units of the light profile `intensity` are the units of the data the light profile was fitted to. For example, \n", + "HST data is often electrons per second, so the intensity is in units of electrons per second per pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "light = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=2.0, # in units of e- pix^-1 s^-1, representative of HST data in units of electrons per second\n", + " effective_radius=0.1,\n", + " sersic_index=3.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The total flux, in units of electrons per second, is computed by summing the image of the light profile over all pixels.\n", + "\n", + "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", + "and high enough resolution that it captures all the light from the light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "image = light.image_2d_from(grid=grid)\n", + "\n", + "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For a spherical Sersic function, there is an analytic expression for the total flux, shown below.\n", + "\n", + "However, because the light profile is in units of pix^-1, the total flux computed via this expression is in slightly\n", + "strange units we need to account for afterwards of e- s^-1 arcsec^2 pix^-1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_flux_strange_units = (\n", + " light.intensity\n", + " * (light.effective_radius**2)\n", + " * 2\n", + " * np.pi\n", + " * light.sersic_index\n", + " * (\n", + " np.exp(light.sersic_constant)\n", + " / (light.sersic_constant ** (2 * light.sersic_index))\n", + " )\n", + " * gamma(2 * light.sersic_index)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To get the total flux in units of e- s^-1, we divide by the total grid area (in arcsec^2) and multiply by the total\n", + "number of pixels, which are provided by the grid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_flux = (total_flux_strange_units / grid.total_area) * grid.total_pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The two calculations come out very close to one another, and become closer the more pixels we use in the grid to \n", + "compute the total flux.\n", + "\n", + "If possible, you should use analytic expressions to compute the total flux of a light profile, as this is exact, \n", + "especially if computing magnitudes precisely is important for your science case.\n", + "\n", + "However, for many light profiles the total flux cannot easily be computed analytically, and the summed image approach\n", + "sufficient.\n", + "\n", + "__Mega Janskys / steradian (MJy/sr): James Webb Space Telescope__\n", + "\n", + "James Webb Space Telescope (JWST) NIRCam data is often provided in units of Mega Janskys per steradian (MJy/sr).\n", + "We therefore show how to convert the intensity of a light profile from MJy/sr to absolute AB magnitudes.\n", + "\n", + "This calculation is well documented in the JWST documentation, and we are following the steps in the following\n", + "webpage:\n", + "\n", + "https://jwst-docs.stsci.edu/jwst-near-infrared-camera/nircam-performance/nircam-absolute-flux-calibration-and-zeropoints#gsc.tab=0\n", + "\n", + "First, we need a light profile, which we'll assume is a Sersic profilee. If you're analyzing real JWST data, you'll\n", + "need to use the light profile that was fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "light = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=2.0, # in units of MJy sr^-1 pix^-1\n", + " effective_radius=0.1,\n", + " sersic_index=3.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "According to the document above, flux density in MJy/sr can be converted to AB magnitude using the following formula:\n", + "\n", + " mag_AB = -6.10 - 2.5 * log10(flux[MJy/sr]*PIXAR_SR[sr/pix] ) = ZP_AB \u2013 2.5 log10(flux[MJy/sr])\n", + "\n", + "Where ZP_AB is the zeropoint: \n", + "\n", + " ZP_AB = \u20136.10 \u2013 2.5 log10(PIXAR_SR[sr/pix]). \n", + "\n", + "For example, ZP_AB = 28.0 for PIXAR_SR = 2.29e-14 (corresponding to pixel size 0.0312\").\n", + "\n", + "For data in units of MJy/sr, computing the total flux that goes into the log10 term is straightforward, it is\n", + "simply the sum of the image of the light profile. \n", + "\n", + "We compute this using a grid, which must be large enough that all light from the light profile is included. Below,\n", + "we use a grid which extends to 10\" from the centre of the light profile, which is sufficient for this example,\n", + "but you may need to increase this size for your own data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "image = light.image_2d_from(grid=grid)\n", + "\n", + "total_flux = np.sum(image) # In units of MJy sr^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now convert this total flux to an AB magnitude using the zero point of the JWST NIRCam filter we are analyzing.\n", + "\n", + "As stated above, the zero point is given by:\n", + "\n", + " ZP_AB = \u20136.10 \u2013 2.5 log10(PIXAR_SR[sr/pix])\n", + " \n", + "Where the value of PIXAR_SR is provided in the JWST documentation for the filter you are analyzing. \n", + "\n", + "The Pixar_SR values for JWST (James Webb Space Telescope) NIRCam filters refer to the pixel scale in steradians (sr) \n", + "for each filter, which is a measure of the solid angle covered by each pixel. These values are important for \n", + "calibrating and understanding how light is captured by the instrument.\n", + "\n", + "For the F444W filter, which we are using in this example, the value is 2.29e-14 (corresponding to a pixel size o\n", + "f 0.0312\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixar_sr = 2.29e-14\n", + "\n", + "zero_point = -6.10 - 2.5 * np.log10(pixar_sr)\n", + "\n", + "magnitude_ab = zero_point - 2.5 * np.log10(total_flux)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "With an absolute magnitude and quantity of light in physical units, you should now be able to convert these values to\n", + "whatever units you need for your science case.\n", + "\n", + "__Electrons Per Second (e s^-1): Hubble Space Telescope__\n", + "\n", + "Hubble Space Telescope (HST) data is often provided in units of electrons per second (e- s^-1). \n", + "\n", + "We therefore show how to convert the intensity of a light profile from electrons per second to an absolute magnitude." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "light = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=2.0, # in units of e- pix^-1 s^-1\n", + " effective_radius=0.1,\n", + " sersic_index=3.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We first compute the total flux in electrons per second by summing the image of the light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "image = light.image_2d_from(grid=grid) # in units e- s^-1 as summed over pixels\n", + "\n", + "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now use the zero point of the HST filter we are analyzing to convert this total flux to an AB magnitude.\n", + "\n", + "The zero point for the F814W filter, which we are using in this example, is 25.943.\n", + "\n", + "Zero points of the HST ACS filter are provided here: https://acszeropoints.stsci.edu, for other filters you should\n", + "consult the HST documentation.\n", + "\n", + "The zero point below is defined in units such that it converts the total flux from input units of electrons per second,\n", + "you should make sure your HST data is in these units and that the zero point you are using follows the same convention." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "zero_point_f814w = 25.943\n", + "\n", + "magnitude_ab = zero_point_f814w - 2.5 * np.log10(total_flux)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "With an absolute magnitude and quantity of light in physical units, you should now be able to convert these values to\n", + "whatever units you need for your science case.\n", + "\n", + "For HST, a few quantitites that may be useful and worth looking into are:\n", + "\n", + "- The HST PHOTFLAM value, in units of erg cm^-2 s^-1 A^-1 e-^-1, which is used to convert to ergs, which radio \n", + " astronomers may be interested in.\n", + " \n", + "- The HST PHOTNU value, in units of Jy (e s^-1), which converts to Janskys, which is often used by SED fitting\n", + " software.\n", + "\n", + "__Latent Variables: Total Flux Directly from the Fit__\n", + "\n", + "The examples above all computed total flux by hand: build a light profile, sample it on a grid, sum the image, then\n", + "apply the zero point. PyAutoLens does exactly this automatically as part of every fit and records the result as a\n", + "latent variable in the `latent/samples.csv` file beside the search output. You can skip the manual recipe entirely\n", + "and just read the column.\n", + "\n", + "Three lensing flux latents ship default-on (they need no instrument inputs and run on every fit unless disabled in\n", + "`config/latent.yaml`):\n", + "\n", + "- `total_lens_flux` \u2014 total integrated flux of the lens galaxy (the sum of `fit.tracer.galaxies[0]`'s model image),\n", + " in the *raw* image units the fit was performed in. For HST data in e- s^-1, this is e- s^-1; for JWST data in\n", + " MJy/sr, this is MJy/sr.\n", + "\n", + "- `total_lensed_source_flux` \u2014 image-plane integrated flux of the source galaxy after lensing. This is the source\n", + " flux as observed: stretched, multiplied, and distorted by the lens. Same raw units.\n", + "\n", + "- `total_source_flux` \u2014 source-plane intrinsic flux of the source galaxy: the flux the source would have if the\n", + " lens weren't there. Computed via `fit.tracer_linear_light_profiles_to_light_profiles` so that linear light\n", + " profiles (MGEs, pixelizations) work correctly. The ratio\n", + " `total_lensed_source_flux / total_source_flux` is the integrated magnification (also recorded directly as the\n", + " `magnification` latent).\n", + "\n", + "To convert any of these to AB magnitudes or microjanskies, apply the same zero-point recipe used above. Suppose\n", + "you have an HST F814W fit with the F814W zero point of 25.943; reading the lens galaxy flux from your result and\n", + "converting goes:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autogalaxy.imaging.model.latent import (\n", + " ab_mag_via_flux_from,\n", + " flux_mujy_via_ab_mag_from,\n", + ")\n", + "\n", + "# Stand-in for what you'd read from `latent.csv` \u2014 in a real script this is one column of one row, e.g.\n", + "# total_lens_flux = pd.read_csv(search.paths.output_path / \"latent\" / \"samples.csv\")[\"total_lens_flux\"].iloc[-1]\n", + "total_lens_flux = 1234.5 # e- s^-1\n", + "\n", + "zero_point_f814w = 25.943\n", + "ab_mag_lens = ab_mag_via_flux_from(flux=total_lens_flux, magzero=zero_point_f814w)\n", + "flux_mujy_lens = flux_mujy_via_ab_mag_from(ab_mag=ab_mag_lens)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The two helpers used above are the same ones the library uses internally to populate the `_mujy` variants of the\n", + "latents (`total_lens_flux_mujy`, `total_lensed_source_flux_mujy`, `total_source_flux_mujy`). Those variants are\n", + "default-off because they need a `magzero` you supply per-instrument. If you have a single fixed zero-point you can\n", + "flip them on by:\n", + "\n", + "1. Setting `total_lens_flux_mujy: true` (and friends) in your project's `config/latent.yaml`.\n", + "2. Passing `magzero=` when constructing the analysis:\n", + " `analysis = al.AnalysisImaging(dataset=dataset, magzero=25.943)`.\n", + "\n", + "The latent dispatcher then writes the converted \u00b5Jy columns into `latent/samples.csv` directly, so you don't have\n", + "to run the conversion in post. If you enable the `_mujy` latents but forget the `magzero` keyword, the columns are\n", + "populated with NaN and a single warning per process notes that the conversion was skipped \u2014 the fit itself is\n", + "unaffected.\n", + "\n", + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/guides/units/mass_to_light_ratio_units.ipynb b/notebooks/guides/units/mass_to_light_ratio_units.ipynb index e33cdde4e..467c70ad4 100644 --- a/notebooks/guides/units/mass_to_light_ratio_units.ipynb +++ b/notebooks/guides/units/mass_to_light_ratio_units.ipynb @@ -1,168 +1,205 @@ { - "cells": [ - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import numpy as np\n", - "import autolens as al\n", - "import autogalaxy as ag\n", - "import astropy.cosmology as cosmology" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Author: Kaihao Wang\n", - "\n", - "This guide is in development and not written with extensive text and descrptions (yet!).\n", - "\n", - "In a nutshell, it shows how to conver the units of a Sersic mass profile from PyAutoLens's internal units ot physical\n", - "units like solar masses. The code therefore also includes conversion of a Sersic light profile from internal units\n", - "(e.g. counts/s/arcsec^2) to physical units (e.g. solar luminosity).\n", - "\n", - "If you read this and anything is unclear please join the SLACK channel and ask a quesiton in the #generla channel,\n", - "we can then work on making this script more clear :)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class mass_profile(al.mp.Sersic):\n", - " \"\"\"\n", - " This is an example light-trace-mass mass profile showing how to calculate signals on each pixel\n", - " from the intensity of the light profile and mass-to-light ratio in M_sun/L_sun.\n", - "\n", - " I'll let the light profile to be a Sersic sphere.\n", - " What I aim to do is to make the units of mass intensity = intensity * mass-to-light ratio\n", - " to be critical_surface_density / (counts/s) correctly.\n", - "\n", - " Note that the units of intensity must be counts / s / arcsec^2, since astropy's output of critical\n", - " surface density is in units of solar mass / arcsec^2\n", - " \"\"\"\n", - "\n", - " def __init__(\n", - " self,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " solar_magnitude: float, # absolute magnitude for a given band\n", - " effective_radius: float,\n", - " sersic_index: float,\n", - " centre=(0.0, 0.0),\n", - " zero_point: float = 25.23,\n", - " intensity: float = 1.0, # units: counts/s/arcsec^2\n", - " mass_to_light_ratio: float = 1.0, # units: M_sun/L_sun\n", - " cosmo=al.cosmo.FlatLambdaCDM(),\n", - " ) -> None:\n", - " # critical surface density in solar mass / arcsec^2\n", - " self.critical_surface_density = (\n", - " cosmo.critical_surface_density_between_redshifts_from(\n", - " redshift_0=redshift_lens, redshift_1=redshift_source\n", - " )\n", - " )\n", - " self.redshift = redshift_lens\n", - " self.cosmo = cosmo\n", - "\n", - " super().__init__(\n", - " centre=centre,\n", - " intensity=intensity,\n", - " effective_radius=effective_radius,\n", - " sersic_index=sersic_index,\n", - " mass_to_light_ratio=mass_to_light_ratio\n", - " * self.unit_mass2light_instrument(solar_magnitude, zero_point),\n", - " )\n", - "\n", - " def magnitude_absolute2apparent(self, mag):\n", - " \"\"\"\n", - " Calculate the apparent magnitude of an object with a given absolute magnitude and on a certain redshift\n", - " \"\"\"\n", - " distance = self.cosmo.luminosity_distance(self.redshift) * 1e6\n", - " return mag + 5 * np.log10(float(distance) / 10)\n", - "\n", - " def mag2counts(self, mag, zero_mag=25.23):\n", - " \"\"\"\n", - " Convert apparent magnitude to counts in e-/s\n", - " \"\"\"\n", - " return 10 ** ((zero_mag - mag) / 2.5)\n", - "\n", - " def unit_mass2light_instrument(self, solar_magnitude, zero_point):\n", - " \"\"\"\n", - " Calculate a factor so that (mass-to-light ratio * factor * intensity) could be equal to kappa\n", - " This is the key of mass-to-light ratio units conversion.\n", - " \"\"\"\n", - "\n", - " solar_magapp = self.magnitude_absolute2apparent(solar_magnitude)\n", - " counts_per_sec_per_solar_luminosity = self.mag2counts(solar_magapp, zero_point)\n", - "\n", - " return 1 / self.critical_surface_density / counts_per_sec_per_solar_luminosity\n", - "\n", - "\n", - "if __name__ == \"__main__\":\n", - " # This is a Sersic light profile at redshift = 0.2, whose absolute magnitude is -15.4, corresponding to an apparent magnitude ~ 24.6\n", - " light = al.lp.Sersic(\n", - " centre=(0, 0),\n", - " intensity=0.892476, # unit: counts/s/arcsec^2\n", - " effective_radius=0.381768,\n", - " sersic_index=1.3,\n", - " )\n", - "\n", - " # Assuming it's mass-to-light ratio is 2 and the absolute magnitude of the sun is 4.83,\n", - " # this turns out a total mass of 2.47 * 10^8 times solar mass\n", - " # Let's check if my light-trace-mass profile can give the right answer.\n", - " mass = mass_profile(\n", - " redshift_lens=0.2,\n", - " redshift_source=1.0,\n", - " solar_magnitude=4.83,\n", - " intensity=0.892476,\n", - " effective_radius=0.381768,\n", - " sersic_index=1.3,\n", - " mass_to_light_ratio=2,\n", - " zero_point=25.23,\n", - " )\n", - "\n", - " from scipy.integrate import quad\n", - "\n", - " def delta_mass(r):\n", - " # the length scale of MassProfile is defined on the unit of arcsec, so that r is in unit of arcsec\n", - " # thus, critical_surface_density should also be defined in per arcsec^2 unit\n", - " return mass.convergence_func(r) * mass.critical_surface_density * 2 * np.pi * r\n", - "\n", - " total_mass, _ = quad(delta_mass, 0, np.inf)\n", - "\n", - " print(mass.mass_to_light_ratio)\n", - " print(\"The total mass of this galaxy should be ~ 2.47 * 10^8\")\n", - " print(f\"While the result of the light-trace-mass profile is {total_mass:.4e}\")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import numpy as np\n", + "import autolens as al\n", + "import autogalaxy as ag\n", + "import astropy.cosmology as cosmology" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Author: Kaihao Wang\n", + "\n", + "This guide is in development and not written with extensive text and descrptions (yet!).\n", + "\n", + "In a nutshell, it shows how to conver the units of a Sersic mass profile from PyAutoLens's internal units ot physical\n", + "units like solar masses. The code therefore also includes conversion of a Sersic light profile from internal units\n", + "(e.g. counts/s/arcsec^2) to physical units (e.g. solar luminosity).\n", + "\n", + "If you read this and anything is unclear please join the SLACK channel and ask a quesiton in the #generla channel,\n", + "we can then work on making this script more clear :)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class mass_profile(al.mp.Sersic):\n", + " \"\"\"\n", + " This is an example light-trace-mass mass profile showing how to calculate signals on each pixel\n", + " from the intensity of the light profile and mass-to-light ratio in M_sun/L_sun.\n", + "\n", + " I'll let the light profile to be a Sersic sphere.\n", + " What I aim to do is to make the units of mass intensity = intensity * mass-to-light ratio\n", + " to be critical_surface_density / (counts/s) correctly.\n", + "\n", + " Note that the units of intensity must be counts / s / arcsec^2, since astropy's output of critical\n", + " surface density is in units of solar mass / arcsec^2\n", + " \"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " solar_magnitude: float, # absolute magnitude for a given band\n", + " effective_radius: float,\n", + " sersic_index: float,\n", + " centre=(0.0, 0.0),\n", + " zero_point: float = 25.23,\n", + " intensity: float = 1.0, # units: counts/s/arcsec^2\n", + " mass_to_light_ratio: float = 1.0, # units: M_sun/L_sun\n", + " cosmo=al.cosmo.FlatLambdaCDM(),\n", + " ) -> None:\n", + " # critical surface density in solar mass / arcsec^2\n", + " self.critical_surface_density = (\n", + " cosmo.critical_surface_density_between_redshifts_from(\n", + " redshift_0=redshift_lens, redshift_1=redshift_source\n", + " )\n", + " )\n", + " self.redshift = redshift_lens\n", + " self.cosmo = cosmo\n", + "\n", + " super().__init__(\n", + " centre=centre,\n", + " intensity=intensity,\n", + " effective_radius=effective_radius,\n", + " sersic_index=sersic_index,\n", + " mass_to_light_ratio=mass_to_light_ratio\n", + " * self.unit_mass2light_instrument(solar_magnitude, zero_point),\n", + " )\n", + "\n", + " def magnitude_absolute2apparent(self, mag):\n", + " \"\"\"\n", + " Calculate the apparent magnitude of an object with a given absolute magnitude and on a certain redshift\n", + " \"\"\"\n", + " distance = self.cosmo.luminosity_distance(self.redshift) * 1e6\n", + " return mag + 5 * np.log10(float(distance) / 10)\n", + "\n", + " def mag2counts(self, mag, zero_mag=25.23):\n", + " \"\"\"\n", + " Convert apparent magnitude to counts in e-/s\n", + " \"\"\"\n", + " return 10 ** ((zero_mag - mag) / 2.5)\n", + "\n", + " def unit_mass2light_instrument(self, solar_magnitude, zero_point):\n", + " \"\"\"\n", + " Calculate a factor so that (mass-to-light ratio * factor * intensity) could be equal to kappa\n", + " This is the key of mass-to-light ratio units conversion.\n", + " \"\"\"\n", + "\n", + " solar_magapp = self.magnitude_absolute2apparent(solar_magnitude)\n", + " counts_per_sec_per_solar_luminosity = self.mag2counts(solar_magapp, zero_point)\n", + "\n", + " return 1 / self.critical_surface_density / counts_per_sec_per_solar_luminosity\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " # This is a Sersic light profile at redshift = 0.2, whose absolute magnitude is -15.4, corresponding to an apparent magnitude ~ 24.6\n", + " light = al.lp.Sersic(\n", + " centre=(0, 0),\n", + " intensity=0.892476, # unit: counts/s/arcsec^2\n", + " effective_radius=0.381768,\n", + " sersic_index=1.3,\n", + " )\n", + "\n", + " # Assuming it's mass-to-light ratio is 2 and the absolute magnitude of the sun is 4.83,\n", + " # this turns out a total mass of 2.47 * 10^8 times solar mass\n", + " # Let's check if my light-trace-mass profile can give the right answer.\n", + " mass = mass_profile(\n", + " redshift_lens=0.2,\n", + " redshift_source=1.0,\n", + " solar_magnitude=4.83,\n", + " intensity=0.892476,\n", + " effective_radius=0.381768,\n", + " sersic_index=1.3,\n", + " mass_to_light_ratio=2,\n", + " zero_point=25.23,\n", + " )\n", + "\n", + " from scipy.integrate import quad\n", + "\n", + " def delta_mass(r):\n", + " # the length scale of MassProfile is defined on the unit of arcsec, so that r is in unit of arcsec\n", + " # thus, critical_surface_density should also be defined in per arcsec^2 unit\n", + " return mass.convergence_func(r) * mass.critical_surface_density * 2 * np.pi * r\n", + "\n", + " total_mass, _ = quad(delta_mass, 0, np.inf)\n", + "\n", + " print(mass.mass_to_light_ratio)\n", + " print(\"The total mass of this galaxy should be ~ 2.47 * 10^8\")\n", + " print(f\"While the result of the light-trace-mass profile is {total_mass:.4e}\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/data.ipynb b/notebooks/imaging/data_preparation/examples/data.ipynb index 03a158d2c..ba3bad3d2 100644 --- a/notebooks/imaging/data_preparation/examples/data.ipynb +++ b/notebooks/imaging/data_preparation/examples/data.ipynb @@ -1,350 +1,387 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Image\n", - "=======================\n", - "\n", - "The image is the image of your galaxy, which comes from a telescope like the Hubble Space telescope (HST).\n", - "\n", - "This tutorial describes preprocessing your dataset`s image to adhere to the units and formats required by PyAutoLens.\n", - "\n", - "__Contents__\n", - "\n", - "- **Pixel Scale:** Overview of the pixel-to-arcsecond conversion factor for common telescopes.\n", - "- **Loading Data From Individual Fits Files:** Loading an image from FITS files and inspecting its standards.\n", - "- **Converting Data To Electrons Per Second:** Converting image flux units between electrons per second, counts and ADUs.\n", - "- **Resizing Data:** Trimming or padding a large postage stamp to an appropriate size.\n", - "- **Background Subtraction:** Overview of background sky subtraction tools and modeling approaches.\n", - "\n", - "__Pixel Scale__\n", - "\n", - "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", - "your telescope. You should look up now if you are unsure of the value.\n", - "\n", - "The pixel scale of some common telescopes is as follows:\n", - "\n", - " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", - " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", - " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", - " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", - " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", - "\n", - "It is absolutely vital you use the correct pixel scale, so double check this value!\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Data From Individual Fits Files__\n", - "\n", - "Load an image from .fits files (a format commonly used by Astronomers) via the `Array2D` object.\n", - "\n", - "This image represents a good data-reduction that conforms **PyAutoLens** formatting standards!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This image conforms to **PyAutoLens** standards for the following reasons.\n", - "\n", - " - Units: The image flux is in units of electrons per second (as opposed to electrons, counts, ADU`s etc.).\n", - " Internal **PyAutoLens** functions for computing quantities like a galaxy magnitude assume the data and model\n", - " light profiles are in electrons per second.\n", - "\n", - " - Centering: The lens galaxy is at the centre of the image (as opposed to in a corner). Default **PyAutoLens**\n", - " parameter priors assume the galaxy is at the centre of the image.\n", - "\n", - " - Stamp Size: The image is a postage stamp cut-out of the galaxy, but does not include many pixels around the edge of\n", - " the galaxy. It is advisible to cut out a postage stamp of the galaxy, as opposed to the entire image, as this reduces\n", - " the amount of memory **PyAutoLens** uses, speeds up the analysis and ensures visualization zooms around the galaxy.\n", - " Conforming to this standard is not necessary to ensure an accurate analsyis.\n", - "\n", - " - Background Sky Subtraction: The image has had its background sky subtracted.\n", - "\n", - "If your image conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you noise-map and PSF conform to standards first!).\n", - "\n", - "If it does not conform to standards, this script illustrates **PyAutoLens** functionality which can be used to\n", - "convert it to standards.\n", - "\n", - "__Converting Data To Electrons Per Second__\n", - "\n", - "Brightness units: the image`s flux values should be in units of electrons per second (as opposed to electrons,\n", - "counts, ADU`s etc.).\n", - "\n", - "Although **PyAutoLens** can technically perform an analysis using other units, the default setup assumes electrons per\n", - "second (e.g. the priors on `LightProfile` intensity and `Regularization` parameters). Thus, images not in electrons per\n", - "second should be converted!\n", - "\n", - "The data loaded above is in units of electrons per second, lets convert it to counts to illustrate how this is done.\n", - "\n", - "Converting from electrons per second to counts (and visa versa) means we must know the exposure time of our observation,\n", - "which will either be in the .fits header information of your data or be an output of your data reduction pipeline.\n", - "\n", - "We create an `Array2D` of the exposure time map, which is the exposure time of each pixel in the image assuming that\n", - "all pixels have the same exposure time. This is a good approximation for most HST observations, but not for all." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "exposure_time = 1000.0\n", - "\n", - "exposure_time_map = al.Array2D.full(\n", - " fill_value=exposure_time,\n", - " shape_native=data.shape_native,\n", - " pixel_scales=data.pixel_scales,\n", - ")\n", - "\n", - "data_counts = al.preprocess.array_eps_to_counts(\n", - " array_eps=data, exposure_time_map=exposure_time_map\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the image in counts, we can see that the flux values are now much higher values (e.g. ~1000 or above)\n", - "compared to the data in electrons per second (e.g. ~1 or below)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data_counts, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It is therefore straightforward to convert an image to electrons per second from counts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_eps = al.preprocess.array_counts_to_eps(\n", - " array_counts=data_counts, exposure_time_map=exposure_time_map\n", - ")\n", - "\n", - "aplt.plot_array(array=data_eps, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If the effective exposure-time map is output as part of the data reduction, you can use this to convert the image to\n", - "electrons per second instead.\n", - "\n", - "[The code below is commented out because the simulated data does not have an effective exposure time map in .fits\n", - "format.]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# exposure_time_map = al.Array2D.from_fits(\n", - "# file_path=Path(dataset_path, \"exposure_time_map.fits\"),\n", - "# pixel_scales=data_eps.pixel_scales,\n", - "# )\n", - "#\n", - "# data_eps = al.preprocess.array_counts_to_eps(\n", - "# array_counts=data_counts, exposure_time_map=exposure_time_map\n", - "# )\n", - "#\n", - "# aplt.plot_array(array=data_eps, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**PyAutoLens** can also convert data to / from units of electrons per second to ADUs, which uses both the exposure\n", - "time and instrumental gain of the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_in_adus = al.preprocess.array_eps_to_adus(\n", - " array_eps=data, gain=4.0, exposure_time_map=exposure_time_map\n", - ")\n", - "\n", - "aplt.plot_array(array=data_in_adus, title=\"\")\n", - "\n", - "data_eps = al.preprocess.array_adus_to_eps(\n", - " array_adus=data_in_adus, gain=4.0, exposure_time_map=exposure_time_map\n", - ")\n", - "\n", - "aplt.plot_array(array=data_eps, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In `autolens_workspace/*/data_preparation/noise_map.py` we show that a noise-map must also be in units of\n", - "electrons per second, and that the same functions as above can be used to do this.\n", - "\n", - "__Resizing Data__\n", - "\n", - "The bigger the postage stamp cut-out of the image the more memory it requires to store. Visualization will be less\n", - "ideal too, as the lens galaxy will be a smaller blob in the centre relative to the large surrounding edges of the image. Why\n", - "keep the edges surrounding the lens and source galaxy if they are masked out anyway?\n", - "\n", - "If you have a large postage stamp you can trim it using the preprocess module, which is centered on the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_large_stamp_trimmed = al.preprocess.array_with_new_shape(\n", - " array=data, new_shape=(130, 130)\n", - ")\n", - "\n", - "aplt.plot_array(array=data_large_stamp_trimmed, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Stamps can also be too small, if the mask you input to the analysis is larger than the postage stamp extends.\n", - "\n", - "In this case, you either need to reproduce the data with a larger postage stamp, or use a smaller mask.\n", - "\n", - "__Background Subtraction__\n", - "\n", - "The background of an image is the light that is not associated with the lens or source galaxies we are\n", - "interested in. This is due to light from the sky, zodiacal light, and light from other galaxies in the\n", - "field of view. The background should have been subtracted from the image before it was reduced, but\n", - "sometimes this is not the case.\n", - "\n", - "It is recommend you use data processing tools outside of **PyAutoLens** to subtract the background from your image,\n", - "as these have been optimized for this task. However, if you do not have access to these tools, **PyAutoLens** has\n", - "functions in the `preprocess` module that can estimate and subtract the background of an image.\n", - "\n", - "The preprocess module is found here:\n", - "\n", - "https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/dataset/preprocess.py\n", - "\n", - "Functions related to background subtraction are:\n", - "\n", - "- `background_sky_level_via_edges_from`\n", - "- `background_noise_map_via_edges_from`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Image\n", + "=======================\n", + "\n", + "The image is the image of your galaxy, which comes from a telescope like the Hubble Space telescope (HST).\n", + "\n", + "This tutorial describes preprocessing your dataset`s image to adhere to the units and formats required by PyAutoLens.\n", + "\n", + "__Contents__\n", + "\n", + "- **Pixel Scale:** Overview of the pixel-to-arcsecond conversion factor for common telescopes.\n", + "- **Loading Data From Individual Fits Files:** Loading an image from FITS files and inspecting its standards.\n", + "- **Converting Data To Electrons Per Second:** Converting image flux units between electrons per second, counts and ADUs.\n", + "- **Resizing Data:** Trimming or padding a large postage stamp to an appropriate size.\n", + "- **Background Subtraction:** Overview of background sky subtraction tools and modeling approaches.\n", + "\n", + "__Pixel Scale__\n", + "\n", + "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", + "your telescope. You should look up now if you are unsure of the value.\n", + "\n", + "The pixel scale of some common telescopes is as follows:\n", + "\n", + " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", + " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", + " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", + " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", + " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", + "\n", + "It is absolutely vital you use the correct pixel scale, so double check this value!\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Data From Individual Fits Files__\n", + "\n", + "Load an image from .fits files (a format commonly used by Astronomers) via the `Array2D` object.\n", + "\n", + "This image represents a good data-reduction that conforms **PyAutoLens** formatting standards!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This image conforms to **PyAutoLens** standards for the following reasons.\n", + "\n", + " - Units: The image flux is in units of electrons per second (as opposed to electrons, counts, ADU`s etc.).\n", + " Internal **PyAutoLens** functions for computing quantities like a galaxy magnitude assume the data and model\n", + " light profiles are in electrons per second.\n", + "\n", + " - Centering: The lens galaxy is at the centre of the image (as opposed to in a corner). Default **PyAutoLens**\n", + " parameter priors assume the galaxy is at the centre of the image.\n", + "\n", + " - Stamp Size: The image is a postage stamp cut-out of the galaxy, but does not include many pixels around the edge of\n", + " the galaxy. It is advisible to cut out a postage stamp of the galaxy, as opposed to the entire image, as this reduces\n", + " the amount of memory **PyAutoLens** uses, speeds up the analysis and ensures visualization zooms around the galaxy.\n", + " Conforming to this standard is not necessary to ensure an accurate analsyis.\n", + "\n", + " - Background Sky Subtraction: The image has had its background sky subtracted.\n", + "\n", + "If your image conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you noise-map and PSF conform to standards first!).\n", + "\n", + "If it does not conform to standards, this script illustrates **PyAutoLens** functionality which can be used to\n", + "convert it to standards.\n", + "\n", + "__Converting Data To Electrons Per Second__\n", + "\n", + "Brightness units: the image`s flux values should be in units of electrons per second (as opposed to electrons,\n", + "counts, ADU`s etc.).\n", + "\n", + "Although **PyAutoLens** can technically perform an analysis using other units, the default setup assumes electrons per\n", + "second (e.g. the priors on `LightProfile` intensity and `Regularization` parameters). Thus, images not in electrons per\n", + "second should be converted!\n", + "\n", + "The data loaded above is in units of electrons per second, lets convert it to counts to illustrate how this is done.\n", + "\n", + "Converting from electrons per second to counts (and visa versa) means we must know the exposure time of our observation,\n", + "which will either be in the .fits header information of your data or be an output of your data reduction pipeline.\n", + "\n", + "We create an `Array2D` of the exposure time map, which is the exposure time of each pixel in the image assuming that\n", + "all pixels have the same exposure time. This is a good approximation for most HST observations, but not for all." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "exposure_time = 1000.0\n", + "\n", + "exposure_time_map = al.Array2D.full(\n", + " fill_value=exposure_time,\n", + " shape_native=data.shape_native,\n", + " pixel_scales=data.pixel_scales,\n", + ")\n", + "\n", + "data_counts = al.preprocess.array_eps_to_counts(\n", + " array_eps=data, exposure_time_map=exposure_time_map\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the image in counts, we can see that the flux values are now much higher values (e.g. ~1000 or above)\n", + "compared to the data in electrons per second (e.g. ~1 or below)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data_counts, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It is therefore straightforward to convert an image to electrons per second from counts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_eps = al.preprocess.array_counts_to_eps(\n", + " array_counts=data_counts, exposure_time_map=exposure_time_map\n", + ")\n", + "\n", + "aplt.plot_array(array=data_eps, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If the effective exposure-time map is output as part of the data reduction, you can use this to convert the image to\n", + "electrons per second instead.\n", + "\n", + "[The code below is commented out because the simulated data does not have an effective exposure time map in .fits\n", + "format.]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# exposure_time_map = al.Array2D.from_fits(\n", + "# file_path=Path(dataset_path, \"exposure_time_map.fits\"),\n", + "# pixel_scales=data_eps.pixel_scales,\n", + "# )\n", + "#\n", + "# data_eps = al.preprocess.array_counts_to_eps(\n", + "# array_counts=data_counts, exposure_time_map=exposure_time_map\n", + "# )\n", + "#\n", + "# aplt.plot_array(array=data_eps, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**PyAutoLens** can also convert data to / from units of electrons per second to ADUs, which uses both the exposure\n", + "time and instrumental gain of the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_in_adus = al.preprocess.array_eps_to_adus(\n", + " array_eps=data, gain=4.0, exposure_time_map=exposure_time_map\n", + ")\n", + "\n", + "aplt.plot_array(array=data_in_adus, title=\"\")\n", + "\n", + "data_eps = al.preprocess.array_adus_to_eps(\n", + " array_adus=data_in_adus, gain=4.0, exposure_time_map=exposure_time_map\n", + ")\n", + "\n", + "aplt.plot_array(array=data_eps, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In `autolens_workspace/*/data_preparation/noise_map.py` we show that a noise-map must also be in units of\n", + "electrons per second, and that the same functions as above can be used to do this.\n", + "\n", + "__Resizing Data__\n", + "\n", + "The bigger the postage stamp cut-out of the image the more memory it requires to store. Visualization will be less\n", + "ideal too, as the lens galaxy will be a smaller blob in the centre relative to the large surrounding edges of the image. Why\n", + "keep the edges surrounding the lens and source galaxy if they are masked out anyway?\n", + "\n", + "If you have a large postage stamp you can trim it using the preprocess module, which is centered on the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_large_stamp_trimmed = al.preprocess.array_with_new_shape(\n", + " array=data, new_shape=(130, 130)\n", + ")\n", + "\n", + "aplt.plot_array(array=data_large_stamp_trimmed, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Stamps can also be too small, if the mask you input to the analysis is larger than the postage stamp extends.\n", + "\n", + "In this case, you either need to reproduce the data with a larger postage stamp, or use a smaller mask.\n", + "\n", + "__Background Subtraction__\n", + "\n", + "The background of an image is the light that is not associated with the lens or source galaxies we are\n", + "interested in. This is due to light from the sky, zodiacal light, and light from other galaxies in the\n", + "field of view. The background should have been subtracted from the image before it was reduced, but\n", + "sometimes this is not the case.\n", + "\n", + "It is recommend you use data processing tools outside of **PyAutoLens** to subtract the background from your image,\n", + "as these have been optimized for this task. However, if you do not have access to these tools, **PyAutoLens** has\n", + "functions in the `preprocess` module that can estimate and subtract the background of an image.\n", + "\n", + "The preprocess module is found here:\n", + "\n", + "https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/dataset/preprocess.py\n", + "\n", + "Functions related to background subtraction are:\n", + "\n", + "- `background_sky_level_via_edges_from`\n", + "- `background_noise_map_via_edges_from`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/noise_map.ipynb b/notebooks/imaging/data_preparation/examples/noise_map.ipynb index f41c67165..9c2069d09 100644 --- a/notebooks/imaging/data_preparation/examples/noise_map.ipynb +++ b/notebooks/imaging/data_preparation/examples/noise_map.ipynb @@ -1,201 +1,238 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Noise-map\n", - "===========================\n", - "\n", - "The noise-map defines the uncertainty in every pixel of your lens image, where values are defined as the\n", - "RMS standard deviation in every pixel (not the variances, HST WHT-map values, etc.).\n", - "\n", - "You MUST be certain that the noise-map is the RMS standard deviations or else your analysis will be incorrect!\n", - "\n", - "This tutorial describes preprocessing your dataset`s noise-map to adhere to the units and formats required\n", - "by **PyAutoLens**.\n", - "\n", - "__Contents__\n", - "\n", - "- **Pixel Scale:** Overview of the pixel-to-arcsecond conversion factor for common telescopes.\n", - "- **Loading Data From Individual Fits Files:** Loading a noise-map from FITS files and inspecting its standards.\n", - "- **1) Tools Illustrated In Image:** Overview of unit conversion and resizing tools from the image preparation script.\n", - "- **Noise Conversions:** Functions for computing noise-maps from various input formats.\n", - "\n", - "__Pixel Scale__\n", - "\n", - "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", - "your telescope. You should look up now if you are unsure of the value.\n", - "\n", - "The pixel scale of some common telescopes is as follows:\n", - "\n", - " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", - " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", - " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", - " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", - " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", - "\n", - "It is absolutely vital you use the correct pixel scale, so double check this value!\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Data From Individual Fits Files__\n", - "\n", - "Load a noise-map from .fits files (a format commonly used by Astronomers) via the `Array2D` object.\n", - "\n", - "This noise-map represents a good data-reduction that conforms **PyAutoLens** formatting standards!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "noise_map = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"noise_map.fits\", pixel_scales=0.1\n", - ")\n", - "\n", - "aplt.plot_array(array=noise_map, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This noise-map conforms to **PyAutoLens** standards for the following reasons:\n", - "\n", - " - Units: Like its corresponding image, it is in units of electrons per second (as opposed to electrons, counts,\n", - " ADU`s etc.). Internal **PyAutoLens** functions for computing quantities like a galaxy magnitude assume the data and\n", - " model light profiles are in electrons per second.\n", - "\n", - " - Values: The noise-map values themselves are the RMS standard deviations of the noise in every pixel. When a model\n", - " is fitted to data in **PyAutoLens** and a likelihood is evaluated, this calculation assumes that this is the\n", - " corresponding definition of the noise-map. The noise map therefore should not be the variance of the noise, or\n", - " another definition of noise.\n", - "\n", - " - Poisson: The noise-map includes the Poisson noise contribution of the image (e.g. due to Poisson count statistics\n", - " in the lens and source galaxies), in addition to the contribution of background noise from the sky background.\n", - " Data reduction pipelines often remove the Poisson noise contribution, but this is incorrect and will lead to\n", - " incorrect results.\n", - "\n", - "Given the image should be centred and cut-out around the lens and source galaxies, so should the noise-map.\n", - "\n", - "If your noise-map conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you image and PSF conform to standards first!).\n", - "\n", - "If it does not conform to standards, this script illustrates **PyAutoLens** functionality which can be used to\n", - "convert it to standards.\n", - "\n", - "__1) Tools Illustrated In Image__\n", - "\n", - "The script `data_prepatation/examples/image.ipynb` illustrates the following preparation steps:\n", - "\n", - "1) Converted it from counts / ADUs / other units to electrons per second.\n", - "2) Trimmed / padded the image.\n", - "3) Recentered the image.\n", - "\n", - "You can perform identical operations on your noise-map (assuming it is in the same units and has the dimensions as the\n", - "image.\n", - "\n", - "__Noise Conversions__\n", - "\n", - "There are many different ways the noise-map can be reduced, and it varies depending on the telescope and its\n", - "specific data reduction.\n", - "\n", - "The preprocess module contains example functions for computing noise-maps, which may help you calculate your noise-map\n", - "from the data you currently have (if it is not already RMS values including the Poisson noise contribution and\n", - "background sky contribution).\n", - "\n", - "https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/dataset/preprocess.py\n", - "\n", - "Functions related to the noise map are:\n", - "\n", - "- `noise_map_via_data_eps_and_exposure_time_map_from`\n", - "- `noise_map_via_weight_map_from`\n", - "- `noise_map_via_inverse_noise_map_from`\n", - "- `noise_map_via_data_eps_exposure_time_map_and_background_noise_map_from`\n", - "- `noise_map_via_data_eps_exposure_time_map_and_background_variances_from`\n", - "- `poisson_noise_via_data_eps_from`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Noise-map\n", + "===========================\n", + "\n", + "The noise-map defines the uncertainty in every pixel of your lens image, where values are defined as the\n", + "RMS standard deviation in every pixel (not the variances, HST WHT-map values, etc.).\n", + "\n", + "You MUST be certain that the noise-map is the RMS standard deviations or else your analysis will be incorrect!\n", + "\n", + "This tutorial describes preprocessing your dataset`s noise-map to adhere to the units and formats required\n", + "by **PyAutoLens**.\n", + "\n", + "__Contents__\n", + "\n", + "- **Pixel Scale:** Overview of the pixel-to-arcsecond conversion factor for common telescopes.\n", + "- **Loading Data From Individual Fits Files:** Loading a noise-map from FITS files and inspecting its standards.\n", + "- **1) Tools Illustrated In Image:** Overview of unit conversion and resizing tools from the image preparation script.\n", + "- **Noise Conversions:** Functions for computing noise-maps from various input formats.\n", + "\n", + "__Pixel Scale__\n", + "\n", + "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", + "your telescope. You should look up now if you are unsure of the value.\n", + "\n", + "The pixel scale of some common telescopes is as follows:\n", + "\n", + " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", + " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", + " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", + " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", + " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", + "\n", + "It is absolutely vital you use the correct pixel scale, so double check this value!\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Data From Individual Fits Files__\n", + "\n", + "Load a noise-map from .fits files (a format commonly used by Astronomers) via the `Array2D` object.\n", + "\n", + "This noise-map represents a good data-reduction that conforms **PyAutoLens** formatting standards!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "noise_map = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"noise_map.fits\", pixel_scales=0.1\n", + ")\n", + "\n", + "aplt.plot_array(array=noise_map, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This noise-map conforms to **PyAutoLens** standards for the following reasons:\n", + "\n", + " - Units: Like its corresponding image, it is in units of electrons per second (as opposed to electrons, counts,\n", + " ADU`s etc.). Internal **PyAutoLens** functions for computing quantities like a galaxy magnitude assume the data and\n", + " model light profiles are in electrons per second.\n", + "\n", + " - Values: The noise-map values themselves are the RMS standard deviations of the noise in every pixel. When a model\n", + " is fitted to data in **PyAutoLens** and a likelihood is evaluated, this calculation assumes that this is the\n", + " corresponding definition of the noise-map. The noise map therefore should not be the variance of the noise, or\n", + " another definition of noise.\n", + "\n", + " - Poisson: The noise-map includes the Poisson noise contribution of the image (e.g. due to Poisson count statistics\n", + " in the lens and source galaxies), in addition to the contribution of background noise from the sky background.\n", + " Data reduction pipelines often remove the Poisson noise contribution, but this is incorrect and will lead to\n", + " incorrect results.\n", + "\n", + "Given the image should be centred and cut-out around the lens and source galaxies, so should the noise-map.\n", + "\n", + "If your noise-map conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you image and PSF conform to standards first!).\n", + "\n", + "If it does not conform to standards, this script illustrates **PyAutoLens** functionality which can be used to\n", + "convert it to standards.\n", + "\n", + "__1) Tools Illustrated In Image__\n", + "\n", + "The script `data_prepatation/examples/image.ipynb` illustrates the following preparation steps:\n", + "\n", + "1) Converted it from counts / ADUs / other units to electrons per second.\n", + "2) Trimmed / padded the image.\n", + "3) Recentered the image.\n", + "\n", + "You can perform identical operations on your noise-map (assuming it is in the same units and has the dimensions as the\n", + "image.\n", + "\n", + "__Noise Conversions__\n", + "\n", + "There are many different ways the noise-map can be reduced, and it varies depending on the telescope and its\n", + "specific data reduction.\n", + "\n", + "The preprocess module contains example functions for computing noise-maps, which may help you calculate your noise-map\n", + "from the data you currently have (if it is not already RMS values including the Poisson noise contribution and\n", + "background sky contribution).\n", + "\n", + "https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/dataset/preprocess.py\n", + "\n", + "Functions related to the noise map are:\n", + "\n", + "- `noise_map_via_data_eps_and_exposure_time_map_from`\n", + "- `noise_map_via_weight_map_from`\n", + "- `noise_map_via_inverse_noise_map_from`\n", + "- `noise_map_via_data_eps_exposure_time_map_and_background_noise_map_from`\n", + "- `noise_map_via_data_eps_exposure_time_map_and_background_variances_from`\n", + "- `poisson_noise_via_data_eps_from`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/optional/extra_galaxies_centres.ipynb b/notebooks/imaging/data_preparation/examples/optional/extra_galaxies_centres.ipynb index 9f633c0f2..ff1d98659 100644 --- a/notebooks/imaging/data_preparation/examples/optional/extra_galaxies_centres.ipynb +++ b/notebooks/imaging/data_preparation/examples/optional/extra_galaxies_centres.ipynb @@ -1,256 +1,293 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Extra Galaxies (Optional)\n", - "===========================================\n", - "\n", - "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source\n", - "and whose mass may contribute to the ray-tracing and lens model.\n", - "\n", - "We can include these extra galaxies in the lens model, either as light profiles, mass profiles, or both, using the\n", - "modeling API, where these nearby objects are denoted `extra_galaxies`.\n", - "\n", - "This script marks the (y,x) arcsecond locations of these extra galaxies, so that when they are included in the lens model\n", - "the centre of these extra galaxies light and / or mass profiles are fixed to these values (or their priors are initialized\n", - "surrounding these centres).\n", - "\n", - "This tutorial closely mirrors tutorial 7, `lens_light_centre`, where the main purpose of this script is to mark the\n", - "centres of every object we'll model as an extra galaxy. A GUI is also available to do this.\n", - "\n", - "__Contents__\n", - "\n", - "- **Masking Extra Galaxies:** The example `mask_extra_galaxies.py` masks the regions of an image where extra galaxies are present.\n", - "- **Output:** Save this as a .png image in the dataset folder for easy inspection later.\n", - "\n", - "__Masking Extra Galaxies__\n", - "\n", - "The example `mask_extra_galaxies.py` masks the regions of an image where extra galaxies are present. This mask is used\n", - "to remove their signal from the data and increase their noise to make them not impact the fit. This means their\n", - "luminous emission does not need to be included in the model, reducing the number of free parameters and speeding up the\n", - "analysis. It is still a choice whether their mass is included in the model.\n", - "\n", - "Which approach you use to account for the emission of extra galaxies, modeling or masking, depends on how significant\n", - "the blending of their emission with the lens and source galaxies is and how much it impacts the model-fit.\n", - "\n", - "__Links / Resources__\n", - "\n", - "The script `data_preparation/gui/extra_galaxies_centres.ipynb` shows how to use a Graphical User Interface (GUI) to mark\n", - "the extra galaxy centres in this way.\n", - "\n", - "The script `features/extra_galaxies/modeling` shows how to use extra galaxies in a model-fit, including loading the\n", - "extra galaxy centres created by this script.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the extra galaxy centres are output, which is `dataset/imaging/extra_galaxies`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the `Imaging` dataset, so that the lens light centres can be plotted over the strong lens image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the extra galaxy centres, which is a `Grid2DIrregular` object of (y,x) values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_centres = al.Grid2DIrregular(values=[(1.0, 3.5), (-2.0, -3.5)])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the image and extra galaxy centres, so we can check that the centre overlaps the lens light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Save this as a .png image in the dataset folder for easy inspection later." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the extra galaxy centres to the dataset folder of the lens, so that we can load them from a .json file\n", - "when we model them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=extra_galaxies_centres,\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The workspace also includes a GUI for drawing extra galaxy centres, which can be found at\n", - "`autolens_workspace/*/imaging/data_preparation/gui/extra_galaxies_centres.py`.\n", - "\n", - "This tools allows you `click` on the image where an image of the lensed source is, and it will use the brightest pixel\n", - "within a 5x5 box of pixels to select the coordinate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Extra Galaxies (Optional)\n", + "===========================================\n", + "\n", + "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source\n", + "and whose mass may contribute to the ray-tracing and lens model.\n", + "\n", + "We can include these extra galaxies in the lens model, either as light profiles, mass profiles, or both, using the\n", + "modeling API, where these nearby objects are denoted `extra_galaxies`.\n", + "\n", + "This script marks the (y,x) arcsecond locations of these extra galaxies, so that when they are included in the lens model\n", + "the centre of these extra galaxies light and / or mass profiles are fixed to these values (or their priors are initialized\n", + "surrounding these centres).\n", + "\n", + "This tutorial closely mirrors tutorial 7, `lens_light_centre`, where the main purpose of this script is to mark the\n", + "centres of every object we'll model as an extra galaxy. A GUI is also available to do this.\n", + "\n", + "__Contents__\n", + "\n", + "- **Masking Extra Galaxies:** The example `mask_extra_galaxies.py` masks the regions of an image where extra galaxies are present.\n", + "- **Output:** Save this as a .png image in the dataset folder for easy inspection later.\n", + "\n", + "__Masking Extra Galaxies__\n", + "\n", + "The example `mask_extra_galaxies.py` masks the regions of an image where extra galaxies are present. This mask is used\n", + "to remove their signal from the data and increase their noise to make them not impact the fit. This means their\n", + "luminous emission does not need to be included in the model, reducing the number of free parameters and speeding up the\n", + "analysis. It is still a choice whether their mass is included in the model.\n", + "\n", + "Which approach you use to account for the emission of extra galaxies, modeling or masking, depends on how significant\n", + "the blending of their emission with the lens and source galaxies is and how much it impacts the model-fit.\n", + "\n", + "__Links / Resources__\n", + "\n", + "The script `data_preparation/gui/extra_galaxies_centres.ipynb` shows how to use a Graphical User Interface (GUI) to mark\n", + "the extra galaxy centres in this way.\n", + "\n", + "The script `features/extra_galaxies/modeling` shows how to use extra galaxies in a model-fit, including loading the\n", + "extra galaxy centres created by this script.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the extra galaxy centres are output, which is `dataset/imaging/extra_galaxies`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the `Imaging` dataset, so that the lens light centres can be plotted over the strong lens image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the extra galaxy centres, which is a `Grid2DIrregular` object of (y,x) values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_centres = al.Grid2DIrregular(values=[(1.0, 3.5), (-2.0, -3.5)])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the image and extra galaxy centres, so we can check that the centre overlaps the lens light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Save this as a .png image in the dataset folder for easy inspection later." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the extra galaxy centres to the dataset folder of the lens, so that we can load them from a .json file\n", + "when we model them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=extra_galaxies_centres,\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The workspace also includes a GUI for drawing extra galaxy centres, which can be found at\n", + "`autolens_workspace/*/imaging/data_preparation/gui/extra_galaxies_centres.py`.\n", + "\n", + "This tools allows you `click` on the image where an image of the lensed source is, and it will use the brightest pixel\n", + "within a 5x5 box of pixels to select the coordinate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/optional/info.ipynb b/notebooks/imaging/data_preparation/examples/optional/info.ipynb index 7e8b61c16..fe3a21959 100644 --- a/notebooks/imaging/data_preparation/examples/optional/info.ipynb +++ b/notebooks/imaging/data_preparation/examples/optional/info.ipynb @@ -1,188 +1,225 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Info (Optional)\n", - "=================================\n", - "\n", - "Auxiliary information about a strong lens dataset may used during an analysis or afterwards when interpreting the\n", - " modeling results. For example, the redshifts of the source and lens galaxy.\n", - "\n", - "By storing these as an `info.json` file in the lens's dataset folder, it is straight forward to load the redshifts\n", - "in a modeling script and pass them to a fit, such that **PyAutoLens** can then output results in physical\n", - "units (e.g. kpc instead of arc-seconds).\n", - "\n", - "For analysing large quantities of modeling results, **PyAutoLens** has an sqlite database feature. The info file\n", - "may can also be loaded by the database after a model-fit has completed, such that when one is interpreting\n", - "the results of a model fit additional data on a lens can be used to.\n", - "\n", - "For example, to plot the model-results against other measurements of a lens not made by PyAutoLens. Examples of such\n", - "data might be:\n", - "\n", - "- The velocity dispersion of the lens galaxy.\n", - "- The stellar mass of the lens galaxy.\n", - "- The results of previous strong lens models to the lens performed in previous papers.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where info is output, which is `dataset/imaging/simple`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The info is written as a Python dictionary and can have as many entries as desired added to it. Any information you\n", - "want to include int he interpretation of your lens models should be included here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "info = {\n", - " \"redshift_lens\": 0.5,\n", - " \"redshift_source\": 1.0,\n", - " \"velocity_dispersion\": 250000,\n", - " \"stellar mass\": 1e11,\n", - "}" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The info is stored in the dataset folder as a .json file.\n", - "\n", - "We cannot `dump` a .json file using a string which contains a directory, so we dump it to the location of this\n", - "script and move it to the appropriate dataset folder. We first delete existing info file in the dataset folder." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import os\n", - "import shutil\n", - "import json\n", - "\n", - "info_file = \"info.json\"\n", - "\n", - "with open(info_file, \"w+\") as f:\n", - " json.dump(info, f, indent=4)\n", - "\n", - "if Path(dataset_path, \"info.json\").exists():\n", - " os.remove(Path(dataset_path, \"info.json\"))\n", - "\n", - "shutil.move(\"info.json\", Path(dataset_path, \"info.json\"))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For the info to be available to the results of a model-fit, the modeling script must load the info file from the .json and\n", - "pass it to the search.run() or pipeline.run() function:\n", - "\n", - "info_file = Path(dataset_path, \"info.json\")\n", - "\n", - "with open(info_file, \"r\") as f:\n", - " info = json.load(f)\n", - "\n", - "result = search.fit(model=model, analysis=analysis, info=info)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Info (Optional)\n", + "=================================\n", + "\n", + "Auxiliary information about a strong lens dataset may used during an analysis or afterwards when interpreting the\n", + " modeling results. For example, the redshifts of the source and lens galaxy.\n", + "\n", + "By storing these as an `info.json` file in the lens's dataset folder, it is straight forward to load the redshifts\n", + "in a modeling script and pass them to a fit, such that **PyAutoLens** can then output results in physical\n", + "units (e.g. kpc instead of arc-seconds).\n", + "\n", + "For analysing large quantities of modeling results, **PyAutoLens** has an sqlite database feature. The info file\n", + "may can also be loaded by the database after a model-fit has completed, such that when one is interpreting\n", + "the results of a model fit additional data on a lens can be used to.\n", + "\n", + "For example, to plot the model-results against other measurements of a lens not made by PyAutoLens. Examples of such\n", + "data might be:\n", + "\n", + "- The velocity dispersion of the lens galaxy.\n", + "- The stellar mass of the lens galaxy.\n", + "- The results of previous strong lens models to the lens performed in previous papers.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where info is output, which is `dataset/imaging/simple`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The info is written as a Python dictionary and can have as many entries as desired added to it. Any information you\n", + "want to include int he interpretation of your lens models should be included here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "info = {\n", + " \"redshift_lens\": 0.5,\n", + " \"redshift_source\": 1.0,\n", + " \"velocity_dispersion\": 250000,\n", + " \"stellar mass\": 1e11,\n", + "}" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The info is stored in the dataset folder as a .json file.\n", + "\n", + "We cannot `dump` a .json file using a string which contains a directory, so we dump it to the location of this\n", + "script and move it to the appropriate dataset folder. We first delete existing info file in the dataset folder." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "import shutil\n", + "import json\n", + "\n", + "info_file = \"info.json\"\n", + "\n", + "with open(info_file, \"w+\") as f:\n", + " json.dump(info, f, indent=4)\n", + "\n", + "if Path(dataset_path, \"info.json\").exists():\n", + " os.remove(Path(dataset_path, \"info.json\"))\n", + "\n", + "shutil.move(\"info.json\", Path(dataset_path, \"info.json\"))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For the info to be available to the results of a model-fit, the modeling script must load the info file from the .json and\n", + "pass it to the search.run() or pipeline.run() function:\n", + "\n", + "info_file = Path(dataset_path, \"info.json\")\n", + "\n", + "with open(info_file, \"r\") as f:\n", + " info = json.load(f)\n", + "\n", + "result = search.fit(model=model, analysis=analysis, info=info)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/optional/lens_light_centre.ipynb b/notebooks/imaging/data_preparation/examples/optional/lens_light_centre.ipynb index 68b79e729..7e04cf17d 100644 --- a/notebooks/imaging/data_preparation/examples/optional/lens_light_centre.ipynb +++ b/notebooks/imaging/data_preparation/examples/optional/lens_light_centre.ipynb @@ -1,220 +1,257 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Lens Light Centre (Optional)\n", - "==============================================\n", - "\n", - "This script marks the (y,x) arcsecond locations of the lens galaxy light centre(s) of the strong lens you are\n", - "analysing. These can be used as fixed values for the lens light and mass models in a model-fit.\n", - "\n", - "This reduces the number of free parameters fitted for in a lens model and removes inaccurate solutions where\n", - "the lens mass model centre is unrealistically far from its true centre.\n", - "\n", - "Advanced `chaining` scripts often use these input centres in the early fits to infer an accurate initial lens model,\n", - "amd then make the centres free parameters in later searches to ensure a general and accurate lens model is inferred.\n", - "\n", - "If you create a `light_centre` for your dataset, you must also update your modeling script to use them.\n", - "\n", - "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using a fixed\n", - "lens light and / or mass centre to help the non-linear search find a good lens model.\n", - "\n", - "Links / Resources:\n", - "\n", - "The script `data_preparation/gui/lens_light_centre.ipynb` shows how to use a Graphical User Interface (GUI) to mask the\n", - "lens galaxy light centres.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the lens light centre is output, which is `dataset/imaging/simple`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the `Imaging` dataset, so that the lens light centres can be plotted over the strong lens image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, create a lens light centre, which is a `Grid2DIrregular` object of (y,x) values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "light_centre = al.Grid2DIrregular(values=[(0.0, 0.0)])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now lets plot the image and lens light centre, so we can check that the centre overlaps the lens light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we`re happy with the lens light centre(s), lets output them to the dataset folder of the lens, so that we can\n", - "load them from a .json file in our pipelines!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=light_centre,\n", - " file_path=Path(dataset_path, \"lens_light_centre.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The workspace also includes a GUI for drawing lens light centres, which can be found at\n", - "`autolens_workspace/*/imaging/data_preparation/gui/light_centres.py`.\n", - "\n", - "This tools allows you `click` on the image where the lens light centres are, and it uses the brightest\n", - "pixel within a 5x5 box of pixels to select the coordinate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Lens Light Centre (Optional)\n", + "==============================================\n", + "\n", + "This script marks the (y,x) arcsecond locations of the lens galaxy light centre(s) of the strong lens you are\n", + "analysing. These can be used as fixed values for the lens light and mass models in a model-fit.\n", + "\n", + "This reduces the number of free parameters fitted for in a lens model and removes inaccurate solutions where\n", + "the lens mass model centre is unrealistically far from its true centre.\n", + "\n", + "Advanced `chaining` scripts often use these input centres in the early fits to infer an accurate initial lens model,\n", + "amd then make the centres free parameters in later searches to ensure a general and accurate lens model is inferred.\n", + "\n", + "If you create a `light_centre` for your dataset, you must also update your modeling script to use them.\n", + "\n", + "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using a fixed\n", + "lens light and / or mass centre to help the non-linear search find a good lens model.\n", + "\n", + "Links / Resources:\n", + "\n", + "The script `data_preparation/gui/lens_light_centre.ipynb` shows how to use a Graphical User Interface (GUI) to mask the\n", + "lens galaxy light centres.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the lens light centre is output, which is `dataset/imaging/simple`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the `Imaging` dataset, so that the lens light centres can be plotted over the strong lens image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, create a lens light centre, which is a `Grid2DIrregular` object of (y,x) values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "light_centre = al.Grid2DIrregular(values=[(0.0, 0.0)])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now lets plot the image and lens light centre, so we can check that the centre overlaps the lens light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we`re happy with the lens light centre(s), lets output them to the dataset folder of the lens, so that we can\n", + "load them from a .json file in our pipelines!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=light_centre,\n", + " file_path=Path(dataset_path, \"lens_light_centre.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The workspace also includes a GUI for drawing lens light centres, which can be found at\n", + "`autolens_workspace/*/imaging/data_preparation/gui/light_centres.py`.\n", + "\n", + "This tools allows you `click` on the image where the lens light centres are, and it uses the brightest\n", + "pixel within a 5x5 box of pixels to select the coordinate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/optional/mask.ipynb b/notebooks/imaging/data_preparation/examples/optional/mask.ipynb index 79fd57308..df3b09ef6 100644 --- a/notebooks/imaging/data_preparation/examples/optional/mask.ipynb +++ b/notebooks/imaging/data_preparation/examples/optional/mask.ipynb @@ -1,297 +1,334 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Mask (Optional)\n", - "=================================\n", - "\n", - "The mask removes the regions of the image where the lens and source galaxy are not present, typically the edges of the\n", - "image.\n", - "\n", - "Example modeling scripts internally create a 3.0\" circular mask and therefore do not require that a mask has been\n", - "created externally via a data preparation script.\n", - "\n", - "This script shows how to create customize masked (e.g. annular, ellipses) which are tailored to match the lens or\n", - "lensed source emission.\n", - "\n", - "If you have not analysed your dataset yet and do not know of a specific reason why you need the bespoke masks\n", - "created by this script, it is recommended that you simply use the default ~3.0\" circular mask internally made in each\n", - "script and omit this data preparation tutorial.\n", - "\n", - "Links / Resources:\n", - "\n", - "The `examples/mask.ipynb` scripts shows how to create customize masked (e.g. annular, ellipses)\n", - "which are tailored to match the lens or lensed source emission of your data.\n", - "\n", - "The script `data_preparation/gui/mask.ipynb` shows how to use a Graphical User Interface (GUI) to create an even\n", - "more custom mask.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This tool allows one to mask a bespoke mask for a given image of a strong lens, which is loaded before a\n", - "pipeline is run and passed to that pipeline.\n", - "\n", - "Whereas in the previous 3 tutorials we used the data_raw folder of `autolens/propocess`, the mask is generated from\n", - "the reduced dataset, so we'll example `Imaging` in the `autolens_workspace/dataset` folder where your dataset reduced\n", - "following `data_preparation` tutorials 1-3 should be located.\n", - "\n", - "Setup the path to the autolens_workspace, using the correct path name below.\n", - "\n", - "The `dataset label` is the name of the folder in the `autolens_workspace/dataset` folder and `dataset_name` the\n", - "folder the dataset is stored in, e.g, `/autolens_workspace/dataset/dataset_type/dataset_name`. The mask will be\n", - "output here as `mask.fits`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"simple\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Returns the path where the mask will be output, which in this case is\n", - "`/autolens_workspace/dataset/imaging/simple`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "First, load the image of the dataset, so that the mask can be plotted over the strong lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create a mask for this dataset, using the Mask2D object I`ll use a circular-annular mask here, but I`ve commented\n", - "other options you might want to use (feel free to experiment!)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular_annular(\n", - " shape_native=data.shape_native,\n", - " pixel_scales=data.pixel_scales,\n", - " inner_radius=0.5,\n", - " outer_radius=2.5,\n", - " centre=(0.0, 0.0),\n", - ")\n", - "\n", - "# mask = al.Mask2D.circular(\n", - "# shape_native=data.shape_native,\n", - "# pixel_scales=data.pixel_scales,\n", - "# radius=2.5,\n", - "# centre=(0.0, 0.0),\n", - "# )\n", - "\n", - "# mask = al.Mask2D.elliptical(\n", - "# shape_native=data.shape_native,\n", - "# pixel_scales=data.pixel_scales,\n", - "# major_axis_radius=2.5,\n", - "# axis_ratio=0.7,\n", - "# angle=45.0,\n", - "# centre=(0.0, 0.0),\n", - "# )\n", - "\n", - "# mask = al.Mask2D.elliptical_annular(\n", - "# shape_native=data.shape_native,\n", - "# pixel_scales=data.pixel_scales,\n", - "# inner_major_axis_radius=0.5,\n", - "# inner_axis_ratio=0.7,\n", - "# inner_phi=45.0,\n", - "# outer_major_axis_radius=0.5,\n", - "# outer_axis_ratio=0.7,\n", - "# outer_phi=45.0,\n", - "# centre=(0.0, 0.0),\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now lets plot the image and mask, so we can check that the mask includes the regions of the image we want." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the masked image to clearly show what parts of the source are included." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = data.apply_mask(mask=mask)\n", - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we`re happy with the mask, lets output it to the dataset folder of the lens, so that we can load it from a .fits\n", - "file in our pipelines!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_array(array=mask, file_path=Path(dataset_path, \"mask.fits\"), overwrite=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The workspace also includes a GUI for drawing a mask, which can be found at\n", - "`autolens_workspace/*/imaging/data_preparation/gui/mask.py`. This tools allows you to draw the mask via a `spray paint` mouse\n", - "icon, such that you can draw irregular masks more tailored to the source's light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Mask (Optional)\n", + "=================================\n", + "\n", + "The mask removes the regions of the image where the lens and source galaxy are not present, typically the edges of the\n", + "image.\n", + "\n", + "Example modeling scripts internally create a 3.0\" circular mask and therefore do not require that a mask has been\n", + "created externally via a data preparation script.\n", + "\n", + "This script shows how to create customize masked (e.g. annular, ellipses) which are tailored to match the lens or\n", + "lensed source emission.\n", + "\n", + "If you have not analysed your dataset yet and do not know of a specific reason why you need the bespoke masks\n", + "created by this script, it is recommended that you simply use the default ~3.0\" circular mask internally made in each\n", + "script and omit this data preparation tutorial.\n", + "\n", + "Links / Resources:\n", + "\n", + "The `examples/mask.ipynb` scripts shows how to create customize masked (e.g. annular, ellipses)\n", + "which are tailored to match the lens or lensed source emission of your data.\n", + "\n", + "The script `data_preparation/gui/mask.ipynb` shows how to use a Graphical User Interface (GUI) to create an even\n", + "more custom mask.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This tool allows one to mask a bespoke mask for a given image of a strong lens, which is loaded before a\n", + "pipeline is run and passed to that pipeline.\n", + "\n", + "Whereas in the previous 3 tutorials we used the data_raw folder of `autolens/propocess`, the mask is generated from\n", + "the reduced dataset, so we'll example `Imaging` in the `autolens_workspace/dataset` folder where your dataset reduced\n", + "following `data_preparation` tutorials 1-3 should be located.\n", + "\n", + "Setup the path to the autolens_workspace, using the correct path name below.\n", + "\n", + "The `dataset label` is the name of the folder in the `autolens_workspace/dataset` folder and `dataset_name` the\n", + "folder the dataset is stored in, e.g, `/autolens_workspace/dataset/dataset_type/dataset_name`. The mask will be\n", + "output here as `mask.fits`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"simple\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Returns the path where the mask will be output, which in this case is\n", + "`/autolens_workspace/dataset/imaging/simple`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "First, load the image of the dataset, so that the mask can be plotted over the strong lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create a mask for this dataset, using the Mask2D object I`ll use a circular-annular mask here, but I`ve commented\n", + "other options you might want to use (feel free to experiment!)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular_annular(\n", + " shape_native=data.shape_native,\n", + " pixel_scales=data.pixel_scales,\n", + " inner_radius=0.5,\n", + " outer_radius=2.5,\n", + " centre=(0.0, 0.0),\n", + ")\n", + "\n", + "# mask = al.Mask2D.circular(\n", + "# shape_native=data.shape_native,\n", + "# pixel_scales=data.pixel_scales,\n", + "# radius=2.5,\n", + "# centre=(0.0, 0.0),\n", + "# )\n", + "\n", + "# mask = al.Mask2D.elliptical(\n", + "# shape_native=data.shape_native,\n", + "# pixel_scales=data.pixel_scales,\n", + "# major_axis_radius=2.5,\n", + "# axis_ratio=0.7,\n", + "# angle=45.0,\n", + "# centre=(0.0, 0.0),\n", + "# )\n", + "\n", + "# mask = al.Mask2D.elliptical_annular(\n", + "# shape_native=data.shape_native,\n", + "# pixel_scales=data.pixel_scales,\n", + "# inner_major_axis_radius=0.5,\n", + "# inner_axis_ratio=0.7,\n", + "# inner_phi=45.0,\n", + "# outer_major_axis_radius=0.5,\n", + "# outer_axis_ratio=0.7,\n", + "# outer_phi=45.0,\n", + "# centre=(0.0, 0.0),\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now lets plot the image and mask, so we can check that the mask includes the regions of the image we want." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the masked image to clearly show what parts of the source are included." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = data.apply_mask(mask=mask)\n", + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we`re happy with the mask, lets output it to the dataset folder of the lens, so that we can load it from a .fits\n", + "file in our pipelines!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_array(array=mask, file_path=Path(dataset_path, \"mask.fits\"), overwrite=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The workspace also includes a GUI for drawing a mask, which can be found at\n", + "`autolens_workspace/*/imaging/data_preparation/gui/mask.py`. This tools allows you to draw the mask via a `spray paint` mouse\n", + "icon, such that you can draw irregular masks more tailored to the source's light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/optional/mask_extra_galaxies.ipynb b/notebooks/imaging/data_preparation/examples/optional/mask_extra_galaxies.ipynb index 2aa6e350f..b9bc06864 100644 --- a/notebooks/imaging/data_preparation/examples/optional/mask_extra_galaxies.ipynb +++ b/notebooks/imaging/data_preparation/examples/optional/mask_extra_galaxies.ipynb @@ -1,273 +1,310 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Extra Galaxies Mask (Optional)\n", - "================================================\n", - "\n", - "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated\n", - "with the strong lens we are studying. The emission from these images will impact our model fitting and needs to be\n", - "removed from the analysis.\n", - "\n", - "This script creates a mask of these regions of the image, called the `mask_extra_galaxies`, which can be used to\n", - "prevent them from impacting a fit. This mask may also include emission from objects which are not technically galaxies,\n", - "but blend with the galaxy we are studying in a similar way. Common examples of such objects are foreground stars\n", - "or emission due to the data reduction process.\n", - "\n", - "The mask can be applied in different ways. For example, it could be applied such that the image pixels are discarded\n", - "from the fit entirely, Alternatively the mask could be used to set the image values to (near) zero and increase their\n", - "corresponding noise-map to large values.\n", - "\n", - "The exact method used depends on the nature of the model being fitted. For simple fits like a light profile a mask\n", - "is appropriate, as removing image pixels does not change how the model is fitted. However, for more complex models\n", - "fits, like those using a pixelization, masking regions of the image in a way that removes their image pixels entirely\n", - "from the fit can produce discontinuities in the pixelixation. In this case, scaling the data and noise-map values\n", - "may be a better approach.\n", - "\n", - "This script outputs a `mask_extra_galaxies.fits` file, which can be loaded and used before a model fit, in whatever\n", - "way is appropriate for the model being fitted.\n", - "\n", - "__Links / Resources__\n", - "\n", - "The script `data_preparation/gui/extra_galaxies_mask.ipynb` shows how to use a Graphical User Interface (GUI) to create\n", - "the extra galaxies mask.\n", - "\n", - "__Contents__\n", - "\n", - "- **Output:** Output to a .png file for easy inspection.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import numpy as np\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the extra galaxy mask is output, which is `dataset/imaging/extra_galaxies`.\n", - "\n", - "The corresponding simulator (`scripts/imaging/features/extra_galaxies/simulator.py`) already writes a default\n", - "`mask_extra_galaxies.fits` automatically. This script demonstrates how to override that default with your own\n", - "centres + radii \u2014 useful when working with real data where the extra galaxy locations are not known in advance." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the dataset image, so that the location of galaxies is clear when scaling the noise-map." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Define the extra galaxies mask as the union of circles centred on each extra galaxy.\n", - "\n", - "`Mask2D.circular` honours `PYAUTO_SMALL_DATASETS=1` (caps to 15x15 at 0.6\"/px) and works on the actual loaded\n", - "data shape, so the same code path runs under both normal and small-dataset modes without modification." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_mask = np.zeros(data.shape_native, dtype=bool)\n", - "\n", - "for centre, radius in [\n", - " ((1.0, 3.5), 1.5),\n", - " ((-2.0, -3.5), 2.4),\n", - "]:\n", - " circle = al.Mask2D.circular(\n", - " shape_native=data.shape_native,\n", - " pixel_scales=data.pixel_scales,\n", - " centre=centre,\n", - " radius=radius,\n", - " invert=True, # True inside the circle (i.e. masked region)\n", - " )\n", - " extra_galaxies_mask = np.logical_or(extra_galaxies_mask, circle.native)\n", - "\n", - "mask = al.Mask2D(mask=extra_galaxies_mask, pixel_scales=data.pixel_scales)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Apply the extra galaxies mask to the image, which will remove them from visualization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = data.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the data with the new mask, in order to check that the mask removes the regions of the image corresponding to the\n", - "extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output to a .png file for easy inspection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the extra galaxies mask, which will be load and used before a model fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_array(\n", - " array=mask, file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"), overwrite=True\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The workspace also includes a GUI for image and noise-map scaling, which can be found at\n", - "`autolens_workspace/*/imaging/data_preparation/gui/mask_extra_galaxies.py`.\n", - "\n", - "This tools allows you `spray paint` on the image where an you want to scale, allow irregular patterns (i.e. not\n", - "rectangles) to be scaled." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Extra Galaxies Mask (Optional)\n", + "================================================\n", + "\n", + "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated\n", + "with the strong lens we are studying. The emission from these images will impact our model fitting and needs to be\n", + "removed from the analysis.\n", + "\n", + "This script creates a mask of these regions of the image, called the `mask_extra_galaxies`, which can be used to\n", + "prevent them from impacting a fit. This mask may also include emission from objects which are not technically galaxies,\n", + "but blend with the galaxy we are studying in a similar way. Common examples of such objects are foreground stars\n", + "or emission due to the data reduction process.\n", + "\n", + "The mask can be applied in different ways. For example, it could be applied such that the image pixels are discarded\n", + "from the fit entirely, Alternatively the mask could be used to set the image values to (near) zero and increase their\n", + "corresponding noise-map to large values.\n", + "\n", + "The exact method used depends on the nature of the model being fitted. For simple fits like a light profile a mask\n", + "is appropriate, as removing image pixels does not change how the model is fitted. However, for more complex models\n", + "fits, like those using a pixelization, masking regions of the image in a way that removes their image pixels entirely\n", + "from the fit can produce discontinuities in the pixelixation. In this case, scaling the data and noise-map values\n", + "may be a better approach.\n", + "\n", + "This script outputs a `mask_extra_galaxies.fits` file, which can be loaded and used before a model fit, in whatever\n", + "way is appropriate for the model being fitted.\n", + "\n", + "__Links / Resources__\n", + "\n", + "The script `data_preparation/gui/extra_galaxies_mask.ipynb` shows how to use a Graphical User Interface (GUI) to create\n", + "the extra galaxies mask.\n", + "\n", + "__Contents__\n", + "\n", + "- **Output:** Output to a .png file for easy inspection.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the extra galaxy mask is output, which is `dataset/imaging/extra_galaxies`.\n", + "\n", + "The corresponding simulator (`scripts/imaging/features/extra_galaxies/simulator.py`) already writes a default\n", + "`mask_extra_galaxies.fits` automatically. This script demonstrates how to override that default with your own\n", + "centres + radii \u2014 useful when working with real data where the extra galaxy locations are not known in advance." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the dataset image, so that the location of galaxies is clear when scaling the noise-map." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Define the extra galaxies mask as the union of circles centred on each extra galaxy.\n", + "\n", + "`Mask2D.circular` honours `PYAUTO_SMALL_DATASETS=1` (caps to 15x15 at 0.6\"/px) and works on the actual loaded\n", + "data shape, so the same code path runs under both normal and small-dataset modes without modification." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_mask = np.zeros(data.shape_native, dtype=bool)\n", + "\n", + "for centre, radius in [\n", + " ((1.0, 3.5), 1.5),\n", + " ((-2.0, -3.5), 2.4),\n", + "]:\n", + " circle = al.Mask2D.circular(\n", + " shape_native=data.shape_native,\n", + " pixel_scales=data.pixel_scales,\n", + " centre=centre,\n", + " radius=radius,\n", + " invert=True, # True inside the circle (i.e. masked region)\n", + " )\n", + " extra_galaxies_mask = np.logical_or(extra_galaxies_mask, circle.native)\n", + "\n", + "mask = al.Mask2D(mask=extra_galaxies_mask, pixel_scales=data.pixel_scales)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Apply the extra galaxies mask to the image, which will remove them from visualization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = data.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the data with the new mask, in order to check that the mask removes the regions of the image corresponding to the\n", + "extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output to a .png file for easy inspection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the extra galaxies mask, which will be load and used before a model fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_array(\n", + " array=mask, file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"), overwrite=True\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The workspace also includes a GUI for image and noise-map scaling, which can be found at\n", + "`autolens_workspace/*/imaging/data_preparation/gui/mask_extra_galaxies.py`.\n", + "\n", + "This tools allows you `spray paint` on the image where an you want to scale, allow irregular patterns (i.e. not\n", + "rectangles) to be scaled." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/optional/positions.ipynb b/notebooks/imaging/data_preparation/examples/optional/positions.ipynb index 626d0b7b7..13353253f 100644 --- a/notebooks/imaging/data_preparation/examples/optional/positions.ipynb +++ b/notebooks/imaging/data_preparation/examples/optional/positions.ipynb @@ -1,224 +1,261 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Positions (Optional)\n", - "======================================\n", - "\n", - "The script manually marks the (y,x) arc-second positions of the multiply imaged lensed source galaxy in the image-plane,\n", - "under the assumption that they originate from the same location in the source-plane.\n", - "\n", - "A non-linear search (e.g. Nautilus) can then use these positions to preferentially choose mass models where these\n", - "positions trace close to one another in the source-plane. This speeding up the initial fitting of lens models and\n", - "removes unwanted solutions from parameter space which have too much or too little mass in the lens galaxy.\n", - "\n", - "If you create positions for your dataset, you must also update your modeling script to use them by loading them\n", - "and passing them to the `Analysis` object via a `PositionsLH` object.\n", - "\n", - "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using positions\n", - "to help the non-linear search find a good lens model.\n", - "\n", - "Links / Resources:\n", - "\n", - "Position-based lens model resampling is particularly important for fitting pixelized source models, for the\n", - "reasons disucssed in the following readthedocs\n", - "webapge https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", - "\n", - "The script `data_preparation/gui/positions.ipynb` shows how to use a Graphical User Interface (GUI) to mask the\n", - "positions on the lensed source.\n", - "\n", - "See `autolens_workspace/*/guides/modeling/customize` for an example.of how to use positions in a\n", - "`modeling` script.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where positions are output, which is `dataset/imaging/simple`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the `Imaging` dataset, so that the positions can be plotted over the strong lens image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, create a set of positions, which is a Coordinate of (y,x) values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(\n", - " values=[(0.4, 1.6), (1.58, -0.35), (-0.43, -1.59), (-1.45, 0.2)]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now lets plot the image and positions, so we can check that the positions overlap different regions of the source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we`re happy with the positions, lets output them to the dataset folder of the lens, so that we can load them from a\n", - ".json file in our pipelines!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=positions,\n", - " file_path=Path(dataset_path, \"positions.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Positions (Optional)\n", + "======================================\n", + "\n", + "The script manually marks the (y,x) arc-second positions of the multiply imaged lensed source galaxy in the image-plane,\n", + "under the assumption that they originate from the same location in the source-plane.\n", + "\n", + "A non-linear search (e.g. Nautilus) can then use these positions to preferentially choose mass models where these\n", + "positions trace close to one another in the source-plane. This speeding up the initial fitting of lens models and\n", + "removes unwanted solutions from parameter space which have too much or too little mass in the lens galaxy.\n", + "\n", + "If you create positions for your dataset, you must also update your modeling script to use them by loading them\n", + "and passing them to the `Analysis` object via a `PositionsLH` object.\n", + "\n", + "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using positions\n", + "to help the non-linear search find a good lens model.\n", + "\n", + "Links / Resources:\n", + "\n", + "Position-based lens model resampling is particularly important for fitting pixelized source models, for the\n", + "reasons disucssed in the following readthedocs\n", + "webapge https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", + "\n", + "The script `data_preparation/gui/positions.ipynb` shows how to use a Graphical User Interface (GUI) to mask the\n", + "positions on the lensed source.\n", + "\n", + "See `autolens_workspace/*/guides/modeling/customize` for an example.of how to use positions in a\n", + "`modeling` script.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where positions are output, which is `dataset/imaging/simple`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the `Imaging` dataset, so that the positions can be plotted over the strong lens image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, create a set of positions, which is a Coordinate of (y,x) values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(\n", + " values=[(0.4, 1.6), (1.58, -0.35), (-0.43, -1.59), (-1.45, 0.2)]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now lets plot the image and positions, so we can check that the positions overlap different regions of the source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we`re happy with the positions, lets output them to the dataset folder of the lens, so that we can load them from a\n", + ".json file in our pipelines!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=positions,\n", + " file_path=Path(dataset_path, \"positions.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/examples/psf.ipynb b/notebooks/imaging/data_preparation/examples/psf.ipynb index 997f2a87d..d986ba7e2 100644 --- a/notebooks/imaging/data_preparation/examples/psf.ipynb +++ b/notebooks/imaging/data_preparation/examples/psf.ipynb @@ -1,246 +1,283 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: PSF\n", - "=====================\n", - "\n", - "The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope. It is used by\n", - "PyAutoLens when fitting a dataset to include these effects, such that does not bias the model.\n", - "\n", - "It should be estimated from a stack of stars in the image during data reduction or using a PSF simulator (e.g. TinyTim\n", - "for Hubble).\n", - "\n", - "This tutorial describes preprocessing your dataset`s psf to adhere to the units and formats required by PyAutoLens.\n", - "\n", - "__Contents__\n", - "\n", - "- **Pixel Scale:** The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units.\n", - "- **Loading Data From Individual Fits Files:** Load a PSF from .fits files (a format commonly used by Astronomers) via the `Array2D` object.\n", - "- **PSF Dimensions:** The PSF dimensions must be odd x odd (e.g.\n", - "- **PSF Normalization:** The PSF should also be normalized to unity.\n", - "\n", - "__Pixel Scale__\n", - "\n", - "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", - "your telescope. You should look up now if you are unsure of the value.\n", - "\n", - "The pixel scale of some common telescopes is as follows:\n", - "\n", - " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", - " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", - " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", - " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", - " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", - "\n", - "It is absolutely vital you use the correct pixel scale, so double check this value!\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "# %matplotlib\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Setup the path the datasets we'll use to illustrate preprocessing, which is the folder `dataset/imaging/data_preparation`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Data From Individual Fits Files__\n", - "\n", - "Load a PSF from .fits files (a format commonly used by Astronomers) via the `Array2D` object. \n", - "\n", - "This image represents a good data-reduction that conforms **PyAutoLens** formatting standards!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_fits(\n", - " file_path=dataset_path / \"psf.fits\", hdu=0, pixel_scales=0.1\n", - ")\n", - "\n", - "aplt.plot_array(array=psf.kernel, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This psf conforms to **PyAutoLens** standards for the following reasons.\n", - "\n", - " - Size: The PSF has a shape 21 x 21 pixels, which is large enough to capture the PSF core and thus capture the \n", - " majority of the blurring effect, but not so large that the convolution slows down the analysis. Large \n", - " PSFs (e.g. 51 x 51) are supported, but will lead to much slower run times. The size of the PSF should be carefully \n", - " chosen to ensure it captures the majority of blurring due to the telescope optics, which for most instruments is \n", - " something around 11 x 11 to 21 x 21.\n", - "\n", - " - Oddness: The PSF has dimensions which are odd (an even PSF would for example have shape 20 x 20). The \n", - " convolution of an even PSF introduces a small shift in the modle images and produces an offset in the inferred\n", - " model parameters.\n", - " \n", - " - Normalization: The PSF has been normalized such that all values within the kernel sum to 1 (note how all values in \n", - " the example PSF are below zero with the majority below 0.01). This ensures that flux is conserved when convolution \n", - " is performed, ensuring that quantities like a galaxy's magnitude are computed accurately.\n", - "\n", - " - Centering: The PSF is at the centre of the array (as opposed to in a corner), ensuring that no shift is introduced\n", - " due to PSF blurring on the inferred model parameters.\n", - "\n", - "If your PSF conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you noise-map and image conform to standards first!).\n", - "\n", - "If it does not conform to standards, this script illustrates **PyAutoLens** functionality which can be used to \n", - "convert it to standards. \n", - "\n", - "__1) PSF Size__\n", - "\n", - "The majority of PSF blurring occurs at its central core, which is the most important region for lens modeling. \n", - "\n", - "By default, the size of the PSF kernel in the .fits is used to perform convolution. The larger this stamp, the longer \n", - "this convolution will take to run. Large PSFs (e.g. > 51 x 51) could have significantly slow down on run-time. \n", - "\n", - "In general we recommend the PSF size is 21 x 21. The example below is 11 x 11, which for this simulated data is just \n", - "about acceptable but would be on the small side for many real telescopes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_fits(\n", - " file_path=dataset_path / \"psf.fits\", hdu=0, pixel_scales=0.1\n", - ")\n", - "\n", - "aplt.plot_array(array=psf.kernel, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can resize a psf the same way that we resize an image.\n", - "\n", - "Below, we resize the PSF to 5 x 5 pixels, which is too small for a realistic analysis and just for demonstration \n", - "purposes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "trimmed_psf_kernel = al.preprocess.array_with_new_shape(\n", - " array=psf.kernel, new_shape=(5, 5)\n", - ")\n", - "\n", - "aplt.plot_array(array=trimmed_psf_kernel, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__PSF Dimensions__\n", - "\n", - "The PSF dimensions must be odd x odd (e.g. 21 x 21), because even-sized PSF kernels introduce a half-pixel offset in \n", - "the convolution routine which can lead to systematics in the lens analysis. \n", - "\n", - "The preprocess module contains functions for converting an even-sized PSF to an odd-sized PSF.\n", - "\n", - "https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/dataset/preprocess.py\n", - "\n", - "- `psf_with_odd_dimensions_from`\n", - "\n", - "However, this uses an interpolation routine that will not be perfect. The best way to create an odd-sized PSF is to do \n", - "so via the data reduction procedure. If this is a possibility, do that, this function is only for when you have no\n", - "other choice.\n", - "\n", - "__PSF Normalization__\n", - "\n", - "The PSF should also be normalized to unity. That is, the sum of all values in the kernel \n", - "should sum to 1. This ensures that the PSF convolution does not change the overall normalization of the image.\n", - "\n", - "PyAutoLens automatically normalized PSF when they are passed into a `Imaging` or `SimulatedImaging` object, so you \n", - "do not actually need to normalize your PSF. However, it is better to do it now, just in case.\n", - "\n", - "Below, we show how to normalize a PSF when it is loaded from a .fits file, by simply including the `normalize=True`\n", - "argument (the default value is `False`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_fits(\n", - " file_path=dataset_path / \"psf.fits\",\n", - " hdu=0,\n", - " pixel_scales=0.1,\n", - " normalize=True,\n", - ")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: PSF\n", + "=====================\n", + "\n", + "The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope. It is used by\n", + "PyAutoLens when fitting a dataset to include these effects, such that does not bias the model.\n", + "\n", + "It should be estimated from a stack of stars in the image during data reduction or using a PSF simulator (e.g. TinyTim\n", + "for Hubble).\n", + "\n", + "This tutorial describes preprocessing your dataset`s psf to adhere to the units and formats required by PyAutoLens.\n", + "\n", + "__Contents__\n", + "\n", + "- **Pixel Scale:** The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units.\n", + "- **Loading Data From Individual Fits Files:** Load a PSF from .fits files (a format commonly used by Astronomers) via the `Array2D` object.\n", + "- **PSF Dimensions:** The PSF dimensions must be odd x odd (e.g.\n", + "- **PSF Normalization:** The PSF should also be normalized to unity.\n", + "\n", + "__Pixel Scale__\n", + "\n", + "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", + "your telescope. You should look up now if you are unsure of the value.\n", + "\n", + "The pixel scale of some common telescopes is as follows:\n", + "\n", + " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", + " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", + " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", + " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", + " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", + "\n", + "It is absolutely vital you use the correct pixel scale, so double check this value!\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `data_preparation/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "# %matplotlib\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Setup the path the datasets we'll use to illustrate preprocessing, which is the folder `dataset/imaging/data_preparation`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Data From Individual Fits Files__\n", + "\n", + "Load a PSF from .fits files (a format commonly used by Astronomers) via the `Array2D` object. \n", + "\n", + "This image represents a good data-reduction that conforms **PyAutoLens** formatting standards!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_fits(\n", + " file_path=dataset_path / \"psf.fits\", hdu=0, pixel_scales=0.1\n", + ")\n", + "\n", + "aplt.plot_array(array=psf.kernel, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This psf conforms to **PyAutoLens** standards for the following reasons.\n", + "\n", + " - Size: The PSF has a shape 21 x 21 pixels, which is large enough to capture the PSF core and thus capture the \n", + " majority of the blurring effect, but not so large that the convolution slows down the analysis. Large \n", + " PSFs (e.g. 51 x 51) are supported, but will lead to much slower run times. The size of the PSF should be carefully \n", + " chosen to ensure it captures the majority of blurring due to the telescope optics, which for most instruments is \n", + " something around 11 x 11 to 21 x 21.\n", + "\n", + " - Oddness: The PSF has dimensions which are odd (an even PSF would for example have shape 20 x 20). The \n", + " convolution of an even PSF introduces a small shift in the modle images and produces an offset in the inferred\n", + " model parameters.\n", + " \n", + " - Normalization: The PSF has been normalized such that all values within the kernel sum to 1 (note how all values in \n", + " the example PSF are below zero with the majority below 0.01). This ensures that flux is conserved when convolution \n", + " is performed, ensuring that quantities like a galaxy's magnitude are computed accurately.\n", + "\n", + " - Centering: The PSF is at the centre of the array (as opposed to in a corner), ensuring that no shift is introduced\n", + " due to PSF blurring on the inferred model parameters.\n", + "\n", + "If your PSF conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you noise-map and image conform to standards first!).\n", + "\n", + "If it does not conform to standards, this script illustrates **PyAutoLens** functionality which can be used to \n", + "convert it to standards. \n", + "\n", + "__1) PSF Size__\n", + "\n", + "The majority of PSF blurring occurs at its central core, which is the most important region for lens modeling. \n", + "\n", + "By default, the size of the PSF kernel in the .fits is used to perform convolution. The larger this stamp, the longer \n", + "this convolution will take to run. Large PSFs (e.g. > 51 x 51) could have significantly slow down on run-time. \n", + "\n", + "In general we recommend the PSF size is 21 x 21. The example below is 11 x 11, which for this simulated data is just \n", + "about acceptable but would be on the small side for many real telescopes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_fits(\n", + " file_path=dataset_path / \"psf.fits\", hdu=0, pixel_scales=0.1\n", + ")\n", + "\n", + "aplt.plot_array(array=psf.kernel, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can resize a psf the same way that we resize an image.\n", + "\n", + "Below, we resize the PSF to 5 x 5 pixels, which is too small for a realistic analysis and just for demonstration \n", + "purposes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "trimmed_psf_kernel = al.preprocess.array_with_new_shape(\n", + " array=psf.kernel, new_shape=(5, 5)\n", + ")\n", + "\n", + "aplt.plot_array(array=trimmed_psf_kernel, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__PSF Dimensions__\n", + "\n", + "The PSF dimensions must be odd x odd (e.g. 21 x 21), because even-sized PSF kernels introduce a half-pixel offset in \n", + "the convolution routine which can lead to systematics in the lens analysis. \n", + "\n", + "The preprocess module contains functions for converting an even-sized PSF to an odd-sized PSF.\n", + "\n", + "https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/dataset/preprocess.py\n", + "\n", + "- `psf_with_odd_dimensions_from`\n", + "\n", + "However, this uses an interpolation routine that will not be perfect. The best way to create an odd-sized PSF is to do \n", + "so via the data reduction procedure. If this is a possibility, do that, this function is only for when you have no\n", + "other choice.\n", + "\n", + "__PSF Normalization__\n", + "\n", + "The PSF should also be normalized to unity. That is, the sum of all values in the kernel \n", + "should sum to 1. This ensures that the PSF convolution does not change the overall normalization of the image.\n", + "\n", + "PyAutoLens automatically normalized PSF when they are passed into a `Imaging` or `SimulatedImaging` object, so you \n", + "do not actually need to normalize your PSF. However, it is better to do it now, just in case.\n", + "\n", + "Below, we show how to normalize a PSF when it is loaded from a .fits file, by simply including the `normalize=True`\n", + "argument (the default value is `False`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_fits(\n", + " file_path=dataset_path / \"psf.fits\",\n", + " hdu=0,\n", + " pixel_scales=0.1,\n", + " normalize=True,\n", + ")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/gui/extra_galaxies_centres.ipynb b/notebooks/imaging/data_preparation/gui/extra_galaxies_centres.ipynb index e2c7dfcc5..805835ff7 100644 --- a/notebooks/imaging/data_preparation/gui/extra_galaxies_centres.ipynb +++ b/notebooks/imaging/data_preparation/gui/extra_galaxies_centres.ipynb @@ -1,260 +1,297 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "GUI Preprocessing: Extra Galaxies Centres\n", - "=========================================\n", - "\n", - "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source\n", - "and whose mass may contribute to the ray-tracing and lens model.\n", - "\n", - "The example `imaging/data_preparation/example/optional/extra_galaxies_centres.py` provides a full description of\n", - "what the extra galaxies are and how they are used in the model-fit. You should read this script first before\n", - "using this script.\n", - "\n", - "This script uses a GUI to mark the (y,x) arcsecond locations of these extra galaxies, in contrast to the example\n", - "above which requires you to input these values manually.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Search Box:** When you click on a pixel to mark a position, the search box looks around this click and finds the.\n", - "- **Clicker:** Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order.\n", - "- **Output:** Now lets plot the image and extra galaxy centres, so we can check that the centre overlaps the." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from matplotlib import pyplot as plt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "The path where the extra galaxy centres are output, which is `dataset/imaging/extra_galaxies`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the image which we will use to mark the lens light centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Box__\n", - "\n", - "When you click on a pixel to mark a position, the search box looks around this click and finds the pixel with\n", - "the highest flux to mark the position.\n", - "\n", - "The `search_box_size` is the number of pixels around your click this search takes place." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_box_size = 5" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Clicker__\n", - "\n", - "Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order to determine\n", - "the extra galaxy centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "clicker = al.Clicker(\n", - " image=data, pixel_scales=pixel_scales, search_box_size=search_box_size\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Set up the clicker canvas and load the GUI which you can now click on to mark the extra galaxy centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "n_y, n_x = data.shape_native\n", - "hw = int(n_x / 2) * pixel_scales\n", - "ext = [-hw, hw, -hw, hw]\n", - "fig = plt.figure(figsize=(14, 14))\n", - "plt.imshow(data.native, cmap=\"jet\", extent=ext)\n", - "plt.colorbar()\n", - "cid = fig.canvas.mpl_connect(\"button_press_event\", clicker.onclick)\n", - "plt.show()\n", - "fig.canvas.mpl_disconnect(cid)\n", - "plt.close(fig)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use the results of the Clicker GUI to create the list of extra galaxy centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_centres = al.Grid2DIrregular(values=clicker.click_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Now lets plot the image and extra galaxy centres, so we can check that the centre overlaps the brightest pixels in the\n", - "extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output this image of the extra galaxy centres to a .png file in the dataset folder for future reference." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the extra galaxy centres to the dataset folder of the lens, so that we can load them from a .json file \n", - "when we model them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=extra_galaxies_centres,\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "GUI Preprocessing: Extra Galaxies Centres\n", + "=========================================\n", + "\n", + "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source\n", + "and whose mass may contribute to the ray-tracing and lens model.\n", + "\n", + "The example `imaging/data_preparation/example/optional/extra_galaxies_centres.py` provides a full description of\n", + "what the extra galaxies are and how they are used in the model-fit. You should read this script first before\n", + "using this script.\n", + "\n", + "This script uses a GUI to mark the (y,x) arcsecond locations of these extra galaxies, in contrast to the example\n", + "above which requires you to input these values manually.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Search Box:** When you click on a pixel to mark a position, the search box looks around this click and finds the.\n", + "- **Clicker:** Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order.\n", + "- **Output:** Now lets plot the image and extra galaxy centres, so we can check that the centre overlaps the." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from matplotlib import pyplot as plt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "The path where the extra galaxy centres are output, which is `dataset/imaging/extra_galaxies`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the image which we will use to mark the lens light centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Box__\n", + "\n", + "When you click on a pixel to mark a position, the search box looks around this click and finds the pixel with\n", + "the highest flux to mark the position.\n", + "\n", + "The `search_box_size` is the number of pixels around your click this search takes place." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_box_size = 5" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Clicker__\n", + "\n", + "Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order to determine\n", + "the extra galaxy centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "clicker = al.Clicker(\n", + " image=data, pixel_scales=pixel_scales, search_box_size=search_box_size\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set up the clicker canvas and load the GUI which you can now click on to mark the extra galaxy centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "n_y, n_x = data.shape_native\n", + "hw = int(n_x / 2) * pixel_scales\n", + "ext = [-hw, hw, -hw, hw]\n", + "fig = plt.figure(figsize=(14, 14))\n", + "plt.imshow(data.native, cmap=\"jet\", extent=ext)\n", + "plt.colorbar()\n", + "cid = fig.canvas.mpl_connect(\"button_press_event\", clicker.onclick)\n", + "plt.show()\n", + "fig.canvas.mpl_disconnect(cid)\n", + "plt.close(fig)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use the results of the Clicker GUI to create the list of extra galaxy centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_centres = al.Grid2DIrregular(values=clicker.click_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Now lets plot the image and extra galaxy centres, so we can check that the centre overlaps the brightest pixels in the\n", + "extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output this image of the extra galaxy centres to a .png file in the dataset folder for future reference." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the extra galaxy centres to the dataset folder of the lens, so that we can load them from a .json file \n", + "when we model them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=extra_galaxies_centres,\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/gui/lens_light_centre.ipynb b/notebooks/imaging/data_preparation/gui/lens_light_centre.ipynb index 08435b4db..63b255577 100644 --- a/notebooks/imaging/data_preparation/gui/lens_light_centre.ipynb +++ b/notebooks/imaging/data_preparation/gui/lens_light_centre.ipynb @@ -1,267 +1,304 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "GUI Preprocessing: Lens Light Centre\n", - "====================================\n", - "\n", - "This tool allows one to input the lens light centre(s) of a strong lens(es) via a GUI, which can be used as a fixed\n", - "value in pipelines.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Search Box:** When you click on a pixel to mark a position, the search box looks around this click and finds the.\n", - "- **Clicker:** Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order.\n", - "- **Output:** Now lets plot the image and lens light centres, so we can check that the centre overlaps the." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from matplotlib import pyplot as plt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Setup the path the datasets we'll use to illustrate preprocessing, which is the \n", - "folder `dataset/imaging/lens_sersic`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"lens_sersic\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the image which we will use to mark the lens light centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Box__\n", - "\n", - "When you click on a pixel to mark a position, the search box looks around this click and finds the pixel with\n", - "the highest flux to mark the position.\n", - "\n", - "The `search_box_size` is the number of pixels around your click this search takes place." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_box_size = 5" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Clicker__\n", - "\n", - "Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order to determine\n", - "the lens light centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "clicker = al.Clicker(\n", - " image=data, pixel_scales=pixel_scales, search_box_size=search_box_size\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Set up the clicker canvas and load the GUI which you can now click on to mark the lens light centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "n_y, n_x = data.shape_native\n", - "hw = int(n_x / 2) * pixel_scales\n", - "ext = [-hw, hw, -hw, hw]\n", - "fig = plt.figure(figsize=(14, 14))\n", - "plt.imshow(data.native, cmap=\"jet\", extent=ext)\n", - "plt.colorbar()\n", - "cid = fig.canvas.mpl_connect(\"button_press_event\", clicker.onclick)\n", - "plt.show()\n", - "fig.canvas.mpl_disconnect(cid)\n", - "plt.close(fig)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use the results of the Clicker GUI to create the list of lens light centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_light_centres = al.Grid2DIrregular(values=clicker.click_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Now lets plot the image and lens light centres, so we can check that the centre overlaps the brightest pixel in the\n", - "lens light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output this image of the lens light centres to a .png file in the dataset folder for future reference." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the lens light centres to a .json file in the dataset folder, so we can load them in modeling scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=lens_light_centres,\n", - " file_path=Path(dataset_path, \"lens_light_centre.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "GUI Preprocessing: Lens Light Centre\n", + "====================================\n", + "\n", + "This tool allows one to input the lens light centre(s) of a strong lens(es) via a GUI, which can be used as a fixed\n", + "value in pipelines.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Search Box:** When you click on a pixel to mark a position, the search box looks around this click and finds the.\n", + "- **Clicker:** Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order.\n", + "- **Output:** Now lets plot the image and lens light centres, so we can check that the centre overlaps the." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from matplotlib import pyplot as plt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Setup the path the datasets we'll use to illustrate preprocessing, which is the \n", + "folder `dataset/imaging/lens_sersic`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"lens_sersic\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the image which we will use to mark the lens light centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Box__\n", + "\n", + "When you click on a pixel to mark a position, the search box looks around this click and finds the pixel with\n", + "the highest flux to mark the position.\n", + "\n", + "The `search_box_size` is the number of pixels around your click this search takes place." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_box_size = 5" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Clicker__\n", + "\n", + "Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order to determine\n", + "the lens light centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "clicker = al.Clicker(\n", + " image=data, pixel_scales=pixel_scales, search_box_size=search_box_size\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set up the clicker canvas and load the GUI which you can now click on to mark the lens light centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "n_y, n_x = data.shape_native\n", + "hw = int(n_x / 2) * pixel_scales\n", + "ext = [-hw, hw, -hw, hw]\n", + "fig = plt.figure(figsize=(14, 14))\n", + "plt.imshow(data.native, cmap=\"jet\", extent=ext)\n", + "plt.colorbar()\n", + "cid = fig.canvas.mpl_connect(\"button_press_event\", clicker.onclick)\n", + "plt.show()\n", + "fig.canvas.mpl_disconnect(cid)\n", + "plt.close(fig)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use the results of the Clicker GUI to create the list of lens light centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_light_centres = al.Grid2DIrregular(values=clicker.click_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Now lets plot the image and lens light centres, so we can check that the centre overlaps the brightest pixel in the\n", + "lens light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output this image of the lens light centres to a .png file in the dataset folder for future reference." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the lens light centres to a .json file in the dataset folder, so we can load them in modeling scripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=lens_light_centres,\n", + " file_path=Path(dataset_path, \"lens_light_centre.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/gui/mask.ipynb b/notebooks/imaging/data_preparation/gui/mask.ipynb index 71de507de..7b0e5b21a 100644 --- a/notebooks/imaging/data_preparation/gui/mask.ipynb +++ b/notebooks/imaging/data_preparation/gui/mask.ipynb @@ -1,191 +1,228 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "GUI Preprocessing: Mask\n", - "=======================\n", - "\n", - "This tool allows one to mask a bespoke mask for a given image of a strong lens using an interactive GUI. This mask\n", - "can then be loaded before a pipeline is run and passed to that pipeline so as to become the default masked used by a\n", - "search (if a mask function is not passed to that search).\n", - "\n", - "This GUI is adapted from the following code: https://gist.github.com/brikeats/4f63f867fd8ea0f196c78e9b835150ab\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Scribbler:** Load the Scribbler GUI for drawing the mask.\n", - "- **Output:** Now lets plot the image and mask, so we can check that the mask includes the regions of the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "import numpy as np" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Setup the path the datasets we'll use to illustrate preprocessing, which is the \n", - "folder `dataset/imaging/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the `Imaging` dataset, so that the mask can be plotted over the strong lens image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scribbler__\n", - "\n", - "Load the Scribbler GUI for drawing the mask. \n", - "\n", - "Push Esc when you are finished drawing the mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scribbler = al.Scribbler(image=data.native)\n", - "mask = scribbler.show_mask()\n", - "mask = al.Mask2D(mask=np.invert(mask), pixel_scales=pixel_scales)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Now lets plot the image and mask, so we can check that the mask includes the regions of the image we want." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output this image of the mask to a .png file in the dataset folder for future reference." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output it to the dataset folder of the lens, so that we can load it from a .fits in our modeling scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_array(\n", - " array=mask, file_path=Path(dataset_path, \"mask_gui.fits\"), overwrite=True\n", - ")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "GUI Preprocessing: Mask\n", + "=======================\n", + "\n", + "This tool allows one to mask a bespoke mask for a given image of a strong lens using an interactive GUI. This mask\n", + "can then be loaded before a pipeline is run and passed to that pipeline so as to become the default masked used by a\n", + "search (if a mask function is not passed to that search).\n", + "\n", + "This GUI is adapted from the following code: https://gist.github.com/brikeats/4f63f867fd8ea0f196c78e9b835150ab\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Scribbler:** Load the Scribbler GUI for drawing the mask.\n", + "- **Output:** Now lets plot the image and mask, so we can check that the mask includes the regions of the image." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "import numpy as np" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Setup the path the datasets we'll use to illustrate preprocessing, which is the \n", + "folder `dataset/imaging/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the `Imaging` dataset, so that the mask can be plotted over the strong lens image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scribbler__\n", + "\n", + "Load the Scribbler GUI for drawing the mask. \n", + "\n", + "Push Esc when you are finished drawing the mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scribbler = al.Scribbler(image=data.native)\n", + "mask = scribbler.show_mask()\n", + "mask = al.Mask2D(mask=np.invert(mask), pixel_scales=pixel_scales)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Now lets plot the image and mask, so we can check that the mask includes the regions of the image we want." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output this image of the mask to a .png file in the dataset folder for future reference." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output it to the dataset folder of the lens, so that we can load it from a .fits in our modeling scripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_array(\n", + " array=mask, file_path=Path(dataset_path, \"mask_gui.fits\"), overwrite=True\n", + ")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/gui/mask_extra_galaxies.ipynb b/notebooks/imaging/data_preparation/gui/mask_extra_galaxies.ipynb index e5fa6de83..99b94a27c 100644 --- a/notebooks/imaging/data_preparation/gui/mask_extra_galaxies.ipynb +++ b/notebooks/imaging/data_preparation/gui/mask_extra_galaxies.ipynb @@ -1,260 +1,297 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "GUI Preprocessing: Extra Galaxies Mask (Optional)\n", - "=================================================\n", - "\n", - "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated\n", - "with the strong lenswe are studying. The emission from these images will impact our model fitting and needs to be\n", - "removed from the analysis.\n", - "\n", - "The example `imaging/data_preparation/example/optional/extra_galaxies_mask.py` provides a full description of\n", - "what the extra galaxies are and how they are used in the model-fit. You should read this script first before\n", - "using this script.\n", - "\n", - "This script uses a GUI to mark the regions of the image where these extra galaxies are located, in contrast to the\n", - "example above which requires you to input these values manually.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Scribbler:** Load the Scribbler GUI for spray painting the scaled regions of the dataset.\n", - "- **Output:** The new image is plotted for inspection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "import numpy as np" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "The path where the extra galaxy mask is output, which is `dataset/imaging/extra_galaxies`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the `Imaging` data, where the extra galaxies are visible in the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")\n", - "\n", - "data = al.Array2D(\n", - " values=np.nan_to_num(data, nan=0.0, posinf=0.0, neginf=0.0), mask=data.mask\n", - ")\n", - "\n", - "cmap = \"jet\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Create a 3.0\" mask to plot over the image to guide where extra galaxy light needs its emission removed and noise scaled." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=data.shape_native, pixel_scales=data.pixel_scales, radius=mask_radius\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scribbler__\n", - "\n", - "Load the Scribbler GUI for spray painting the scaled regions of the dataset. \n", - "\n", - "Push Esc when you are finished spray painting." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scribbler = al.Scribbler(image=data.native, cmap=cmap, mask_overlay=mask)\n", - "mask = scribbler.show_mask()\n", - "mask = al.Mask2D(mask=mask, pixel_scales=pixel_scales)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The GUI has now closed and the extra galaxies mask has been created.\n", - "\n", - "Apply the extra galaxies mask to the image, which will remove them from visualization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = data.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "The new image is plotted for inspection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the data with the new mask, in order to check that the mask removes the regions of the image corresponding to the\n", - "extra galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output to a .png file for easy inspection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the extra galaxies mask, which will be load and used before a model fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_array(\n", - " array=mask, file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"), overwrite=True\n", - ")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "GUI Preprocessing: Extra Galaxies Mask (Optional)\n", + "=================================================\n", + "\n", + "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated\n", + "with the strong lenswe are studying. The emission from these images will impact our model fitting and needs to be\n", + "removed from the analysis.\n", + "\n", + "The example `imaging/data_preparation/example/optional/extra_galaxies_mask.py` provides a full description of\n", + "what the extra galaxies are and how they are used in the model-fit. You should read this script first before\n", + "using this script.\n", + "\n", + "This script uses a GUI to mark the regions of the image where these extra galaxies are located, in contrast to the\n", + "example above which requires you to input these values manually.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Scribbler:** Load the Scribbler GUI for spray painting the scaled regions of the dataset.\n", + "- **Output:** The new image is plotted for inspection." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "import numpy as np" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "The path where the extra galaxy mask is output, which is `dataset/imaging/extra_galaxies`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the `Imaging` data, where the extra galaxies are visible in the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")\n", + "\n", + "data = al.Array2D(\n", + " values=np.nan_to_num(data, nan=0.0, posinf=0.0, neginf=0.0), mask=data.mask\n", + ")\n", + "\n", + "cmap = \"jet\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Create a 3.0\" mask to plot over the image to guide where extra galaxy light needs its emission removed and noise scaled." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=data.shape_native, pixel_scales=data.pixel_scales, radius=mask_radius\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scribbler__\n", + "\n", + "Load the Scribbler GUI for spray painting the scaled regions of the dataset. \n", + "\n", + "Push Esc when you are finished spray painting." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scribbler = al.Scribbler(image=data.native, cmap=cmap, mask_overlay=mask)\n", + "mask = scribbler.show_mask()\n", + "mask = al.Mask2D(mask=mask, pixel_scales=pixel_scales)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The GUI has now closed and the extra galaxies mask has been created.\n", + "\n", + "Apply the extra galaxies mask to the image, which will remove them from visualization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = data.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "The new image is plotted for inspection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the data with the new mask, in order to check that the mask removes the regions of the image corresponding to the\n", + "extra galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output to a .png file for easy inspection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the extra galaxies mask, which will be load and used before a model fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_array(\n", + " array=mask, file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"), overwrite=True\n", + ")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/gui/positions.ipynb b/notebooks/imaging/data_preparation/gui/positions.ipynb index 6ba535e2d..d2e73db5e 100644 --- a/notebooks/imaging/data_preparation/gui/positions.ipynb +++ b/notebooks/imaging/data_preparation/gui/positions.ipynb @@ -1,275 +1,312 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "GUI Preprocessing: Positions\n", - "============================\n", - "\n", - "This tool allows one to input the positions of strong lenses via a GUI, which can be used to penalize inaccurate\n", - "mass models during lensing modeling.\n", - "\n", - "This GUI is adapted from the following code: https://gist.github.com/brikeats/4f63f867fd8ea0f196c78e9b835150ab\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Search Box:** When you click on a pixel to mark a position, the search box looks around this click and finds the.\n", - "- **Clicker:** Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order.\n", - "- **Output:** Now lets plot the image and positions,, so we can check that the positions overlap the brightest." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from matplotlib import pyplot as plt\n", - "import numpy as np" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Setup the path the datasets we'll use to illustrate preprocessing, which is the \n", - "folder `dataset/imaging/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The pixel scale of the imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales = 0.1" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Load the image which we will use to mark the positions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search Box__\n", - "\n", - "When you click on a pixel to mark a position, the search box looks around this click and finds the pixel with\n", - "the highest flux to mark the position.\n", - "\n", - "The `search_box_size` is the number of pixels around your click this search takes place." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_box_size = 5" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Clicker__\n", - "\n", - "Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order to determine\n", - "the positions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "clicker = al.Clicker(\n", - " image=data, pixel_scales=pixel_scales, search_box_size=search_box_size\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For lenses with bright lens light emission, it can be difficult to get the source light to show. The normalization\n", - "below uses a log-scale with a capped maximum, which better contrasts the lens and source emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "cmap = \"jet\"\n", - "\n", - "norm = cmap.norm_from(array=None)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Set up the clicker canvas and load the GUI which you can now click on to mark the positionss." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "n_y, n_x = data.shape_native\n", - "hw = int(n_x / 2) * pixel_scales\n", - "ext = [-hw, hw, -hw, hw]\n", - "fig = plt.figure(figsize=(14, 14))\n", - "plt.imshow(data.native, cmap=\"jet\", extent=ext)\n", - "plt.colorbar()\n", - "cid = fig.canvas.mpl_connect(\"button_press_event\", clicker.onclick)\n", - "plt.show()\n", - "fig.canvas.mpl_disconnect(cid)\n", - "plt.close(fig)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use the results of the Clicker GUI to create the list of the positions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(values=clicker.click_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Now lets plot the image and positions,, so we can check that the positions overlap the brightest pixels in the\n", - "lensed source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output this image of the positions to a .png file in the dataset folder for future reference." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output the positions to a .json file in the dataset folder, so we can load them in modeling scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=positions,\n", - " file_path=Path(dataset_path, \"positions.json\"),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "GUI Preprocessing: Positions\n", + "============================\n", + "\n", + "This tool allows one to input the positions of strong lenses via a GUI, which can be used to penalize inaccurate\n", + "mass models during lensing modeling.\n", + "\n", + "This GUI is adapted from the following code: https://gist.github.com/brikeats/4f63f867fd8ea0f196c78e9b835150ab\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Search Box:** When you click on a pixel to mark a position, the search box looks around this click and finds the.\n", + "- **Clicker:** Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order.\n", + "- **Output:** Now lets plot the image and positions,, so we can check that the positions overlap the brightest." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from matplotlib import pyplot as plt\n", + "import numpy as np" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Setup the path the datasets we'll use to illustrate preprocessing, which is the \n", + "folder `dataset/imaging/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The pixel scale of the imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales = 0.1" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Load the image which we will use to mark the positions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"data.fits\", pixel_scales=pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search Box__\n", + "\n", + "When you click on a pixel to mark a position, the search box looks around this click and finds the pixel with\n", + "the highest flux to mark the position.\n", + "\n", + "The `search_box_size` is the number of pixels around your click this search takes place." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_box_size = 5" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Clicker__\n", + "\n", + "Set up the `Clicker` object from the `clicker.py` module, which monitors your mouse clicks in order to determine\n", + "the positions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "clicker = al.Clicker(\n", + " image=data, pixel_scales=pixel_scales, search_box_size=search_box_size\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For lenses with bright lens light emission, it can be difficult to get the source light to show. The normalization\n", + "below uses a log-scale with a capped maximum, which better contrasts the lens and source emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cmap = \"jet\"\n", + "\n", + "norm = cmap.norm_from(array=None)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set up the clicker canvas and load the GUI which you can now click on to mark the positionss." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "n_y, n_x = data.shape_native\n", + "hw = int(n_x / 2) * pixel_scales\n", + "ext = [-hw, hw, -hw, hw]\n", + "fig = plt.figure(figsize=(14, 14))\n", + "plt.imshow(data.native, cmap=\"jet\", extent=ext)\n", + "plt.colorbar()\n", + "cid = fig.canvas.mpl_connect(\"button_press_event\", clicker.onclick)\n", + "plt.show()\n", + "fig.canvas.mpl_disconnect(cid)\n", + "plt.close(fig)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use the results of the Clicker GUI to create the list of the positions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(values=clicker.click_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Now lets plot the image and positions,, so we can check that the positions overlap the brightest pixels in the\n", + "lensed source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output this image of the positions to a .png file in the dataset folder for future reference." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output the positions to a .json file in the dataset folder, so we can load them in modeling scripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=positions,\n", + " file_path=Path(dataset_path, \"positions.json\"),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/manual/mask_irregular.ipynb b/notebooks/imaging/data_preparation/manual/mask_irregular.ipynb index 703fb7296..42fc65668 100644 --- a/notebooks/imaging/data_preparation/manual/mask_irregular.ipynb +++ b/notebooks/imaging/data_preparation/manual/mask_irregular.ipynb @@ -1,198 +1,235 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Manual Preprocessing: Mask Irregular\n", - "====================================" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "import numpy as np" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This tool allows one to mask a bespoke mask for a given image of a strong lens, which can then be loaded\n", - "before a model-fit.\n", - "\n", - "This tool creates an irregular mask, which can form any shape and is not restricted to circles, annuli, ellipses,\n", - "etc. This mask is created as follows:\n", - "\n", - "1) Blur the observed image with a Gaussian kernel of specified FWHM.\n", - "2) Compute the absolute S/N map of that blurred image and the noise-map.\n", - "3) Create the mask for all pixels with a S/N above a theshold value.\n", - "\n", - "For strong lenses without a lens light component this masks create a source-only mask. If the lens light is included\n", - "it includes the lens light and source.\n", - "\n", - "The following parameters determine the behaviour of this function:\n", - "\n", - "The sigma value (e.g. FWHM) of the Gaussian the image is blurred with and the S/N threshold defining above which a \n", - "image-pixel value must be to not be masked." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "blurring_gaussian_sigma = 0.1\n", - "snr_cut = 10.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Setup the path the datasets we'll use to illustrate preprocessing, which is the \n", - "folder `dataset/imaging/data_preparation/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", - "noise_map = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"noise_map.fits\", pixel_scales=0.1\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Returns the 2D Gaussian that the image is blurred with. This blurring smooths over noise in the image, which will \n", - "otherwise lead unmasked values with in individual pixels if not smoothed over correctly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "blurring_gaussian = al.Convolver.from_gaussian(\n", - " shape_native=(31, 31),\n", - " pixel_scales=data.pixel_scales,\n", - " sigma=blurring_gaussian_sigma,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Blur the image with this Gaussian smoothing kernel and plot the resulting image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "blurred_image = blurring_gaussian.convolved_image_from(image=data, blurring_image=None)\n", - "aplt.plot_array(array=blurred_image, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now compute the absolute signal-to-noise map of this blurred image, given the noise-map of the observed dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "blurred_signal_to_noise_map = blurred_image / noise_map\n", - "aplt.plot_array(array=blurred_signal_to_noise_map, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now create the mask in 2ll pixels where the signal to noise is above some threshold value." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = np.where(blurred_signal_to_noise_map.native > snr_cut, False, True)\n", - "mask = al.Mask2D(mask=mask, pixel_scales=data.pixel_scales)\n", - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we`re happy with the mask, lets output it to the dataset folder of the lens, so that we can load it from a .fits\n", - "file in our pipelines!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_array(array=mask, file_path=Path(dataset_path, \"mask.fits\"), overwrite=True)\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Manual Preprocessing: Mask Irregular\n", + "====================================" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "import numpy as np" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This tool allows one to mask a bespoke mask for a given image of a strong lens, which can then be loaded\n", + "before a model-fit.\n", + "\n", + "This tool creates an irregular mask, which can form any shape and is not restricted to circles, annuli, ellipses,\n", + "etc. This mask is created as follows:\n", + "\n", + "1) Blur the observed image with a Gaussian kernel of specified FWHM.\n", + "2) Compute the absolute S/N map of that blurred image and the noise-map.\n", + "3) Create the mask for all pixels with a S/N above a theshold value.\n", + "\n", + "For strong lenses without a lens light component this masks create a source-only mask. If the lens light is included\n", + "it includes the lens light and source.\n", + "\n", + "The following parameters determine the behaviour of this function:\n", + "\n", + "The sigma value (e.g. FWHM) of the Gaussian the image is blurred with and the S/N threshold defining above which a \n", + "image-pixel value must be to not be masked." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "blurring_gaussian_sigma = 0.1\n", + "snr_cut = 10.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Setup the path the datasets we'll use to illustrate preprocessing, which is the \n", + "folder `dataset/imaging/data_preparation/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", + "noise_map = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"noise_map.fits\", pixel_scales=0.1\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Returns the 2D Gaussian that the image is blurred with. This blurring smooths over noise in the image, which will \n", + "otherwise lead unmasked values with in individual pixels if not smoothed over correctly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "blurring_gaussian = al.Convolver.from_gaussian(\n", + " shape_native=(31, 31),\n", + " pixel_scales=data.pixel_scales,\n", + " sigma=blurring_gaussian_sigma,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Blur the image with this Gaussian smoothing kernel and plot the resulting image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "blurred_image = blurring_gaussian.convolved_image_from(image=data, blurring_image=None)\n", + "aplt.plot_array(array=blurred_image, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now compute the absolute signal-to-noise map of this blurred image, given the noise-map of the observed dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "blurred_signal_to_noise_map = blurred_image / noise_map\n", + "aplt.plot_array(array=blurred_signal_to_noise_map, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now create the mask in 2ll pixels where the signal to noise is above some threshold value." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = np.where(blurred_signal_to_noise_map.native > snr_cut, False, True)\n", + "mask = al.Mask2D(mask=mask, pixel_scales=data.pixel_scales)\n", + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we`re happy with the mask, lets output it to the dataset folder of the lens, so that we can load it from a .fits\n", + "file in our pipelines!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_array(array=mask, file_path=Path(dataset_path, \"mask.fits\"), overwrite=True)\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/data_preparation/start_here.ipynb b/notebooks/imaging/data_preparation/start_here.ipynb index fc17d9828..4dfe3a52d 100644 --- a/notebooks/imaging/data_preparation/start_here.ipynb +++ b/notebooks/imaging/data_preparation/start_here.ipynb @@ -1,393 +1,430 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Imaging: Data Preparation\n", - "=========================\n", - "\n", - "When a CCD imaging dataset is analysed, it must conform to certain standards in order for the analysis\n", - "to be performed correctly. This tutorial describes these standards and links to more detailed scripts which will help\n", - "you prepare your dataset to adhere to them if it does not already.\n", - "\n", - "__Contents__\n", - "\n", - "- **Pixel Scale:** The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units.\n", - "- **Image:** The image is the image of your strong lens, which comes from a telescope like the Hubble Space.\n", - "- **Noise Map:** The noise-map defines the uncertainty in every pixel of your strong lens image, where values are.\n", - "- **PSF:** The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope.\n", - "- **Data Processing Complete:** If your image, noise-map and PSF conform the standards above, you are ready to analyse your dataset!\n", - "\n", - "__Pixel Scale__\n", - "\n", - "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", - "your telescope. You should look up now if you are unsure of the value.\n", - "\n", - "The pixel scale of some common telescopes is as follows:\n", - "\n", - " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", - " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", - " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", - " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", - " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", - "\n", - "It is absolutely vital you use the correct pixel scale, so double check this value!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image__\n", - "\n", - "The image is the image of your strong lens, which comes from a telescope like the Hubble Space telescope (HST).\n", - "\n", - "Lets inspect an image which conforms to **PyAutoLens** standards:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\") / \"imaging\" / \"simple\"\n", - "\n", - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This image conforms to **PyAutoLens** standards for the following reasons.\n", - "\n", - " - Units: The image flux is in units of electrons per second (as opposed to electrons, counts, ADU`s etc.). \n", - " Internal **PyAutoLens** functions for computing quantities like galaxy magnitudes assume the data and model\n", - " light profiles are in electrons per second.\n", - " \n", - " - Centering: The lens galaxy is at the centre of the image (as opposed to in a corner). Default **PyAutoLens**\n", - " parameter priors assume the lens galaxy is at the centre of the image.\n", - " \n", - " - Stamp Size: The image is a postage stamp cut-out of the lens, but does not include many pixels around the edge of\n", - " the lens. It is advised to cut out a postage stamp of the lens, as opposed to the entire image, as this reduces\n", - " the amount of memory **PyAutoLens** uses, speeds up the analysis and ensures visualization zooms around the lens\n", - " and source. However, conforming to this standard is not necessary to ensure an accurate **PyAutoLens** analysis.\n", - " \n", - "If your image conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you noise-map and PSF conform to standards first!).\n", - "\n", - "**Links / Resources:**\n", - "\n", - " - `imaging/data_preparation/examples/data.ipynb`: tools to process the data to conform to these standards.\n", - "\n", - "__Noise Map__\n", - "\n", - "The noise-map defines the uncertainty in every pixel of your strong lens image, where values are defined as the \n", - "RMS standard deviation in every pixel (not the variances, HST WHT-map values, etc.). \n", - "\n", - "Lets inspect a noise-map which conforms to **PyAutoLens** standards:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_map = al.Array2D.from_fits(\n", - " file_path=dataset_path / \"noise_map.fits\", pixel_scales=0.1\n", - ")\n", - "\n", - "aplt.plot_array(array=noise_map, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This noise-map conforms to **PyAutoLens** standards for the following reasons:\n", - "\n", - " - Units: Like its corresponding image, it is in units of electrons per second (as opposed to electrons, counts, \n", - " ADU`s etc.). Internal **PyAutoLens** functions for computing quantities like a galaxy magnitude assume the data and \n", - " model light profiles are in electrons per second.\n", - "\n", - " - Values: The noise-map values themselves are the RMS standard deviations of the noise in every pixel. When a model \n", - " is fitted to data in **PyAutoLens** and a likelihood is evaluated, this calculation assumes that this is the\n", - " corresponding definition of the noise-map. The noise map therefore should not be the variance of the noise, or \n", - " another definition of noise.\n", - "\n", - "If you are not certain what the definition of the noise-map you have available to you is, or do not know how to\n", - "compute a noise-map at all, you should refer to the instrument handbook of the telescope your data is from. It is\n", - "absolutely vital that the noise-map is correct, as it is the only way **PyAutoLens** can quantify the goodness-of-fit.\n", - "\n", - "A sanity check for a reliable noise map is that the signal-to-noise of the lens galaxy is somewhere between a value of \n", - "10 - 300 and source around 5 - 50, however this is not a definitive test.\n", - " \n", - "Given the image should be centred and cut-out around the lens galaxy, so should the noise-map.\n", - "\n", - "If your noise-map conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you image and PSF conform to standards first!).\n", - "\n", - "**Links / Resources:**\n", - "\n", - " - `imaging/data_preparation/examples/noise_map.ipynb`: tools to process the noise-map to conform to these standards.\n", - "\n", - "__PSF__\n", - "\n", - "The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope. It is used when fitting \n", - "a dataset to include these effects, such that does not bias the model.\n", - "\n", - "It should be estimated from a stack of stars in the image during data reduction or using a PSF simulator (e.g. TinyTim\n", - "for Hubble).\n", - "\n", - "Lets inspect a PSF which conforms to **PyAutoLens** standards:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_fits(\n", - " file_path=dataset_path / \"psf.fits\", hdu=0, pixel_scales=0.1\n", - ")\n", - "\n", - "aplt.plot_array(array=psf.kernel, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This psf conforms to **PyAutoLens** standards for the following reasons.\n", - "\n", - " - Size: The PSF has a shape 21 x 21 pixels, which is large enough to capture the PSF core and thus capture the \n", - " majority of the blurring effect, but not so large that the convolution slows down the analysis. Large \n", - " PSFs (e.g. 51 x 51) are supported, but will lead to much slower run times. The size of the PSF should be carefully \n", - " chosen to ensure it captures the majority of blurring due to the telescope optics, which for most instruments is \n", - " something around 11 x 11 to 21 x 21.\n", - "\n", - " - Oddness: The PSF has dimensions which are odd (an even PSF would for example have shape 20 x 20). The \n", - " convolution of an even PSF introduces a small shift in the model images and produces an offset in the inferred\n", - " lens model parameters. Inputting an even PSF will lead **PyAutoLens** to raise an error.\n", - "\n", - " - Normalization: The PSF has been normalized such that all values within the kernel sum to 1 (note how all values in \n", - " the example PSF are below zero with the majority below 0.01). This ensures that flux is conserved when convolution \n", - " is performed, ensuring that quantities like a galaxy's magnitude are computed accurately.\n", - "\n", - " - Centering: The PSF is at the centre of the array (as opposed to in a corner), ensuring that no shift is introduced\n", - " due to PSF blurring on the inferred model parameters.\n", - "\n", - "If your PSF conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", - "you noise-map and image conform to standards first!).\n", - "\n", - "**Links / Resources:**\n", - "\n", - " - `imaging/data_preparation/examples/psf.ipynb`: tools to process the PSF to conform to these standards.\n", - "\n", - "__Data Processing Complete__\n", - "\n", - "If your image, noise-map and PSF conform the standards above, you are ready to analyse your dataset!\n", - "\n", - "Below, we provide an overview of optional data preparation steos which prepare other aspects of the analysis. \n", - "\n", - "New users are recommended to skim-read the optional steps below so they are aware of them, but to not perform them \n", - "and instead analyse their dataset now. You can come back to the data preparation scripts below if it becomes necessary.\n", - "\n", - "__Mask (Optional)__\n", - "\n", - "The mask removes the regions of the image where the lens and source galaxy are not present, typically the edges of the \n", - "image.\n", - "\n", - "Example modeling scripts internally create a 3.0\" circular mask and therefore do not require that a mask has been \n", - "created externally via a data preparation script. \n", - "\n", - "This script shows how to create customize masked (e.g. annular, ellipses) which are tailored to match the lens or\n", - "lensed source emission. \n", - "\n", - "If you have not analysed your dataset yet and do not know of a specific reason why you need the bespoke masks \n", - "created by this script, it is recommended that you simply use the default ~3.0\" circular mask internally made in each\n", - "script and omit this data preparation tutorial.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/mask.ipynb`: tools to create a bespoke mask for your dataset.\n", - "- `data_preparation/examples/gui/mask.ipynb`: use a Graphical User Interface (GUI) to create a bespoke mask.\n", - "\n", - "__Positions (Optional)__\n", - "\n", - "The script allows you to mark the (y,x) arc-second positions of the multiply imaged lensed source galaxy in \n", - "the image-plane, under the assumption that they originate from the same location in the source-plane.\n", - "\n", - "A non-linear search (e.g. Nautilus) can then use these positions to preferentially choose mass models where these \n", - "positions trace close to one another in the source-plane. This speeding up the initial fitting of lens models and \n", - "removes unwanted solutions from parameter space which have too much or too little mass in the lens galaxy.\n", - "\n", - "If you create positions for your dataset, you must also update your modeling script to use them by loading them \n", - "and passing them to the `Analysis` object via a `PositionsLH` object. \n", - "\n", - "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using positions\n", - "to help the non-linear search find a good lens model.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "Position-based lens model resampling is particularly important for fitting pixelized source models, for the\n", - "reasons disucssed in the following readthedocs \n", - "webapge https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", - "\n", - "- `data_preparation/examples/optional/positions.ipynb`: input the positions manually into a Python script.\n", - "- `data_preparation/gui/positions.ipynb` use a Graphical User Interface (GUI) to mark the positions.\n", - "- `guides/modeling/customize` for an example.of how to use positions in a `modeling` script.\n", - "\n", - "\n", - "__Lens Light Centre (Optional)__\n", - "\n", - "This script allows you to mark the (y,x) arcsecond locations of the lens galaxy light centre(s) of the strong lens\n", - "you are analysing. These can be used as fixed values for the lens light and mass models in a model-fit.\n", - "\n", - "This reduces the number of free parameters fitted for in a lens model and removes inaccurate solutions where\n", - "the lens mass model centre is unrealistically far from its true centre.\n", - "\n", - "Advanced `chaining` scripts often use these input centres in the early fits to infer an accurate initial lens model,\n", - "amd then make the centres free parameters in later searches to ensure a general and accurate lens model is inferred.\n", - "\n", - "If you create a `light_centre` for your dataset, you must also update your modeling script to use them.\n", - "\n", - "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using a fixed\n", - "lens light and / or mass centre to help the non-linear search find a good lens model.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/lens_light_centre.py`: input the lens galaxy light centre manually into a Python script.\n", - "- `data_preparation/gui/lens_light_centre.ipynb` use a Graphical User Interface (GUI) to mask the lens galaxy light centre.\n", - "\n", - "\n", - "__Extra Galaxies (Optional)__\n", - "\n", - "There may be galaxies nearby the lens and source galaxies, whose emission blends with that of the lens and source\n", - "and whose mass may contribute to the ray-tracing and lens model.\n", - "\n", - "We can include these galaxies in the lens model, either as light profiles, mass profiles, or both, using the\n", - "modeling API, where these nearby objects are denoted `extra_galaxies`.\n", - "\n", - "The script `extra_galaxies_centres.py` marks the (y,x) arcsecond locations of these extra galaxies, so that when they \n", - "are included in the lens model the centre of these extra galaxies light and / or mass profiles are fixed to these \n", - "values (or their priors are initialized surrounding these centres).\n", - "\n", - "The example `mask_extra_galaxies.py` (see below) masks the regions of an image where extra galaxies are present. \n", - "This mask is used to remove their signal from the data and increase their noise to make them not impact the fit. \n", - "This means their luminous emission does not need to be included in the model, reducing the number of free parameters \n", - "and speeding up the analysis. It is still a choice whether their mass is included in the model.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/extra_galaxies_centres.py`: input the extra galaxy centres manually into a \n", - " Python script.\n", - "- `data_preparation/gui/extra_galaxies_centres.ipynb`: use a Graphical User Interface (GUI) to mark the extra galaxy centres.\n", - "- `features/extra_galaxies.py` how to use extra galaxies in a model-fit, including loading the extra galaxy centres.\n", - "\n", - "\n", - "__Mask Extra Galaxies (Optional)__\n", - "\n", - "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated \n", - "with the strong lens we are studying. The emission from these images will impact our model fitting and needs to be \n", - "removed from the analysis.\n", - "\n", - "This script creates a mask of these regions of the image, called the `mask_extra_galaxies`, which can be used to\n", - "prevent them from impacting a fit. This mask may also include emission from objects which are not technically galaxies,\n", - "but blend with the galaxy we are studying in a similar way. Common examples of such objects are foreground stars\n", - "or emission due to the data reduction process.\n", - "\n", - "The mask can be applied in different ways. For example, it could be applied such that the image pixels are discarded\n", - "from the fit entirely, Alternatively the mask could be used to set the image values to (near) zero and increase their\n", - "corresponding noise-map to large values.\n", - "\n", - "The exact method used depends on the nature of the model being fitted. For simple fits like a light profile a mask\n", - "is appropriate, as removing image pixels does not change how the model is fitted. However, for more complex models\n", - "fits, like those using a pixelization, masking regions of the image in a way that removes their image pixels entirely\n", - "from the fit can produce discontinuities in the pixelixation. In this case, scaling the data and noise-map values\n", - "may be a better approach.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/mask_extra_galaxies.py`: create the extra galaxies mask manually via a Python script.\n", - "- `data_preparation/gui/extra_galaxies_mask.ipynb` use a Graphical User Interface (GUI) to create the extra galaxies mask.\n", - "- `features/extra_galaxies.py` how to use the extra galaxies mask in a model-fit.\n", - "\n", - "__Info (Optional)__\n", - "\n", - "Auxiliary information about a strong lens dataset may used during an analysis or afterwards when interpreting the \n", - " modeling results. For example, the redshifts of the source and lens galaxy. \n", - "\n", - "By storing these as an `info.json` file in the lens's dataset folder, it is straight forward to load the redshifts \n", - "in a modeling script and pass them to a fit, such that **PyAutoLens** can then output results in physical \n", - "units (e.g. kpc instead of arc-seconds).\n", - "\n", - "For analysing large quantities of modeling results, **PyAutoLens** has an sqlite database feature. The info file \n", - "may can also be loaded by the database after a model-fit has completed, such that when one is interpreting\n", - "the results of a model fit additional data on a lens can be used to. \n", - "\n", - "For example, to plot the model-results against other measurements of a lens not made by PyAutoLens. Examples of such \n", - "data might be:\n", - "\n", - "- The velocity dispersion of the lens galaxy.\n", - "- The stellar mass of the lens galaxy.\n", - "- The results of previous strong lens models to the lens performed in previous papers.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/info.py`: create the info file manually via a Python script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Imaging: Data Preparation\n", + "=========================\n", + "\n", + "When a CCD imaging dataset is analysed, it must conform to certain standards in order for the analysis\n", + "to be performed correctly. This tutorial describes these standards and links to more detailed scripts which will help\n", + "you prepare your dataset to adhere to them if it does not already.\n", + "\n", + "__Contents__\n", + "\n", + "- **Pixel Scale:** The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units.\n", + "- **Image:** The image is the image of your strong lens, which comes from a telescope like the Hubble Space.\n", + "- **Noise Map:** The noise-map defines the uncertainty in every pixel of your strong lens image, where values are.\n", + "- **PSF:** The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope.\n", + "- **Data Processing Complete:** If your image, noise-map and PSF conform the standards above, you are ready to analyse your dataset!\n", + "\n", + "__Pixel Scale__\n", + "\n", + "The \"pixel_scale\" of the image (and the data in general) is pixel-units to arcsecond-units conversion factor of\n", + "your telescope. You should look up now if you are unsure of the value.\n", + "\n", + "The pixel scale of some common telescopes is as follows:\n", + "\n", + " - Hubble Space telescope 0.04\" - 0.1\" (depends on the instrument and wavelength).\n", + " - James Webb Space telescope 0.06\" - 0.1\" (depends on the instrument and wavelength).\n", + " - Euclid 0.1\" (Optical VIS instrument) and 0.2\" (NIR NISP instrument).\n", + " - VRO / LSST 0.2\" - 0.3\" (depends on the instrument and wavelength).\n", + " - Keck Adaptive Optics 0.01\" - 0.03\" (depends on the instrument and wavelength).\n", + "\n", + "It is absolutely vital you use the correct pixel scale, so double check this value!" + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image__\n", + "\n", + "The image is the image of your strong lens, which comes from a telescope like the Hubble Space telescope (HST).\n", + "\n", + "Lets inspect an image which conforms to **PyAutoLens** standards:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / \"imaging\" / \"simple\"\n", + "\n", + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.1)\n", + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This image conforms to **PyAutoLens** standards for the following reasons.\n", + "\n", + " - Units: The image flux is in units of electrons per second (as opposed to electrons, counts, ADU`s etc.). \n", + " Internal **PyAutoLens** functions for computing quantities like galaxy magnitudes assume the data and model\n", + " light profiles are in electrons per second.\n", + " \n", + " - Centering: The lens galaxy is at the centre of the image (as opposed to in a corner). Default **PyAutoLens**\n", + " parameter priors assume the lens galaxy is at the centre of the image.\n", + " \n", + " - Stamp Size: The image is a postage stamp cut-out of the lens, but does not include many pixels around the edge of\n", + " the lens. It is advised to cut out a postage stamp of the lens, as opposed to the entire image, as this reduces\n", + " the amount of memory **PyAutoLens** uses, speeds up the analysis and ensures visualization zooms around the lens\n", + " and source. However, conforming to this standard is not necessary to ensure an accurate **PyAutoLens** analysis.\n", + " \n", + "If your image conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you noise-map and PSF conform to standards first!).\n", + "\n", + "**Links / Resources:**\n", + "\n", + " - `imaging/data_preparation/examples/data.ipynb`: tools to process the data to conform to these standards.\n", + "\n", + "__Noise Map__\n", + "\n", + "The noise-map defines the uncertainty in every pixel of your strong lens image, where values are defined as the \n", + "RMS standard deviation in every pixel (not the variances, HST WHT-map values, etc.). \n", + "\n", + "Lets inspect a noise-map which conforms to **PyAutoLens** standards:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_map = al.Array2D.from_fits(\n", + " file_path=dataset_path / \"noise_map.fits\", pixel_scales=0.1\n", + ")\n", + "\n", + "aplt.plot_array(array=noise_map, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This noise-map conforms to **PyAutoLens** standards for the following reasons:\n", + "\n", + " - Units: Like its corresponding image, it is in units of electrons per second (as opposed to electrons, counts, \n", + " ADU`s etc.). Internal **PyAutoLens** functions for computing quantities like a galaxy magnitude assume the data and \n", + " model light profiles are in electrons per second.\n", + "\n", + " - Values: The noise-map values themselves are the RMS standard deviations of the noise in every pixel. When a model \n", + " is fitted to data in **PyAutoLens** and a likelihood is evaluated, this calculation assumes that this is the\n", + " corresponding definition of the noise-map. The noise map therefore should not be the variance of the noise, or \n", + " another definition of noise.\n", + "\n", + "If you are not certain what the definition of the noise-map you have available to you is, or do not know how to\n", + "compute a noise-map at all, you should refer to the instrument handbook of the telescope your data is from. It is\n", + "absolutely vital that the noise-map is correct, as it is the only way **PyAutoLens** can quantify the goodness-of-fit.\n", + "\n", + "A sanity check for a reliable noise map is that the signal-to-noise of the lens galaxy is somewhere between a value of \n", + "10 - 300 and source around 5 - 50, however this is not a definitive test.\n", + " \n", + "Given the image should be centred and cut-out around the lens galaxy, so should the noise-map.\n", + "\n", + "If your noise-map conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you image and PSF conform to standards first!).\n", + "\n", + "**Links / Resources:**\n", + "\n", + " - `imaging/data_preparation/examples/noise_map.ipynb`: tools to process the noise-map to conform to these standards.\n", + "\n", + "__PSF__\n", + "\n", + "The Point Spread Function (PSF) describes blurring due the optics of your dataset`s telescope. It is used when fitting \n", + "a dataset to include these effects, such that does not bias the model.\n", + "\n", + "It should be estimated from a stack of stars in the image during data reduction or using a PSF simulator (e.g. TinyTim\n", + "for Hubble).\n", + "\n", + "Lets inspect a PSF which conforms to **PyAutoLens** standards:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_fits(\n", + " file_path=dataset_path / \"psf.fits\", hdu=0, pixel_scales=0.1\n", + ")\n", + "\n", + "aplt.plot_array(array=psf.kernel, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This psf conforms to **PyAutoLens** standards for the following reasons.\n", + "\n", + " - Size: The PSF has a shape 21 x 21 pixels, which is large enough to capture the PSF core and thus capture the \n", + " majority of the blurring effect, but not so large that the convolution slows down the analysis. Large \n", + " PSFs (e.g. 51 x 51) are supported, but will lead to much slower run times. The size of the PSF should be carefully \n", + " chosen to ensure it captures the majority of blurring due to the telescope optics, which for most instruments is \n", + " something around 11 x 11 to 21 x 21.\n", + "\n", + " - Oddness: The PSF has dimensions which are odd (an even PSF would for example have shape 20 x 20). The \n", + " convolution of an even PSF introduces a small shift in the model images and produces an offset in the inferred\n", + " lens model parameters. Inputting an even PSF will lead **PyAutoLens** to raise an error.\n", + "\n", + " - Normalization: The PSF has been normalized such that all values within the kernel sum to 1 (note how all values in \n", + " the example PSF are below zero with the majority below 0.01). This ensures that flux is conserved when convolution \n", + " is performed, ensuring that quantities like a galaxy's magnitude are computed accurately.\n", + "\n", + " - Centering: The PSF is at the centre of the array (as opposed to in a corner), ensuring that no shift is introduced\n", + " due to PSF blurring on the inferred model parameters.\n", + "\n", + "If your PSF conforms to all of the above standards, you are good to use it for an analysis (but must also check\n", + "you noise-map and image conform to standards first!).\n", + "\n", + "**Links / Resources:**\n", + "\n", + " - `imaging/data_preparation/examples/psf.ipynb`: tools to process the PSF to conform to these standards.\n", + "\n", + "__Data Processing Complete__\n", + "\n", + "If your image, noise-map and PSF conform the standards above, you are ready to analyse your dataset!\n", + "\n", + "Below, we provide an overview of optional data preparation steos which prepare other aspects of the analysis. \n", + "\n", + "New users are recommended to skim-read the optional steps below so they are aware of them, but to not perform them \n", + "and instead analyse their dataset now. You can come back to the data preparation scripts below if it becomes necessary.\n", + "\n", + "__Mask (Optional)__\n", + "\n", + "The mask removes the regions of the image where the lens and source galaxy are not present, typically the edges of the \n", + "image.\n", + "\n", + "Example modeling scripts internally create a 3.0\" circular mask and therefore do not require that a mask has been \n", + "created externally via a data preparation script. \n", + "\n", + "This script shows how to create customize masked (e.g. annular, ellipses) which are tailored to match the lens or\n", + "lensed source emission. \n", + "\n", + "If you have not analysed your dataset yet and do not know of a specific reason why you need the bespoke masks \n", + "created by this script, it is recommended that you simply use the default ~3.0\" circular mask internally made in each\n", + "script and omit this data preparation tutorial.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/mask.ipynb`: tools to create a bespoke mask for your dataset.\n", + "- `data_preparation/examples/gui/mask.ipynb`: use a Graphical User Interface (GUI) to create a bespoke mask.\n", + "\n", + "__Positions (Optional)__\n", + "\n", + "The script allows you to mark the (y,x) arc-second positions of the multiply imaged lensed source galaxy in \n", + "the image-plane, under the assumption that they originate from the same location in the source-plane.\n", + "\n", + "A non-linear search (e.g. Nautilus) can then use these positions to preferentially choose mass models where these \n", + "positions trace close to one another in the source-plane. This speeding up the initial fitting of lens models and \n", + "removes unwanted solutions from parameter space which have too much or too little mass in the lens galaxy.\n", + "\n", + "If you create positions for your dataset, you must also update your modeling script to use them by loading them \n", + "and passing them to the `Analysis` object via a `PositionsLH` object. \n", + "\n", + "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using positions\n", + "to help the non-linear search find a good lens model.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "Position-based lens model resampling is particularly important for fitting pixelized source models, for the\n", + "reasons disucssed in the following readthedocs \n", + "webapge https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", + "\n", + "- `data_preparation/examples/optional/positions.ipynb`: input the positions manually into a Python script.\n", + "- `data_preparation/gui/positions.ipynb` use a Graphical User Interface (GUI) to mark the positions.\n", + "- `guides/modeling/customize` for an example.of how to use positions in a `modeling` script.\n", + "\n", + "\n", + "__Lens Light Centre (Optional)__\n", + "\n", + "This script allows you to mark the (y,x) arcsecond locations of the lens galaxy light centre(s) of the strong lens\n", + "you are analysing. These can be used as fixed values for the lens light and mass models in a model-fit.\n", + "\n", + "This reduces the number of free parameters fitted for in a lens model and removes inaccurate solutions where\n", + "the lens mass model centre is unrealistically far from its true centre.\n", + "\n", + "Advanced `chaining` scripts often use these input centres in the early fits to infer an accurate initial lens model,\n", + "amd then make the centres free parameters in later searches to ensure a general and accurate lens model is inferred.\n", + "\n", + "If you create a `light_centre` for your dataset, you must also update your modeling script to use them.\n", + "\n", + "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using a fixed\n", + "lens light and / or mass centre to help the non-linear search find a good lens model.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/lens_light_centre.py`: input the lens galaxy light centre manually into a Python script.\n", + "- `data_preparation/gui/lens_light_centre.ipynb` use a Graphical User Interface (GUI) to mask the lens galaxy light centre.\n", + "\n", + "\n", + "__Extra Galaxies (Optional)__\n", + "\n", + "There may be galaxies nearby the lens and source galaxies, whose emission blends with that of the lens and source\n", + "and whose mass may contribute to the ray-tracing and lens model.\n", + "\n", + "We can include these galaxies in the lens model, either as light profiles, mass profiles, or both, using the\n", + "modeling API, where these nearby objects are denoted `extra_galaxies`.\n", + "\n", + "The script `extra_galaxies_centres.py` marks the (y,x) arcsecond locations of these extra galaxies, so that when they \n", + "are included in the lens model the centre of these extra galaxies light and / or mass profiles are fixed to these \n", + "values (or their priors are initialized surrounding these centres).\n", + "\n", + "The example `mask_extra_galaxies.py` (see below) masks the regions of an image where extra galaxies are present. \n", + "This mask is used to remove their signal from the data and increase their noise to make them not impact the fit. \n", + "This means their luminous emission does not need to be included in the model, reducing the number of free parameters \n", + "and speeding up the analysis. It is still a choice whether their mass is included in the model.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/extra_galaxies_centres.py`: input the extra galaxy centres manually into a \n", + " Python script.\n", + "- `data_preparation/gui/extra_galaxies_centres.ipynb`: use a Graphical User Interface (GUI) to mark the extra galaxy centres.\n", + "- `features/extra_galaxies.py` how to use extra galaxies in a model-fit, including loading the extra galaxy centres.\n", + "\n", + "\n", + "__Mask Extra Galaxies (Optional)__\n", + "\n", + "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated \n", + "with the strong lens we are studying. The emission from these images will impact our model fitting and needs to be \n", + "removed from the analysis.\n", + "\n", + "This script creates a mask of these regions of the image, called the `mask_extra_galaxies`, which can be used to\n", + "prevent them from impacting a fit. This mask may also include emission from objects which are not technically galaxies,\n", + "but blend with the galaxy we are studying in a similar way. Common examples of such objects are foreground stars\n", + "or emission due to the data reduction process.\n", + "\n", + "The mask can be applied in different ways. For example, it could be applied such that the image pixels are discarded\n", + "from the fit entirely, Alternatively the mask could be used to set the image values to (near) zero and increase their\n", + "corresponding noise-map to large values.\n", + "\n", + "The exact method used depends on the nature of the model being fitted. For simple fits like a light profile a mask\n", + "is appropriate, as removing image pixels does not change how the model is fitted. However, for more complex models\n", + "fits, like those using a pixelization, masking regions of the image in a way that removes their image pixels entirely\n", + "from the fit can produce discontinuities in the pixelixation. In this case, scaling the data and noise-map values\n", + "may be a better approach.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/mask_extra_galaxies.py`: create the extra galaxies mask manually via a Python script.\n", + "- `data_preparation/gui/extra_galaxies_mask.ipynb` use a Graphical User Interface (GUI) to create the extra galaxies mask.\n", + "- `features/extra_galaxies.py` how to use the extra galaxies mask in a model-fit.\n", + "\n", + "__Info (Optional)__\n", + "\n", + "Auxiliary information about a strong lens dataset may used during an analysis or afterwards when interpreting the \n", + " modeling results. For example, the redshifts of the source and lens galaxy. \n", + "\n", + "By storing these as an `info.json` file in the lens's dataset folder, it is straight forward to load the redshifts \n", + "in a modeling script and pass them to a fit, such that **PyAutoLens** can then output results in physical \n", + "units (e.g. kpc instead of arc-seconds).\n", + "\n", + "For analysing large quantities of modeling results, **PyAutoLens** has an sqlite database feature. The info file \n", + "may can also be loaded by the database after a model-fit has completed, such that when one is interpreting\n", + "the results of a model fit additional data on a lens can be used to. \n", + "\n", + "For example, to plot the model-results against other measurements of a lens not made by PyAutoLens. Examples of such \n", + "data might be:\n", + "\n", + "- The velocity dispersion of the lens galaxy.\n", + "- The stellar mass of the lens galaxy.\n", + "- The results of previous strong lens models to the lens performed in previous papers.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/info.py`: create the info file manually via a Python script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/double_einstein_ring/chaining.ipynb b/notebooks/imaging/features/advanced/double_einstein_ring/chaining.ipynb index 2f0dfe4a5..699579c82 100644 --- a/notebooks/imaging/features/advanced/double_einstein_ring/chaining.ipynb +++ b/notebooks/imaging/features/advanced/double_einstein_ring/chaining.ipynb @@ -1,1128 +1,1165 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Chaining: Double Einstein Ring\n", - "==============================\n", - "\n", - "This script chains two searches to fit `Imaging` data of a 'galaxy-scale' strong lens which has two source galaxies at\n", - "two different redshifts, forming a double Einstein ring system. This fits a model where:\n", - "\n", - " - The lens galaxy's light is omitted.\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The first source galaxy's mass is a `IsothermalSph` and its light an MGE.\n", - " - The second source galaxy's light is an MGE.\n", - "\n", - "The two searches break down as follows:\n", - "\n", - " 1) Model the lens galaxy mass as an `Isothermal` and first source galaxy using an MGE.\n", - " 2) Model the lens, first and second source galaxies, where the first source's mass is an `IsothermalSph` and second\n", - " source is an MGE.\n", - "\n", - "__Why Chain?__\n", - "\n", - "Systems with two (or more) strongly lensed sources are a great example of the benefits of search chaining. The lens\n", - "model can quickly have many parameters (e.g. N > 20), but many of the components being fitted are only mildly covariant\n", - "with one another.\n", - "\n", - "Most importantly, ray-tracing of the first source galaxy does not depend on the properties of the second source galaxy\n", - "at all, meaning it can be used to initialize the lens mass model before the second source is fitted. For the simulated\n", - "data fitted in this example, we'll see that the first search successfully initializes the lens mass model and first\n", - "source model without issue, such that fitting of the second source can be done efficiently.\n", - "\n", - "The only problem is that the light of the second source is included in the data we fit in the first search, and thus\n", - "could bias or impact its model fit. To circumvent this, the first search uses a smaller mask which removes the light\n", - "of the second source from the model-fit. A larger mask included both sources is then used in the second search.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load and plot the strong lens `Imaging` dataset.\n", - "- **Paths:** The path the results of all chained searches are output.\n", - "- **Masking (Search 1):** We apply a smaller circular mask, the radius of which is chosen to remove the light of the second source galaxy.\n", - "- **Model (Search 1):** Search 1 fits a lens model where the lens galaxy's total mass distribution is an `Isothermal` and the first source galaxy's light is an MGE.\n", - "- **Search + Analysis + Model-Fit (Search 1):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "- **Result (Search 1):** The results which are used for prior passing are summarised in the `info` attribute.\n", - "- **Masking (Search 2):** We apply a larger circular mask which includes the light of both source galaxies.\n", - "- **Model (Search 2):** Search 2 fits the lens, first and second source galaxies, where the first source's mass is an `IsothermalSph` and the second source is an MGE.\n", - "- **Search + Analysis + Model-Fit (Search 2):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "- **Result (Search 2):** The final results of the chained model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "- **Pipelines:** Advanced search chaining uses `pipelines` that chain together multiple searches to perform complex lens modeling in a robust and efficient way.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__ \n", - "\n", - "Load and plot the `Imaging` data. N\n", - "\n", - "ote that we use different masks for searches 1 and 2." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__\n", - "\n", - "The path the results of all chained searches are output:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"imaging\") / \"chaining\" / \"double_einstein_ring\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masking (Search 1)__\n", - "\n", - "We apply a smaller circular mask, the radius of which is chosen to remove the light of the second source galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 2.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Search 1 fits a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", - " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=12.\n", - "\n", - "We therefore omit the second source from the model entirely." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source_0 = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "source_1 = af.Model(al.Galaxy, redshift=2.0)\n", - "\n", - "model_1 = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", - "provides a reasonably accurate lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__source_0_parametric\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 1)__\n", - "\n", - "The results which are used for prior passing are summarised in the `info` attribute." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masking (Search 2)__\n", - "\n", - "We apply a larger circular mask, which includes the second source galaxy now that it is included in the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "We use the results of search 1 to create the lens model fitted in search 2, where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [parameters fixed to results of search 1].\n", - " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [parameters fixed to results of search 1].\n", - " - The first source galaxy's mass is a `IsothermalSph` [3 parameters].\n", - " - The second source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=10.\n", - "\n", - "The galaxies are assigned redshifts of 0.5, 1.0 and 2.0. This ensures the multi-plane ray-tracing necessary for the \n", - "double Einstein ring lens system is performed correctly.\n", - "\n", - "The lens galaxy's mass and first source galaxy's light are passed as an `instance` (as opposed to the `model` which \n", - "was used in the API tutorial). By passing these objects as an `instance` it passes the maximum log likelihood parameters \n", - "inferred by search 1 as fixed values that are not free parameters fitted for by the non-linear search of search 2." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(al.Galaxy, redshift=0.5, mass=result_1.instance.galaxies.lens.mass)\n", - "source_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " bulge=result_1.instance.galaxies.source_0.bulge,\n", - " mass=al.mp.IsothermalSph,\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=bulge)\n", - "source_1.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.5)\n", - "source_1.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.5)\n", - "\n", - "model_2 = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 2)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the `model.info` file of the search 2 model-fit to ensure the priors were passed correctly, as \n", - "well as the checkout the results to ensure an accurate power-law mass model is inferred." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__source_1_parametric\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 2)__\n", - "\n", - "The final results can be summarised via printing `info`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "In this example, we used prior passing to initialize a model fit to a double Einstein ring. We exploited the fact that \n", - "ray-tracing of the first source is fully independent of the source behind it, such that we could use it to initialize \n", - "the lens model before fitting the second source.\n", - "\n", - "For certain double Einstein ring systems, it is possible that the light of the first and second sources are harder to\n", - "deblend than the simple masking we used in this example. Manual masks drawn using a GUI which removes the second \n", - "source's light will nevertheless always be possible, but more care may be required.\n", - "\n", - "__Pipelines__\n", - "\n", - "Advanced search chaining uses `pipelines` that chain together multiple searches to perform complex lens modeling \n", - "in a robust and efficient way. \n", - "\n", - "There are currently no pipelines written for double Einstein ring systems, albeit one can craft them by learning the\n", - "API and concepts from existing template pipelines. We are still figuring out the most effective way to model double\n", - "Einstein ring systems, which is why pipeline templates are not yet written.\n", - "\n", - "__SLaM (Source, Light and Mass)__\n", - "\n", - "An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens modeling \n", - "processing into a series of fits that first perfect the source model, then the lens light model and finally the lens\n", - "mass model. \n", - "\n", - "The SLaM pipelines begin with a linear Source pipeline, which then switches to an inversion Source pipeline, \n", - "exploiting the chaining technique demonstrated in this example.\n", - "\n", - "Pipeline: Double Einstein Ring\n", - "==============================\n", - "\n", - "By chaining together four searches this script fits `Imaging` dataset of a 'galaxy-scale' strong lens, which has two source galaxies\n", - "at two different redshifts, forming a double Einstein ring system. In the final model:\n", - "\n", - " - The lens galaxy's light is an MGE.\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The first source galaxy's mass is a `IsothermalSph` and its light is modeled using an `Inversion`.\n", - " - The second source galaxy's light is modeled using an `Inversion`.\n", - "\n", - "The three searches break down as follows:\n", - "\n", - " 1) Model the lens galaxy using an MGE with 2 x 30 Gaussians to subtract its emission.\n", - " 2) Model the lens galaxy mass as an `Isothermal` and first source galaxy using an MGE.\n", - " 3) Model the lens, first and second source galaxies, where the first source's mass is an `IsothermalSph` and second\n", - " source is an MGE.\n", - " 4) Model the first and second source galaxy simultaneously using an `Inversion` and lens galaxy mass as an\n", - " `Isothermal`.\n", - "\n", - "The approach used in this pipeline and benefits of using chaining searching to fit double einstein ring systems are\n", - "described in the script `notebooks/imaging/chaining/double_einstein_ring.ipynb`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__ \n", - "\n", - "Load, plot and mask the `Imaging` data.\n", - "\n", - "ote that we use different masks for each search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__\n", - "\n", - "The path the results of all chained searches are output:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"imaging\") / \"chaining\" / \"double_einstein_ring\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masking (Search 1 & 2)__\n", - "\n", - "We apply a smaller circular mask, the radius of which is chosen to remove the light of the second source galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 2.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Search 1 fits a lens model where:\n", - "\n", - " - The lens galaxy's light is an MGE bulge with 2 x 30 Gaussians [6 parameters].\n", - " - The lens galaxy's mass and both source galaxies are omitted.\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=7." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(lens=lens))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__lens_light\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "We use the results of search 1 to create the lens model fitted in search 2, where:\n", - "\n", - " - The lens galaxy's light is an MGE bulge [Parameters fixed to results of search 1].\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", - " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=12.\n", - "\n", - "We therefore omit the second source from the model entirely." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "mass.centre = result_1.model.galaxies.lens.bulge.profile_list[0].centre\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, bulge=result_1.instance.galaxies.lens.bulge, mass=mass\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source_0 = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "model_2 = af.Collection(galaxies=af.Collection(lens=lens, source_0=source_0))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 2)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__parametric_source_0\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masking (Search 3)__\n", - "\n", - "We apply a larger circular mask, which includes the second source galaxy now that it is included in the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 3)__\n", - "\n", - "We use the results of searches 1 & 2 to create the lens model fitted in search 3, where:\n", - "\n", - " - The lens galaxy's light is an MGE bulge [Parameters fixed to results of search 1].\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [Parameters fixed to results of search 2].\n", - " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [Parameters fixed to results of search 2].\n", - " - The first source galaxy's mass is a `IsothermalSph` [3 parameters].\n", - " - The second source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=10.\n", - "\n", - "The galaxies are assigned redshifts of 0.5, 1.0 and 2.0. This ensures the multi-plane ray-tracing necessary for the \n", - "double Einstein ring lens system is performed correctly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=result_1.instance.galaxies.lens.bulge,\n", - " mass=result_2.model.galaxies.lens.mass,\n", - ")\n", - "source_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " bulge=result_2.model.galaxies.source_0.bulge,\n", - " mass=al.mp.IsothermalSph,\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=bulge)\n", - "\n", - "model_3 = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 3)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_3 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[3]__source_2_parametric\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_3 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_3 = search_3.fit(model=model_3, analysis=analysis_3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 4)__\n", - "\n", - "We use the results of searches 1, 2 & 3 to create the lens model fitted in search 4, where:\n", - "\n", - " - The lens galaxy's light is an MGE with 2 x 30 Gaussians bulge [7 Parameters: we do not use the results of search 1 to \n", - " initialize priors].\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [5 Parameters: priors initialized from search 2].\n", - " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters: priors initialized from search 2].\n", - " - The first source galaxy's mass is a `IsothermalSph` [3 parameters: priors initialized from search 3].\n", - " - The second source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters: priors initialized from search 3].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=29.\n", - "\n", - "The galaxies are assigned redshifts of 0.5, 1.0 and 2.0. This ensures the multi-plane ray-tracing necessary for the \n", - "double Einstein ring lens system is performed correctly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=result_2.model.galaxies.lens.mass,\n", - ")\n", - "source_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " bulge=result_2.model.galaxies.source_0.bulge,\n", - " mass=result_3.model.galaxies.source_0.mass,\n", - ")\n", - "source_1 = af.Model(\n", - " al.Galaxy, redshift=2.0, bulge=result_3.model.galaxies.source_1.bulge\n", - ")\n", - "\n", - "model_4 = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 4)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_4 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[4]__parametric_all\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_4 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_4 = search_4.fit(model=model_4, analysis=analysis_4)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 5)__\n", - "\n", - "We use the results of search 4 to create the lens model fitted in search 5, where:\n", - "\n", - " - The lens galaxy's light is an MGE bulge [Parameters fixed to results of search 4].\n", - " - The lens galaxy's total mass distribution is again an `Isothermal` [Parameters fixed to results of search 4].\n", - " - The first source galaxy's mass is a `IsothermalSph` [Parameters fixed to results of search 4].\n", - " - The first source-galaxy's light uses an `Overlay` image-mesh, `RectangularAdaptDensity` mesh and `Constant` regularization \n", - " scheme [3 parameters].\n", - " - The second source-galaxy's light uses an `Overlay` image-mesh, `RectangularAdaptDensity` mesh and `Constant` regularization \n", - " scheme [3 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=6." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = result_4.instance.galaxies.lens\n", - "source_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " mass=result_4.instance.galaxies.source_0.mass,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", - " regularization=al.reg.Constant,\n", - " ),\n", - ")\n", - "source_1 = af.Model(\n", - " al.Galaxy,\n", - " redshift=2.0,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", - " regularization=al.reg.Constant,\n", - " ),\n", - ")\n", - "model_5 = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_5 = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Model-Fit__\n", - "\n", - "We now create the non-linear search and perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_5 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[5]__sources_pixelization\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "result_5 = search_5.fit(model=model_5, analysis=analysis_5)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "In this example, we passed used prior passing to initialize a model fit to a double Einstein ring using \n", - "two `Inversion`'s.\n", - "\n", - "Fitting just the `Inversion` by itself for a double Einstein ring system is practically impossible, due to the \n", - "unphysical solutions which reconstruct its light as a demagnified version of each source. Furthermore, it helped to \n", - "ensure that the model-fit ran efficiently.\n", - "\n", - "__Pipelines__\n", - "\n", - "Advanced search chaining uses `pipelines` that chain together multiple searches to perform complex lens modeling \n", - "in a robust and efficient way. \n", - "\n", - "There are currently no pipelines written for double Einstein ring systems, albeit one can craft them by learning the\n", - "API and concepts from existing template pipelines. We are still figuring out the most effective way to model double\n", - "Einstein ring systems, which is why pipeline templates are not yet written.\n", - "\n", - "__SLaM (Source, Light and Mass)__\n", - "\n", - "An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens modeling \n", - "processing into a series of fits that first perfect the source model, then the lens light model and finally the lens\n", - "mass model. \n", - "\n", - "The SLaM pipelines begin with a parametric Source pipeline, which then switches to an inversion Source pipeline, \n", - "exploiting the chaining technique demonstrated in this example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Chaining: Double Einstein Ring\n", + "==============================\n", + "\n", + "This script chains two searches to fit `Imaging` data of a 'galaxy-scale' strong lens which has two source galaxies at\n", + "two different redshifts, forming a double Einstein ring system. This fits a model where:\n", + "\n", + " - The lens galaxy's light is omitted.\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The first source galaxy's mass is a `IsothermalSph` and its light an MGE.\n", + " - The second source galaxy's light is an MGE.\n", + "\n", + "The two searches break down as follows:\n", + "\n", + " 1) Model the lens galaxy mass as an `Isothermal` and first source galaxy using an MGE.\n", + " 2) Model the lens, first and second source galaxies, where the first source's mass is an `IsothermalSph` and second\n", + " source is an MGE.\n", + "\n", + "__Why Chain?__\n", + "\n", + "Systems with two (or more) strongly lensed sources are a great example of the benefits of search chaining. The lens\n", + "model can quickly have many parameters (e.g. N > 20), but many of the components being fitted are only mildly covariant\n", + "with one another.\n", + "\n", + "Most importantly, ray-tracing of the first source galaxy does not depend on the properties of the second source galaxy\n", + "at all, meaning it can be used to initialize the lens mass model before the second source is fitted. For the simulated\n", + "data fitted in this example, we'll see that the first search successfully initializes the lens mass model and first\n", + "source model without issue, such that fitting of the second source can be done efficiently.\n", + "\n", + "The only problem is that the light of the second source is included in the data we fit in the first search, and thus\n", + "could bias or impact its model fit. To circumvent this, the first search uses a smaller mask which removes the light\n", + "of the second source from the model-fit. A larger mask included both sources is then used in the second search.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load and plot the strong lens `Imaging` dataset.\n", + "- **Paths:** The path the results of all chained searches are output.\n", + "- **Masking (Search 1):** We apply a smaller circular mask, the radius of which is chosen to remove the light of the second source galaxy.\n", + "- **Model (Search 1):** Search 1 fits a lens model where the lens galaxy's total mass distribution is an `Isothermal` and the first source galaxy's light is an MGE.\n", + "- **Search + Analysis + Model-Fit (Search 1):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "- **Result (Search 1):** The results which are used for prior passing are summarised in the `info` attribute.\n", + "- **Masking (Search 2):** We apply a larger circular mask which includes the light of both source galaxies.\n", + "- **Model (Search 2):** Search 2 fits the lens, first and second source galaxies, where the first source's mass is an `IsothermalSph` and the second source is an MGE.\n", + "- **Search + Analysis + Model-Fit (Search 2):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "- **Result (Search 2):** The final results of the chained model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "- **Pipelines:** Advanced search chaining uses `pipelines` that chain together multiple searches to perform complex lens modeling in a robust and efficient way.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__ \n", + "\n", + "Load and plot the `Imaging` data. N\n", + "\n", + "ote that we use different masks for searches 1 and 2." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__\n", + "\n", + "The path the results of all chained searches are output:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"imaging\") / \"chaining\" / \"double_einstein_ring\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masking (Search 1)__\n", + "\n", + "We apply a smaller circular mask, the radius of which is chosen to remove the light of the second source galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 2.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Search 1 fits a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=12.\n", + "\n", + "We therefore omit the second source from the model entirely." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source_0 = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "source_1 = af.Model(al.Galaxy, redshift=2.0)\n", + "\n", + "model_1 = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", + "provides a reasonably accurate lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__source_0_parametric\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 1)__\n", + "\n", + "The results which are used for prior passing are summarised in the `info` attribute." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masking (Search 2)__\n", + "\n", + "We apply a larger circular mask, which includes the second source galaxy now that it is included in the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "We use the results of search 1 to create the lens model fitted in search 2, where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [parameters fixed to results of search 1].\n", + " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [parameters fixed to results of search 1].\n", + " - The first source galaxy's mass is a `IsothermalSph` [3 parameters].\n", + " - The second source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=10.\n", + "\n", + "The galaxies are assigned redshifts of 0.5, 1.0 and 2.0. This ensures the multi-plane ray-tracing necessary for the \n", + "double Einstein ring lens system is performed correctly.\n", + "\n", + "The lens galaxy's mass and first source galaxy's light are passed as an `instance` (as opposed to the `model` which \n", + "was used in the API tutorial). By passing these objects as an `instance` it passes the maximum log likelihood parameters \n", + "inferred by search 1 as fixed values that are not free parameters fitted for by the non-linear search of search 2." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(al.Galaxy, redshift=0.5, mass=result_1.instance.galaxies.lens.mass)\n", + "source_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " bulge=result_1.instance.galaxies.source_0.bulge,\n", + " mass=al.mp.IsothermalSph,\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=bulge)\n", + "source_1.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.5)\n", + "source_1.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.5)\n", + "\n", + "model_2 = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 2)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the `model.info` file of the search 2 model-fit to ensure the priors were passed correctly, as \n", + "well as the checkout the results to ensure an accurate power-law mass model is inferred." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__source_1_parametric\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 2)__\n", + "\n", + "The final results can be summarised via printing `info`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "In this example, we used prior passing to initialize a model fit to a double Einstein ring. We exploited the fact that \n", + "ray-tracing of the first source is fully independent of the source behind it, such that we could use it to initialize \n", + "the lens model before fitting the second source.\n", + "\n", + "For certain double Einstein ring systems, it is possible that the light of the first and second sources are harder to\n", + "deblend than the simple masking we used in this example. Manual masks drawn using a GUI which removes the second \n", + "source's light will nevertheless always be possible, but more care may be required.\n", + "\n", + "__Pipelines__\n", + "\n", + "Advanced search chaining uses `pipelines` that chain together multiple searches to perform complex lens modeling \n", + "in a robust and efficient way. \n", + "\n", + "There are currently no pipelines written for double Einstein ring systems, albeit one can craft them by learning the\n", + "API and concepts from existing template pipelines. We are still figuring out the most effective way to model double\n", + "Einstein ring systems, which is why pipeline templates are not yet written.\n", + "\n", + "__SLaM (Source, Light and Mass)__\n", + "\n", + "An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens modeling \n", + "processing into a series of fits that first perfect the source model, then the lens light model and finally the lens\n", + "mass model. \n", + "\n", + "The SLaM pipelines begin with a linear Source pipeline, which then switches to an inversion Source pipeline, \n", + "exploiting the chaining technique demonstrated in this example.\n", + "\n", + "Pipeline: Double Einstein Ring\n", + "==============================\n", + "\n", + "By chaining together four searches this script fits `Imaging` dataset of a 'galaxy-scale' strong lens, which has two source galaxies\n", + "at two different redshifts, forming a double Einstein ring system. In the final model:\n", + "\n", + " - The lens galaxy's light is an MGE.\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The first source galaxy's mass is a `IsothermalSph` and its light is modeled using an `Inversion`.\n", + " - The second source galaxy's light is modeled using an `Inversion`.\n", + "\n", + "The three searches break down as follows:\n", + "\n", + " 1) Model the lens galaxy using an MGE with 2 x 30 Gaussians to subtract its emission.\n", + " 2) Model the lens galaxy mass as an `Isothermal` and first source galaxy using an MGE.\n", + " 3) Model the lens, first and second source galaxies, where the first source's mass is an `IsothermalSph` and second\n", + " source is an MGE.\n", + " 4) Model the first and second source galaxy simultaneously using an `Inversion` and lens galaxy mass as an\n", + " `Isothermal`.\n", + "\n", + "The approach used in this pipeline and benefits of using chaining searching to fit double einstein ring systems are\n", + "described in the script `notebooks/imaging/chaining/double_einstein_ring.ipynb`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__ \n", + "\n", + "Load, plot and mask the `Imaging` data.\n", + "\n", + "ote that we use different masks for each search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__\n", + "\n", + "The path the results of all chained searches are output:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"imaging\") / \"chaining\" / \"double_einstein_ring\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masking (Search 1 & 2)__\n", + "\n", + "We apply a smaller circular mask, the radius of which is chosen to remove the light of the second source galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 2.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Search 1 fits a lens model where:\n", + "\n", + " - The lens galaxy's light is an MGE bulge with 2 x 30 Gaussians [6 parameters].\n", + " - The lens galaxy's mass and both source galaxies are omitted.\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=7." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(lens=lens))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__lens_light\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "We use the results of search 1 to create the lens model fitted in search 2, where:\n", + "\n", + " - The lens galaxy's light is an MGE bulge [Parameters fixed to results of search 1].\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=12.\n", + "\n", + "We therefore omit the second source from the model entirely." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "mass.centre = result_1.model.galaxies.lens.bulge.profile_list[0].centre\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, bulge=result_1.instance.galaxies.lens.bulge, mass=mass\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source_0 = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "model_2 = af.Collection(galaxies=af.Collection(lens=lens, source_0=source_0))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 2)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__parametric_source_0\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masking (Search 3)__\n", + "\n", + "We apply a larger circular mask, which includes the second source galaxy now that it is included in the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 3)__\n", + "\n", + "We use the results of searches 1 & 2 to create the lens model fitted in search 3, where:\n", + "\n", + " - The lens galaxy's light is an MGE bulge [Parameters fixed to results of search 1].\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [Parameters fixed to results of search 2].\n", + " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [Parameters fixed to results of search 2].\n", + " - The first source galaxy's mass is a `IsothermalSph` [3 parameters].\n", + " - The second source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=10.\n", + "\n", + "The galaxies are assigned redshifts of 0.5, 1.0 and 2.0. This ensures the multi-plane ray-tracing necessary for the \n", + "double Einstein ring lens system is performed correctly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=result_1.instance.galaxies.lens.bulge,\n", + " mass=result_2.model.galaxies.lens.mass,\n", + ")\n", + "source_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " bulge=result_2.model.galaxies.source_0.bulge,\n", + " mass=al.mp.IsothermalSph,\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=bulge)\n", + "\n", + "model_3 = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 3)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_3 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[3]__source_2_parametric\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_3 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_3 = search_3.fit(model=model_3, analysis=analysis_3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 4)__\n", + "\n", + "We use the results of searches 1, 2 & 3 to create the lens model fitted in search 4, where:\n", + "\n", + " - The lens galaxy's light is an MGE with 2 x 30 Gaussians bulge [7 Parameters: we do not use the results of search 1 to \n", + " initialize priors].\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 Parameters: priors initialized from search 2].\n", + " - The first source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters: priors initialized from search 2].\n", + " - The first source galaxy's mass is a `IsothermalSph` [3 parameters: priors initialized from search 3].\n", + " - The second source galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters: priors initialized from search 3].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=29.\n", + "\n", + "The galaxies are assigned redshifts of 0.5, 1.0 and 2.0. This ensures the multi-plane ray-tracing necessary for the \n", + "double Einstein ring lens system is performed correctly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=result_2.model.galaxies.lens.mass,\n", + ")\n", + "source_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " bulge=result_2.model.galaxies.source_0.bulge,\n", + " mass=result_3.model.galaxies.source_0.mass,\n", + ")\n", + "source_1 = af.Model(\n", + " al.Galaxy, redshift=2.0, bulge=result_3.model.galaxies.source_1.bulge\n", + ")\n", + "\n", + "model_4 = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 4)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_4 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[4]__parametric_all\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_4 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_4 = search_4.fit(model=model_4, analysis=analysis_4)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 5)__\n", + "\n", + "We use the results of search 4 to create the lens model fitted in search 5, where:\n", + "\n", + " - The lens galaxy's light is an MGE bulge [Parameters fixed to results of search 4].\n", + " - The lens galaxy's total mass distribution is again an `Isothermal` [Parameters fixed to results of search 4].\n", + " - The first source galaxy's mass is a `IsothermalSph` [Parameters fixed to results of search 4].\n", + " - The first source-galaxy's light uses an `Overlay` image-mesh, `RectangularAdaptDensity` mesh and `Constant` regularization \n", + " scheme [3 parameters].\n", + " - The second source-galaxy's light uses an `Overlay` image-mesh, `RectangularAdaptDensity` mesh and `Constant` regularization \n", + " scheme [3 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=6." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = result_4.instance.galaxies.lens\n", + "source_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " mass=result_4.instance.galaxies.source_0.mass,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", + " regularization=al.reg.Constant,\n", + " ),\n", + ")\n", + "source_1 = af.Model(\n", + " al.Galaxy,\n", + " redshift=2.0,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", + " regularization=al.reg.Constant,\n", + " ),\n", + ")\n", + "model_5 = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_5 = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Model-Fit__\n", + "\n", + "We now create the non-linear search and perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_5 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[5]__sources_pixelization\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "result_5 = search_5.fit(model=model_5, analysis=analysis_5)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "In this example, we passed used prior passing to initialize a model fit to a double Einstein ring using \n", + "two `Inversion`'s.\n", + "\n", + "Fitting just the `Inversion` by itself for a double Einstein ring system is practically impossible, due to the \n", + "unphysical solutions which reconstruct its light as a demagnified version of each source. Furthermore, it helped to \n", + "ensure that the model-fit ran efficiently.\n", + "\n", + "__Pipelines__\n", + "\n", + "Advanced search chaining uses `pipelines` that chain together multiple searches to perform complex lens modeling \n", + "in a robust and efficient way. \n", + "\n", + "There are currently no pipelines written for double Einstein ring systems, albeit one can craft them by learning the\n", + "API and concepts from existing template pipelines. We are still figuring out the most effective way to model double\n", + "Einstein ring systems, which is why pipeline templates are not yet written.\n", + "\n", + "__SLaM (Source, Light and Mass)__\n", + "\n", + "An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens modeling \n", + "processing into a series of fits that first perfect the source model, then the lens light model and finally the lens\n", + "mass model. \n", + "\n", + "The SLaM pipelines begin with a parametric Source pipeline, which then switches to an inversion Source pipeline, \n", + "exploiting the chaining technique demonstrated in this example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/double_einstein_ring/fit.ipynb b/notebooks/imaging/features/advanced/double_einstein_ring/fit.ipynb index 965201571..6867625a3 100644 --- a/notebooks/imaging/features/advanced/double_einstein_ring/fit.ipynb +++ b/notebooks/imaging/features/advanced/double_einstein_ring/fit.ipynb @@ -1,473 +1,510 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Double Einstein Ring Fit\n", - "==================================\n", - "\n", - "A double Einstein ring lens is a strong lens system where two source galaxies at different redshifts are lensed\n", - "by the foreground lens galaxy. They appear as two distinct Einstein rings in the image-plane.\n", - "\n", - "This script illustrates the API for performing a fit to a double Einstein ring lens via the standard `Tracer`\n", - "and `FitImaging` objects, without invoking a non-linear search. It is intended to make the multi-plane\n", - "ray-tracing API concrete before the reader moves on to `modeling.py` (search-based) or `chaining.py` / `slam.py`\n", - "(realistic, robust modeling).\n", - "\n", - "The source galaxies are both modelled with a Multi Gaussian Expansion (MGE), the same source parameterization\n", - "used in `chaining.py` and `slam.py`. The MGE is built from a `Basis` of linear `Gaussian` light profiles, whose\n", - "`intensity` values are solved for via linear algebra at fit time.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **MGE Basis:** Build a `Basis` of linear Gaussians, used for both source galaxies.\n", - "- **Galaxies:** Compose the lens galaxy plus two source galaxies at different redshifts.\n", - "- **Tracer:** Build the three-plane `Tracer` that performs the multi-plane ray-tracing.\n", - "- **Fit:** Create a `FitImaging` and inspect the fit.\n", - "- **Multi-Plane Ray-Tracing:** A short tour of how the second source-plane sees deflection from both the lens\n", - " galaxy AND the first source galaxy's mass.\n", - "- **Intensities:** The solved-for linear light profile `intensity` values for each Gaussian, per source.\n", - "- **Wrap Up:** Summary and next steps.\n", - "\n", - "__Prerequisites__\n", - "\n", - "This script focuses on the API specific to a double Einstein ring fit. For background on the underlying single-plane\n", - "fit API and the MGE source parameterization, you should read first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/fit.py` \u2014 the standard single-plane fit.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py` \u2014 the MGE fit API and `Basis` of\n", - " linear Gaussians.\n", - "\n", - "The redshifts (`lens=0.5`, `source_0=1.0`, `source_1=2.0`) match those used by the simulator and modeling examples." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `double_einstein_ring` via .fits files.\n", - "\n", - "This dataset has a double Einstein ring, due to the two source galaxies at different redshifts behind the lens galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes both source-galaxy Einstein rings." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling, with finer sub-pixelization at the centre where the lens galaxy's mass is most\n", - "strongly deflecting light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE Basis__\n", - "\n", - "We build a single `Basis` of linear Gaussians which we will use as the source-galaxy light model for both\n", - "`source_0` and `source_1`.\n", - "\n", - "The Gaussians share a common centre and elliptical components (set to spherical here for simplicity), and have\n", - "`sigma` values spaced in log10 increments from 0.01\" up to a reasonable size cap. The `intensity` of each\n", - "Gaussian is a linear parameter, solved for by linear algebra at fit time \u2014 no non-linear search is required.\n", - "\n", - "Two `Basis` objects are constructed, with different sphericity centres, so that `source_0` and `source_1` each\n", - "have their own light component centred on the (known) position from the simulator.\n", - "\n", - "For background on the MGE `Basis` API, see\n", - "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "# Centres match the simulator: source_0 is offset from the lens, source_1 is on the far side of the lens.\n", - "\n", - "source_0_bulge = build_source_basis(centre=(-0.15, -0.15))\n", - "source_1_bulge = build_source_basis(centre=(-0.45, 0.45))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The Gaussians of each basis cannot be plotted yet because their `intensity` values have not been solved for \u2014\n", - "linear light profiles only acquire an `intensity` once a `FitImaging` runs its linear algebra step. After the\n", - "fit below, we visualise each source's MGE basis with its solved-for intensities.\n", - "\n", - "We set up the plotting grid we will use post-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "We now compose the three galaxies that form the double Einstein ring system:\n", - "\n", - " - `lens` (z=0.5): an `Isothermal` mass profile, matching the simulator. The lens galaxy has no light in this\n", - " simulated dataset.\n", - " - `source_0` (z=1.0): the MGE basis above as a light component, AND an `IsothermalSph` mass profile. `source_0`\n", - " acts as both a light source AND a deflector \u2014 its mass distribution bends the light from the higher-redshift\n", - " `source_1`, contributing to the second Einstein ring.\n", - " - `source_1` (z=2.0): only an MGE light component. It contributes light, but its own mass is negligible at\n", - " this redshift in the simulated data.\n", - "\n", - "The mass profile values for `lens` and `source_0` are set to the simulator's true values, so the fit visibly\n", - "recovers both Einstein rings without a search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.5,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=source_0_bulge,\n", - " mass=al.mp.IsothermalSph(centre=(-0.15, -0.15), einstein_radius=0.3),\n", - ")\n", - "\n", - "source_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " bulge=source_1_bulge,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "The `Tracer` performs the multi-plane ray-tracing. PyAutoLens orders the galaxies internally by redshift, so\n", - "the tracer first deflects image-plane (y,x) coordinates through the lens-plane (z=0.5) onto `source_0`'s plane\n", - "(z=1.0), then continues to deflect through `source_0`'s mass to reach `source_1`'s plane (z=2.0).\n", - "\n", - "Both the lens galaxy's mass AND `source_0`'s mass contribute to the deflection map applied to coordinates\n", - "arriving at `source_1`'s plane. This is what makes double Einstein ring systems sensitive to angular diameter\n", - "distance ratios, and therefore to cosmological parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens, source_0, source_1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We pass the `Tracer` to a `FitImaging` to fit the dataset. The fit performs the multi-plane ray-tracing, evaluates\n", - "each source galaxy's light at its own source-plane, sums the resulting image-plane contributions, convolves with\n", - "the PSF, and computes the residuals against the data.\n", - "\n", - "The `linear_light_profile_intensity_dict` of the fit will hold a solved-for `intensity` for every Gaussian in\n", - "both source MGE bases." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multi-Plane Ray-Tracing__\n", - "\n", - "The tracer exposes per-plane ray-traced grids via `traced_grid_2d_list_from`. The list returned has one grid per\n", - "plane (image-plane, `source_0`-plane, `source_1`-plane).\n", - "\n", - "Inspecting these grids confirms the chained deflection: the grid arriving at `source_1`'s plane has been\n", - "deflected first by the lens galaxy and then by `source_0`'s mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.grid)\n", - "\n", - "print(f\"Number of planes traced through: {len(traced_grids)}\")\n", - "print(f\"Plane 0 (image-plane) \u2014 first 3 coordinates: {traced_grids[0][:3]}\")\n", - "print(f\"Plane 1 (source_0 at z=1.0) \u2014 first 3 coordinates: {traced_grids[1][:3]}\")\n", - "print(f\"Plane 2 (source_1 at z=2.0) \u2014 first 3 coordinates: {traced_grids[2][:3]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "After the fit, every linear Gaussian in each source's MGE basis has been assigned an `intensity` via linear\n", - "algebra. These are available via the fit's `linear_light_profile_intensity_dict`, keyed by light profile object.\n", - "\n", - "We print the intensity of the first Gaussian in each source's basis to confirm both sources have been\n", - "reconstructed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"\\nFirst Gaussian intensity, source_0 = \"\n", - " f\"{fit.linear_light_profile_intensity_dict[source_0_bulge.profile_list[0]]}\"\n", - ")\n", - "print(\n", - " f\"First Gaussian intensity, source_1 = \"\n", - " f\"{fit.linear_light_profile_intensity_dict[source_1_bulge.profile_list[0]]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A `Tracer` where every linear light profile has been replaced with an ordinary light profile carrying its\n", - "solved-for `intensity` is also accessible from the fit, which is useful for visualising each MGE basis with\n", - "its actual reconstructed amplitude." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer_fitted.galaxies[1].bulge, grid=plot_grid)\n", - "subplot_basis_image(basis=tracer_fitted.galaxies[2].bulge, grid=plot_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated the multi-plane ray-tracing and MGE source API for a double Einstein ring lens, without\n", - "invoking a non-linear search.\n", - "\n", - "In a real modeling workflow:\n", - "\n", - " - `modeling.py` shows how to fit the same system using `Nautilus`, but \"cheats\" by initialising priors at the\n", - " true values. It is therefore only useful as a tutorial.\n", - " - `chaining.py` is the practical workflow \u2014 two chained searches that initialise the lens and `source_0` first,\n", - " then introduce `source_1`. This is the script you'll actually use to fit data.\n", - " - `slam.py` is the most robust pipeline for production-quality DSPL modeling, ending in a pixelized source\n", - " reconstruction.\n", - "\n", - "The key takeaway from this script is that double Einstein rings are fit with the same `Tracer` + `FitImaging`\n", - "objects as any other lens; the only difference is that the `Tracer` contains three (or more) galaxies at\n", - "different redshifts and `source_0` carries both light AND mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Double Einstein Ring Fit\n", + "==================================\n", + "\n", + "A double Einstein ring lens is a strong lens system where two source galaxies at different redshifts are lensed\n", + "by the foreground lens galaxy. They appear as two distinct Einstein rings in the image-plane.\n", + "\n", + "This script illustrates the API for performing a fit to a double Einstein ring lens via the standard `Tracer`\n", + "and `FitImaging` objects, without invoking a non-linear search. It is intended to make the multi-plane\n", + "ray-tracing API concrete before the reader moves on to `modeling.py` (search-based) or `chaining.py` / `slam.py`\n", + "(realistic, robust modeling).\n", + "\n", + "The source galaxies are both modelled with a Multi Gaussian Expansion (MGE), the same source parameterization\n", + "used in `chaining.py` and `slam.py`. The MGE is built from a `Basis` of linear `Gaussian` light profiles, whose\n", + "`intensity` values are solved for via linear algebra at fit time.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **MGE Basis:** Build a `Basis` of linear Gaussians, used for both source galaxies.\n", + "- **Galaxies:** Compose the lens galaxy plus two source galaxies at different redshifts.\n", + "- **Tracer:** Build the three-plane `Tracer` that performs the multi-plane ray-tracing.\n", + "- **Fit:** Create a `FitImaging` and inspect the fit.\n", + "- **Multi-Plane Ray-Tracing:** A short tour of how the second source-plane sees deflection from both the lens\n", + " galaxy AND the first source galaxy's mass.\n", + "- **Intensities:** The solved-for linear light profile `intensity` values for each Gaussian, per source.\n", + "- **Wrap Up:** Summary and next steps.\n", + "\n", + "__Prerequisites__\n", + "\n", + "This script focuses on the API specific to a double Einstein ring fit. For background on the underlying single-plane\n", + "fit API and the MGE source parameterization, you should read first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/fit.py` \u2014 the standard single-plane fit.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py` \u2014 the MGE fit API and `Basis` of\n", + " linear Gaussians.\n", + "\n", + "The redshifts (`lens=0.5`, `source_0=1.0`, `source_1=2.0`) match those used by the simulator and modeling examples." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `double_einstein_ring` via .fits files.\n", + "\n", + "This dataset has a double Einstein ring, due to the two source galaxies at different redshifts behind the lens galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes both source-galaxy Einstein rings." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling, with finer sub-pixelization at the centre where the lens galaxy's mass is most\n", + "strongly deflecting light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE Basis__\n", + "\n", + "We build a single `Basis` of linear Gaussians which we will use as the source-galaxy light model for both\n", + "`source_0` and `source_1`.\n", + "\n", + "The Gaussians share a common centre and elliptical components (set to spherical here for simplicity), and have\n", + "`sigma` values spaced in log10 increments from 0.01\" up to a reasonable size cap. The `intensity` of each\n", + "Gaussian is a linear parameter, solved for by linear algebra at fit time \u2014 no non-linear search is required.\n", + "\n", + "Two `Basis` objects are constructed, with different sphericity centres, so that `source_0` and `source_1` each\n", + "have their own light component centred on the (known) position from the simulator.\n", + "\n", + "For background on the MGE `Basis` API, see\n", + "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "# Centres match the simulator: source_0 is offset from the lens, source_1 is on the far side of the lens.\n", + "\n", + "source_0_bulge = build_source_basis(centre=(-0.15, -0.15))\n", + "source_1_bulge = build_source_basis(centre=(-0.45, 0.45))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The Gaussians of each basis cannot be plotted yet because their `intensity` values have not been solved for \u2014\n", + "linear light profiles only acquire an `intensity` once a `FitImaging` runs its linear algebra step. After the\n", + "fit below, we visualise each source's MGE basis with its solved-for intensities.\n", + "\n", + "We set up the plotting grid we will use post-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "We now compose the three galaxies that form the double Einstein ring system:\n", + "\n", + " - `lens` (z=0.5): an `Isothermal` mass profile, matching the simulator. The lens galaxy has no light in this\n", + " simulated dataset.\n", + " - `source_0` (z=1.0): the MGE basis above as a light component, AND an `IsothermalSph` mass profile. `source_0`\n", + " acts as both a light source AND a deflector \u2014 its mass distribution bends the light from the higher-redshift\n", + " `source_1`, contributing to the second Einstein ring.\n", + " - `source_1` (z=2.0): only an MGE light component. It contributes light, but its own mass is negligible at\n", + " this redshift in the simulated data.\n", + "\n", + "The mass profile values for `lens` and `source_0` are set to the simulator's true values, so the fit visibly\n", + "recovers both Einstein rings without a search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.5,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=source_0_bulge,\n", + " mass=al.mp.IsothermalSph(centre=(-0.15, -0.15), einstein_radius=0.3),\n", + ")\n", + "\n", + "source_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " bulge=source_1_bulge,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "The `Tracer` performs the multi-plane ray-tracing. PyAutoLens orders the galaxies internally by redshift, so\n", + "the tracer first deflects image-plane (y,x) coordinates through the lens-plane (z=0.5) onto `source_0`'s plane\n", + "(z=1.0), then continues to deflect through `source_0`'s mass to reach `source_1`'s plane (z=2.0).\n", + "\n", + "Both the lens galaxy's mass AND `source_0`'s mass contribute to the deflection map applied to coordinates\n", + "arriving at `source_1`'s plane. This is what makes double Einstein ring systems sensitive to angular diameter\n", + "distance ratios, and therefore to cosmological parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens, source_0, source_1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We pass the `Tracer` to a `FitImaging` to fit the dataset. The fit performs the multi-plane ray-tracing, evaluates\n", + "each source galaxy's light at its own source-plane, sums the resulting image-plane contributions, convolves with\n", + "the PSF, and computes the residuals against the data.\n", + "\n", + "The `linear_light_profile_intensity_dict` of the fit will hold a solved-for `intensity` for every Gaussian in\n", + "both source MGE bases." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multi-Plane Ray-Tracing__\n", + "\n", + "The tracer exposes per-plane ray-traced grids via `traced_grid_2d_list_from`. The list returned has one grid per\n", + "plane (image-plane, `source_0`-plane, `source_1`-plane).\n", + "\n", + "Inspecting these grids confirms the chained deflection: the grid arriving at `source_1`'s plane has been\n", + "deflected first by the lens galaxy and then by `source_0`'s mass." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.grid)\n", + "\n", + "print(f\"Number of planes traced through: {len(traced_grids)}\")\n", + "print(f\"Plane 0 (image-plane) \u2014 first 3 coordinates: {traced_grids[0][:3]}\")\n", + "print(f\"Plane 1 (source_0 at z=1.0) \u2014 first 3 coordinates: {traced_grids[1][:3]}\")\n", + "print(f\"Plane 2 (source_1 at z=2.0) \u2014 first 3 coordinates: {traced_grids[2][:3]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "After the fit, every linear Gaussian in each source's MGE basis has been assigned an `intensity` via linear\n", + "algebra. These are available via the fit's `linear_light_profile_intensity_dict`, keyed by light profile object.\n", + "\n", + "We print the intensity of the first Gaussian in each source's basis to confirm both sources have been\n", + "reconstructed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"\\nFirst Gaussian intensity, source_0 = \"\n", + " f\"{fit.linear_light_profile_intensity_dict[source_0_bulge.profile_list[0]]}\"\n", + ")\n", + "print(\n", + " f\"First Gaussian intensity, source_1 = \"\n", + " f\"{fit.linear_light_profile_intensity_dict[source_1_bulge.profile_list[0]]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Tracer` where every linear light profile has been replaced with an ordinary light profile carrying its\n", + "solved-for `intensity` is also accessible from the fit, which is useful for visualising each MGE basis with\n", + "its actual reconstructed amplitude." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer_fitted.galaxies[1].bulge, grid=plot_grid)\n", + "subplot_basis_image(basis=tracer_fitted.galaxies[2].bulge, grid=plot_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated the multi-plane ray-tracing and MGE source API for a double Einstein ring lens, without\n", + "invoking a non-linear search.\n", + "\n", + "In a real modeling workflow:\n", + "\n", + " - `modeling.py` shows how to fit the same system using `Nautilus`, but \"cheats\" by initialising priors at the\n", + " true values. It is therefore only useful as a tutorial.\n", + " - `chaining.py` is the practical workflow \u2014 two chained searches that initialise the lens and `source_0` first,\n", + " then introduce `source_1`. This is the script you'll actually use to fit data.\n", + " - `slam.py` is the most robust pipeline for production-quality DSPL modeling, ending in a pixelized source\n", + " reconstruction.\n", + "\n", + "The key takeaway from this script is that double Einstein rings are fit with the same `Tracer` + `FitImaging`\n", + "objects as any other lens; the only difference is that the `Tracer` contains three (or more) galaxies at\n", + "different redshifts and `source_0` carries both light AND mass." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/double_einstein_ring/likelihood_function.ipynb b/notebooks/imaging/features/advanced/double_einstein_ring/likelihood_function.ipynb index 3505fb22e..7ab823cd6 100644 --- a/notebooks/imaging/features/advanced/double_einstein_ring/likelihood_function.ipynb +++ b/notebooks/imaging/features/advanced/double_einstein_ring/likelihood_function.ipynb @@ -1,370 +1,407 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Double Einstein Ring__\n", - "\n", - "This script describes the additional steps required to compute the `log_likelihood` for a double Einstein ring\n", - "lens \u2014 a strong lens system with two source galaxies at different redshifts behind the foreground lens.\n", - "\n", - "This script does NOT repeat the steps shared with single-plane lensing (mask, image-plane grid, convolution,\n", - "chi-squared, noise normalization). It documents only the parts of the likelihood function which are specific\n", - "to double Einstein ring multi-plane ray-tracing.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function below builds directly on standard imaging and MGE likelihood functions. You should read\n", - "these notebooks first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 the canonical single-plane log likelihood\n", - " walkthrough, covering image-plane grids, ray-tracing, source-plane evaluation, PSF convolution, chi-squared\n", - " and the noise normalization term.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a `Basis`\n", - " of linear Gaussians is solved for via linear algebra.\n", - "\n", - "Sections covered in those scripts (e.g. \"Chi Squared\", \"Noise Normalization Term\", \"Calculate The Log\n", - "Likelihood\") are not repeated here; this script focuses entirely on what changes for a double Einstein ring.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script (see above).\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Galaxies:** Three galaxies at three redshifts \u2014 lens, source_0 (light + mass), source_1 (light only).\n", - "- **Multi-Plane Ray-Tracing:** The deflection chain that produces three ray-traced grids, one per plane.\n", - "- **Source-Plane Images:** Each source galaxy's light is evaluated at its own ray-traced grid.\n", - "- **Model Image:** Sum of both source-plane contributions, then PSF convolution.\n", - "- **Likelihood:** Reference up to the canonical imaging likelihood for chi-squared, noise normalization, and\n", - " the final log likelihood expression.\n", - "- **Fit Check:** Confirm the manual reconstruction matches `FitImaging.log_likelihood`.\n", - "- **Wrap Up.**\n", - "\n", - "__What Changes For A Double Einstein Ring__\n", - "\n", - "For a single-plane lens, ray-tracing maps image-plane (y,x) coordinates onto a single source-plane via the lens\n", - "galaxy's deflection map alpha_lens(theta). The source galaxy's light is then evaluated at the source-plane\n", - "coordinates and projected back into the image plane.\n", - "\n", - "For a double Einstein ring, there are TWO source-planes at different redshifts, and the first source galaxy\n", - "acts as a deflector for the second. The deflection chain is:\n", - "\n", - " Plane 0 (image-plane) : theta\n", - " Plane 1 (source_0, z=1.0) : theta - alpha_lens(theta)\n", - " Plane 2 (source_1, z=2.0) : theta - alpha_lens(theta) - beta_01 * alpha_source_0(plane_1_grid)\n", - "\n", - "where beta_01 is a scaling factor derived from the angular diameter distances between the lens, source_0 and\n", - "source_1 \u2014 this is what makes double Einstein ring systems sensitive to cosmology. The factor is computed\n", - "internally by `Tracer.traced_grid_2d_list_from` based on the redshifts and cosmology.\n", - "\n", - "The model image is then the sum of the two source-plane reconstructions projected back to the image plane,\n", - "PSF-convolved, and compared to the data exactly as in the single-plane case." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the double Einstein ring dataset. The auto-simulation block mirrors the other example scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "The three galaxies that participate in the multi-plane ray-tracing:\n", - "\n", - " - `lens` (z=0.5): an `Isothermal` mass profile. No light, matching the simulator.\n", - " - `source_0` (z=1.0): an MGE light component (a simple basis of 10 linear Gaussians) AND an `IsothermalSph`\n", - " mass profile. `source_0` deflects the light from `source_1`.\n", - " - `source_1` (z=2.0): an MGE light component only.\n", - "\n", - "The mass-profile parameters and source centres are set to the simulator's true values so the manual likelihood\n", - "computation below produces a sensible-looking model image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 10\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.5,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=build_source_basis(centre=(-0.15, -0.15)),\n", - " mass=al.mp.IsothermalSph(centre=(-0.15, -0.15), einstein_radius=0.3),\n", - ")\n", - "\n", - "source_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " bulge=build_source_basis(centre=(-0.45, 0.45)),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source_0, source_1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multi-Plane Ray-Tracing__\n", - "\n", - "The single call below performs the full deflection chain. `traced_grid_2d_list_from` returns one grid per plane,\n", - "in redshift order: image-plane grid (no deflection), source_0-plane grid (deflected by the lens), source_1-plane\n", - "grid (deflected by the lens AND source_0).\n", - "\n", - "Note: the cosmological scaling factor `beta_01` mentioned in the header is applied internally to the\n", - "source_0-plane deflection contribution to the source_1-plane grid. PyAutoLens uses `Planck18` by default; to\n", - "make `Om0` a free parameter, see `modeling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=dataset.grid)\n", - "\n", - "grid_image_plane = traced_grid_list[0]\n", - "grid_source_0 = traced_grid_list[1]\n", - "grid_source_1 = traced_grid_list[2]\n", - "\n", - "print(f\"Number of planes traced: {len(traced_grid_list)}\")\n", - "print(f\"Plane 1 (source_0) first coord: {grid_source_0[0]}\")\n", - "print(f\"Plane 2 (source_1) first coord: {grid_source_1[0]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot the image-plane grid after it has been deflected to each source-plane. This visualises how the\n", - "caustics of the lens galaxy carve up image-plane (y,x) coordinates and route them to different source-plane\n", - "positions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=grid_source_0, title=\"Ray-traced grid at source_0 plane (z=1.0)\")\n", - "aplt.plot_grid(grid=grid_source_1, title=\"Ray-traced grid at source_1 plane (z=2.0)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Images__\n", - "\n", - "Each source galaxy's light component is evaluated at the ray-traced grid that arrives at its own redshift plane.\n", - "\n", - "For linear light profiles (Gaussians in our MGE), `image_2d_from` returns an `Array2D` of image-plane (y,x)\n", - "pixel values per profile, with a placeholder `intensity=1.0`. The true `intensity` is solved for at the linear\n", - "algebra step (see the MGE likelihood prerequisite).\n", - "\n", - "For this manual walkthrough we use the convenience method `image_2d_from` on the `Tracer`, which evaluates\n", - "every galaxy's light at the correct plane and sums them into a single model image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image_unconvolved = tracer.image_2d_from(grid=dataset.grid)\n", - "\n", - "aplt.plot_array(\n", - " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "What `image_2d_from` does internally for our double Einstein ring:\n", - "\n", - " 1. Ray-traces the image-plane grid to obtain `grid_source_0` and `grid_source_1`.\n", - " 2. Evaluates `source_0`'s MGE at `grid_source_0`, producing its image-plane contribution.\n", - " 3. Evaluates `source_1`'s MGE at `grid_source_1`, producing its image-plane contribution.\n", - " 4. Sums all source contributions into the model image.\n", - "\n", - "For a single-plane lens there is only one source-plane grid and one such evaluation; for the double Einstein\n", - "ring there are two.\n", - "\n", - "__Model Image__\n", - "\n", - "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model image\n", - "above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", - "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`.\n", - "\n", - "We delegate the remaining steps to `FitImaging`, which handles the linear-algebra step that solves for each\n", - "Gaussian's `intensity` and assembles the full `log_likelihood`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(f\"\\nLog likelihood of the manual double Einstein ring fit: {fit.log_likelihood}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood__\n", - "\n", - "The final `log_likelihood` combines:\n", - "\n", - " - The chi-squared term, computed from the residuals between the PSF-convolved model image and the data,\n", - " weighted by the noise map.\n", - " - The noise normalization term, the standard Gaussian normalization over all unmasked pixels.\n", - " - The linear algebra terms (regularization and curvature determinants) introduced by the MGE\n", - " `Basis` of linear Gaussians.\n", - "\n", - "The first two are documented in `imaging/likelihood_function.py`; the third in\n", - "`imaging/features/multi_gaussian_expansion/likelihood_function.py`. No new terms are introduced by the multi-plane\n", - "ray-tracing \u2014 the only change is the source-plane evaluation step described above.\n", - "\n", - "__Wrap Up__\n", - "\n", - "The double Einstein ring `log_likelihood` differs from the single-plane case in exactly one place: the source\n", - "galaxies are evaluated on different ray-traced grids, one per source-plane redshift. Every other step (PSF\n", - "convolution, chi-squared, noise normalization, linear algebra) is shared with the single-plane likelihood and\n", - "documented in the prerequisite scripts.\n", - "\n", - "The deflection scaling factor `beta_01` between source_0 and source_1 is the physical reason double Einstein\n", - "rings constrain cosmology: it depends on the ratio of angular diameter distances `D_{ls_0 -> s_1} / D_{s_0 ->\n", - "s_1}`, which in turn depends on the cosmological parameters. This is why the `modeling.py` example exposes the\n", - "option to make `Om0` a free parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Double Einstein Ring__\n", + "\n", + "This script describes the additional steps required to compute the `log_likelihood` for a double Einstein ring\n", + "lens \u2014 a strong lens system with two source galaxies at different redshifts behind the foreground lens.\n", + "\n", + "This script does NOT repeat the steps shared with single-plane lensing (mask, image-plane grid, convolution,\n", + "chi-squared, noise normalization). It documents only the parts of the likelihood function which are specific\n", + "to double Einstein ring multi-plane ray-tracing.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function below builds directly on standard imaging and MGE likelihood functions. You should read\n", + "these notebooks first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 the canonical single-plane log likelihood\n", + " walkthrough, covering image-plane grids, ray-tracing, source-plane evaluation, PSF convolution, chi-squared\n", + " and the noise normalization term.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a `Basis`\n", + " of linear Gaussians is solved for via linear algebra.\n", + "\n", + "Sections covered in those scripts (e.g. \"Chi Squared\", \"Noise Normalization Term\", \"Calculate The Log\n", + "Likelihood\") are not repeated here; this script focuses entirely on what changes for a double Einstein ring.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script (see above).\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Galaxies:** Three galaxies at three redshifts \u2014 lens, source_0 (light + mass), source_1 (light only).\n", + "- **Multi-Plane Ray-Tracing:** The deflection chain that produces three ray-traced grids, one per plane.\n", + "- **Source-Plane Images:** Each source galaxy's light is evaluated at its own ray-traced grid.\n", + "- **Model Image:** Sum of both source-plane contributions, then PSF convolution.\n", + "- **Likelihood:** Reference up to the canonical imaging likelihood for chi-squared, noise normalization, and\n", + " the final log likelihood expression.\n", + "- **Fit Check:** Confirm the manual reconstruction matches `FitImaging.log_likelihood`.\n", + "- **Wrap Up.**\n", + "\n", + "__What Changes For A Double Einstein Ring__\n", + "\n", + "For a single-plane lens, ray-tracing maps image-plane (y,x) coordinates onto a single source-plane via the lens\n", + "galaxy's deflection map alpha_lens(theta). The source galaxy's light is then evaluated at the source-plane\n", + "coordinates and projected back into the image plane.\n", + "\n", + "For a double Einstein ring, there are TWO source-planes at different redshifts, and the first source galaxy\n", + "acts as a deflector for the second. The deflection chain is:\n", + "\n", + " Plane 0 (image-plane) : theta\n", + " Plane 1 (source_0, z=1.0) : theta - alpha_lens(theta)\n", + " Plane 2 (source_1, z=2.0) : theta - alpha_lens(theta) - beta_01 * alpha_source_0(plane_1_grid)\n", + "\n", + "where beta_01 is a scaling factor derived from the angular diameter distances between the lens, source_0 and\n", + "source_1 \u2014 this is what makes double Einstein ring systems sensitive to cosmology. The factor is computed\n", + "internally by `Tracer.traced_grid_2d_list_from` based on the redshifts and cosmology.\n", + "\n", + "The model image is then the sum of the two source-plane reconstructions projected back to the image plane,\n", + "PSF-convolved, and compared to the data exactly as in the single-plane case." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the double Einstein ring dataset. The auto-simulation block mirrors the other example scripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "The three galaxies that participate in the multi-plane ray-tracing:\n", + "\n", + " - `lens` (z=0.5): an `Isothermal` mass profile. No light, matching the simulator.\n", + " - `source_0` (z=1.0): an MGE light component (a simple basis of 10 linear Gaussians) AND an `IsothermalSph`\n", + " mass profile. `source_0` deflects the light from `source_1`.\n", + " - `source_1` (z=2.0): an MGE light component only.\n", + "\n", + "The mass-profile parameters and source centres are set to the simulator's true values so the manual likelihood\n", + "computation below produces a sensible-looking model image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 10\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.5,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=build_source_basis(centre=(-0.15, -0.15)),\n", + " mass=al.mp.IsothermalSph(centre=(-0.15, -0.15), einstein_radius=0.3),\n", + ")\n", + "\n", + "source_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " bulge=build_source_basis(centre=(-0.45, 0.45)),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source_0, source_1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multi-Plane Ray-Tracing__\n", + "\n", + "The single call below performs the full deflection chain. `traced_grid_2d_list_from` returns one grid per plane,\n", + "in redshift order: image-plane grid (no deflection), source_0-plane grid (deflected by the lens), source_1-plane\n", + "grid (deflected by the lens AND source_0).\n", + "\n", + "Note: the cosmological scaling factor `beta_01` mentioned in the header is applied internally to the\n", + "source_0-plane deflection contribution to the source_1-plane grid. PyAutoLens uses `Planck18` by default; to\n", + "make `Om0` a free parameter, see `modeling.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=dataset.grid)\n", + "\n", + "grid_image_plane = traced_grid_list[0]\n", + "grid_source_0 = traced_grid_list[1]\n", + "grid_source_1 = traced_grid_list[2]\n", + "\n", + "print(f\"Number of planes traced: {len(traced_grid_list)}\")\n", + "print(f\"Plane 1 (source_0) first coord: {grid_source_0[0]}\")\n", + "print(f\"Plane 2 (source_1) first coord: {grid_source_1[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot the image-plane grid after it has been deflected to each source-plane. This visualises how the\n", + "caustics of the lens galaxy carve up image-plane (y,x) coordinates and route them to different source-plane\n", + "positions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=grid_source_0, title=\"Ray-traced grid at source_0 plane (z=1.0)\")\n", + "aplt.plot_grid(grid=grid_source_1, title=\"Ray-traced grid at source_1 plane (z=2.0)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Images__\n", + "\n", + "Each source galaxy's light component is evaluated at the ray-traced grid that arrives at its own redshift plane.\n", + "\n", + "For linear light profiles (Gaussians in our MGE), `image_2d_from` returns an `Array2D` of image-plane (y,x)\n", + "pixel values per profile, with a placeholder `intensity=1.0`. The true `intensity` is solved for at the linear\n", + "algebra step (see the MGE likelihood prerequisite).\n", + "\n", + "For this manual walkthrough we use the convenience method `image_2d_from` on the `Tracer`, which evaluates\n", + "every galaxy's light at the correct plane and sums them into a single model image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image_unconvolved = tracer.image_2d_from(grid=dataset.grid)\n", + "\n", + "aplt.plot_array(\n", + " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "What `image_2d_from` does internally for our double Einstein ring:\n", + "\n", + " 1. Ray-traces the image-plane grid to obtain `grid_source_0` and `grid_source_1`.\n", + " 2. Evaluates `source_0`'s MGE at `grid_source_0`, producing its image-plane contribution.\n", + " 3. Evaluates `source_1`'s MGE at `grid_source_1`, producing its image-plane contribution.\n", + " 4. Sums all source contributions into the model image.\n", + "\n", + "For a single-plane lens there is only one source-plane grid and one such evaluation; for the double Einstein\n", + "ring there are two.\n", + "\n", + "__Model Image__\n", + "\n", + "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model image\n", + "above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", + "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`.\n", + "\n", + "We delegate the remaining steps to `FitImaging`, which handles the linear-algebra step that solves for each\n", + "Gaussian's `intensity` and assembles the full `log_likelihood`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(f\"\\nLog likelihood of the manual double Einstein ring fit: {fit.log_likelihood}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood__\n", + "\n", + "The final `log_likelihood` combines:\n", + "\n", + " - The chi-squared term, computed from the residuals between the PSF-convolved model image and the data,\n", + " weighted by the noise map.\n", + " - The noise normalization term, the standard Gaussian normalization over all unmasked pixels.\n", + " - The linear algebra terms (regularization and curvature determinants) introduced by the MGE\n", + " `Basis` of linear Gaussians.\n", + "\n", + "The first two are documented in `imaging/likelihood_function.py`; the third in\n", + "`imaging/features/multi_gaussian_expansion/likelihood_function.py`. No new terms are introduced by the multi-plane\n", + "ray-tracing \u2014 the only change is the source-plane evaluation step described above.\n", + "\n", + "__Wrap Up__\n", + "\n", + "The double Einstein ring `log_likelihood` differs from the single-plane case in exactly one place: the source\n", + "galaxies are evaluated on different ray-traced grids, one per source-plane redshift. Every other step (PSF\n", + "convolution, chi-squared, noise normalization, linear algebra) is shared with the single-plane likelihood and\n", + "documented in the prerequisite scripts.\n", + "\n", + "The deflection scaling factor `beta_01` between source_0 and source_1 is the physical reason double Einstein\n", + "rings constrain cosmology: it depends on the ratio of angular diameter distances `D_{ls_0 -> s_1} / D_{s_0 ->\n", + "s_1}`, which in turn depends on the cosmological parameters. This is why the `modeling.py` example exposes the\n", + "option to make `Om0` a free parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/double_einstein_ring/modeling.ipynb b/notebooks/imaging/features/advanced/double_einstein_ring/modeling.ipynb index be3b60f1d..eb14a9439 100644 --- a/notebooks/imaging/features/advanced/double_einstein_ring/modeling.ipynb +++ b/notebooks/imaging/features/advanced/double_einstein_ring/modeling.ipynb @@ -1,571 +1,608 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Double Einstein Ring\n", - "=======================================\n", - "\n", - "A double Einstein ring lens is a strong lens system where there are two source galaxies at different redshifts\n", - "behind the lens galaxy. They appear as two distinct Einstein rings in the image-plane, and can constrain\n", - "Cosmological parameters in a way single Einstein ring lenses cannot.\n", - "\n", - "To analyse these systems correctly the mass of the lens galaxy and the first source galaxy must be modeled\n", - "simultaneously, and the emission of both source galaxies must be modeled simultaneously.\n", - "\n", - "This script illustrates the PyAutoLens API for modeling a double Einstein ring lens.\n", - "\n", - "__Practical Use: Read This First__\n", - "\n", - "This script is a tutorial. It produces a working fit by \"cheating\" \u2014 every prior is initialised at the true\n", - "simulator value, narrowed by a small Gaussian. On real data this is impossible, and a single Nautilus search on\n", - "a 16-parameter double Einstein ring model would almost certainly converge to a local maximum.\n", - "\n", - "The script you will actually use to fit a double Einstein ring on real data is\n", - "`autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/chaining.py`, which runs two chained\n", - "non-linear searches: the first initialises the lens mass and `source_0` using a smaller mask that excludes\n", - "`source_1`, the second introduces `source_1` and frees `source_0`'s mass. This is also significantly more\n", - "computationally efficient than the single-search approach below.\n", - "\n", - "For production-quality modeling, see `slam.py` in the same directory.\n", - "\n", - "Read this script to understand the model composition API, then jump to `chaining.py`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Cheating:** Initializing a double Einstein ring lens model is difficult, due to the complexity of parameter.\n", - "- **Cosmology:** Double Einstein rings allow cosmological parameters to be constrained \u2014 `Om0` is fixed at\n", - " Planck18 here, with a commented-out snippet showing how to make it free.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a double Einstein ring where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The first lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The second lens galaxy / first source galaxy's light is a linear `ExponentialSph` and its mass a `IsothermalSph`.\n", - " - The second source galaxy's light is a linear `ExponentialSph`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `double_einstein_ring` via .fits files.\n", - "\n", - "This dataset has a double Einstien ring, due to the two source galaxies at different redshifts behind the lens galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Visualization of this dataset shows two distinct Einstein rings, which are the two source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of both of the lensed source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The first lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", - "\n", - " - The second lens / first source galaxy's light are MGE models [8 parameters].\n", - "\n", - " - The second source galaxy's light is a linear `ExponentialSph` [3 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=16.\n", - "\n", - "Note that the galaxies are assigned redshifts of 0.5, 1.0 and 2.0. This ensures the multi-plane ray-tracing necessary\n", - "for the double Einstein ring lens system is performed correctly.\n", - "\n", - "The `centre` values input into `mge_model_from` are explained below.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " centre_prior_is_uniform=True,\n", - " centre=(0.0, 0.0),\n", - ")\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - "# Source 0:\n", - "\n", - "bulge = af.Model(al.lp_linear.ExponentialCoreSph)\n", - "mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - "source_0 = af.Model(\n", - " al.Galaxy, redshift=1.0, bulge=bulge, mass=mass, centre=(0.15, 0.15)\n", - ")\n", - "\n", - "# Source 1:\n", - "\n", - "bulge = af.Model(al.lp_linear.ExponentialCoreSph)\n", - "\n", - "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=bulge, centre=(0.0, 0.0))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cheating__\n", - "\n", - "Initializing a double Einstein ring lens model is difficult, due to the complexity of parameter space. It is common to \n", - "infer local maxima, which this script does if default broad priors on every model parameter are assumed.\n", - "\n", - "To infer the correct model, we \"cheat\" and overwrite all of the priors of the model parameters to start centred on \n", - "their true values. This is why the true `centre` values were input into the `mge_model_from` functions above.\n", - "\n", - "For real data, we obviously do not know the true parameters and therefore cannot cheat in this way. Readers should\n", - "checkout the **PyAutoLens**'s advanced feature `chaining`, which chains together multiple non-linear searches. \n", - "\n", - "This feature is described in HowToLens chapter 3 and specific examples for a double Einstein ring are given in\n", - "the script `guides/modeling/chaining/double_einstein_ring.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens.mass.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", - "lens.mass.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", - "\n", - "source_0.mass.centre_0 = af.GaussianPrior(mean=-0.15, sigma=0.2)\n", - "source_0.mass.centre_1 = af.GaussianPrior(mean=-0.15, sigma=0.2)\n", - "source_0.mass.einstein_radius = af.GaussianPrior(mean=0.4, sigma=0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cosmology__\n", - "\n", - "Double Einstein rings allow cosmological parameters to be constrained, because they provide information on the\n", - "different angular diameter distances between the lens, `source_0` and `source_1`. The deflection scaling factor\n", - "between the two source-planes (sometimes written `beta_01`) depends on those distances and therefore on the\n", - "cosmology.\n", - "\n", - "For this tutorial, we use a fixed `Planck18` cosmology and treat no cosmological parameters as free. This keeps\n", - "the model dimensionality and the parameter space tractable for the single-search \"cheating\" workflow below; a\n", - "realistic cosmological constraint requires both the chained-search workflow in `chaining.py` and significantly\n", - "more data than this one simulated system.\n", - "\n", - "To make `Om0` (Omega_m) a free parameter in your own fit, uncomment the three lines below. They construct a\n", - "`FlatLambdaCDM` cosmology as a free model, override the prior on `Om0`, and include the cosmology in the overall\n", - "`Collection`. The remaining cosmological parameters (`H0`, `Tcmb0`, etc.) stay fixed at their Planck18 values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# cosmology = af.Model(al.cosmo.FlatLambdaCDM)\n", - "# cosmology.Om0 = af.GaussianPrior(mean=0.3, sigma=0.1)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", - " # cosmology=cosmology,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "This confirms the model is composed of three galaxies, two of which are lensed source galaxies, and a fixed\n", - "Planck18 cosmology." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"double_einstein_ring\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=2000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", - "required by a model.\n", - "\n", - "Double source plane lenses can use a lot of VRAM, because the multi-plane ray-tracing and creation of multiple\n", - "images for different source planes can require all the additional data to be stored in VRAM. This will\n", - "at least double the VRAM requirements compared to a single lens plane model, but often more than this.\n", - "\n", - "Given VRAM use is an important consideration, we print out the estimated VRAM required for this\n", - "model-fit and advise you do this for your own double source plane lens model-fits.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The likelihood evaluation time for analysing double Einstein ring lens is quite a lot longer than single lens plane\n", - "lenses. This is because multi-plane ray-tracing calculations are computationally expensive. \n", - "\n", - "However, the real hit on run-time is the large number of free parameters in the model, which is often 10+ parameters\n", - "more than a single lens plane model. This means that the non-linear search takes longer to converge on a solution.\n", - "In this example, we cheated by initializing the priors on the model close to the correct solution. \n", - "\n", - "Combining pixelized source analyses with double Einstein ring lenses is very computationally expensive, because the\n", - "linear algebra calculations become significantly more expensive. This is not shown in this script, but is worth\n", - "baring in mind.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this):" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "These plots show that the lens and both sources of the double Einstein ring were fitted successfully." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "These examples show how the results API can be extended to investigate double Einstein ring results.\n", - "\n", - "__Wrap Up__\n", - "\n", - "Double Einstein ring systems can be fitted in **PyAutoLens**, however this script bypass the most difficult aspect\n", - "of fitting these systems by \"cheating\", and manually adjusting the priors to be near their true values.\n", - "\n", - "Modeling real observations of double Einstein rings is one of the hardest lens modeling tasks, and requires an high\n", - "degree of lens modeling expertise to make a success.\n", - "\n", - "If you have not already, I recommend you familiarize yourself with and use all of the following **PyAutoLens features\n", - "to model a real double Einstein ring:\n", - "\n", - " - Basis based light profiles (e.g. ``shapelets.ipynb` / `multi_gaussian_expansion.ipynb`): these allow one to fit\n", - " complex lens and source morphologies whilst keeping the dimensionality of the problem low.\n", - "\n", - " - Search chaining (e.g. `guides/modeling/chaining` and HowToLens chapter 3): by breaking the model-fit into a series\n", - " of Nautilus searches models of gradually increasing complexity can be fitted.\n", - "\n", - " - pixelizations (e.g. `pixelization.ipynb` and HowToLens chapter 4): to infer the cosmological parameters reliably\n", - " the source must be reconstructed on an adaptive mesh to capture a irregular morphological features." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Double Einstein Ring\n", + "=======================================\n", + "\n", + "A double Einstein ring lens is a strong lens system where there are two source galaxies at different redshifts\n", + "behind the lens galaxy. They appear as two distinct Einstein rings in the image-plane, and can constrain\n", + "Cosmological parameters in a way single Einstein ring lenses cannot.\n", + "\n", + "To analyse these systems correctly the mass of the lens galaxy and the first source galaxy must be modeled\n", + "simultaneously, and the emission of both source galaxies must be modeled simultaneously.\n", + "\n", + "This script illustrates the PyAutoLens API for modeling a double Einstein ring lens.\n", + "\n", + "__Practical Use: Read This First__\n", + "\n", + "This script is a tutorial. It produces a working fit by \"cheating\" \u2014 every prior is initialised at the true\n", + "simulator value, narrowed by a small Gaussian. On real data this is impossible, and a single Nautilus search on\n", + "a 16-parameter double Einstein ring model would almost certainly converge to a local maximum.\n", + "\n", + "The script you will actually use to fit a double Einstein ring on real data is\n", + "`autolens_workspace/scripts/imaging/features/advanced/double_einstein_ring/chaining.py`, which runs two chained\n", + "non-linear searches: the first initialises the lens mass and `source_0` using a smaller mask that excludes\n", + "`source_1`, the second introduces `source_1` and frees `source_0`'s mass. This is also significantly more\n", + "computationally efficient than the single-search approach below.\n", + "\n", + "For production-quality modeling, see `slam.py` in the same directory.\n", + "\n", + "Read this script to understand the model composition API, then jump to `chaining.py`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Cheating:** Initializing a double Einstein ring lens model is difficult, due to the complexity of parameter.\n", + "- **Cosmology:** Double Einstein rings allow cosmological parameters to be constrained \u2014 `Om0` is fixed at\n", + " Planck18 here, with a commented-out snippet showing how to make it free.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a double Einstein ring where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The first lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The second lens galaxy / first source galaxy's light is a linear `ExponentialSph` and its mass a `IsothermalSph`.\n", + " - The second source galaxy's light is a linear `ExponentialSph`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `double_einstein_ring` via .fits files.\n", + "\n", + "This dataset has a double Einstien ring, due to the two source galaxies at different redshifts behind the lens galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Visualization of this dataset shows two distinct Einstein rings, which are the two source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of both of the lensed source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The first lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + "\n", + " - The second lens / first source galaxy's light are MGE models [8 parameters].\n", + "\n", + " - The second source galaxy's light is a linear `ExponentialSph` [3 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=16.\n", + "\n", + "Note that the galaxies are assigned redshifts of 0.5, 1.0 and 2.0. This ensures the multi-plane ray-tracing necessary\n", + "for the double Einstein ring lens system is performed correctly.\n", + "\n", + "The `centre` values input into `mge_model_from` are explained below.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " centre_prior_is_uniform=True,\n", + " centre=(0.0, 0.0),\n", + ")\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + "# Source 0:\n", + "\n", + "bulge = af.Model(al.lp_linear.ExponentialCoreSph)\n", + "mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + "source_0 = af.Model(\n", + " al.Galaxy, redshift=1.0, bulge=bulge, mass=mass, centre=(0.15, 0.15)\n", + ")\n", + "\n", + "# Source 1:\n", + "\n", + "bulge = af.Model(al.lp_linear.ExponentialCoreSph)\n", + "\n", + "source_1 = af.Model(al.Galaxy, redshift=2.0, bulge=bulge, centre=(0.0, 0.0))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cheating__\n", + "\n", + "Initializing a double Einstein ring lens model is difficult, due to the complexity of parameter space. It is common to \n", + "infer local maxima, which this script does if default broad priors on every model parameter are assumed.\n", + "\n", + "To infer the correct model, we \"cheat\" and overwrite all of the priors of the model parameters to start centred on \n", + "their true values. This is why the true `centre` values were input into the `mge_model_from` functions above.\n", + "\n", + "For real data, we obviously do not know the true parameters and therefore cannot cheat in this way. Readers should\n", + "checkout the **PyAutoLens**'s advanced feature `chaining`, which chains together multiple non-linear searches. \n", + "\n", + "This feature is described in HowToLens chapter 3 and specific examples for a double Einstein ring are given in\n", + "the script `guides/modeling/chaining/double_einstein_ring.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens.mass.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", + "lens.mass.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", + "\n", + "source_0.mass.centre_0 = af.GaussianPrior(mean=-0.15, sigma=0.2)\n", + "source_0.mass.centre_1 = af.GaussianPrior(mean=-0.15, sigma=0.2)\n", + "source_0.mass.einstein_radius = af.GaussianPrior(mean=0.4, sigma=0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cosmology__\n", + "\n", + "Double Einstein rings allow cosmological parameters to be constrained, because they provide information on the\n", + "different angular diameter distances between the lens, `source_0` and `source_1`. The deflection scaling factor\n", + "between the two source-planes (sometimes written `beta_01`) depends on those distances and therefore on the\n", + "cosmology.\n", + "\n", + "For this tutorial, we use a fixed `Planck18` cosmology and treat no cosmological parameters as free. This keeps\n", + "the model dimensionality and the parameter space tractable for the single-search \"cheating\" workflow below; a\n", + "realistic cosmological constraint requires both the chained-search workflow in `chaining.py` and significantly\n", + "more data than this one simulated system.\n", + "\n", + "To make `Om0` (Omega_m) a free parameter in your own fit, uncomment the three lines below. They construct a\n", + "`FlatLambdaCDM` cosmology as a free model, override the prior on `Om0`, and include the cosmology in the overall\n", + "`Collection`. The remaining cosmological parameters (`H0`, `Tcmb0`, etc.) stay fixed at their Planck18 values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# cosmology = af.Model(al.cosmo.FlatLambdaCDM)\n", + "# cosmology.Om0 = af.GaussianPrior(mean=0.3, sigma=0.1)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1),\n", + " # cosmology=cosmology,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "This confirms the model is composed of three galaxies, two of which are lensed source galaxies, and a fixed\n", + "Planck18 cosmology." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"double_einstein_ring\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=2000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", + "required by a model.\n", + "\n", + "Double source plane lenses can use a lot of VRAM, because the multi-plane ray-tracing and creation of multiple\n", + "images for different source planes can require all the additional data to be stored in VRAM. This will\n", + "at least double the VRAM requirements compared to a single lens plane model, but often more than this.\n", + "\n", + "Given VRAM use is an important consideration, we print out the estimated VRAM required for this\n", + "model-fit and advise you do this for your own double source plane lens model-fits.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The likelihood evaluation time for analysing double Einstein ring lens is quite a lot longer than single lens plane\n", + "lenses. This is because multi-plane ray-tracing calculations are computationally expensive. \n", + "\n", + "However, the real hit on run-time is the large number of free parameters in the model, which is often 10+ parameters\n", + "more than a single lens plane model. This means that the non-linear search takes longer to converge on a solution.\n", + "In this example, we cheated by initializing the priors on the model close to the correct solution. \n", + "\n", + "Combining pixelized source analyses with double Einstein ring lenses is very computationally expensive, because the\n", + "linear algebra calculations become significantly more expensive. This is not shown in this script, but is worth\n", + "baring in mind.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this):" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "These plots show that the lens and both sources of the double Einstein ring were fitted successfully." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "These examples show how the results API can be extended to investigate double Einstein ring results.\n", + "\n", + "__Wrap Up__\n", + "\n", + "Double Einstein ring systems can be fitted in **PyAutoLens**, however this script bypass the most difficult aspect\n", + "of fitting these systems by \"cheating\", and manually adjusting the priors to be near their true values.\n", + "\n", + "Modeling real observations of double Einstein rings is one of the hardest lens modeling tasks, and requires an high\n", + "degree of lens modeling expertise to make a success.\n", + "\n", + "If you have not already, I recommend you familiarize yourself with and use all of the following **PyAutoLens features\n", + "to model a real double Einstein ring:\n", + "\n", + " - Basis based light profiles (e.g. ``shapelets.ipynb` / `multi_gaussian_expansion.ipynb`): these allow one to fit\n", + " complex lens and source morphologies whilst keeping the dimensionality of the problem low.\n", + "\n", + " - Search chaining (e.g. `guides/modeling/chaining` and HowToLens chapter 3): by breaking the model-fit into a series\n", + " of Nautilus searches models of gradually increasing complexity can be fitted.\n", + "\n", + " - pixelizations (e.g. `pixelization.ipynb` and HowToLens chapter 4): to infer the cosmological parameters reliably\n", + " the source must be reconstructed on an adaptive mesh to capture a irregular morphological features." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/double_einstein_ring/simulator.ipynb b/notebooks/imaging/features/advanced/double_einstein_ring/simulator.ipynb index 2a41832b6..8f8068b1d 100644 --- a/notebooks/imaging/features/advanced/double_einstein_ring/simulator.ipynb +++ b/notebooks/imaging/features/advanced/double_einstein_ring/simulator.ipynb @@ -1,373 +1,410 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Double Einstein Ring\n", - "===============================\n", - "\n", - "A double Einstein ring lens is a strong lens system where there are two source galaxies at different redshifts\n", - "behind the lens galaxy. They appear as two distinct Einstein rings in the image-plane, and can constrain\n", - "Cosmological parameters in a way single Einstein ring lenses cannot.\n", - "\n", - "This script simulates a double Einstein ring lens, which is used in `modeling/features/double_einstein_ring.py`\n", - "and other script to illustrate how to model these systems correctly.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where there are two source's at different different\n", - "redshifts, creating a double einstein ring. Specifically:\n", - "\n", - " - The first lens galaxy's light is an `Sersic` and total mass distribution is an `Isothermal`.\n", - " - The second galaxy, which is the first source galaxy, total mass distribution is an `Isothermal` and its light is\n", - " an `Sersic`.\n", - " - The third galaxy, which is the second source galaxy, light profile is an `Sersic`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "\n", - "The `Tracer` below is composed of three galaxies, a lens and two sources, where the first source also acts as a\n", - "lens for the second source.\n", - "\n", - "The redshifts of the galaxies are used to determine how many images are formed of the source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.5,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.ExponentialCoreSph(\n", - " centre=(-0.15, -0.15), intensity=1.2, effective_radius=0.1\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-0.15, -0.15), einstein_radius=0.3),\n", - ")\n", - "\n", - "source_galaxy_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " bulge=al.lp.ExponentialCoreSph(\n", - " centre=(-0.45, 0.45), intensity=0.6, effective_radius=0.07\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the \n", - "folder `autolens_workspace/imaging/ double_einstein_ring`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Double Einstein Ring\n", + "===============================\n", + "\n", + "A double Einstein ring lens is a strong lens system where there are two source galaxies at different redshifts\n", + "behind the lens galaxy. They appear as two distinct Einstein rings in the image-plane, and can constrain\n", + "Cosmological parameters in a way single Einstein ring lenses cannot.\n", + "\n", + "This script simulates a double Einstein ring lens, which is used in `modeling/features/double_einstein_ring.py`\n", + "and other script to illustrate how to model these systems correctly.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where there are two source's at different different\n", + "redshifts, creating a double einstein ring. Specifically:\n", + "\n", + " - The first lens galaxy's light is an `Sersic` and total mass distribution is an `Isothermal`.\n", + " - The second galaxy, which is the first source galaxy, total mass distribution is an `Isothermal` and its light is\n", + " an `Sersic`.\n", + " - The third galaxy, which is the second source galaxy, light profile is an `Sersic`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "\n", + "The `Tracer` below is composed of three galaxies, a lens and two sources, where the first source also acts as a\n", + "lens for the second source.\n", + "\n", + "The redshifts of the galaxies are used to determine how many images are formed of the source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.5,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.ExponentialCoreSph(\n", + " centre=(-0.15, -0.15), intensity=1.2, effective_radius=0.1\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-0.15, -0.15), einstein_radius=0.3),\n", + ")\n", + "\n", + "source_galaxy_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " bulge=al.lp.ExponentialCoreSph(\n", + " centre=(-0.45, 0.45), intensity=0.6, effective_radius=0.07\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the \n", + "folder `autolens_workspace/imaging/ double_einstein_ring`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/double_einstein_ring/slam.ipynb b/notebooks/imaging/features/advanced/double_einstein_ring/slam.ipynb index 14963fb2a..499931165 100644 --- a/notebooks/imaging/features/advanced/double_einstein_ring/slam.ipynb +++ b/notebooks/imaging/features/advanced/double_einstein_ring/slam.ipynb @@ -1,752 +1,789 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM (Source, Light and Mass): Double Einstein Ring\n", - "===================================================\n", - "\n", - "This script adapts the SLaM (Source, Light and Mass) pipelines to a Double Source Plane Lens (DSPL) system, where\n", - "there are two source galaxies at different redshifts behind the lens galaxy. Both source galaxies must be modeled\n", - "simultaneously because the first source (at the intermediate redshift) acts as both a background source of the lens\n", - "galaxy and as an additional deflector for the second (higher-redshift) source.\n", - "\n", - "This script is the DSPL analogue of `guides/modeling/slam_start_here.py`. It follows the same API conventions \u2014\n", - "each pipeline stage is a plain inline Python function, priors are chained via `al.util.chaining.mass_from`, image\n", - "positions are derived automatically via `positions_likelihood_from`, and MGE light profiles are constructed via\n", - "`al.model_util.mge_model_from`.\n", - "\n", - "__DSPL-Specific Differences From Standard SLaM__\n", - "\n", - " - There are two source galaxies (`source_0` at redshift 1.0, `source_1` at redshift 2.0). `source_0` is a light\n", - " source AND a mass deflector for `source_1`.\n", - " - The SOURCE LP PIPELINE is split into two searches. The first fits lens + `source_0` only, providing a stable\n", - " starting point. The second frees `source_0`'s mass and adds `source_1`'s light.\n", - " - The SOURCE PIX PIPELINE has an extra search: one pixelizes `source_0` while `source_1` is a bare ray-tracing\n", - " galaxy, and the next pixelizes `source_1` with `source_0`'s mass fixed from the previous search.\n", - " - Two `PositionsLH` likelihoods are used once both sources are active, one per source-plane redshift.\n", - " - Adapt images are stitched across pipeline stages (e.g. `source_0`'s adapt image comes from a different result\n", - " than `source_1`'s).\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE and SOURCE PIX PIPELINE, this DSPL SLaM modeling script fits an `Imaging` dataset of a\n", - "double Einstein ring system where in the final model:\n", - "\n", - " - The lens galaxy's light is a bulge with MGE light profile.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` plus an `ExternalShear`.\n", - " - The first source galaxy's light is a `Pixelization` and its mass is an `Isothermal`.\n", - " - The second source galaxy's light is a `Pixelization`.\n", - "\n", - "Optional LIGHT LP and MASS TOTAL stages are left as a follow-up exercise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 1__\n", - "\n", - "The first SOURCE LP PIPELINE search initializes a model where `source_1` is ignored and only the lens galaxy and\n", - "`source_0` are fit. This single-plane fit provides robust initial priors for the lens light, lens mass, shear and\n", - "`source_0` light before the more complex DSPL model is introduced.\n", - "\n", - "The model:\n", - " - Lens light: MGE with 2 x 20 Gaussians.\n", - " - Lens mass: `Isothermal` + `ExternalShear`.\n", - " - `source_0` light: MGE with 1 x 20 Gaussians." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source_0: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_0_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source_0,\n", - " bulge=source_0_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE 2__\n", - "\n", - "The second SOURCE LP PIPELINE search introduces the second source galaxy. The lens bulge / mass / shear and\n", - "`source_0` light are fixed to the instance values from search 1, and we free:\n", - "\n", - " - `source_0`'s mass: `Isothermal` with a prior tightly centred on the origin (the first source typically sits\n", - " close to the lens centre).\n", - " - `source_1`'s light: MGE with 1 x 20 Gaussians.\n", - "\n", - "A `PositionsLH` for `source_0` is attached to the analysis \u2014 automatically derived from the search-1 result via\n", - "`positions_likelihood_from` \u2014 to prevent unphysical mass models during the DSPL fit. No `source_1` positions are\n", - "available yet (this is the first search that fits `source_1`); they are introduced from the pixelized results\n", - "further down the pipeline." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_lp_result_1: af.Result,\n", - " redshift_source_1: float,\n", - " n_batch: int = 30,\n", - ") -> af.Result:\n", - " positions_likelihood_source_0 = source_lp_result_1.positions_likelihood_from(\n", - " factor=3.0,\n", - " minimum_threshold=0.3,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[positions_likelihood_source_0],\n", - " )\n", - "\n", - " source_0_mass = af.Model(al.mp.Isothermal)\n", - " source_0_mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - " source_0_mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - " source_0_mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - " source_1_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result_1.instance.galaxies.lens.bulge,\n", - " mass=source_lp_result_1.instance.galaxies.lens.mass,\n", - " shear=source_lp_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_1.instance.galaxies.source_0.redshift,\n", - " bulge=source_lp_result_1.instance.galaxies.source_0.bulge,\n", - " mass=source_0_mass,\n", - " ),\n", - " source_1=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source_1,\n", - " bulge=source_1_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1 \u2014 source_0__\n", - "\n", - "Pixelizes `source_0` while `source_1` is present only as a bare galaxy for ray-tracing purposes. The lens mass is\n", - "freed with priors initialized from the SOURCE LP PIPELINE result, and `source_1` is not fit (no light, no mass)\n", - "so this search constrains the lens mass via `source_0` alone.\n", - "\n", - "Adapt images come from the SOURCE LP PIPELINE result 1 (the single-plane fit)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1_source_0(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result_1: af.Result,\n", - " source_lp_result_2: af.Result,\n", - " redshift_source_1: float,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result_1\n", - " )\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " positions_likelihood_source_0 = source_lp_result_1.positions_likelihood_from(\n", - " factor=3.0,\n", - " minimum_threshold=0.2,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[positions_likelihood_source_0],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.Isothermal),\n", - " mass_result=source_lp_result_2.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result_2.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result_2.instance.galaxies.lens.bulge,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " source_1=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source_1,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]_source_0\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1 \u2014 source_1__\n", - "\n", - "Pixelizes `source_1`. `source_0`'s mass is freed with priors initialized from the previous pixelized search, so\n", - "the second source plane constrains the lensing by the first source. The lens mass is fixed from\n", - "`source_pix_result_1_source_0`.\n", - "\n", - "Two `PositionsLH` are attached \u2014 one per source plane \u2014 to prevent unphysical reconstructions.\n", - "\n", - "Adapt images are stitched: the lens adapt image comes from the LP pipeline result 2; `source_0`'s adapt image\n", - "comes from the pixelized search above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1_source_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result_2: af.Result,\n", - " source_pix_result_1_source_0: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " lp2_dict = al.galaxy_name_image_dict_via_result_from(result=source_lp_result_2)\n", - " pix_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1_source_0\n", - " )\n", - "\n", - " galaxy_name_image_dict = {\n", - " \"('galaxies', 'lens')\": lp2_dict[\"('galaxies', 'lens')\"],\n", - " \"('galaxies', 'source_0')\": pix_dict[\"('galaxies', 'source_0')\"],\n", - " }\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", - "\n", - " positions_likelihood_source_0 = (\n", - " source_pix_result_1_source_0.positions_likelihood_from(\n", - " factor=3.0,\n", - " minimum_threshold=0.2,\n", - " plane_redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", - " )\n", - " )\n", - " positions_likelihood_source_1 = source_lp_result_2.positions_likelihood_from(\n", - " factor=3.0,\n", - " minimum_threshold=0.2,\n", - " plane_redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " positions_likelihood_source_0,\n", - " positions_likelihood_source_1,\n", - " ],\n", - " )\n", - "\n", - " source_0_mass = al.util.chaining.mass_from(\n", - " mass=af.Model(al.mp.Isothermal),\n", - " mass_result=source_lp_result_2.model.galaxies.source_0.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result_2.instance.galaxies.lens.bulge,\n", - " mass=source_pix_result_1_source_0.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1_source_0.instance.galaxies.lens.shear,\n", - " ),\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", - " mass=source_0_mass,\n", - " ),\n", - " source_1=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]_source_1\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "The final SOURCE PIX PIPELINE search fits both source galaxies simultaneously with adaptive pixelizations.\n", - "Lens mass, shear and `source_0`'s mass are all fixed to the maximum-likelihood instances of the previous\n", - "pixelized searches; only the pixelization regularization parameters are free.\n", - "\n", - "The `RectangularAdaptImage` (or equivalent) mesh uses the high-quality adapt images built up over the earlier\n", - "pipeline stages to adapt each source-plane pixelization to its reconstructed morphology." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result_2: af.Result,\n", - " source_pix_result_1_source_0: af.Result,\n", - " source_pix_result_1_source_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " lp2_dict = al.galaxy_name_image_dict_via_result_from(result=source_lp_result_2)\n", - " pix0_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1_source_0\n", - " )\n", - " pix1_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1_source_1\n", - " )\n", - "\n", - " galaxy_name_image_dict = {\n", - " \"('galaxies', 'lens')\": lp2_dict[\"('galaxies', 'lens')\"],\n", - " \"('galaxies', 'source_0')\": pix0_dict[\"('galaxies', 'source_0')\"],\n", - " \"('galaxies', 'source_1')\": pix1_dict[\"('galaxies', 'source_1')\"],\n", - " }\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result_2.instance.galaxies.lens.bulge,\n", - " mass=source_pix_result_1_source_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1_source_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source_0=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", - " mass=source_pix_result_1_source_1.instance.galaxies.source_0.mass,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " source_1=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"double_einstein_ring\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam_dspl\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens galaxy and the two source galaxies. Multi-plane ray-tracing uses these to compute the\n", - "correct deflection angles between each plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source_0 = 1.0\n", - "redshift_source_1 = 2.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The pixelization mesh shape is fixed before modeling; see `features/pixelization/modeling` for details." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below runs the full DSPL SLaM pipeline. See the docstring above each function for a description of\n", - "each stage." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result_1 = source_lp_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source_0=redshift_source_0,\n", - ")\n", - "\n", - "source_lp_result_2 = source_lp_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_lp_result_1=source_lp_result_1,\n", - " redshift_source_1=redshift_source_1,\n", - ")\n", - "\n", - "source_pix_result_1_source_0 = source_pix_1_source_0(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result_1=source_lp_result_1,\n", - " source_lp_result_2=source_lp_result_2,\n", - " redshift_source_1=redshift_source_1,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "source_pix_result_1_source_1 = source_pix_1_source_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result_2=source_lp_result_2,\n", - " source_pix_result_1_source_0=source_pix_result_1_source_0,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result_2=source_lp_result_2,\n", - " source_pix_result_1_source_0=source_pix_result_1_source_0,\n", - " source_pix_result_1_source_1=source_pix_result_1_source_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM (Source, Light and Mass): Double Einstein Ring\n", + "===================================================\n", + "\n", + "This script adapts the SLaM (Source, Light and Mass) pipelines to a Double Source Plane Lens (DSPL) system, where\n", + "there are two source galaxies at different redshifts behind the lens galaxy. Both source galaxies must be modeled\n", + "simultaneously because the first source (at the intermediate redshift) acts as both a background source of the lens\n", + "galaxy and as an additional deflector for the second (higher-redshift) source.\n", + "\n", + "This script is the DSPL analogue of `guides/modeling/slam_start_here.py`. It follows the same API conventions \u2014\n", + "each pipeline stage is a plain inline Python function, priors are chained via `al.util.chaining.mass_from`, image\n", + "positions are derived automatically via `positions_likelihood_from`, and MGE light profiles are constructed via\n", + "`al.model_util.mge_model_from`.\n", + "\n", + "__DSPL-Specific Differences From Standard SLaM__\n", + "\n", + " - There are two source galaxies (`source_0` at redshift 1.0, `source_1` at redshift 2.0). `source_0` is a light\n", + " source AND a mass deflector for `source_1`.\n", + " - The SOURCE LP PIPELINE is split into two searches. The first fits lens + `source_0` only, providing a stable\n", + " starting point. The second frees `source_0`'s mass and adds `source_1`'s light.\n", + " - The SOURCE PIX PIPELINE has an extra search: one pixelizes `source_0` while `source_1` is a bare ray-tracing\n", + " galaxy, and the next pixelizes `source_1` with `source_0`'s mass fixed from the previous search.\n", + " - Two `PositionsLH` likelihoods are used once both sources are active, one per source-plane redshift.\n", + " - Adapt images are stitched across pipeline stages (e.g. `source_0`'s adapt image comes from a different result\n", + " than `source_1`'s).\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE and SOURCE PIX PIPELINE, this DSPL SLaM modeling script fits an `Imaging` dataset of a\n", + "double Einstein ring system where in the final model:\n", + "\n", + " - The lens galaxy's light is a bulge with MGE light profile.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` plus an `ExternalShear`.\n", + " - The first source galaxy's light is a `Pixelization` and its mass is an `Isothermal`.\n", + " - The second source galaxy's light is a `Pixelization`.\n", + "\n", + "Optional LIGHT LP and MASS TOTAL stages are left as a follow-up exercise." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 1__\n", + "\n", + "The first SOURCE LP PIPELINE search initializes a model where `source_1` is ignored and only the lens galaxy and\n", + "`source_0` are fit. This single-plane fit provides robust initial priors for the lens light, lens mass, shear and\n", + "`source_0` light before the more complex DSPL model is introduced.\n", + "\n", + "The model:\n", + " - Lens light: MGE with 2 x 20 Gaussians.\n", + " - Lens mass: `Isothermal` + `ExternalShear`.\n", + " - `source_0` light: MGE with 1 x 20 Gaussians." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source_0: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_0_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source_0,\n", + " bulge=source_0_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE 2__\n", + "\n", + "The second SOURCE LP PIPELINE search introduces the second source galaxy. The lens bulge / mass / shear and\n", + "`source_0` light are fixed to the instance values from search 1, and we free:\n", + "\n", + " - `source_0`'s mass: `Isothermal` with a prior tightly centred on the origin (the first source typically sits\n", + " close to the lens centre).\n", + " - `source_1`'s light: MGE with 1 x 20 Gaussians.\n", + "\n", + "A `PositionsLH` for `source_0` is attached to the analysis \u2014 automatically derived from the search-1 result via\n", + "`positions_likelihood_from` \u2014 to prevent unphysical mass models during the DSPL fit. No `source_1` positions are\n", + "available yet (this is the first search that fits `source_1`); they are introduced from the pixelized results\n", + "further down the pipeline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_lp_result_1: af.Result,\n", + " redshift_source_1: float,\n", + " n_batch: int = 30,\n", + ") -> af.Result:\n", + " positions_likelihood_source_0 = source_lp_result_1.positions_likelihood_from(\n", + " factor=3.0,\n", + " minimum_threshold=0.3,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[positions_likelihood_source_0],\n", + " )\n", + "\n", + " source_0_mass = af.Model(al.mp.Isothermal)\n", + " source_0_mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + " source_0_mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + " source_0_mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + " source_1_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result_1.instance.galaxies.lens.bulge,\n", + " mass=source_lp_result_1.instance.galaxies.lens.mass,\n", + " shear=source_lp_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_1.instance.galaxies.source_0.redshift,\n", + " bulge=source_lp_result_1.instance.galaxies.source_0.bulge,\n", + " mass=source_0_mass,\n", + " ),\n", + " source_1=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source_1,\n", + " bulge=source_1_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1 \u2014 source_0__\n", + "\n", + "Pixelizes `source_0` while `source_1` is present only as a bare galaxy for ray-tracing purposes. The lens mass is\n", + "freed with priors initialized from the SOURCE LP PIPELINE result, and `source_1` is not fit (no light, no mass)\n", + "so this search constrains the lens mass via `source_0` alone.\n", + "\n", + "Adapt images come from the SOURCE LP PIPELINE result 1 (the single-plane fit)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1_source_0(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result_1: af.Result,\n", + " source_lp_result_2: af.Result,\n", + " redshift_source_1: float,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result_1\n", + " )\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " positions_likelihood_source_0 = source_lp_result_1.positions_likelihood_from(\n", + " factor=3.0,\n", + " minimum_threshold=0.2,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[positions_likelihood_source_0],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.Isothermal),\n", + " mass_result=source_lp_result_2.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result_2.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result_2.instance.galaxies.lens.bulge,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " source_1=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source_1,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]_source_0\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1 \u2014 source_1__\n", + "\n", + "Pixelizes `source_1`. `source_0`'s mass is freed with priors initialized from the previous pixelized search, so\n", + "the second source plane constrains the lensing by the first source. The lens mass is fixed from\n", + "`source_pix_result_1_source_0`.\n", + "\n", + "Two `PositionsLH` are attached \u2014 one per source plane \u2014 to prevent unphysical reconstructions.\n", + "\n", + "Adapt images are stitched: the lens adapt image comes from the LP pipeline result 2; `source_0`'s adapt image\n", + "comes from the pixelized search above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1_source_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result_2: af.Result,\n", + " source_pix_result_1_source_0: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " lp2_dict = al.galaxy_name_image_dict_via_result_from(result=source_lp_result_2)\n", + " pix_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1_source_0\n", + " )\n", + "\n", + " galaxy_name_image_dict = {\n", + " \"('galaxies', 'lens')\": lp2_dict[\"('galaxies', 'lens')\"],\n", + " \"('galaxies', 'source_0')\": pix_dict[\"('galaxies', 'source_0')\"],\n", + " }\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", + "\n", + " positions_likelihood_source_0 = (\n", + " source_pix_result_1_source_0.positions_likelihood_from(\n", + " factor=3.0,\n", + " minimum_threshold=0.2,\n", + " plane_redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", + " )\n", + " )\n", + " positions_likelihood_source_1 = source_lp_result_2.positions_likelihood_from(\n", + " factor=3.0,\n", + " minimum_threshold=0.2,\n", + " plane_redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " positions_likelihood_source_0,\n", + " positions_likelihood_source_1,\n", + " ],\n", + " )\n", + "\n", + " source_0_mass = al.util.chaining.mass_from(\n", + " mass=af.Model(al.mp.Isothermal),\n", + " mass_result=source_lp_result_2.model.galaxies.source_0.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result_2.instance.galaxies.lens.bulge,\n", + " mass=source_pix_result_1_source_0.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1_source_0.instance.galaxies.lens.shear,\n", + " ),\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", + " mass=source_0_mass,\n", + " ),\n", + " source_1=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]_source_1\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "The final SOURCE PIX PIPELINE search fits both source galaxies simultaneously with adaptive pixelizations.\n", + "Lens mass, shear and `source_0`'s mass are all fixed to the maximum-likelihood instances of the previous\n", + "pixelized searches; only the pixelization regularization parameters are free.\n", + "\n", + "The `RectangularAdaptImage` (or equivalent) mesh uses the high-quality adapt images built up over the earlier\n", + "pipeline stages to adapt each source-plane pixelization to its reconstructed morphology." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result_2: af.Result,\n", + " source_pix_result_1_source_0: af.Result,\n", + " source_pix_result_1_source_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " lp2_dict = al.galaxy_name_image_dict_via_result_from(result=source_lp_result_2)\n", + " pix0_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1_source_0\n", + " )\n", + " pix1_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1_source_1\n", + " )\n", + "\n", + " galaxy_name_image_dict = {\n", + " \"('galaxies', 'lens')\": lp2_dict[\"('galaxies', 'lens')\"],\n", + " \"('galaxies', 'source_0')\": pix0_dict[\"('galaxies', 'source_0')\"],\n", + " \"('galaxies', 'source_1')\": pix1_dict[\"('galaxies', 'source_1')\"],\n", + " }\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_name_image_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result_2.instance.galaxies.lens.bulge,\n", + " mass=source_pix_result_1_source_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1_source_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source_0=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_0.redshift,\n", + " mass=source_pix_result_1_source_1.instance.galaxies.source_0.mass,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " source_1=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result_2.instance.galaxies.source_1.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"double_einstein_ring\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/double_einstein_ring/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam_dspl\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens galaxy and the two source galaxies. Multi-plane ray-tracing uses these to compute the\n", + "correct deflection angles between each plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source_0 = 1.0\n", + "redshift_source_1 = 2.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The pixelization mesh shape is fixed before modeling; see `features/pixelization/modeling` for details." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below runs the full DSPL SLaM pipeline. See the docstring above each function for a description of\n", + "each stage." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result_1 = source_lp_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source_0=redshift_source_0,\n", + ")\n", + "\n", + "source_lp_result_2 = source_lp_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_lp_result_1=source_lp_result_1,\n", + " redshift_source_1=redshift_source_1,\n", + ")\n", + "\n", + "source_pix_result_1_source_0 = source_pix_1_source_0(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result_1=source_lp_result_1,\n", + " source_lp_result_2=source_lp_result_2,\n", + " redshift_source_1=redshift_source_1,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "source_pix_result_1_source_1 = source_pix_1_source_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result_2=source_lp_result_2,\n", + " source_pix_result_1_source_0=source_pix_result_1_source_0,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result_2=source_lp_result_2,\n", + " source_pix_result_1_source_0=source_pix_result_1_source_0,\n", + " source_pix_result_1_source_1=source_pix_result_1_source_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/los_halos/simulator.ipynb b/notebooks/imaging/features/advanced/los_halos/simulator.ipynb index 4e264d7e4..72ca0da75 100644 --- a/notebooks/imaging/features/advanced/los_halos/simulator.ipynb +++ b/notebooks/imaging/features/advanced/los_halos/simulator.ipynb @@ -1,579 +1,616 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Line-of-Sight Halos\n", - "==============================\n", - "\n", - "This script simulates a strong gravitational lens including line-of-sight (LOS) dark matter halos\n", - "that perturb the lensed images via multi-plane ray tracing.\n", - "\n", - "LOS halos are sampled from a cosmological halo mass function within a light-cone geometry, converted\n", - "to truncated NFW profiles, and placed on multiple redshift planes between the observer and the source.\n", - "A compensatory negative convergence (kappa) sheet is added to each plane to maintain mass conservation,\n", - "following the methodology of He et al. (2022, MNRAS 511, 3046).\n", - "\n", - "Without the negative kappa sheets, LOS halos systematically over-lens the images because the total\n", - "convergence is not conserved. The negative sheets account for the smooth average contribution of\n", - "all halos, so that only the *fluctuations* above the mean affect the lensing.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The simulated dataset is output to the ``dataset/imaging/los_halos`` folder.\n", - "- **Grid:** Define the 2D grid on which the image is evaluated and simulated.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **PSF:** A Gaussian PSF models the telescope optics.\n", - "- **Simulator:** The simulator defines the exposure time, PSF, background sky level, and noise properties.\n", - "- **LOS Configuration:** Parameters controlling the line-of-sight halo population.\n", - "- **Sample LOS Halos:** The ``LOSSampler`` handles the full pipeline.\n", - "- **Ray Tracing:** Define the main lens galaxy and source galaxy, then combine with the LOS galaxies to create a.\n", - "- **Output:** Output the simulated dataset to .fits files.\n", - "- **Visualize:** Output subplots and summary images as .png files for quick inspection.\n", - "- **Tracer json:** Save the tracer as a .json file for reproducibility.\n", - "- **LOS Diagnostics:** Save the LOS halo sample list and negative sheet values for post-simulation analysis.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates ``Imaging`` of a galaxy-scale strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is a ``PowerLaw`` and ``ExternalShear``.\n", - " - The source galaxy's light is a ``SersicCore``.\n", - " - Line-of-sight halos are ``NFWTruncatedSph`` profiles on multiple redshift planes.\n", - " - Each redshift plane includes a ``MassSheet`` with negative kappa." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import numpy as np\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "from autolens.lens.los import LOSSampler, los_planes_from" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The simulated dataset is output to the ``dataset/imaging/los_halos`` folder." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"los_halos\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2D grid on which the image is evaluated and simulated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(161, 161),\n", - " pixel_scales=0.05,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive oversampling ensures the central bright regions are evaluated at higher resolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__PSF__\n", - "\n", - "A Gaussian PSF models the telescope optics." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(13, 13), sigma=0.05, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulator__\n", - "\n", - "The simulator defines the exposure time, PSF, background sky level, and noise properties." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=8000.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - " noise_seed=666,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LOS Configuration__\n", - "\n", - "Parameters controlling the line-of-sight halo population.\n", - "\n", - " - ``z_lens``: Redshift of the main lens galaxy.\n", - " - ``z_source``: Redshift of the background source galaxy.\n", - " - ``planes_before_lens`` / ``planes_after_lens``: Number of LOS planes in front of and behind the main lens.\n", - " With [4, 4] we get 9 planes total (4 in front, the lens plane, and 4 behind).\n", - " - ``m_min`` / ``m_max``: Halo mass range in solar masses (M_sun).\n", - " - ``cone_radius_arcsec``: Angular radius of the light cone in arcseconds.\n", - " - ``c_scatter``: Log-normal scatter in the concentration-mass relation (dex).\n", - " - ``truncation_factor``: Overdensity factor defining the truncation radius (e.g. 100 for r_100).\n", - " - ``seed``: Random seed for reproducibility." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "z_lens = 0.5\n", - "z_source = 1.0\n", - "\n", - "planes_before_lens = 4\n", - "planes_after_lens = 4\n", - "\n", - "m_min = 1e7\n", - "m_max = 1e10\n", - "\n", - "cone_radius_arcsec = 5.0\n", - "c_scatter = 0.15\n", - "truncation_factor = 100.0\n", - "\n", - "seed = 42" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mass Function and Mass-Concentration Coefficients__\n", - "\n", - "The ``LOSSampler`` can compute these from the ``hmf`` and ``colossus`` libraries, but here we\n", - "provide pre-computed values for each plane to avoid those dependencies.\n", - "\n", - "Each row contains ``[A, B]`` where:\n", - "\n", - " - For the mass function: ``log10(dn/dm) = A * log10(m) + B``\n", - " - For the mass-concentration relation: ``c(m) = A * log10(m) + B``\n", - "\n", - "These coefficients were computed using the Sheth-Mo-Tormen mass function and the Ludlow+16\n", - "concentration-mass relation for Planck15 cosmology at the redshift of each plane centre.\n", - "\n", - "If you have ``hmf`` and ``colossus`` installed, you can pass ``cosmology_astropy`` and\n", - "``cosmology_name_colossus`` to ``LOSSampler`` instead, and it will compute these automatically." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "_, plane_centres = los_planes_from(\n", - " z_lens=z_lens,\n", - " z_source=z_source,\n", - " planes_before_lens=planes_before_lens,\n", - " planes_after_lens=planes_after_lens,\n", - ")\n", - "\n", - "n_planes = len(plane_centres)\n", - "\n", - "print(f\"Number of LOS planes: {n_planes}\")\n", - "print(f\"Plane centres: {plane_centres}\")\n", - "\n", - "try:\n", - " from autolens.lens.los import mass_function_ab_from, mass_concentration_ab_from\n", - " from astropy.cosmology import Planck15 as astropy_planck15\n", - "\n", - " mass_function_coefficients = np.zeros((n_planes, 2))\n", - " mass_concentration_coefficients = np.zeros((n_planes, 2))\n", - "\n", - " for i, z in enumerate(plane_centres):\n", - " mass_function_coefficients[i] = mass_function_ab_from(\n", - " redshift=z, cosmology_astropy=astropy_planck15\n", - " )\n", - " mass_concentration_coefficients[i] = mass_concentration_ab_from(redshift=z)\n", - "\n", - " print(\n", - " \"Computed mass function and mass-concentration coefficients from hmf/colossus.\"\n", - " )\n", - "\n", - "except ImportError:\n", - " print(\"hmf/colossus not available, using approximate pre-computed coefficients.\")\n", - "\n", - " mass_function_coefficients = np.tile([-1.9, 8.0], (n_planes, 1))\n", - " mass_concentration_coefficients = np.tile([-3.0, 40.0], (n_planes, 1))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sample LOS Halos__\n", - "\n", - "The ``LOSSampler`` handles the full pipeline:\n", - "\n", - " 1. Slices the light cone into redshift planes.\n", - " 2. For each plane, samples halo masses, positions, and concentrations.\n", - " 3. Converts each halo to an ``NFWTruncatedSph`` via physical-to-lensing unit conversion.\n", - " 4. Computes the negative kappa sheet for each plane.\n", - " 5. Returns a list of ``Galaxy`` objects ready for the ``Tracer``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autogalaxy.cosmology import Planck15\n", - "\n", - "cosmology = Planck15()\n", - "\n", - "sampler = LOSSampler(\n", - " z_lens=z_lens,\n", - " z_source=z_source,\n", - " planes_before_lens=planes_before_lens,\n", - " planes_after_lens=planes_after_lens,\n", - " m_min=m_min,\n", - " m_max=m_max,\n", - " cone_radius_arcsec=cone_radius_arcsec,\n", - " c_scatter=c_scatter,\n", - " truncation_factor=truncation_factor,\n", - " cosmology=cosmology,\n", - " mass_function_coefficients=mass_function_coefficients,\n", - " mass_concentration_coefficients=mass_concentration_coefficients,\n", - " seed=seed,\n", - ")\n", - "\n", - "los_galaxies = sampler.galaxies_from()\n", - "\n", - "n_halos = sum(\n", - " 1\n", - " for g in los_galaxies\n", - " if hasattr(g, \"mass\") and isinstance(g.mass, al.mp.NFWTruncatedSph)\n", - ")\n", - "n_sheets = sum(\n", - " 1\n", - " for g in los_galaxies\n", - " if hasattr(g, \"mass_sheet\") and isinstance(g.mass_sheet, al.mp.MassSheet)\n", - ")\n", - "\n", - "print(f\"Sampled {n_halos} LOS halos across {n_sheets} planes.\")\n", - "\n", - "for g in los_galaxies:\n", - " if hasattr(g, \"mass_sheet\") and isinstance(g.mass_sheet, al.mp.MassSheet):\n", - " print(f\" Plane z={g.redshift:.4f}: kappa_neg = {g.mass_sheet.kappa:.6e}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Define the main lens galaxy and source galaxy, then combine with the LOS galaxies\n", - "to create a multi-plane ``Tracer``.\n", - "\n", - "The ``Tracer`` automatically groups galaxies by redshift into planes and performs\n", - "multi-plane ray tracing through all of them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=z_lens,\n", - " mass=al.mp.PowerLaw(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.059, -0.027),\n", - " slope=2.264,\n", - " einstein_radius=1.6,\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.0, gamma_2=0.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=z_source,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.02, -0.03),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=75.0),\n", - " intensity=1.5,\n", - " effective_radius=0.15,\n", - " sersic_index=3.5,\n", - " radius_break=0.025,\n", - " ),\n", - ")\n", - "\n", - "all_galaxies = los_galaxies + [lens_galaxy, source_galaxy]\n", - "\n", - "tracer = al.Tracer(galaxies=all_galaxies)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot the tracer's image to see the combined lensing effect of the main lens\n", - "and all LOS halos." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image with LOS Halos\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By passing the tracer and grid to the simulator, we create the simulated CCD imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output subplots and summary images as .png files for quick inspection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the tracer as a .json file for reproducibility." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LOS Diagnostics__\n", - "\n", - "Save the LOS halo sample list and negative sheet values for post-simulation analysis.\n", - "\n", - "A useful validation check is to compare the deflection field of the full tracer (with LOS halos)\n", - "against a smooth tracer (without LOS halos). If the negative kappa sheets are correct, the\n", - "*difference* in deflection angles should not systematically point toward the centre of the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "halo_info = []\n", - "sheet_info = []\n", - "\n", - "for g in los_galaxies:\n", - " if hasattr(g, \"mass\") and isinstance(g.mass, al.mp.NFWTruncatedSph):\n", - " halo_info.append(\n", - " [\n", - " g.redshift,\n", - " g.mass.centre[0],\n", - " g.mass.centre[1],\n", - " g.mass.kappa_s,\n", - " g.mass.scale_radius,\n", - " g.mass.truncation_radius,\n", - " ]\n", - " )\n", - " elif hasattr(g, \"mass_sheet\") and isinstance(g.mass_sheet, al.mp.MassSheet):\n", - " sheet_info.append([g.redshift, g.mass_sheet.kappa])\n", - "\n", - "if len(halo_info) > 0:\n", - " np.save(dataset_path / \"los_halo_list.npy\", np.array(halo_info))\n", - " print(f\"Saved {len(halo_info)} halo parameters to los_halo_list.npy\")\n", - "\n", - "if len(sheet_info) > 0:\n", - " np.save(dataset_path / \"los_sheet_values.npy\", np.array(sheet_info))\n", - " print(f\"Saved {len(sheet_info)} sheet values to los_sheet_values.npy\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder ``autolens_workspace/dataset/imaging/los_halos``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Line-of-Sight Halos\n", + "==============================\n", + "\n", + "This script simulates a strong gravitational lens including line-of-sight (LOS) dark matter halos\n", + "that perturb the lensed images via multi-plane ray tracing.\n", + "\n", + "LOS halos are sampled from a cosmological halo mass function within a light-cone geometry, converted\n", + "to truncated NFW profiles, and placed on multiple redshift planes between the observer and the source.\n", + "A compensatory negative convergence (kappa) sheet is added to each plane to maintain mass conservation,\n", + "following the methodology of He et al. (2022, MNRAS 511, 3046).\n", + "\n", + "Without the negative kappa sheets, LOS halos systematically over-lens the images because the total\n", + "convergence is not conserved. The negative sheets account for the smooth average contribution of\n", + "all halos, so that only the *fluctuations* above the mean affect the lensing.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The simulated dataset is output to the ``dataset/imaging/los_halos`` folder.\n", + "- **Grid:** Define the 2D grid on which the image is evaluated and simulated.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **PSF:** A Gaussian PSF models the telescope optics.\n", + "- **Simulator:** The simulator defines the exposure time, PSF, background sky level, and noise properties.\n", + "- **LOS Configuration:** Parameters controlling the line-of-sight halo population.\n", + "- **Sample LOS Halos:** The ``LOSSampler`` handles the full pipeline.\n", + "- **Ray Tracing:** Define the main lens galaxy and source galaxy, then combine with the LOS galaxies to create a.\n", + "- **Output:** Output the simulated dataset to .fits files.\n", + "- **Visualize:** Output subplots and summary images as .png files for quick inspection.\n", + "- **Tracer json:** Save the tracer as a .json file for reproducibility.\n", + "- **LOS Diagnostics:** Save the LOS halo sample list and negative sheet values for post-simulation analysis.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates ``Imaging`` of a galaxy-scale strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is a ``PowerLaw`` and ``ExternalShear``.\n", + " - The source galaxy's light is a ``SersicCore``.\n", + " - Line-of-sight halos are ``NFWTruncatedSph`` profiles on multiple redshift planes.\n", + " - Each redshift plane includes a ``MassSheet`` with negative kappa." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import numpy as np\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "from autolens.lens.los import LOSSampler, los_planes_from" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The simulated dataset is output to the ``dataset/imaging/los_halos`` folder." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"los_halos\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2D grid on which the image is evaluated and simulated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(161, 161),\n", + " pixel_scales=0.05,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive oversampling ensures the central bright regions are evaluated at higher resolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__PSF__\n", + "\n", + "A Gaussian PSF models the telescope optics." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(13, 13), sigma=0.05, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulator__\n", + "\n", + "The simulator defines the exposure time, PSF, background sky level, and noise properties." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=8000.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + " noise_seed=666,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LOS Configuration__\n", + "\n", + "Parameters controlling the line-of-sight halo population.\n", + "\n", + " - ``z_lens``: Redshift of the main lens galaxy.\n", + " - ``z_source``: Redshift of the background source galaxy.\n", + " - ``planes_before_lens`` / ``planes_after_lens``: Number of LOS planes in front of and behind the main lens.\n", + " With [4, 4] we get 9 planes total (4 in front, the lens plane, and 4 behind).\n", + " - ``m_min`` / ``m_max``: Halo mass range in solar masses (M_sun).\n", + " - ``cone_radius_arcsec``: Angular radius of the light cone in arcseconds.\n", + " - ``c_scatter``: Log-normal scatter in the concentration-mass relation (dex).\n", + " - ``truncation_factor``: Overdensity factor defining the truncation radius (e.g. 100 for r_100).\n", + " - ``seed``: Random seed for reproducibility." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "z_lens = 0.5\n", + "z_source = 1.0\n", + "\n", + "planes_before_lens = 4\n", + "planes_after_lens = 4\n", + "\n", + "m_min = 1e7\n", + "m_max = 1e10\n", + "\n", + "cone_radius_arcsec = 5.0\n", + "c_scatter = 0.15\n", + "truncation_factor = 100.0\n", + "\n", + "seed = 42" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mass Function and Mass-Concentration Coefficients__\n", + "\n", + "The ``LOSSampler`` can compute these from the ``hmf`` and ``colossus`` libraries, but here we\n", + "provide pre-computed values for each plane to avoid those dependencies.\n", + "\n", + "Each row contains ``[A, B]`` where:\n", + "\n", + " - For the mass function: ``log10(dn/dm) = A * log10(m) + B``\n", + " - For the mass-concentration relation: ``c(m) = A * log10(m) + B``\n", + "\n", + "These coefficients were computed using the Sheth-Mo-Tormen mass function and the Ludlow+16\n", + "concentration-mass relation for Planck15 cosmology at the redshift of each plane centre.\n", + "\n", + "If you have ``hmf`` and ``colossus`` installed, you can pass ``cosmology_astropy`` and\n", + "``cosmology_name_colossus`` to ``LOSSampler`` instead, and it will compute these automatically." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "_, plane_centres = los_planes_from(\n", + " z_lens=z_lens,\n", + " z_source=z_source,\n", + " planes_before_lens=planes_before_lens,\n", + " planes_after_lens=planes_after_lens,\n", + ")\n", + "\n", + "n_planes = len(plane_centres)\n", + "\n", + "print(f\"Number of LOS planes: {n_planes}\")\n", + "print(f\"Plane centres: {plane_centres}\")\n", + "\n", + "try:\n", + " from autolens.lens.los import mass_function_ab_from, mass_concentration_ab_from\n", + " from astropy.cosmology import Planck15 as astropy_planck15\n", + "\n", + " mass_function_coefficients = np.zeros((n_planes, 2))\n", + " mass_concentration_coefficients = np.zeros((n_planes, 2))\n", + "\n", + " for i, z in enumerate(plane_centres):\n", + " mass_function_coefficients[i] = mass_function_ab_from(\n", + " redshift=z, cosmology_astropy=astropy_planck15\n", + " )\n", + " mass_concentration_coefficients[i] = mass_concentration_ab_from(redshift=z)\n", + "\n", + " print(\n", + " \"Computed mass function and mass-concentration coefficients from hmf/colossus.\"\n", + " )\n", + "\n", + "except ImportError:\n", + " print(\"hmf/colossus not available, using approximate pre-computed coefficients.\")\n", + "\n", + " mass_function_coefficients = np.tile([-1.9, 8.0], (n_planes, 1))\n", + " mass_concentration_coefficients = np.tile([-3.0, 40.0], (n_planes, 1))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sample LOS Halos__\n", + "\n", + "The ``LOSSampler`` handles the full pipeline:\n", + "\n", + " 1. Slices the light cone into redshift planes.\n", + " 2. For each plane, samples halo masses, positions, and concentrations.\n", + " 3. Converts each halo to an ``NFWTruncatedSph`` via physical-to-lensing unit conversion.\n", + " 4. Computes the negative kappa sheet for each plane.\n", + " 5. Returns a list of ``Galaxy`` objects ready for the ``Tracer``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autogalaxy.cosmology import Planck15\n", + "\n", + "cosmology = Planck15()\n", + "\n", + "sampler = LOSSampler(\n", + " z_lens=z_lens,\n", + " z_source=z_source,\n", + " planes_before_lens=planes_before_lens,\n", + " planes_after_lens=planes_after_lens,\n", + " m_min=m_min,\n", + " m_max=m_max,\n", + " cone_radius_arcsec=cone_radius_arcsec,\n", + " c_scatter=c_scatter,\n", + " truncation_factor=truncation_factor,\n", + " cosmology=cosmology,\n", + " mass_function_coefficients=mass_function_coefficients,\n", + " mass_concentration_coefficients=mass_concentration_coefficients,\n", + " seed=seed,\n", + ")\n", + "\n", + "los_galaxies = sampler.galaxies_from()\n", + "\n", + "n_halos = sum(\n", + " 1\n", + " for g in los_galaxies\n", + " if hasattr(g, \"mass\") and isinstance(g.mass, al.mp.NFWTruncatedSph)\n", + ")\n", + "n_sheets = sum(\n", + " 1\n", + " for g in los_galaxies\n", + " if hasattr(g, \"mass_sheet\") and isinstance(g.mass_sheet, al.mp.MassSheet)\n", + ")\n", + "\n", + "print(f\"Sampled {n_halos} LOS halos across {n_sheets} planes.\")\n", + "\n", + "for g in los_galaxies:\n", + " if hasattr(g, \"mass_sheet\") and isinstance(g.mass_sheet, al.mp.MassSheet):\n", + " print(f\" Plane z={g.redshift:.4f}: kappa_neg = {g.mass_sheet.kappa:.6e}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Define the main lens galaxy and source galaxy, then combine with the LOS galaxies\n", + "to create a multi-plane ``Tracer``.\n", + "\n", + "The ``Tracer`` automatically groups galaxies by redshift into planes and performs\n", + "multi-plane ray tracing through all of them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=z_lens,\n", + " mass=al.mp.PowerLaw(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.059, -0.027),\n", + " slope=2.264,\n", + " einstein_radius=1.6,\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.0, gamma_2=0.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=z_source,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.02, -0.03),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=75.0),\n", + " intensity=1.5,\n", + " effective_radius=0.15,\n", + " sersic_index=3.5,\n", + " radius_break=0.025,\n", + " ),\n", + ")\n", + "\n", + "all_galaxies = los_galaxies + [lens_galaxy, source_galaxy]\n", + "\n", + "tracer = al.Tracer(galaxies=all_galaxies)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot the tracer's image to see the combined lensing effect of the main lens\n", + "and all LOS halos." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image with LOS Halos\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By passing the tracer and grid to the simulator, we create the simulated CCD imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output subplots and summary images as .png files for quick inspection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the tracer as a .json file for reproducibility." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LOS Diagnostics__\n", + "\n", + "Save the LOS halo sample list and negative sheet values for post-simulation analysis.\n", + "\n", + "A useful validation check is to compare the deflection field of the full tracer (with LOS halos)\n", + "against a smooth tracer (without LOS halos). If the negative kappa sheets are correct, the\n", + "*difference* in deflection angles should not systematically point toward the centre of the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "halo_info = []\n", + "sheet_info = []\n", + "\n", + "for g in los_galaxies:\n", + " if hasattr(g, \"mass\") and isinstance(g.mass, al.mp.NFWTruncatedSph):\n", + " halo_info.append(\n", + " [\n", + " g.redshift,\n", + " g.mass.centre[0],\n", + " g.mass.centre[1],\n", + " g.mass.kappa_s,\n", + " g.mass.scale_radius,\n", + " g.mass.truncation_radius,\n", + " ]\n", + " )\n", + " elif hasattr(g, \"mass_sheet\") and isinstance(g.mass_sheet, al.mp.MassSheet):\n", + " sheet_info.append([g.redshift, g.mass_sheet.kappa])\n", + "\n", + "if len(halo_info) > 0:\n", + " np.save(dataset_path / \"los_halo_list.npy\", np.array(halo_info))\n", + " print(f\"Saved {len(halo_info)} halo parameters to los_halo_list.npy\")\n", + "\n", + "if len(sheet_info) > 0:\n", + " np.save(dataset_path / \"los_sheet_values.npy\", np.array(sheet_info))\n", + " print(f\"Saved {len(sheet_info)} sheet values to los_sheet_values.npy\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder ``autolens_workspace/dataset/imaging/los_halos``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/los_halos/simulator_jax.ipynb b/notebooks/imaging/features/advanced/los_halos/simulator_jax.ipynb index 46302b1c6..496740a21 100644 --- a/notebooks/imaging/features/advanced/los_halos/simulator_jax.ipynb +++ b/notebooks/imaging/features/advanced/los_halos/simulator_jax.ipynb @@ -1,596 +1,633 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Line-of-Sight Halos (JAX)\n", - "=====================================\n", - "\n", - "This script simulates the same strong lens configuration as ``simulator.py`` \u2014 a galaxy-scale\n", - "lens with line-of-sight (LOS) dark matter halos \u2014 but uses a JAX-accelerated path that compiles\n", - "the full simulation pipeline under ``jax.jit``.\n", - "\n", - "The standard ``Tracer``-based simulator in ``simulator.py`` uses Python loops to iterate over\n", - "redshift planes and sum deflections from each halo. This works well for typical lens models with\n", - "a few galaxies, but becomes a bottleneck when the number of halos grows into the hundreds or\n", - "thousands (as in substructure forward models where the halo population is drawn from a mass\n", - "function and varies between realisations).\n", - "\n", - "The JAX path replaces these Python loops with two JAX primitives:\n", - "\n", - " - ``jax.vmap`` vectorises the deflection computation across all halos on a given plane, so\n", - " that a single GPU kernel evaluates every halo simultaneously.\n", - "\n", - " - ``jax.lax.scan`` iterates over redshift planes with a fixed-structure loop that compiles\n", - " to a single XLA operation, regardless of how many planes or halos are involved.\n", - "\n", - "Together, these allow the full ``theta -> noisy image`` pipeline to compile once under\n", - "``jax.jit`` and then be reused without recompilation for different halo populations. The\n", - "compiled function can also be batched via ``jax.vmap`` to produce many realisations per GPU\n", - "launch.\n", - "\n", - "__Contents__\n", - "\n", - "- **LOS Configuration:** Redshifts, mass range, light-cone parameters (same as ``simulator.py``).\n", - "- **Sample LOS Halos:** Use ``LOSSampler`` to draw a halo population (runs in NumPy).\n", - "- **Grid:** Define the 2D grid on which the image is evaluated and simulated.\n", - "- **PSF:** A Gaussian PSF kernel for convolution.\n", - "- **Lens Galaxy and Source Galaxy:** The main lens (``PowerLaw`` + ``ExternalShear``) and\n", - " the background source (``SersicCore``), identical to ``simulator.py``.\n", - "- **Convert to Padded Arrays:** Pack the ``Galaxy`` list from ``LOSSampler`` into fixed-shape\n", - " JAX arrays with boolean masks for unused slots.\n", - "- **Scaling Matrix:** Precompute the cosmological scaling factors between all redshift planes.\n", - "- **Parameterized Functions:** Define the lens mass, source light and (optionally) lens light\n", - " as functions of ``(grid, params)`` so their parameters are dynamic inputs to ``jax.jit``.\n", - "- **Single Realisation:** ``simulate_substructure`` produces one noisy lensed image.\n", - "- **Batched Realisations:** ``batched_simulate_substructure`` uses ``jax.vmap`` to produce\n", - " many images at once.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates ``Imaging`` of a galaxy-scale strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is a ``PowerLaw`` and ``ExternalShear``.\n", - " - The source galaxy's light is a ``SersicCore``.\n", - " - Line-of-sight halos are ``NFWTruncatedSph`` profiles on multiple redshift planes.\n", - " - Each redshift plane includes a ``MassSheet`` with negative kappa." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import jax\n", - "import jax.numpy as jnp\n", - "\n", - "import autoarray as aa\n", - "import autogalaxy as ag\n", - "import autolens as al\n", - "from autolens.lens import substructure_util\n", - "from autolens.lens.los import LOSSampler, los_planes_from" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LOS Configuration__\n", - "\n", - "Parameters controlling the line-of-sight halo population. These are identical to those in\n", - "``simulator.py`` and are described in detail there." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "z_lens = 0.5\n", - "z_source = 1.0\n", - "\n", - "planes_before_lens = 4\n", - "planes_after_lens = 4\n", - "\n", - "m_min = 1e7\n", - "m_max = 1e10\n", - "\n", - "cone_radius_arcsec = 5.0\n", - "c_scatter = 0.15\n", - "truncation_factor = 100.0\n", - "\n", - "seed = 42" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sample LOS Halos__\n", - "\n", - "The ``LOSSampler`` draws halo masses, positions and concentrations from a cosmological mass\n", - "function, converts each halo to an ``NFWTruncatedSph`` profile, and adds a compensatory\n", - "negative kappa ``MassSheet`` to each plane. See ``simulator.py`` for the full explanation\n", - "of the sampling pipeline and the mass function coefficients." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "_, plane_centres = los_planes_from(\n", - " z_lens=z_lens,\n", - " z_source=z_source,\n", - " planes_before_lens=planes_before_lens,\n", - " planes_after_lens=planes_after_lens,\n", - ")\n", - "\n", - "n_planes = len(plane_centres)\n", - "\n", - "try:\n", - " from autolens.lens.los import mass_function_ab_from, mass_concentration_ab_from\n", - " from astropy.cosmology import Planck15 as astropy_planck15\n", - "\n", - " mass_function_coefficients = np.zeros((n_planes, 2))\n", - " mass_concentration_coefficients = np.zeros((n_planes, 2))\n", - "\n", - " for i, z in enumerate(plane_centres):\n", - " mass_function_coefficients[i] = mass_function_ab_from(\n", - " redshift=z, cosmology_astropy=astropy_planck15\n", - " )\n", - " mass_concentration_coefficients[i] = mass_concentration_ab_from(redshift=z)\n", - "\n", - "except ImportError:\n", - " mass_function_coefficients = np.tile([-1.9, 8.0], (n_planes, 1))\n", - " mass_concentration_coefficients = np.tile([-3.0, 40.0], (n_planes, 1))\n", - "\n", - "from autogalaxy.cosmology import Planck15\n", - "\n", - "cosmology = Planck15()\n", - "\n", - "sampler = LOSSampler(\n", - " z_lens=z_lens,\n", - " z_source=z_source,\n", - " planes_before_lens=planes_before_lens,\n", - " planes_after_lens=planes_after_lens,\n", - " m_min=m_min,\n", - " m_max=m_max,\n", - " cone_radius_arcsec=cone_radius_arcsec,\n", - " c_scatter=c_scatter,\n", - " truncation_factor=truncation_factor,\n", - " cosmology=cosmology,\n", - " mass_function_coefficients=mass_function_coefficients,\n", - " mass_concentration_coefficients=mass_concentration_coefficients,\n", - " seed=seed,\n", - ")\n", - "\n", - "los_galaxies = sampler.galaxies_from()\n", - "\n", - "n_halos = sum(\n", - " 1\n", - " for g in los_galaxies\n", - " if hasattr(g, \"mass\") and isinstance(g.mass, al.mp.NFWTruncatedSph)\n", - ")\n", - "\n", - "print(f\"Sampled {n_halos} LOS halos across {n_planes} planes.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2D grid on which the image is evaluated and simulated. For the JAX path, we also\n", - "extract the raw ``(M, 2)`` coordinate array, because the ``simulate_substructure`` function\n", - "operates on plain JAX arrays rather than autoarray grid objects." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(101, 101),\n", - " pixel_scales=0.05,\n", - ")\n", - "\n", - "grid_array = jnp.array(grid.array)\n", - "image_shape = grid.shape_native" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__PSF__\n", - "\n", - "A Gaussian PSF kernel, normalised to unit sum. For the JAX path this is a plain 2D array\n", - "rather than an ``al.Convolver`` object, because ``simulate_substructure`` uses\n", - "``jax.scipy.signal.fftconvolve`` directly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf_sigma = 0.05\n", - "psf_size = 13\n", - "half = psf_size // 2\n", - "y = np.arange(psf_size) - half\n", - "x = np.arange(psf_size) - half\n", - "yy, xx = np.meshgrid(y, x, indexing=\"ij\")\n", - "psf_kernel_np = np.exp(-(yy**2 + xx**2) / (2 * (psf_sigma / grid.pixel_scales[0]) ** 2))\n", - "psf_kernel_np /= psf_kernel_np.sum()\n", - "psf_kernel = jnp.array(psf_kernel_np)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy and Source Galaxy__\n", - "\n", - "The lens galaxy and source galaxy are identical to those in ``simulator.py``.\n", - "\n", - "The lens galaxy's total mass distribution is a ``PowerLaw`` (the dominant smooth mass component)\n", - "plus an ``ExternalShear`` that accounts for tidal perturbations from the large-scale environment.\n", - "\n", - "The source galaxy's light distribution is a ``SersicCore``, which is a Sersic profile with a\n", - "flattened central core." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=z_lens,\n", - " mass=al.mp.PowerLaw(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.059, -0.027),\n", - " slope=2.264,\n", - " einstein_radius=1.6,\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.0, gamma_2=0.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=z_source,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.02, -0.03),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=75.0),\n", - " intensity=1.5,\n", - " effective_radius=0.15,\n", - " sersic_index=3.5,\n", - " radius_break=0.025,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Convert to Padded Arrays__\n", - "\n", - "The ``LOSSampler`` returns a list of ``Galaxy`` objects. The JAX path needs fixed-shape arrays\n", - "so that ``jax.jit`` can compile the simulation once and reuse it for different halo populations\n", - "without recompiling.\n", - "\n", - "``galaxies_to_halo_arrays`` extracts the lensing parameters (centre, kappa_s, scale_radius,\n", - "truncation_radius) from every ``NFWTruncatedSph`` halo and the negative kappa from every\n", - "``MassSheet``, then pads each plane's halo list to a fixed maximum ``max_n``. Unused slots\n", - "have their mask set to ``False`` so they contribute zero deflection.\n", - "\n", - "The full list of redshift planes includes the LOS plane centres plus the source redshift. If\n", - "the lens redshift coincides with a LOS plane centre it is included automatically; otherwise it\n", - "is added so that the lens galaxy's deflections are applied at the correct redshift." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plane_redshifts = sorted(set(list(plane_centres) + [z_lens, z_source]))\n", - "\n", - "max_n = 200\n", - "\n", - "halo_params, halo_mask, sheet_kappas = substructure_util.galaxies_to_halo_arrays(\n", - " galaxies=los_galaxies,\n", - " plane_redshifts=plane_redshifts,\n", - " max_n=max_n,\n", - " profile_cls=al.mp.NFWTruncatedSph,\n", - ")\n", - "\n", - "n_active = int(halo_mask.sum())\n", - "print(f\"Padded to max_n={max_n} per plane ({n_active} active slots across all planes).\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling Matrix__\n", - "\n", - "The cosmological scaling factors between every pair of redshift planes are precomputed once,\n", - "outside ``jax.jit``. This ``(n_planes, n_planes)`` matrix encodes how deflections at one plane\n", - "propagate to subsequent planes via the angular diameter distances, and is a constant input to\n", - "the compiled simulation function." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scaling_matrix = substructure_util.precompute_scaling_matrix(\n", - " plane_redshifts=plane_redshifts,\n", - " cosmology=cosmology,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parameterized Functions__\n", - "\n", - "The lens mass model, source light, and (optionally) lens light are each wrapped as a function\n", - "of ``(grid, params)`` where ``params`` is a 1-D JAX array of the profile parameters. The\n", - "function body constructs the profile objects from those parameters and evaluates them on the\n", - "grid.\n", - "\n", - "Because the ``params`` array is a dynamic input (not a constant captured in a closure), JAX\n", - "traces through the profile construction with traced parameter values. This means the lens mass\n", - "model, source light, and lens light can all be varied between realisations without triggering\n", - "recompilation \u2014 which is essential for inference where these parameters are being fitted.\n", - "\n", - "The function itself (which profile classes to use, how many components) is fixed at trace time.\n", - "Only the parameter *values* are dynamic." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def lens_mass_fn(grid_raw, params):\n", - " power_law = al.mp.PowerLaw(\n", - " centre=(params[0], params[1]),\n", - " ell_comps=(params[2], params[3]),\n", - " slope=params[4],\n", - " einstein_radius=params[5],\n", - " )\n", - " shear = al.mp.ExternalShear(gamma_1=params[6], gamma_2=params[7])\n", - " galaxy = al.Galaxy(redshift=z_lens, mass=power_law, shear=shear)\n", - " g = aa.Grid2DIrregular(values=grid_raw, xp=jnp)\n", - " return galaxy.deflections_yx_2d_from(grid=g, xp=jnp).array\n", - "\n", - "\n", - "lens_mass_params = jnp.array([0.0, 0.0, 0.059, -0.027, 2.264, 1.6, 0.0, 0.0])\n", - "\n", - "\n", - "def source_light_fn(grid_raw, params):\n", - " bulge = al.lp.SersicCore(\n", - " centre=(params[0], params[1]),\n", - " ell_comps=(params[2], params[3]),\n", - " intensity=params[4],\n", - " effective_radius=params[5],\n", - " sersic_index=params[6],\n", - " radius_break=params[7],\n", - " )\n", - " galaxy = al.Galaxy(redshift=z_source, bulge=bulge)\n", - " g = aa.Grid2DIrregular(values=grid_raw, xp=jnp)\n", - " return galaxy.image_2d_from(grid=g, xp=jnp).array\n", - "\n", - "\n", - "ell_comps_source = al.convert.ell_comps_from(axis_ratio=0.8, angle=75.0)\n", - "source_light_params = jnp.array(\n", - " [\n", - " 0.02,\n", - " -0.03,\n", - " ell_comps_source[0],\n", - " ell_comps_source[1],\n", - " 1.5,\n", - " 0.15,\n", - " 3.5,\n", - " 0.025,\n", - " ]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The ``lens_plane_mask`` is a float array with ``1.0`` at the lens-galaxy plane and ``0.0``\n", - "elsewhere. Inside the scan, the lens galaxy's deflections are computed at every plane step but\n", - "multiplied by this mask so they only contribute at the correct redshift." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_plane_mask = jnp.array(\n", - " [1.0 if abs(z - z_lens) < 1e-6 else 0.0 for z in plane_redshifts]\n", - ")\n", - "\n", - "lens_plane_idx = int(jnp.argmax(lens_plane_mask))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Single Realisation__\n", - "\n", - "``simulate_substructure`` compiles the full pipeline under ``jax.jit``:\n", - "\n", - " 1. Multi-plane ray tracing via ``jax.lax.scan`` over redshift planes, with halo deflections\n", - " computed by ``jax.vmap`` across all halos on each plane.\n", - " 2. Source light evaluation on the final (source-plane) traced grid.\n", - " 3. PSF convolution via ``jax.scipy.signal.fftconvolve``.\n", - " 4. Poisson noise using a ``jax.random.PRNGKey``.\n", - "\n", - "Passing ``prng_key=None`` skips the noise step and returns the clean lensed-and-convolved image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "key = jax.random.PRNGKey(seed)\n", - "\n", - "image = substructure_util.simulate_substructure(\n", - " grid=grid_array,\n", - " image_shape=image_shape,\n", - " halo_params=halo_params,\n", - " halo_mask=halo_mask,\n", - " scaling_matrix=scaling_matrix,\n", - " lens_mass_fn=lens_mass_fn,\n", - " lens_mass_params=lens_mass_params,\n", - " lens_plane_mask=lens_plane_mask,\n", - " sheet_kappas=sheet_kappas,\n", - " source_light_fn=source_light_fn,\n", - " source_light_params=source_light_params,\n", - " psf_kernel=psf_kernel,\n", - " exposure_time=8000.0,\n", - " background_sky_level=0.1,\n", - " prng_key=key,\n", - " halo_profile_cls=al.mp.NFWTruncatedSph,\n", - ")\n", - "\n", - "print(f\"Single image shape: {image.shape}, max intensity: {float(jnp.max(image)):.4f}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Batched Realisations__\n", - "\n", - "``batched_simulate_substructure`` uses ``jax.vmap`` to produce many images at once. Each\n", - "realisation in the batch has a different halo population (drawn by running ``LOSSampler`` with\n", - "a different seed) and a different noise realisation (from a different ``PRNGKey``).\n", - "\n", - "The grid, PSF, lens galaxy, source galaxy and scaling matrix are shared across the batch \u2014 only\n", - "the halo parameters, masks, sheet kappas and noise keys vary.\n", - "\n", - "``los_realizations_to_arrays`` is a convenience helper that runs ``galaxies_to_halo_arrays``\n", - "on each realisation and stacks the results into batch-dimensioned arrays." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "batch_size = 8\n", - "\n", - "realization_galaxies = []\n", - "for i in range(batch_size):\n", - " sampler_i = LOSSampler(\n", - " z_lens=z_lens,\n", - " z_source=z_source,\n", - " planes_before_lens=planes_before_lens,\n", - " planes_after_lens=planes_after_lens,\n", - " m_min=m_min,\n", - " m_max=m_max,\n", - " cone_radius_arcsec=cone_radius_arcsec,\n", - " c_scatter=c_scatter,\n", - " truncation_factor=truncation_factor,\n", - " cosmology=cosmology,\n", - " mass_function_coefficients=mass_function_coefficients,\n", - " mass_concentration_coefficients=mass_concentration_coefficients,\n", - " seed=100 + i,\n", - " )\n", - " realization_galaxies.append(sampler_i.galaxies_from())\n", - "\n", - "hp_batch, hm_batch, sk_batch = substructure_util.los_realizations_to_arrays(\n", - " realization_galaxies=realization_galaxies,\n", - " plane_redshifts=plane_redshifts,\n", - " max_n=max_n,\n", - " profile_cls=al.mp.NFWTruncatedSph,\n", - ")\n", - "\n", - "keys = jax.random.split(jax.random.PRNGKey(0), batch_size)\n", - "\n", - "lens_mass_params_batch = jnp.tile(lens_mass_params, (batch_size, 1))\n", - "source_light_params_batch = jnp.tile(source_light_params, (batch_size, 1))\n", - "\n", - "images_batch = substructure_util.batched_simulate_substructure(\n", - " grid=grid_array,\n", - " image_shape=image_shape,\n", - " halo_params_batch=hp_batch,\n", - " halo_mask_batch=hm_batch,\n", - " scaling_matrix=scaling_matrix,\n", - " lens_mass_fn=lens_mass_fn,\n", - " lens_mass_params_batch=lens_mass_params_batch,\n", - " lens_plane_mask=lens_plane_mask,\n", - " sheet_kappas_batch=sk_batch,\n", - " source_light_fn=source_light_fn,\n", - " source_light_params_batch=source_light_params_batch,\n", - " psf_kernel=psf_kernel,\n", - " exposure_time=8000.0,\n", - " background_sky_level=0.1,\n", - " prng_keys=keys,\n", - " halo_profile_cls=al.mp.NFWTruncatedSph,\n", - ")\n", - "\n", - "print(f\"Batch of {batch_size} images, shape: {images_batch.shape}\")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Line-of-Sight Halos (JAX)\n", + "=====================================\n", + "\n", + "This script simulates the same strong lens configuration as ``simulator.py`` \u2014 a galaxy-scale\n", + "lens with line-of-sight (LOS) dark matter halos \u2014 but uses a JAX-accelerated path that compiles\n", + "the full simulation pipeline under ``jax.jit``.\n", + "\n", + "The standard ``Tracer``-based simulator in ``simulator.py`` uses Python loops to iterate over\n", + "redshift planes and sum deflections from each halo. This works well for typical lens models with\n", + "a few galaxies, but becomes a bottleneck when the number of halos grows into the hundreds or\n", + "thousands (as in substructure forward models where the halo population is drawn from a mass\n", + "function and varies between realisations).\n", + "\n", + "The JAX path replaces these Python loops with two JAX primitives:\n", + "\n", + " - ``jax.vmap`` vectorises the deflection computation across all halos on a given plane, so\n", + " that a single GPU kernel evaluates every halo simultaneously.\n", + "\n", + " - ``jax.lax.scan`` iterates over redshift planes with a fixed-structure loop that compiles\n", + " to a single XLA operation, regardless of how many planes or halos are involved.\n", + "\n", + "Together, these allow the full ``theta -> noisy image`` pipeline to compile once under\n", + "``jax.jit`` and then be reused without recompilation for different halo populations. The\n", + "compiled function can also be batched via ``jax.vmap`` to produce many realisations per GPU\n", + "launch.\n", + "\n", + "__Contents__\n", + "\n", + "- **LOS Configuration:** Redshifts, mass range, light-cone parameters (same as ``simulator.py``).\n", + "- **Sample LOS Halos:** Use ``LOSSampler`` to draw a halo population (runs in NumPy).\n", + "- **Grid:** Define the 2D grid on which the image is evaluated and simulated.\n", + "- **PSF:** A Gaussian PSF kernel for convolution.\n", + "- **Lens Galaxy and Source Galaxy:** The main lens (``PowerLaw`` + ``ExternalShear``) and\n", + " the background source (``SersicCore``), identical to ``simulator.py``.\n", + "- **Convert to Padded Arrays:** Pack the ``Galaxy`` list from ``LOSSampler`` into fixed-shape\n", + " JAX arrays with boolean masks for unused slots.\n", + "- **Scaling Matrix:** Precompute the cosmological scaling factors between all redshift planes.\n", + "- **Parameterized Functions:** Define the lens mass, source light and (optionally) lens light\n", + " as functions of ``(grid, params)`` so their parameters are dynamic inputs to ``jax.jit``.\n", + "- **Single Realisation:** ``simulate_substructure`` produces one noisy lensed image.\n", + "- **Batched Realisations:** ``batched_simulate_substructure`` uses ``jax.vmap`` to produce\n", + " many images at once.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates ``Imaging`` of a galaxy-scale strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is a ``PowerLaw`` and ``ExternalShear``.\n", + " - The source galaxy's light is a ``SersicCore``.\n", + " - Line-of-sight halos are ``NFWTruncatedSph`` profiles on multiple redshift planes.\n", + " - Each redshift plane includes a ``MassSheet`` with negative kappa." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import jax\n", + "import jax.numpy as jnp\n", + "\n", + "import autoarray as aa\n", + "import autogalaxy as ag\n", + "import autolens as al\n", + "from autolens.lens import substructure_util\n", + "from autolens.lens.los import LOSSampler, los_planes_from" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LOS Configuration__\n", + "\n", + "Parameters controlling the line-of-sight halo population. These are identical to those in\n", + "``simulator.py`` and are described in detail there." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "z_lens = 0.5\n", + "z_source = 1.0\n", + "\n", + "planes_before_lens = 4\n", + "planes_after_lens = 4\n", + "\n", + "m_min = 1e7\n", + "m_max = 1e10\n", + "\n", + "cone_radius_arcsec = 5.0\n", + "c_scatter = 0.15\n", + "truncation_factor = 100.0\n", + "\n", + "seed = 42" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sample LOS Halos__\n", + "\n", + "The ``LOSSampler`` draws halo masses, positions and concentrations from a cosmological mass\n", + "function, converts each halo to an ``NFWTruncatedSph`` profile, and adds a compensatory\n", + "negative kappa ``MassSheet`` to each plane. See ``simulator.py`` for the full explanation\n", + "of the sampling pipeline and the mass function coefficients." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "_, plane_centres = los_planes_from(\n", + " z_lens=z_lens,\n", + " z_source=z_source,\n", + " planes_before_lens=planes_before_lens,\n", + " planes_after_lens=planes_after_lens,\n", + ")\n", + "\n", + "n_planes = len(plane_centres)\n", + "\n", + "try:\n", + " from autolens.lens.los import mass_function_ab_from, mass_concentration_ab_from\n", + " from astropy.cosmology import Planck15 as astropy_planck15\n", + "\n", + " mass_function_coefficients = np.zeros((n_planes, 2))\n", + " mass_concentration_coefficients = np.zeros((n_planes, 2))\n", + "\n", + " for i, z in enumerate(plane_centres):\n", + " mass_function_coefficients[i] = mass_function_ab_from(\n", + " redshift=z, cosmology_astropy=astropy_planck15\n", + " )\n", + " mass_concentration_coefficients[i] = mass_concentration_ab_from(redshift=z)\n", + "\n", + "except ImportError:\n", + " mass_function_coefficients = np.tile([-1.9, 8.0], (n_planes, 1))\n", + " mass_concentration_coefficients = np.tile([-3.0, 40.0], (n_planes, 1))\n", + "\n", + "from autogalaxy.cosmology import Planck15\n", + "\n", + "cosmology = Planck15()\n", + "\n", + "sampler = LOSSampler(\n", + " z_lens=z_lens,\n", + " z_source=z_source,\n", + " planes_before_lens=planes_before_lens,\n", + " planes_after_lens=planes_after_lens,\n", + " m_min=m_min,\n", + " m_max=m_max,\n", + " cone_radius_arcsec=cone_radius_arcsec,\n", + " c_scatter=c_scatter,\n", + " truncation_factor=truncation_factor,\n", + " cosmology=cosmology,\n", + " mass_function_coefficients=mass_function_coefficients,\n", + " mass_concentration_coefficients=mass_concentration_coefficients,\n", + " seed=seed,\n", + ")\n", + "\n", + "los_galaxies = sampler.galaxies_from()\n", + "\n", + "n_halos = sum(\n", + " 1\n", + " for g in los_galaxies\n", + " if hasattr(g, \"mass\") and isinstance(g.mass, al.mp.NFWTruncatedSph)\n", + ")\n", + "\n", + "print(f\"Sampled {n_halos} LOS halos across {n_planes} planes.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2D grid on which the image is evaluated and simulated. For the JAX path, we also\n", + "extract the raw ``(M, 2)`` coordinate array, because the ``simulate_substructure`` function\n", + "operates on plain JAX arrays rather than autoarray grid objects." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(101, 101),\n", + " pixel_scales=0.05,\n", + ")\n", + "\n", + "grid_array = jnp.array(grid.array)\n", + "image_shape = grid.shape_native" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__PSF__\n", + "\n", + "A Gaussian PSF kernel, normalised to unit sum. For the JAX path this is a plain 2D array\n", + "rather than an ``al.Convolver`` object, because ``simulate_substructure`` uses\n", + "``jax.scipy.signal.fftconvolve`` directly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf_sigma = 0.05\n", + "psf_size = 13\n", + "half = psf_size // 2\n", + "y = np.arange(psf_size) - half\n", + "x = np.arange(psf_size) - half\n", + "yy, xx = np.meshgrid(y, x, indexing=\"ij\")\n", + "psf_kernel_np = np.exp(-(yy**2 + xx**2) / (2 * (psf_sigma / grid.pixel_scales[0]) ** 2))\n", + "psf_kernel_np /= psf_kernel_np.sum()\n", + "psf_kernel = jnp.array(psf_kernel_np)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy and Source Galaxy__\n", + "\n", + "The lens galaxy and source galaxy are identical to those in ``simulator.py``.\n", + "\n", + "The lens galaxy's total mass distribution is a ``PowerLaw`` (the dominant smooth mass component)\n", + "plus an ``ExternalShear`` that accounts for tidal perturbations from the large-scale environment.\n", + "\n", + "The source galaxy's light distribution is a ``SersicCore``, which is a Sersic profile with a\n", + "flattened central core." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=z_lens,\n", + " mass=al.mp.PowerLaw(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.059, -0.027),\n", + " slope=2.264,\n", + " einstein_radius=1.6,\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.0, gamma_2=0.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=z_source,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.02, -0.03),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=75.0),\n", + " intensity=1.5,\n", + " effective_radius=0.15,\n", + " sersic_index=3.5,\n", + " radius_break=0.025,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Convert to Padded Arrays__\n", + "\n", + "The ``LOSSampler`` returns a list of ``Galaxy`` objects. The JAX path needs fixed-shape arrays\n", + "so that ``jax.jit`` can compile the simulation once and reuse it for different halo populations\n", + "without recompiling.\n", + "\n", + "``galaxies_to_halo_arrays`` extracts the lensing parameters (centre, kappa_s, scale_radius,\n", + "truncation_radius) from every ``NFWTruncatedSph`` halo and the negative kappa from every\n", + "``MassSheet``, then pads each plane's halo list to a fixed maximum ``max_n``. Unused slots\n", + "have their mask set to ``False`` so they contribute zero deflection.\n", + "\n", + "The full list of redshift planes includes the LOS plane centres plus the source redshift. If\n", + "the lens redshift coincides with a LOS plane centre it is included automatically; otherwise it\n", + "is added so that the lens galaxy's deflections are applied at the correct redshift." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plane_redshifts = sorted(set(list(plane_centres) + [z_lens, z_source]))\n", + "\n", + "max_n = 200\n", + "\n", + "halo_params, halo_mask, sheet_kappas = substructure_util.galaxies_to_halo_arrays(\n", + " galaxies=los_galaxies,\n", + " plane_redshifts=plane_redshifts,\n", + " max_n=max_n,\n", + " profile_cls=al.mp.NFWTruncatedSph,\n", + ")\n", + "\n", + "n_active = int(halo_mask.sum())\n", + "print(f\"Padded to max_n={max_n} per plane ({n_active} active slots across all planes).\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling Matrix__\n", + "\n", + "The cosmological scaling factors between every pair of redshift planes are precomputed once,\n", + "outside ``jax.jit``. This ``(n_planes, n_planes)`` matrix encodes how deflections at one plane\n", + "propagate to subsequent planes via the angular diameter distances, and is a constant input to\n", + "the compiled simulation function." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scaling_matrix = substructure_util.precompute_scaling_matrix(\n", + " plane_redshifts=plane_redshifts,\n", + " cosmology=cosmology,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parameterized Functions__\n", + "\n", + "The lens mass model, source light, and (optionally) lens light are each wrapped as a function\n", + "of ``(grid, params)`` where ``params`` is a 1-D JAX array of the profile parameters. The\n", + "function body constructs the profile objects from those parameters and evaluates them on the\n", + "grid.\n", + "\n", + "Because the ``params`` array is a dynamic input (not a constant captured in a closure), JAX\n", + "traces through the profile construction with traced parameter values. This means the lens mass\n", + "model, source light, and lens light can all be varied between realisations without triggering\n", + "recompilation \u2014 which is essential for inference where these parameters are being fitted.\n", + "\n", + "The function itself (which profile classes to use, how many components) is fixed at trace time.\n", + "Only the parameter *values* are dynamic." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def lens_mass_fn(grid_raw, params):\n", + " power_law = al.mp.PowerLaw(\n", + " centre=(params[0], params[1]),\n", + " ell_comps=(params[2], params[3]),\n", + " slope=params[4],\n", + " einstein_radius=params[5],\n", + " )\n", + " shear = al.mp.ExternalShear(gamma_1=params[6], gamma_2=params[7])\n", + " galaxy = al.Galaxy(redshift=z_lens, mass=power_law, shear=shear)\n", + " g = aa.Grid2DIrregular(values=grid_raw, xp=jnp)\n", + " return galaxy.deflections_yx_2d_from(grid=g, xp=jnp).array\n", + "\n", + "\n", + "lens_mass_params = jnp.array([0.0, 0.0, 0.059, -0.027, 2.264, 1.6, 0.0, 0.0])\n", + "\n", + "\n", + "def source_light_fn(grid_raw, params):\n", + " bulge = al.lp.SersicCore(\n", + " centre=(params[0], params[1]),\n", + " ell_comps=(params[2], params[3]),\n", + " intensity=params[4],\n", + " effective_radius=params[5],\n", + " sersic_index=params[6],\n", + " radius_break=params[7],\n", + " )\n", + " galaxy = al.Galaxy(redshift=z_source, bulge=bulge)\n", + " g = aa.Grid2DIrregular(values=grid_raw, xp=jnp)\n", + " return galaxy.image_2d_from(grid=g, xp=jnp).array\n", + "\n", + "\n", + "ell_comps_source = al.convert.ell_comps_from(axis_ratio=0.8, angle=75.0)\n", + "source_light_params = jnp.array(\n", + " [\n", + " 0.02,\n", + " -0.03,\n", + " ell_comps_source[0],\n", + " ell_comps_source[1],\n", + " 1.5,\n", + " 0.15,\n", + " 3.5,\n", + " 0.025,\n", + " ]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The ``lens_plane_mask`` is a float array with ``1.0`` at the lens-galaxy plane and ``0.0``\n", + "elsewhere. Inside the scan, the lens galaxy's deflections are computed at every plane step but\n", + "multiplied by this mask so they only contribute at the correct redshift." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_plane_mask = jnp.array(\n", + " [1.0 if abs(z - z_lens) < 1e-6 else 0.0 for z in plane_redshifts]\n", + ")\n", + "\n", + "lens_plane_idx = int(jnp.argmax(lens_plane_mask))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Single Realisation__\n", + "\n", + "``simulate_substructure`` compiles the full pipeline under ``jax.jit``:\n", + "\n", + " 1. Multi-plane ray tracing via ``jax.lax.scan`` over redshift planes, with halo deflections\n", + " computed by ``jax.vmap`` across all halos on each plane.\n", + " 2. Source light evaluation on the final (source-plane) traced grid.\n", + " 3. PSF convolution via ``jax.scipy.signal.fftconvolve``.\n", + " 4. Poisson noise using a ``jax.random.PRNGKey``.\n", + "\n", + "Passing ``prng_key=None`` skips the noise step and returns the clean lensed-and-convolved image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "key = jax.random.PRNGKey(seed)\n", + "\n", + "image = substructure_util.simulate_substructure(\n", + " grid=grid_array,\n", + " image_shape=image_shape,\n", + " halo_params=halo_params,\n", + " halo_mask=halo_mask,\n", + " scaling_matrix=scaling_matrix,\n", + " lens_mass_fn=lens_mass_fn,\n", + " lens_mass_params=lens_mass_params,\n", + " lens_plane_mask=lens_plane_mask,\n", + " sheet_kappas=sheet_kappas,\n", + " source_light_fn=source_light_fn,\n", + " source_light_params=source_light_params,\n", + " psf_kernel=psf_kernel,\n", + " exposure_time=8000.0,\n", + " background_sky_level=0.1,\n", + " prng_key=key,\n", + " halo_profile_cls=al.mp.NFWTruncatedSph,\n", + ")\n", + "\n", + "print(f\"Single image shape: {image.shape}, max intensity: {float(jnp.max(image)):.4f}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Batched Realisations__\n", + "\n", + "``batched_simulate_substructure`` uses ``jax.vmap`` to produce many images at once. Each\n", + "realisation in the batch has a different halo population (drawn by running ``LOSSampler`` with\n", + "a different seed) and a different noise realisation (from a different ``PRNGKey``).\n", + "\n", + "The grid, PSF, lens galaxy, source galaxy and scaling matrix are shared across the batch \u2014 only\n", + "the halo parameters, masks, sheet kappas and noise keys vary.\n", + "\n", + "``los_realizations_to_arrays`` is a convenience helper that runs ``galaxies_to_halo_arrays``\n", + "on each realisation and stacks the results into batch-dimensioned arrays." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "batch_size = 8\n", + "\n", + "realization_galaxies = []\n", + "for i in range(batch_size):\n", + " sampler_i = LOSSampler(\n", + " z_lens=z_lens,\n", + " z_source=z_source,\n", + " planes_before_lens=planes_before_lens,\n", + " planes_after_lens=planes_after_lens,\n", + " m_min=m_min,\n", + " m_max=m_max,\n", + " cone_radius_arcsec=cone_radius_arcsec,\n", + " c_scatter=c_scatter,\n", + " truncation_factor=truncation_factor,\n", + " cosmology=cosmology,\n", + " mass_function_coefficients=mass_function_coefficients,\n", + " mass_concentration_coefficients=mass_concentration_coefficients,\n", + " seed=100 + i,\n", + " )\n", + " realization_galaxies.append(sampler_i.galaxies_from())\n", + "\n", + "hp_batch, hm_batch, sk_batch = substructure_util.los_realizations_to_arrays(\n", + " realization_galaxies=realization_galaxies,\n", + " plane_redshifts=plane_redshifts,\n", + " max_n=max_n,\n", + " profile_cls=al.mp.NFWTruncatedSph,\n", + ")\n", + "\n", + "keys = jax.random.split(jax.random.PRNGKey(0), batch_size)\n", + "\n", + "lens_mass_params_batch = jnp.tile(lens_mass_params, (batch_size, 1))\n", + "source_light_params_batch = jnp.tile(source_light_params, (batch_size, 1))\n", + "\n", + "images_batch = substructure_util.batched_simulate_substructure(\n", + " grid=grid_array,\n", + " image_shape=image_shape,\n", + " halo_params_batch=hp_batch,\n", + " halo_mask_batch=hm_batch,\n", + " scaling_matrix=scaling_matrix,\n", + " lens_mass_fn=lens_mass_fn,\n", + " lens_mass_params_batch=lens_mass_params_batch,\n", + " lens_plane_mask=lens_plane_mask,\n", + " sheet_kappas_batch=sk_batch,\n", + " source_light_fn=source_light_fn,\n", + " source_light_params_batch=source_light_params_batch,\n", + " psf_kernel=psf_kernel,\n", + " exposure_time=8000.0,\n", + " background_sky_level=0.1,\n", + " prng_keys=keys,\n", + " halo_profile_cls=al.mp.NFWTruncatedSph,\n", + ")\n", + "\n", + "print(f\"Batch of {batch_size} images, shape: {images_batch.shape}\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/mass_stellar_dark/chaining.ipynb b/notebooks/imaging/features/advanced/mass_stellar_dark/chaining.ipynb index ed05affbe..3f5235588 100644 --- a/notebooks/imaging/features/advanced/mass_stellar_dark/chaining.ipynb +++ b/notebooks/imaging/features/advanced/mass_stellar_dark/chaining.ipynb @@ -1,424 +1,461 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Chaining: Chaining Lens Light To Mass\n", - "=====================================\n", - "\n", - "This script chains two searches to fit `Imaging` data of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a bulge with an MGE.\n", - " - The lens galaxy's stellar mass distribution is a bulge tied to the light model above.\n", - " - The lens galaxy's dark matter mass distribution is a `NFWSph`.\n", - " - The source galaxy's light is an MGE.\n", - "\n", - "The two searches break down as follows:\n", - "\n", - " 1) Models the lens galaxy's light using an MGE bulge. The source is present in the image, but modeling it is\n", - " omitted.\n", - "\n", - " 2) Models the lens galaxy's mass using a stellar mass distriubtion which is initialized using the bulge light\n", - " models inferred by search 1, alongside a dark matter profile. The source is again modeled using an MGE\n", - "\n", - "__Why Chain?__\n", - "\n", - "For many strong lenses the lens galaxy's light is distinct from the source galaxy's light, and it is therefore a valid\n", - "approach to first subtract the lens's light and then focus on fitting the lens mass model and source's light. This\n", - "provides the following benefits:\n", - "\n", - " - The non-linear parameter space defined by a bulge (N=7), stellar and dark mass (N=5) and parametric source (N=7)\n", - " has N=27 dimensions. By splitting the model-fit into two searches, we fit parameter spaces of dimensions N=11 and then\n", - " N=27, with many priors initialized. These are more efficient to sample and less like to infer a local maxima or\n", - " unphysical solution.\n", - "\n", - " - The lens galaxy's light traces its mass, so we can use the lens light model inferred in search 1 to initialize\n", - " sampling of the stellar mass model in search 2.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset + Masking:** Load, plot and mask the `Imaging` data.\n", - "- **Paths:** The path the results of all chained searches are output.\n", - "- **Model (Search 1):** Search 1 fits a lens model where the lens galaxy's light is an MGE bulge and the lens galaxy's mass and source galaxy are omitted.\n", - "- **Search + Analysis + Model-Fit (Search 1):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "- **Result (Search 1):** The results which are used for prior passing are summarised in the `info` attribute.\n", - "- **Model (Search 2):** Search 2 fits the lens galaxy's mass using a stellar mass distribution initialized from the bulge light model inferred by search 1, alongside a dark matter profile.\n", - "- **Search + Analysis + Model-Fit (Search 2):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "- **Result (Search 2):** The final results of the chained model-fit.\n", - "- **Wrap Up:** In this example, we passed a bulge lens light model to a decomposed stellar + dark matter mass model.\n", - "- **SLaM (Source, Light and Mass):** An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens modeling processing into a series of fits.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__ \n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__\n", - "\n", - "The path the results of all chained searches are output:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"imaging\") / \"chaining\" / \"lens_light_to_light_dark_mass\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Search 1 fits a lens model where:\n", - "\n", - " - The lens galaxy's light is an MGE bulge with 2 x 30 Gaussians [6 parameters].\n", - " - The lens galaxy's mass and source galaxy are omitted.\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=11." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(lens=lens))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__\n", - "\n", - "We use the results of search 1 to create the lens model fitted in search 2, where:\n", - "\n", - "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", - "provides a reasonably accurate lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__lens_light\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 1)__\n", - "\n", - "The results which are used for prior passing are summarised in the `info` attribute." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "We use the results of search 1 to create the lens model fitted in search 2, where:\n", - "\n", - " - The lens galaxy's light and stellar mass is a MGE bulge with 2 x 30 Gaussians [6 parameters: priors initialized from \n", - " search 1].\n", - " - The lens galaxy's dark matter mass distribution is a `NFW` whose centre is aligned with the \n", - " MGE bulge and stellar mass model above [5 parameters].\n", - " - The lens mass model also includes an `ExternalShear` [2 parameters].\n", - " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=22.\n", - "\n", - "We use the `take_attributes` method to pass the priors of the bulge. The reason we use this method is because\n", - "the bulge above use a `LightProfile` (e.g. via `al.lp`), whereas the model below gives a `LightAndMassProfile` \n", - "(e.g. via `al.lmp`). \n", - "\n", - "The `take_attributes` method is used when we pass parameters from two different models. In the example below it finds\n", - "all parameters in the MGE and MGE light models that share the same names\n", - "as parameters in the `MGE and MGE light and mass models and passes their priors \n", - "(in this case, the `centre`, `ell_comps`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = af.Model(al.lmp.Sersic)\n", - "bulge.take_attributes(source=result_1.model)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, dark=af.Model(al.mp.NFW))\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "model_2 = af.Collection(galaxies=af.Collection(lens=lens))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 2)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the `model.info` file of the search 2 model-fit to ensure the priors were passed correctly, as \n", - "well as the checkout the results to ensure an accurate power-law mass model is inferred." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__light_dark_mass\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 2)__\n", - "\n", - "The final results can be summarised via printing `info`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "In this example, we passed a bulge lens light model to a decomposed stellar + dark matter mass model. Thus, we\n", - "use an initial fit of the lens galaxy's light to better constrained our lens mass model! \n", - "\n", - "__SLaM (Source, Light and Mass)__\n", - "\n", - "An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens modeling \n", - "processing into a series of fits that first perfect the source model, then the lens light model and finally the lens\n", - "mass model. \n", - "\n", - "The SLaM Mass pipelines include pipelines which specifically decomposed stellar light + dark matter mass models. These\n", - "follow the Light pipelines and pass the priors of the light model in an identical fashion to this example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Chaining: Chaining Lens Light To Mass\n", + "=====================================\n", + "\n", + "This script chains two searches to fit `Imaging` data of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a bulge with an MGE.\n", + " - The lens galaxy's stellar mass distribution is a bulge tied to the light model above.\n", + " - The lens galaxy's dark matter mass distribution is a `NFWSph`.\n", + " - The source galaxy's light is an MGE.\n", + "\n", + "The two searches break down as follows:\n", + "\n", + " 1) Models the lens galaxy's light using an MGE bulge. The source is present in the image, but modeling it is\n", + " omitted.\n", + "\n", + " 2) Models the lens galaxy's mass using a stellar mass distriubtion which is initialized using the bulge light\n", + " models inferred by search 1, alongside a dark matter profile. The source is again modeled using an MGE\n", + "\n", + "__Why Chain?__\n", + "\n", + "For many strong lenses the lens galaxy's light is distinct from the source galaxy's light, and it is therefore a valid\n", + "approach to first subtract the lens's light and then focus on fitting the lens mass model and source's light. This\n", + "provides the following benefits:\n", + "\n", + " - The non-linear parameter space defined by a bulge (N=7), stellar and dark mass (N=5) and parametric source (N=7)\n", + " has N=27 dimensions. By splitting the model-fit into two searches, we fit parameter spaces of dimensions N=11 and then\n", + " N=27, with many priors initialized. These are more efficient to sample and less like to infer a local maxima or\n", + " unphysical solution.\n", + "\n", + " - The lens galaxy's light traces its mass, so we can use the lens light model inferred in search 1 to initialize\n", + " sampling of the stellar mass model in search 2.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset + Masking:** Load, plot and mask the `Imaging` data.\n", + "- **Paths:** The path the results of all chained searches are output.\n", + "- **Model (Search 1):** Search 1 fits a lens model where the lens galaxy's light is an MGE bulge and the lens galaxy's mass and source galaxy are omitted.\n", + "- **Search + Analysis + Model-Fit (Search 1):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "- **Result (Search 1):** The results which are used for prior passing are summarised in the `info` attribute.\n", + "- **Model (Search 2):** Search 2 fits the lens galaxy's mass using a stellar mass distribution initialized from the bulge light model inferred by search 1, alongside a dark matter profile.\n", + "- **Search + Analysis + Model-Fit (Search 2):** We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "- **Result (Search 2):** The final results of the chained model-fit.\n", + "- **Wrap Up:** In this example, we passed a bulge lens light model to a decomposed stellar + dark matter mass model.\n", + "- **SLaM (Source, Light and Mass):** An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens modeling processing into a series of fits.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__ \n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__\n", + "\n", + "The path the results of all chained searches are output:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"imaging\") / \"chaining\" / \"lens_light_to_light_dark_mass\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Search 1 fits a lens model where:\n", + "\n", + " - The lens galaxy's light is an MGE bulge with 2 x 30 Gaussians [6 parameters].\n", + " - The lens galaxy's mass and source galaxy are omitted.\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=11." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(lens=lens))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__\n", + "\n", + "We use the results of search 1 to create the lens model fitted in search 2, where:\n", + "\n", + "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", + "provides a reasonably accurate lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__lens_light\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 1)__\n", + "\n", + "The results which are used for prior passing are summarised in the `info` attribute." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "We use the results of search 1 to create the lens model fitted in search 2, where:\n", + "\n", + " - The lens galaxy's light and stellar mass is a MGE bulge with 2 x 30 Gaussians [6 parameters: priors initialized from \n", + " search 1].\n", + " - The lens galaxy's dark matter mass distribution is a `NFW` whose centre is aligned with the \n", + " MGE bulge and stellar mass model above [5 parameters].\n", + " - The lens mass model also includes an `ExternalShear` [2 parameters].\n", + " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=22.\n", + "\n", + "We use the `take_attributes` method to pass the priors of the bulge. The reason we use this method is because\n", + "the bulge above use a `LightProfile` (e.g. via `al.lp`), whereas the model below gives a `LightAndMassProfile` \n", + "(e.g. via `al.lmp`). \n", + "\n", + "The `take_attributes` method is used when we pass parameters from two different models. In the example below it finds\n", + "all parameters in the MGE and MGE light models that share the same names\n", + "as parameters in the `MGE and MGE light and mass models and passes their priors \n", + "(in this case, the `centre`, `ell_comps`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = af.Model(al.lmp.Sersic)\n", + "bulge.take_attributes(source=result_1.model)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, dark=af.Model(al.mp.NFW))\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "model_2 = af.Collection(galaxies=af.Collection(lens=lens))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 2)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the `model.info` file of the search 2 model-fit to ensure the priors were passed correctly, as \n", + "well as the checkout the results to ensure an accurate power-law mass model is inferred." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__light_dark_mass\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 2)__\n", + "\n", + "The final results can be summarised via printing `info`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "In this example, we passed a bulge lens light model to a decomposed stellar + dark matter mass model. Thus, we\n", + "use an initial fit of the lens galaxy's light to better constrained our lens mass model! \n", + "\n", + "__SLaM (Source, Light and Mass)__\n", + "\n", + "An even more advanced approach which uses search chaining are the SLaM pipelines, which break the lens modeling \n", + "processing into a series of fits that first perfect the source model, then the lens light model and finally the lens\n", + "mass model. \n", + "\n", + "The SLaM Mass pipelines include pipelines which specifically decomposed stellar light + dark matter mass models. These\n", + "follow the Light pipelines and pass the priors of the light model in an identical fashion to this example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/mass_stellar_dark/fit.ipynb b/notebooks/imaging/features/advanced/mass_stellar_dark/fit.ipynb index a8032bc44..ac723b7a7 100644 --- a/notebooks/imaging/features/advanced/mass_stellar_dark/fit.ipynb +++ b/notebooks/imaging/features/advanced/mass_stellar_dark/fit.ipynb @@ -1,507 +1,544 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Mass Stellar Dark Fit\n", - "===============================\n", - "\n", - "A decomposed mass model splits the lens galaxy's total mass into a stellar component (tied to its observed light\n", - "via a mass-to-light ratio) and a dark matter component (typically an NFW halo). The total deflection at every\n", - "image-plane coordinate is the sum of the deflections produced by each component, plus any external shear.\n", - "\n", - "This script illustrates the API for performing a fit to a decomposed-mass lens via the standard `Tracer` and\n", - "`FitImaging` objects, without invoking a non-linear search. It is intended to make the decomposed-mass\n", - "deflection composition concrete before the reader moves on to `modeling.py` (search-based) or `chaining.py` /\n", - "`slam.py` (realistic, robust modeling).\n", - "\n", - "The lens galaxy uses a linear `lmp.Sersic` light-and-mass profile (a Sersic that simultaneously acts as a light\n", - "profile and a stellar mass profile, coupled by a single `mass_to_light_ratio` parameter), plus a spherical NFW\n", - "dark matter halo and an external shear. The source galaxy is modelled with a Multi Gaussian Expansion (MGE), the\n", - "same source parameterization used in `chaining.py` and `slam.py`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **MGE Basis:** Build a `Basis` of linear Gaussians, used for the source galaxy.\n", - "- **Galaxies:** Compose the decomposed-mass lens galaxy plus the MGE source.\n", - "- **Tracer:** Build the two-plane `Tracer` that performs the ray-tracing.\n", - "- **Fit:** Create a `FitImaging` and inspect the fit.\n", - "- **Decomposed Deflection:** A short tour of how the total lens deflection is the sum of the stellar, dark and\n", - " shear contributions, and how the same composition shows up in convergence.\n", - "- **Intensities:** The solved-for linear light profile `intensity` values for each MGE Gaussian.\n", - "- **Wrap Up:** Summary and next steps.\n", - "\n", - "__Prerequisites__\n", - "\n", - "This script focuses on the API specific to a decomposed-mass fit. For background on the underlying single-plane\n", - "fit API and the MGE source parameterization, you should read first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/fit.py` \u2014 the standard single-plane fit.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py` \u2014 the MGE fit API and `Basis` of\n", - " linear Gaussians.\n", - "\n", - "The galaxy redshifts (`lens=0.5`, `source=1.0`), the lens `lmp.Sersic` mass-to-light ratio (0.2), and the dark\n", - "`NFWSph` parameters (`kappa_s=0.1`, `scale_radius=20.0`) match those used by the simulator and modeling examples." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `mass_stellar_dark` via .fits files.\n", - "\n", - "This dataset shows a single Einstein ring, simulated from a lens galaxy whose total mass is the sum of a Sersic\n", - "stellar component (with mass tied to its light via a constant mass-to-light ratio) and a spherical NFW dark\n", - "matter halo." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling, with finer sub-pixelization at the centre where the lens galaxy's light and\n", - "stellar mass are both most strongly peaked." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE Basis__\n", - "\n", - "We build a `Basis` of 30 linear Gaussians as the source-galaxy light model.\n", - "\n", - "The Gaussians share a common centre (kept spherical here for simplicity) and have `sigma` values spaced in\n", - "log10 increments from 0.01\" up to a reasonable size cap. The `intensity` of each Gaussian is a linear parameter,\n", - "solved for by linear algebra at fit time \u2014 no non-linear search is required.\n", - "\n", - "The MGE centre matches the simulated source position from `simulator.py`.\n", - "\n", - "For background on the MGE `Basis` API, see\n", - "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "source_bulge = build_source_basis(centre=(0.0, 0.0))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The Gaussians cannot be plotted yet because their `intensity` values have not been solved for \u2014 linear light\n", - "profiles only acquire an `intensity` once a `FitImaging` runs its linear algebra step. After the fit below, we\n", - "visualise the source MGE basis with its solved-for intensities.\n", - "\n", - "We set up the plotting grid we will use post-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "We now compose the two galaxies that form the lens system:\n", - "\n", - " - `lens` (z=0.5): a linear `lmp.Sersic` `bulge` which acts as BOTH the lens light AND the stellar mass\n", - " component, coupled by `mass_to_light_ratio`. A spherical `NFWSph` `dark` matter halo aligned with the bulge.\n", - " An `ExternalShear`.\n", - " - `source` (z=1.0): the MGE basis above.\n", - "\n", - "All non-linear parameters are set to the simulator's true values, so the fit visibly recovers the Einstein ring\n", - "without a search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lmp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " mass_to_light_ratio=0.2,\n", - " ),\n", - " dark=al.mp.NFWSph(centre=(0.0, 0.0), kappa_s=0.1, scale_radius=20.0),\n", - " shear=al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=source_bulge,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "The `Tracer` performs the ray-tracing. Internally it queries every mass profile attached to every galaxy in the\n", - "lens plane and sums their deflections. For our lens galaxy, this means the `bulge` contributes a stellar mass\n", - "deflection (its `lmp.Sersic` deflection scaled by `mass_to_light_ratio`), the `dark` halo contributes the\n", - "spherical NFW deflection, and `shear` contributes the external shear deflection \u2014 all summed before mapping\n", - "image-plane coordinates onto the source-plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens, source])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We pass the `Tracer` to a `FitImaging` to fit the dataset. The fit performs the ray-tracing (using the summed\n", - "stellar + dark + shear deflection), evaluates the source MGE at the source-plane, projects back to the image\n", - "plane, convolves with the PSF, and computes the residuals against the data.\n", - "\n", - "The `linear_light_profile_intensity_dict` of the fit will hold a solved-for `intensity` for every Gaussian in\n", - "the source MGE basis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Decomposed Deflection__\n", - "\n", - "This is the section that makes the decomposed-mass fit conceptually distinct. The lens galaxy's total deflection\n", - "map is the SUM of three independent contributions:\n", - "\n", - " alpha_lens(theta) = alpha_stellar(theta) + alpha_dark(theta) + alpha_shear(theta)\n", - "\n", - "where `alpha_stellar` is the `lmp.Sersic` bulge deflection (scaled internally by `mass_to_light_ratio`),\n", - "`alpha_dark` is the spherical NFW deflection, and `alpha_shear` is the external shear. Every individual\n", - "deflection is a public method on the corresponding profile.\n", - "\n", - "We verify this by computing each contribution explicitly and confirming the sum equals what the full lens\n", - "galaxy returns." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = dataset.grid\n", - "\n", - "deflections_stellar = lens.bulge.deflections_yx_2d_from(grid=grid)\n", - "deflections_dark = lens.dark.deflections_yx_2d_from(grid=grid)\n", - "deflections_shear = lens.shear.deflections_yx_2d_from(grid=grid)\n", - "\n", - "deflections_total_summed = deflections_stellar + deflections_dark + deflections_shear\n", - "deflections_total_lens = lens.deflections_yx_2d_from(grid=grid)\n", - "\n", - "print(f\"Stellar deflection (first 3): {deflections_stellar[:3]}\")\n", - "print(f\"Dark deflection (first 3): {deflections_dark[:3]}\")\n", - "print(f\"Shear deflection (first 3): {deflections_shear[:3]}\")\n", - "print(f\"Summed deflection (first 3): {deflections_total_summed[:3]}\")\n", - "print(f\"Lens deflection (first 3): {deflections_total_lens[:3]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The same component-wise decomposition shows up in the convergence (kappa) map. Convergence is what is plotted\n", - "on log-scale \"mass maps\" in the literature, and the decomposed-mass approach lets us inspect the stellar and\n", - "dark contributions separately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "kappa_stellar = lens.bulge.convergence_2d_from(grid=plot_grid)\n", - "kappa_dark = lens.dark.convergence_2d_from(grid=plot_grid)\n", - "\n", - "aplt.plot_array(array=kappa_stellar, title=\"Stellar convergence (M/L * light)\")\n", - "aplt.plot_array(array=kappa_dark, title=\"Dark matter convergence (NFW)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The lensed source's Einstein ring location is set by where `alpha_total` traces image-plane coordinates onto\n", - "the source position. Because the stellar contribution dominates inside the bulge and the dark contribution\n", - "takes over outside, the radial profile of the total deflection differs from any pure isothermal or power-law\n", - "model \u2014 this is the physical reason decomposed-mass fits constrain the M/L ratio and dark-halo concentration.\n", - "\n", - "__Intensities__\n", - "\n", - "After the fit, every linear Gaussian in the source MGE basis has been assigned an `intensity` via linear\n", - "algebra. These are available via the fit's `linear_light_profile_intensity_dict`, keyed by light profile object.\n", - "\n", - "We print the intensity of the first Gaussian in the basis to confirm the source has been reconstructed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"\\nFirst Gaussian intensity, source = \"\n", - " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A `Tracer` where every linear light profile has been replaced with an ordinary light profile carrying its\n", - "solved-for `intensity` is also accessible from the fit, which is useful for visualising the MGE basis with its\n", - "actual reconstructed amplitude." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer_fitted.galaxies[1].bulge, grid=plot_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated the decomposed-mass API and the deflection decomposition, without invoking a non-linear\n", - "search. The lens galaxy's `bulge` simultaneously acts as a light profile and a stellar mass profile (coupled by\n", - "`mass_to_light_ratio`), and the separately-parameterized `dark` NFW halo adds the second mass contribution.\n", - "External shear is included for completeness; it contributes a small additional deflection.\n", - "\n", - "In a real modeling workflow:\n", - "\n", - " - `modeling.py` shows how to fit the same system using `Nautilus`, but \"cheats\" by initialising priors at the\n", - " true values. It is therefore only useful as a tutorial.\n", - " - `chaining.py` is the practical workflow \u2014 two chained searches that fit the lens light first (treating it as\n", - " a pure light profile), then reintroduce the stellar-mass coupling by re-using the bulge result as a\n", - " `lmp.Sersic`. This is the script you'll actually use to fit data.\n", - " - `slam.py` is the most robust pipeline for production-quality decomposed-mass modeling, chaining through\n", - " SOURCE LP, SOURCE PIX, LIGHT LP, and MASS_LIGHT_DARK pipelines and ending in a pixelized source\n", - " reconstruction.\n", - "\n", - "The key takeaway from this script is that decomposed-mass lenses are fit with the same `Tracer` + `FitImaging`\n", - "objects as any other lens; the only difference is that the lens galaxy carries multiple independent mass\n", - "components (here: stellar via `lmp.Sersic`, dark via `NFWSph`, plus external shear) whose deflections sum into\n", - "the total lens-plane deflection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Mass Stellar Dark Fit\n", + "===============================\n", + "\n", + "A decomposed mass model splits the lens galaxy's total mass into a stellar component (tied to its observed light\n", + "via a mass-to-light ratio) and a dark matter component (typically an NFW halo). The total deflection at every\n", + "image-plane coordinate is the sum of the deflections produced by each component, plus any external shear.\n", + "\n", + "This script illustrates the API for performing a fit to a decomposed-mass lens via the standard `Tracer` and\n", + "`FitImaging` objects, without invoking a non-linear search. It is intended to make the decomposed-mass\n", + "deflection composition concrete before the reader moves on to `modeling.py` (search-based) or `chaining.py` /\n", + "`slam.py` (realistic, robust modeling).\n", + "\n", + "The lens galaxy uses a linear `lmp.Sersic` light-and-mass profile (a Sersic that simultaneously acts as a light\n", + "profile and a stellar mass profile, coupled by a single `mass_to_light_ratio` parameter), plus a spherical NFW\n", + "dark matter halo and an external shear. The source galaxy is modelled with a Multi Gaussian Expansion (MGE), the\n", + "same source parameterization used in `chaining.py` and `slam.py`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **MGE Basis:** Build a `Basis` of linear Gaussians, used for the source galaxy.\n", + "- **Galaxies:** Compose the decomposed-mass lens galaxy plus the MGE source.\n", + "- **Tracer:** Build the two-plane `Tracer` that performs the ray-tracing.\n", + "- **Fit:** Create a `FitImaging` and inspect the fit.\n", + "- **Decomposed Deflection:** A short tour of how the total lens deflection is the sum of the stellar, dark and\n", + " shear contributions, and how the same composition shows up in convergence.\n", + "- **Intensities:** The solved-for linear light profile `intensity` values for each MGE Gaussian.\n", + "- **Wrap Up:** Summary and next steps.\n", + "\n", + "__Prerequisites__\n", + "\n", + "This script focuses on the API specific to a decomposed-mass fit. For background on the underlying single-plane\n", + "fit API and the MGE source parameterization, you should read first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/fit.py` \u2014 the standard single-plane fit.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py` \u2014 the MGE fit API and `Basis` of\n", + " linear Gaussians.\n", + "\n", + "The galaxy redshifts (`lens=0.5`, `source=1.0`), the lens `lmp.Sersic` mass-to-light ratio (0.2), and the dark\n", + "`NFWSph` parameters (`kappa_s=0.1`, `scale_radius=20.0`) match those used by the simulator and modeling examples." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `mass_stellar_dark` via .fits files.\n", + "\n", + "This dataset shows a single Einstein ring, simulated from a lens galaxy whose total mass is the sum of a Sersic\n", + "stellar component (with mass tied to its light via a constant mass-to-light ratio) and a spherical NFW dark\n", + "matter halo." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling, with finer sub-pixelization at the centre where the lens galaxy's light and\n", + "stellar mass are both most strongly peaked." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE Basis__\n", + "\n", + "We build a `Basis` of 30 linear Gaussians as the source-galaxy light model.\n", + "\n", + "The Gaussians share a common centre (kept spherical here for simplicity) and have `sigma` values spaced in\n", + "log10 increments from 0.01\" up to a reasonable size cap. The `intensity` of each Gaussian is a linear parameter,\n", + "solved for by linear algebra at fit time \u2014 no non-linear search is required.\n", + "\n", + "The MGE centre matches the simulated source position from `simulator.py`.\n", + "\n", + "For background on the MGE `Basis` API, see\n", + "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "source_bulge = build_source_basis(centre=(0.0, 0.0))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The Gaussians cannot be plotted yet because their `intensity` values have not been solved for \u2014 linear light\n", + "profiles only acquire an `intensity` once a `FitImaging` runs its linear algebra step. After the fit below, we\n", + "visualise the source MGE basis with its solved-for intensities.\n", + "\n", + "We set up the plotting grid we will use post-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "We now compose the two galaxies that form the lens system:\n", + "\n", + " - `lens` (z=0.5): a linear `lmp.Sersic` `bulge` which acts as BOTH the lens light AND the stellar mass\n", + " component, coupled by `mass_to_light_ratio`. A spherical `NFWSph` `dark` matter halo aligned with the bulge.\n", + " An `ExternalShear`.\n", + " - `source` (z=1.0): the MGE basis above.\n", + "\n", + "All non-linear parameters are set to the simulator's true values, so the fit visibly recovers the Einstein ring\n", + "without a search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lmp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " mass_to_light_ratio=0.2,\n", + " ),\n", + " dark=al.mp.NFWSph(centre=(0.0, 0.0), kappa_s=0.1, scale_radius=20.0),\n", + " shear=al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=source_bulge,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "The `Tracer` performs the ray-tracing. Internally it queries every mass profile attached to every galaxy in the\n", + "lens plane and sums their deflections. For our lens galaxy, this means the `bulge` contributes a stellar mass\n", + "deflection (its `lmp.Sersic` deflection scaled by `mass_to_light_ratio`), the `dark` halo contributes the\n", + "spherical NFW deflection, and `shear` contributes the external shear deflection \u2014 all summed before mapping\n", + "image-plane coordinates onto the source-plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens, source])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We pass the `Tracer` to a `FitImaging` to fit the dataset. The fit performs the ray-tracing (using the summed\n", + "stellar + dark + shear deflection), evaluates the source MGE at the source-plane, projects back to the image\n", + "plane, convolves with the PSF, and computes the residuals against the data.\n", + "\n", + "The `linear_light_profile_intensity_dict` of the fit will hold a solved-for `intensity` for every Gaussian in\n", + "the source MGE basis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Decomposed Deflection__\n", + "\n", + "This is the section that makes the decomposed-mass fit conceptually distinct. The lens galaxy's total deflection\n", + "map is the SUM of three independent contributions:\n", + "\n", + " alpha_lens(theta) = alpha_stellar(theta) + alpha_dark(theta) + alpha_shear(theta)\n", + "\n", + "where `alpha_stellar` is the `lmp.Sersic` bulge deflection (scaled internally by `mass_to_light_ratio`),\n", + "`alpha_dark` is the spherical NFW deflection, and `alpha_shear` is the external shear. Every individual\n", + "deflection is a public method on the corresponding profile.\n", + "\n", + "We verify this by computing each contribution explicitly and confirming the sum equals what the full lens\n", + "galaxy returns." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = dataset.grid\n", + "\n", + "deflections_stellar = lens.bulge.deflections_yx_2d_from(grid=grid)\n", + "deflections_dark = lens.dark.deflections_yx_2d_from(grid=grid)\n", + "deflections_shear = lens.shear.deflections_yx_2d_from(grid=grid)\n", + "\n", + "deflections_total_summed = deflections_stellar + deflections_dark + deflections_shear\n", + "deflections_total_lens = lens.deflections_yx_2d_from(grid=grid)\n", + "\n", + "print(f\"Stellar deflection (first 3): {deflections_stellar[:3]}\")\n", + "print(f\"Dark deflection (first 3): {deflections_dark[:3]}\")\n", + "print(f\"Shear deflection (first 3): {deflections_shear[:3]}\")\n", + "print(f\"Summed deflection (first 3): {deflections_total_summed[:3]}\")\n", + "print(f\"Lens deflection (first 3): {deflections_total_lens[:3]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The same component-wise decomposition shows up in the convergence (kappa) map. Convergence is what is plotted\n", + "on log-scale \"mass maps\" in the literature, and the decomposed-mass approach lets us inspect the stellar and\n", + "dark contributions separately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "kappa_stellar = lens.bulge.convergence_2d_from(grid=plot_grid)\n", + "kappa_dark = lens.dark.convergence_2d_from(grid=plot_grid)\n", + "\n", + "aplt.plot_array(array=kappa_stellar, title=\"Stellar convergence (M/L * light)\")\n", + "aplt.plot_array(array=kappa_dark, title=\"Dark matter convergence (NFW)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The lensed source's Einstein ring location is set by where `alpha_total` traces image-plane coordinates onto\n", + "the source position. Because the stellar contribution dominates inside the bulge and the dark contribution\n", + "takes over outside, the radial profile of the total deflection differs from any pure isothermal or power-law\n", + "model \u2014 this is the physical reason decomposed-mass fits constrain the M/L ratio and dark-halo concentration.\n", + "\n", + "__Intensities__\n", + "\n", + "After the fit, every linear Gaussian in the source MGE basis has been assigned an `intensity` via linear\n", + "algebra. These are available via the fit's `linear_light_profile_intensity_dict`, keyed by light profile object.\n", + "\n", + "We print the intensity of the first Gaussian in the basis to confirm the source has been reconstructed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"\\nFirst Gaussian intensity, source = \"\n", + " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Tracer` where every linear light profile has been replaced with an ordinary light profile carrying its\n", + "solved-for `intensity` is also accessible from the fit, which is useful for visualising the MGE basis with its\n", + "actual reconstructed amplitude." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer_fitted.galaxies[1].bulge, grid=plot_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated the decomposed-mass API and the deflection decomposition, without invoking a non-linear\n", + "search. The lens galaxy's `bulge` simultaneously acts as a light profile and a stellar mass profile (coupled by\n", + "`mass_to_light_ratio`), and the separately-parameterized `dark` NFW halo adds the second mass contribution.\n", + "External shear is included for completeness; it contributes a small additional deflection.\n", + "\n", + "In a real modeling workflow:\n", + "\n", + " - `modeling.py` shows how to fit the same system using `Nautilus`, but \"cheats\" by initialising priors at the\n", + " true values. It is therefore only useful as a tutorial.\n", + " - `chaining.py` is the practical workflow \u2014 two chained searches that fit the lens light first (treating it as\n", + " a pure light profile), then reintroduce the stellar-mass coupling by re-using the bulge result as a\n", + " `lmp.Sersic`. This is the script you'll actually use to fit data.\n", + " - `slam.py` is the most robust pipeline for production-quality decomposed-mass modeling, chaining through\n", + " SOURCE LP, SOURCE PIX, LIGHT LP, and MASS_LIGHT_DARK pipelines and ending in a pixelized source\n", + " reconstruction.\n", + "\n", + "The key takeaway from this script is that decomposed-mass lenses are fit with the same `Tracer` + `FitImaging`\n", + "objects as any other lens; the only difference is that the lens galaxy carries multiple independent mass\n", + "components (here: stellar via `lmp.Sersic`, dark via `NFWSph`, plus external shear) whose deflections sum into\n", + "the total lens-plane deflection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/mass_stellar_dark/likelihood_function.ipynb b/notebooks/imaging/features/advanced/mass_stellar_dark/likelihood_function.ipynb index bdb75338b..0fe55d519 100644 --- a/notebooks/imaging/features/advanced/mass_stellar_dark/likelihood_function.ipynb +++ b/notebooks/imaging/features/advanced/mass_stellar_dark/likelihood_function.ipynb @@ -1,381 +1,418 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Mass Stellar Dark__\n", - "\n", - "This script describes the additional steps required to compute the `log_likelihood` for a strong lens whose\n", - "mass model decomposes the lens galaxy's mass into a stellar component (tied to its light via a mass-to-light\n", - "ratio) and a separately-parameterized dark matter halo.\n", - "\n", - "This script does NOT repeat the steps shared with single-plane lensing (mask, image-plane grid, PSF\n", - "convolution, chi-squared, noise normalization, linear-algebra solver for MGE source intensities). It documents\n", - "only the part of the likelihood function which is specific to a decomposed-mass lens: the lens-plane\n", - "deflection composition.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function below builds directly on the standard imaging and MGE likelihood functions. You should\n", - "read these notebooks first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 the canonical single-plane log likelihood\n", - " walkthrough, covering image-plane grids, ray-tracing, source-plane evaluation, PSF convolution, chi-squared\n", - " and the noise normalization term.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a `Basis`\n", - " of linear Gaussians is solved for via linear algebra.\n", - "\n", - "Sections covered in those scripts (e.g. \"Chi Squared\", \"Noise Normalization Term\", \"Mapping Matrix\") are not\n", - "repeated here; this script focuses entirely on what changes for a decomposed mass model.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script (see above).\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Galaxies:** A decomposed-mass lens (stellar + dark + shear) and an MGE source.\n", - "- **Decomposed Deflection:** The lens-plane deflection sum that produces the ray-traced source-plane grid.\n", - "- **Source-Plane Image:** Source MGE evaluated at the ray-traced grid.\n", - "- **Model Image:** PSF convolution and reference up to the canonical chi-squared / noise normalization.\n", - "- **Fit Check:** Confirm the manual ray-tracing matches `Tracer.traced_grid_2d_list_from` and the model fit\n", - " produces a finite `log_likelihood`.\n", - "- **Wrap Up.**\n", - "\n", - "__What Changes For A Decomposed Mass Model__\n", - "\n", - "For a single-plane lens with a total-mass profile such as `Isothermal` or `PowerLaw`, the lens-plane deflection\n", - "at every image-plane coordinate is produced by ONE mass profile:\n", - "\n", - " alpha_lens(theta) = alpha_total(theta ; Isothermal parameters)\n", - "\n", - "For a decomposed mass model, the lens galaxy carries multiple independent mass components, and the lens-plane\n", - "deflection is their SUM:\n", - "\n", - " alpha_lens(theta) = alpha_stellar(theta) + alpha_dark(theta) + alpha_shear(theta)\n", - " = (M/L) * alpha_light(theta ; bulge parameters)\n", - " + alpha_NFW(theta ; kappa_s, scale_radius)\n", - " + alpha_shear(theta ; gamma_1, gamma_2)\n", - "\n", - "The stellar contribution comes from the lens galaxy's light profile, scaled by the `mass_to_light_ratio`\n", - "parameter \u2014 i.e. the lens light is converted into a stellar surface density before being turned into a\n", - "deflection. The dark contribution is the NFW deflection of the dark matter halo. External shear contributes\n", - "a small uniform deflection set by the two shear components.\n", - "\n", - "Every other step of the likelihood (PSF convolution, chi-squared, noise normalization, MGE linear-algebra\n", - "solver) is unchanged." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the mass_stellar_dark dataset. The auto-simulation block mirrors the other example scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "The two galaxies that participate in the ray-tracing:\n", - "\n", - " - `lens` (z=0.5): a linear `lmp.Sersic` bulge (acting as light + stellar mass via a single `mass_to_light_ratio`),\n", - " an `NFWSph` dark matter halo aligned with the bulge, and an `ExternalShear`.\n", - " - `source` (z=1.0): an MGE light component (a simple basis of 10 linear Gaussians).\n", - "\n", - "The mass-profile parameters are set to the simulator's true values so the manual likelihood computation below\n", - "produces a sensible-looking model image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 10\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lmp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " mass_to_light_ratio=0.2,\n", - " ),\n", - " dark=al.mp.NFWSph(centre=(0.0, 0.0), kappa_s=0.1, scale_radius=20.0),\n", - " shear=al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=build_source_basis(centre=(0.0, 0.0)),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Decomposed Deflection__\n", - "\n", - "The single call below performs the standard image-plane \u2192 source-plane ray-tracing.\n", - "`traced_grid_2d_list_from` returns one grid per plane: the image-plane grid (no deflection) and the\n", - "source-plane grid (deflected by every mass profile in the lens plane, summed).\n", - "\n", - "To make the decomposition concrete we re-compute the same source-plane grid by hand. Each mass profile exposes\n", - "its own `deflections_yx_2d_from` method; the SUM of those three deflection maps is what the tracer applies\n", - "internally to produce the source-plane grid:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_grid = dataset.grid\n", - "\n", - "deflections_stellar = lens.bulge.deflections_yx_2d_from(grid=masked_grid)\n", - "deflections_dark = lens.dark.deflections_yx_2d_from(grid=masked_grid)\n", - "deflections_shear = lens.shear.deflections_yx_2d_from(grid=masked_grid)\n", - "\n", - "deflections_total = deflections_stellar + deflections_dark + deflections_shear\n", - "\n", - "grid_source_manual = masked_grid - deflections_total\n", - "\n", - "print(f\"deflections_stellar (first coord): {deflections_stellar[0]}\")\n", - "print(f\"deflections_dark (first coord): {deflections_dark[0]}\")\n", - "print(f\"deflections_shear (first coord): {deflections_shear[0]}\")\n", - "print(f\"deflections_total (first coord): {deflections_total[0]}\")\n", - "print(f\"source-plane grid (first coord, manual): {grid_source_manual[0]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We compare the hand-summed source-plane grid to the one produced by the `Tracer`, confirming the deflection\n", - "decomposition reproduces the internal ray-tracing exactly:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=masked_grid)\n", - "grid_source_tracer = traced_grid_list[1]\n", - "\n", - "print(f\"source-plane grid (first coord, tracer): {grid_source_tracer[0]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Image__\n", - "\n", - "The source galaxy's MGE basis is evaluated at the ray-traced source-plane grid, producing image-plane (y,x)\n", - "pixel values per Gaussian with a placeholder `intensity=1.0`. The true `intensity` of each Gaussian is solved\n", - "for at the linear-algebra step (see the MGE likelihood prerequisite).\n", - "\n", - "For this manual walkthrough we use the convenience method `image_2d_from` on the `Tracer`, which evaluates the\n", - "source MGE at the correct (ray-traced) plane and projects it into the image plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image_unconvolved = tracer.image_2d_from(grid=masked_grid)\n", - "\n", - "aplt.plot_array(\n", - " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "What `image_2d_from` does internally for our decomposed-mass lens:\n", - "\n", - " 1. Computes `alpha_lens(theta) = alpha_stellar + alpha_dark + alpha_shear` (the decomposition above).\n", - " 2. Ray-traces the image-plane grid to obtain `grid_source = grid - alpha_lens`.\n", - " 3. Evaluates the source MGE at `grid_source`, producing its image-plane contribution.\n", - " 4. (If the lens has a light component, also evaluates it at the image-plane grid and adds to the model image.)\n", - "\n", - "For a single-component total-mass lens there is just one profile contributing to step 1; for the decomposed\n", - "mass model there are three.\n", - "\n", - "__Model Image__\n", - "\n", - "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model image\n", - "above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", - "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`. The MGE source's per-Gaussian\n", - "intensities are solved for via the linear-algebra step documented in\n", - "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py`.\n", - "\n", - "We delegate the remaining steps to `FitImaging`, which handles the linear-algebra step that solves for each\n", - "Gaussian's `intensity` and assembles the full `log_likelihood`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(f\"\\nLog likelihood of the manual mass-stellar-dark fit: {fit.log_likelihood}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood__\n", - "\n", - "The final `log_likelihood` combines:\n", - "\n", - " - The chi-squared term, computed from the residuals between the PSF-convolved model image and the data,\n", - " weighted by the noise map.\n", - " - The noise normalization term, the standard Gaussian normalization over all unmasked pixels.\n", - " - The linear algebra terms (regularization and curvature determinants) introduced by the MGE\n", - " `Basis` of linear Gaussians.\n", - "\n", - "The first two are documented in `imaging/likelihood_function.py`; the third in\n", - "`imaging/features/multi_gaussian_expansion/likelihood_function.py`. No new terms are introduced by the\n", - "decomposed mass model \u2014 the only change is the lens-plane deflection composition described above.\n", - "\n", - "__Wrap Up__\n", - "\n", - "The decomposed-mass `log_likelihood` differs from a single-component total-mass case in exactly one place: the\n", - "lens-plane deflection is a sum of three independent contributions (stellar via `(M/L) * alpha_light`, dark via\n", - "`alpha_NFW`, and shear via `alpha_shear`) rather than a single `alpha_total`. Every other step (ray-tracing,\n", - "PSF convolution, chi-squared, noise normalization, linear algebra) is shared with the standard imaging\n", - "likelihood and documented in the prerequisite scripts.\n", - "\n", - "The mass-to-light coupling between the stellar component and the lens light is what makes decomposed-mass fits\n", - "informative about the lens's stellar mass: the same `bulge` parameters that determine the observed light profile\n", - "also determine the stellar deflection, so the data constrains the `mass_to_light_ratio` directly. The dark NFW\n", - "contribution then fills in whatever deflection the stellar component cannot account for \u2014 which is what the\n", - "science measurement of \"stellar vs dark matter contribution to the lens\" rests on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Mass Stellar Dark__\n", + "\n", + "This script describes the additional steps required to compute the `log_likelihood` for a strong lens whose\n", + "mass model decomposes the lens galaxy's mass into a stellar component (tied to its light via a mass-to-light\n", + "ratio) and a separately-parameterized dark matter halo.\n", + "\n", + "This script does NOT repeat the steps shared with single-plane lensing (mask, image-plane grid, PSF\n", + "convolution, chi-squared, noise normalization, linear-algebra solver for MGE source intensities). It documents\n", + "only the part of the likelihood function which is specific to a decomposed-mass lens: the lens-plane\n", + "deflection composition.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function below builds directly on the standard imaging and MGE likelihood functions. You should\n", + "read these notebooks first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 the canonical single-plane log likelihood\n", + " walkthrough, covering image-plane grids, ray-tracing, source-plane evaluation, PSF convolution, chi-squared\n", + " and the noise normalization term.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a `Basis`\n", + " of linear Gaussians is solved for via linear algebra.\n", + "\n", + "Sections covered in those scripts (e.g. \"Chi Squared\", \"Noise Normalization Term\", \"Mapping Matrix\") are not\n", + "repeated here; this script focuses entirely on what changes for a decomposed mass model.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script (see above).\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Galaxies:** A decomposed-mass lens (stellar + dark + shear) and an MGE source.\n", + "- **Decomposed Deflection:** The lens-plane deflection sum that produces the ray-traced source-plane grid.\n", + "- **Source-Plane Image:** Source MGE evaluated at the ray-traced grid.\n", + "- **Model Image:** PSF convolution and reference up to the canonical chi-squared / noise normalization.\n", + "- **Fit Check:** Confirm the manual ray-tracing matches `Tracer.traced_grid_2d_list_from` and the model fit\n", + " produces a finite `log_likelihood`.\n", + "- **Wrap Up.**\n", + "\n", + "__What Changes For A Decomposed Mass Model__\n", + "\n", + "For a single-plane lens with a total-mass profile such as `Isothermal` or `PowerLaw`, the lens-plane deflection\n", + "at every image-plane coordinate is produced by ONE mass profile:\n", + "\n", + " alpha_lens(theta) = alpha_total(theta ; Isothermal parameters)\n", + "\n", + "For a decomposed mass model, the lens galaxy carries multiple independent mass components, and the lens-plane\n", + "deflection is their SUM:\n", + "\n", + " alpha_lens(theta) = alpha_stellar(theta) + alpha_dark(theta) + alpha_shear(theta)\n", + " = (M/L) * alpha_light(theta ; bulge parameters)\n", + " + alpha_NFW(theta ; kappa_s, scale_radius)\n", + " + alpha_shear(theta ; gamma_1, gamma_2)\n", + "\n", + "The stellar contribution comes from the lens galaxy's light profile, scaled by the `mass_to_light_ratio`\n", + "parameter \u2014 i.e. the lens light is converted into a stellar surface density before being turned into a\n", + "deflection. The dark contribution is the NFW deflection of the dark matter halo. External shear contributes\n", + "a small uniform deflection set by the two shear components.\n", + "\n", + "Every other step of the likelihood (PSF convolution, chi-squared, noise normalization, MGE linear-algebra\n", + "solver) is unchanged." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the mass_stellar_dark dataset. The auto-simulation block mirrors the other example scripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "The two galaxies that participate in the ray-tracing:\n", + "\n", + " - `lens` (z=0.5): a linear `lmp.Sersic` bulge (acting as light + stellar mass via a single `mass_to_light_ratio`),\n", + " an `NFWSph` dark matter halo aligned with the bulge, and an `ExternalShear`.\n", + " - `source` (z=1.0): an MGE light component (a simple basis of 10 linear Gaussians).\n", + "\n", + "The mass-profile parameters are set to the simulator's true values so the manual likelihood computation below\n", + "produces a sensible-looking model image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 10\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lmp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " mass_to_light_ratio=0.2,\n", + " ),\n", + " dark=al.mp.NFWSph(centre=(0.0, 0.0), kappa_s=0.1, scale_radius=20.0),\n", + " shear=al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=build_source_basis(centre=(0.0, 0.0)),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Decomposed Deflection__\n", + "\n", + "The single call below performs the standard image-plane \u2192 source-plane ray-tracing.\n", + "`traced_grid_2d_list_from` returns one grid per plane: the image-plane grid (no deflection) and the\n", + "source-plane grid (deflected by every mass profile in the lens plane, summed).\n", + "\n", + "To make the decomposition concrete we re-compute the same source-plane grid by hand. Each mass profile exposes\n", + "its own `deflections_yx_2d_from` method; the SUM of those three deflection maps is what the tracer applies\n", + "internally to produce the source-plane grid:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_grid = dataset.grid\n", + "\n", + "deflections_stellar = lens.bulge.deflections_yx_2d_from(grid=masked_grid)\n", + "deflections_dark = lens.dark.deflections_yx_2d_from(grid=masked_grid)\n", + "deflections_shear = lens.shear.deflections_yx_2d_from(grid=masked_grid)\n", + "\n", + "deflections_total = deflections_stellar + deflections_dark + deflections_shear\n", + "\n", + "grid_source_manual = masked_grid - deflections_total\n", + "\n", + "print(f\"deflections_stellar (first coord): {deflections_stellar[0]}\")\n", + "print(f\"deflections_dark (first coord): {deflections_dark[0]}\")\n", + "print(f\"deflections_shear (first coord): {deflections_shear[0]}\")\n", + "print(f\"deflections_total (first coord): {deflections_total[0]}\")\n", + "print(f\"source-plane grid (first coord, manual): {grid_source_manual[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We compare the hand-summed source-plane grid to the one produced by the `Tracer`, confirming the deflection\n", + "decomposition reproduces the internal ray-tracing exactly:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=masked_grid)\n", + "grid_source_tracer = traced_grid_list[1]\n", + "\n", + "print(f\"source-plane grid (first coord, tracer): {grid_source_tracer[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Image__\n", + "\n", + "The source galaxy's MGE basis is evaluated at the ray-traced source-plane grid, producing image-plane (y,x)\n", + "pixel values per Gaussian with a placeholder `intensity=1.0`. The true `intensity` of each Gaussian is solved\n", + "for at the linear-algebra step (see the MGE likelihood prerequisite).\n", + "\n", + "For this manual walkthrough we use the convenience method `image_2d_from` on the `Tracer`, which evaluates the\n", + "source MGE at the correct (ray-traced) plane and projects it into the image plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image_unconvolved = tracer.image_2d_from(grid=masked_grid)\n", + "\n", + "aplt.plot_array(\n", + " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "What `image_2d_from` does internally for our decomposed-mass lens:\n", + "\n", + " 1. Computes `alpha_lens(theta) = alpha_stellar + alpha_dark + alpha_shear` (the decomposition above).\n", + " 2. Ray-traces the image-plane grid to obtain `grid_source = grid - alpha_lens`.\n", + " 3. Evaluates the source MGE at `grid_source`, producing its image-plane contribution.\n", + " 4. (If the lens has a light component, also evaluates it at the image-plane grid and adds to the model image.)\n", + "\n", + "For a single-component total-mass lens there is just one profile contributing to step 1; for the decomposed\n", + "mass model there are three.\n", + "\n", + "__Model Image__\n", + "\n", + "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model image\n", + "above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", + "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`. The MGE source's per-Gaussian\n", + "intensities are solved for via the linear-algebra step documented in\n", + "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py`.\n", + "\n", + "We delegate the remaining steps to `FitImaging`, which handles the linear-algebra step that solves for each\n", + "Gaussian's `intensity` and assembles the full `log_likelihood`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(f\"\\nLog likelihood of the manual mass-stellar-dark fit: {fit.log_likelihood}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood__\n", + "\n", + "The final `log_likelihood` combines:\n", + "\n", + " - The chi-squared term, computed from the residuals between the PSF-convolved model image and the data,\n", + " weighted by the noise map.\n", + " - The noise normalization term, the standard Gaussian normalization over all unmasked pixels.\n", + " - The linear algebra terms (regularization and curvature determinants) introduced by the MGE\n", + " `Basis` of linear Gaussians.\n", + "\n", + "The first two are documented in `imaging/likelihood_function.py`; the third in\n", + "`imaging/features/multi_gaussian_expansion/likelihood_function.py`. No new terms are introduced by the\n", + "decomposed mass model \u2014 the only change is the lens-plane deflection composition described above.\n", + "\n", + "__Wrap Up__\n", + "\n", + "The decomposed-mass `log_likelihood` differs from a single-component total-mass case in exactly one place: the\n", + "lens-plane deflection is a sum of three independent contributions (stellar via `(M/L) * alpha_light`, dark via\n", + "`alpha_NFW`, and shear via `alpha_shear`) rather than a single `alpha_total`. Every other step (ray-tracing,\n", + "PSF convolution, chi-squared, noise normalization, linear algebra) is shared with the standard imaging\n", + "likelihood and documented in the prerequisite scripts.\n", + "\n", + "The mass-to-light coupling between the stellar component and the lens light is what makes decomposed-mass fits\n", + "informative about the lens's stellar mass: the same `bulge` parameters that determine the observed light profile\n", + "also determine the stellar deflection, so the data constrains the `mass_to_light_ratio` directly. The dark NFW\n", + "contribution then fills in whatever deflection the stellar component cannot account for \u2014 which is what the\n", + "science measurement of \"stellar vs dark matter contribution to the lens\" rests on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/mass_stellar_dark/modeling.ipynb b/notebooks/imaging/features/advanced/mass_stellar_dark/modeling.ipynb index 991b28301..e4b1520c2 100644 --- a/notebooks/imaging/features/advanced/mass_stellar_dark/modeling.ipynb +++ b/notebooks/imaging/features/advanced/mass_stellar_dark/modeling.ipynb @@ -1,498 +1,535 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Mass Stellar Dark\n", - "====================================\n", - "\n", - "The majority of example scripts fit a mass profile which represents the _total_ mass of the lens galaxy (its stars,\n", - "dark matter and other components combined). This typically uses an `Isothermal` or `PowerLaw` mass profile.\n", - "\n", - "This script fits a mass model which decomposes the lens galaxy's mass into its stars and dark matter.\n", - "\n", - "__Practical Use: Read This First__\n", - "\n", - "This script is a tutorial. It produces a working fit by \"cheating\" \u2014 every prior is initialised at the true\n", - "simulator value, narrowed by a small Gaussian. On real data this is impossible, and a single Nautilus search on\n", - "a decomposed stellar + dark matter mass model would almost certainly converge to a local maximum. The mass-to-light\n", - "coupling between the lens light and its stellar mass component makes the parameter degeneracies more severe than\n", - "in total-mass fits: when the M/L ratio drifts, the stellar mass deflection drifts with it, and the dark NFW\n", - "component has to absorb the slack \u2014 these directions are hard for a non-linear search to disentangle without a\n", - "good starting point.\n", - "\n", - "The script you will actually use to fit a decomposed-mass model on real data is\n", - "`autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/chaining.py`, which runs two chained\n", - "non-linear searches: the first fits the lens light alone (treating it as a pure light profile), the second\n", - "introduces the stellar-mass coupling by re-using the bulge result as a `lmp.Sersic` and adds the dark NFW. This\n", - "gives the second search a tight prior on the bulge geometry, which is the parameter set most strongly tied to\n", - "the stellar mass deflection.\n", - "\n", - "For production-quality modeling, see `slam.py` in the same directory, which uses the `MASS_LIGHT_DARK` SLaM\n", - "pipeline to chain through SOURCE LP \u2192 SOURCE PIX \u2192 LIGHT LP \u2192 MASS_LIGHT_DARK and ends with a pixelized source\n", - "reconstruction.\n", - "\n", - "Read this script to understand the model composition API, then jump to `chaining.py`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Decomposed mass models measure direct properties of the stars and dark matter, for example the.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "Decomposed mass models measure direct properties of the stars and dark matter, for example the lens's stellar mass,\n", - "dark matter mass and the relative distribution between the two. Total mass profiles only inform us about the\n", - "superposition of these two components.\n", - "\n", - "Decomposed mass models couple the lens galaxy's light profile to its stellar mass distribution, meaning that\n", - "additional information in the lens galaxy emission is used to constrain the mass model. Whilst total mass models\n", - "also fit the lens light, they do not couple it to the mass model and thus do not exploit this extra information.\n", - "\n", - "Total mass models like the `Isothermal` and `PowerLaw` assume that the overall mass distribution of the lens galaxy\n", - "can be described using a single elliptical coordinate system. The stellar and dark components of a decomposed mass\n", - "model each have their own elliptical coordinate system, meaning that the mass model can be more complex and accurate.\n", - "\n", - "__Disadvantages__\n", - "\n", - "Assumptions must be made about how light and mass are coupled. This script assumes a constant mass-to-light raito,\n", - "however it is not clear this is a reliable assumption in many lens galaxies.\n", - "\n", - "**PyAutoLens** supports more complex mass models which introduce a radial gradient into the mass-to-light ratio.\n", - "However, these are more complex and therefore are difficult to fit robustly. Furthermore, it is still not clear\n", - "whether the way they couple light to mass is a reliable assumption.\n", - "\n", - "Performing ray-tracing with decomposed mass models is also more computationally expensive, meaning that the run times\n", - "of model-fits using these models is typically longer than total mass models.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a linear `Sersic`.\n", - " - The lens galaxy's stellar mass distribution is tied to the light model above.\n", - " - The lens galaxy's dark matter mass distribution is a `NFW`.\n", - " - The source galaxy's light is a Multi Gaussian Expansion.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `mass_stellar_dark` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's light and stellar mass is a linear `Sersic` [7 parameters].\n", - "\n", - " - The lens galaxy's dark matter mass distribution is a `NFW` whose centre is aligned with the \n", - " `Sersic` bulge of the light and stellar mass model above [5 parameters].\n", - "\n", - " - The lens mass model also includes an `ExternalShear` [2 parameters].\n", - "\n", - " - The source galaxy's light is a Multi Gaussian Expansion [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=22.\n", - "\n", - "Note that for the stellar light and mass, we are using a \"light and mass profile\" via the `.lmp` package. This\n", - "profiles simultaneously acts like a light and mass profile.\n", - "\n", - "For the dark matter, we use an `NFW`, which is a common mass profile to represent dark matter.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lmp.Sersic)\n", - "dark = af.Model(al.mp.NFW)\n", - "bulge.centre = dark.centre\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, dark=dark, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the lens model has both a `Sersic` light and mass profile and `NFW` dark matter profile, which \n", - "are aligned." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a full \n", - "description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"mass_stellar_dark\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", - "required by a model.\n", - "\n", - "Deflection angle calculations of stellar mass models and dark matter mass models can use techniques whichs\n", - "store more data in VRAM than other methods. \n", - "\n", - "Given VRAM use is an important consideration, we print out the estimated VRAM required for this\n", - "model-fit and advise you do this for your own double source plane lens model-fits.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The likelihood evaluation time for analysing stellar and dark matter mass models is longer than for total mass models\n", - "like the isothermal or power-law. This is because the deflection angles of these mass profiles are more expensive to\n", - "compute, requiring a Gaussian expansion or numerical calculation.\n", - "\n", - "However, they have far fewer parameters than total mass models, when those models are also modeling the lens light. \n", - "This is because many of the light and mass profile parameters are shared and fitted for simultaneously, reducing the\n", - "overall dimensionality of non-linear parameter space.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "These plots show that a decomposed stars and dark matter model is still able to produce ray-tracing and\n", - "the lensed source's emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "These examples include a results API with specific tools for visualizing and analysing decomposed mass model,\n", - "for example 1D plots which separately show the density of stars and dark matter as a function of radius.\n", - "\n", - "__Wrap Up__\n", - "\n", - "Decomposed mass models have advantages and disavantages compared to total mass models.\n", - "\n", - "The model which is best suited to your needs depends on the science you are hoping to undertake and the quality of the\n", - "data you are fitting.\n", - "\n", - "In general, it is recommended that you first get fits going using total mass models, because they are simpler and make\n", - "fewer assumptions regarding how light is tied to mass. Once you have robust results, decomposed mass models can then\n", - "be fitted and compared in order to gain deeper insight." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Mass Stellar Dark\n", + "====================================\n", + "\n", + "The majority of example scripts fit a mass profile which represents the _total_ mass of the lens galaxy (its stars,\n", + "dark matter and other components combined). This typically uses an `Isothermal` or `PowerLaw` mass profile.\n", + "\n", + "This script fits a mass model which decomposes the lens galaxy's mass into its stars and dark matter.\n", + "\n", + "__Practical Use: Read This First__\n", + "\n", + "This script is a tutorial. It produces a working fit by \"cheating\" \u2014 every prior is initialised at the true\n", + "simulator value, narrowed by a small Gaussian. On real data this is impossible, and a single Nautilus search on\n", + "a decomposed stellar + dark matter mass model would almost certainly converge to a local maximum. The mass-to-light\n", + "coupling between the lens light and its stellar mass component makes the parameter degeneracies more severe than\n", + "in total-mass fits: when the M/L ratio drifts, the stellar mass deflection drifts with it, and the dark NFW\n", + "component has to absorb the slack \u2014 these directions are hard for a non-linear search to disentangle without a\n", + "good starting point.\n", + "\n", + "The script you will actually use to fit a decomposed-mass model on real data is\n", + "`autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/chaining.py`, which runs two chained\n", + "non-linear searches: the first fits the lens light alone (treating it as a pure light profile), the second\n", + "introduces the stellar-mass coupling by re-using the bulge result as a `lmp.Sersic` and adds the dark NFW. This\n", + "gives the second search a tight prior on the bulge geometry, which is the parameter set most strongly tied to\n", + "the stellar mass deflection.\n", + "\n", + "For production-quality modeling, see `slam.py` in the same directory, which uses the `MASS_LIGHT_DARK` SLaM\n", + "pipeline to chain through SOURCE LP \u2192 SOURCE PIX \u2192 LIGHT LP \u2192 MASS_LIGHT_DARK and ends with a pixelized source\n", + "reconstruction.\n", + "\n", + "Read this script to understand the model composition API, then jump to `chaining.py`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Decomposed mass models measure direct properties of the stars and dark matter, for example the.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "Decomposed mass models measure direct properties of the stars and dark matter, for example the lens's stellar mass,\n", + "dark matter mass and the relative distribution between the two. Total mass profiles only inform us about the\n", + "superposition of these two components.\n", + "\n", + "Decomposed mass models couple the lens galaxy's light profile to its stellar mass distribution, meaning that\n", + "additional information in the lens galaxy emission is used to constrain the mass model. Whilst total mass models\n", + "also fit the lens light, they do not couple it to the mass model and thus do not exploit this extra information.\n", + "\n", + "Total mass models like the `Isothermal` and `PowerLaw` assume that the overall mass distribution of the lens galaxy\n", + "can be described using a single elliptical coordinate system. The stellar and dark components of a decomposed mass\n", + "model each have their own elliptical coordinate system, meaning that the mass model can be more complex and accurate.\n", + "\n", + "__Disadvantages__\n", + "\n", + "Assumptions must be made about how light and mass are coupled. This script assumes a constant mass-to-light raito,\n", + "however it is not clear this is a reliable assumption in many lens galaxies.\n", + "\n", + "**PyAutoLens** supports more complex mass models which introduce a radial gradient into the mass-to-light ratio.\n", + "However, these are more complex and therefore are difficult to fit robustly. Furthermore, it is still not clear\n", + "whether the way they couple light to mass is a reliable assumption.\n", + "\n", + "Performing ray-tracing with decomposed mass models is also more computationally expensive, meaning that the run times\n", + "of model-fits using these models is typically longer than total mass models.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a linear `Sersic`.\n", + " - The lens galaxy's stellar mass distribution is tied to the light model above.\n", + " - The lens galaxy's dark matter mass distribution is a `NFW`.\n", + " - The source galaxy's light is a Multi Gaussian Expansion.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `mass_stellar_dark` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's light and stellar mass is a linear `Sersic` [7 parameters].\n", + "\n", + " - The lens galaxy's dark matter mass distribution is a `NFW` whose centre is aligned with the \n", + " `Sersic` bulge of the light and stellar mass model above [5 parameters].\n", + "\n", + " - The lens mass model also includes an `ExternalShear` [2 parameters].\n", + "\n", + " - The source galaxy's light is a Multi Gaussian Expansion [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=22.\n", + "\n", + "Note that for the stellar light and mass, we are using a \"light and mass profile\" via the `.lmp` package. This\n", + "profiles simultaneously acts like a light and mass profile.\n", + "\n", + "For the dark matter, we use an `NFW`, which is a common mass profile to represent dark matter.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lmp.Sersic)\n", + "dark = af.Model(al.mp.NFW)\n", + "bulge.centre = dark.centre\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, dark=dark, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the lens model has both a `Sersic` light and mass profile and `NFW` dark matter profile, which \n", + "are aligned." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a full \n", + "description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"mass_stellar_dark\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", + "required by a model.\n", + "\n", + "Deflection angle calculations of stellar mass models and dark matter mass models can use techniques whichs\n", + "store more data in VRAM than other methods. \n", + "\n", + "Given VRAM use is an important consideration, we print out the estimated VRAM required for this\n", + "model-fit and advise you do this for your own double source plane lens model-fits.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The likelihood evaluation time for analysing stellar and dark matter mass models is longer than for total mass models\n", + "like the isothermal or power-law. This is because the deflection angles of these mass profiles are more expensive to\n", + "compute, requiring a Gaussian expansion or numerical calculation.\n", + "\n", + "However, they have far fewer parameters than total mass models, when those models are also modeling the lens light. \n", + "This is because many of the light and mass profile parameters are shared and fitted for simultaneously, reducing the\n", + "overall dimensionality of non-linear parameter space.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "These plots show that a decomposed stars and dark matter model is still able to produce ray-tracing and\n", + "the lensed source's emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "These examples include a results API with specific tools for visualizing and analysing decomposed mass model,\n", + "for example 1D plots which separately show the density of stars and dark matter as a function of radius.\n", + "\n", + "__Wrap Up__\n", + "\n", + "Decomposed mass models have advantages and disavantages compared to total mass models.\n", + "\n", + "The model which is best suited to your needs depends on the science you are hoping to undertake and the quality of the\n", + "data you are fitting.\n", + "\n", + "In general, it is recommended that you first get fits going using total mass models, because they are simpler and make\n", + "fewer assumptions regarding how light is tied to mass. Once you have robust results, decomposed mass models can then\n", + "be fitted and compared in order to gain deeper insight." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/mass_stellar_dark/simulator.ipynb b/notebooks/imaging/features/advanced/mass_stellar_dark/simulator.ipynb index 8654cee8a..adcff596d 100644 --- a/notebooks/imaging/features/advanced/mass_stellar_dark/simulator.ipynb +++ b/notebooks/imaging/features/advanced/mass_stellar_dark/simulator.ipynb @@ -1,353 +1,390 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Mass Stellar Dark\n", - "============================\n", - "\n", - "The majority of example scripts simulate a mass profile which represents the _total_ mass of the lens galaxy (its stars,\n", - "dark matter and other components combined). This typically uses an `Isothermal` or `PowerLaw` mass profile.\n", - "\n", - "This script simulates a strong lens where the lens mass model decomposes the lens galaxy's mass into its stars and\n", - "dark matter.\n", - "\n", - "This dataset is modeled in the example script `autolens_workspace/scripts/modeling/features/mass_stellar_dark.py`,\n", - "where a discussion of the advantages and disadvantages of fitting decomposed mass models is also given.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens using decomposed light and dark matter profiles where:\n", - "\n", - " - The lens galaxy's light matter mass distribution is an `Sersic`.\n", - " - The lens galaxy's dark `MassProfile` is a `NFWSph`.\n", - " - The source galaxy's light is an `Sersic`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "\n", - "The `lens_galaxy` uses a `bulge` component which has a light and mass profile, and there is also a dark matter \n", - "component included." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lmp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " mass_to_light_ratio=0.2,\n", - " ),\n", - " dark=al.mp.NFWSph(centre=(0.0, 0.0), kappa_s=0.1, scale_radius=20.0),\n", - " shear=al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can then pass this simulator a tracer, which uses the tracer to create a ray-traced image which is simulated as\n", - "imaging dataset following the setup of the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/mass_stellar_dark`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Mass Stellar Dark\n", + "============================\n", + "\n", + "The majority of example scripts simulate a mass profile which represents the _total_ mass of the lens galaxy (its stars,\n", + "dark matter and other components combined). This typically uses an `Isothermal` or `PowerLaw` mass profile.\n", + "\n", + "This script simulates a strong lens where the lens mass model decomposes the lens galaxy's mass into its stars and\n", + "dark matter.\n", + "\n", + "This dataset is modeled in the example script `autolens_workspace/scripts/modeling/features/mass_stellar_dark.py`,\n", + "where a discussion of the advantages and disadvantages of fitting decomposed mass models is also given.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens using decomposed light and dark matter profiles where:\n", + "\n", + " - The lens galaxy's light matter mass distribution is an `Sersic`.\n", + " - The lens galaxy's dark `MassProfile` is a `NFWSph`.\n", + " - The source galaxy's light is an `Sersic`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "\n", + "The `lens_galaxy` uses a `bulge` component which has a light and mass profile, and there is also a dark matter \n", + "component included." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lmp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " mass_to_light_ratio=0.2,\n", + " ),\n", + " dark=al.mp.NFWSph(centre=(0.0, 0.0), kappa_s=0.1, scale_radius=20.0),\n", + " shear=al.mp.ExternalShear(gamma_1=-0.02, gamma_2=0.005),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can then pass this simulator a tracer, which uses the tracer to create a ray-traced image which is simulated as\n", + "imaging dataset following the setup of the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/mass_stellar_dark`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/mass_stellar_dark/slam.ipynb b/notebooks/imaging/features/advanced/mass_stellar_dark/slam.ipynb index 562572f19..489b8b5a4 100644 --- a/notebooks/imaging/features/advanced/mass_stellar_dark/slam.ipynb +++ b/notebooks/imaging/features/advanced/mass_stellar_dark/slam.ipynb @@ -1,716 +1,753 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM (Source, Light and Mass): Mass Light Dark\n", - "===============================================\n", - "\n", - "This example shows how to use the SLaM pipelines to end with a mass model which decomposes the lens into its\n", - "stars and dark matter, using a light plus dark matter mass model.\n", - "\n", - "Unlike other example SLaM pipelines, which end with the MASS TOTAL PIPELINE, this script ends with the\n", - "MASS LIGHT DARK PIPELINE.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light.\n", - "- **MASS LIGHT DARK PIPELINE:** The MASS LIGHT DARK PIPELINE fits a mass model where the stellar mass is tied to the lens light.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Model__\n", - "\n", - "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and a MASS LIGHT DARK PIPELINE this SLaM script\n", - "fits `Imaging` dataset of a strong lens system, where in the final model:\n", - "\n", - " - The lens galaxy's light is a `Sersic` linear light profile.\n", - " - The lens galaxy's stellar mass distribution is a `Sersic` tied to the light model above.\n", - " - The lens galaxy's dark matter mass distribution is modeled as a `NFWMCRLudlow`.\n", - " - The source galaxy's light is a `Pixelization`.\n", - "\n", - "Each SLaM pipeline is implemented as a Python function below, with a documentation string above each function\n", - "describing the pipeline in detail. The full pipeline is run at the bottom of the script.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `slam_start_here` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except:\n", - "\n", - " - The lens galaxy's bulge uses a linear `Sersic` light profile (`al.lp_linear.Sersic`) instead of an MGE.\n", - " - The source galaxy's MGE uses `gaussian_per_basis=1`.\n", - "\n", - "The linear `Sersic` profile is used here because the MASS LIGHT DARK PIPELINE requires a `LightMassProfile`\n", - "(`al.lmp.Sersic`) for the lens stellar mass, which shares the same profile type as the light model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = af.Model(al.lp_linear.Sersic)\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`.\n", - "\n", - "Note that between SOURCE PIX PIPELINE 1 and this search, the calling section applies adaptive over-sampling to\n", - "the dataset using the pixelized source reconstruction from search 1. This improves the accuracy of the\n", - "pixelization by ensuring the over-sampling is adapted to the source morphology." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light profile\n", - "(`al.lp_linear.Sersic`) instead of an MGE. This ensures the light model is consistent with the\n", - "MASS LIGHT DARK PIPELINE, which links stellar mass to a `Sersic` light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " )\n", - "\n", - " lens_bulge = af.Model(al.lp_linear.Sersic)\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS LIGHT DARK PIPELINE__\n", - "\n", - "The MASS LIGHT DARK PIPELINE fits a mass model where the stellar mass is tied to the lens light model and a\n", - "separate dark matter halo is included.\n", - "\n", - "The lens bulge is modeled as a `LightMassProfile` (`al.lmp.Sersic`) whose parameters are initialized from the\n", - "LIGHT LP PIPELINE result via `al.util.chaining.mass_light_dark_from`. The dark matter halo (`NFWMCRLudlow`) centre\n", - "is aligned with the stellar bulge centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_light_dark(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Whether to use the gradient of the mass-to-light ratio profile.\n", - " use_gradient = False\n", - " # Whether to link the mass-to-light ratios of the bulge and disk to the same value.\n", - " link_mass_to_light_ratios = True\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " lp_chain_tracer = al.util.chaining.lp_chain_tracer_from(\n", - " light_result=light_result, settings_search=settings_search\n", - " )\n", - "\n", - " lens_bulge = al.util.chaining.mass_light_dark_from(\n", - " light_result=light_result,\n", - " lp_chain_tracer=lp_chain_tracer,\n", - " name=\"bulge\",\n", - " use_gradient=use_gradient,\n", - " )\n", - " lens_disk = al.util.chaining.mass_light_dark_from(\n", - " light_result=light_result,\n", - " lp_chain_tracer=lp_chain_tracer,\n", - " name=\"disk\",\n", - " use_gradient=use_gradient,\n", - " )\n", - "\n", - " lens_bulge, lens_disk = al.util.chaining.link_ratios(\n", - " link_mass_to_light_ratios=link_mass_to_light_ratios,\n", - " light_result=light_result,\n", - " bulge=lens_bulge,\n", - " disk=lens_disk,\n", - " )\n", - "\n", - " dark = af.Model(al.mp.NFWMCRLudlow)\n", - "\n", - " try:\n", - " dark.centre = lens_bulge.centre\n", - " except AttributeError:\n", - " dark.centre = lens_bulge.profile_list[0].centre\n", - "\n", - " dark.mass_at_200 = af.LogUniformPrior(lower_limit=1e10, upper_limit=1e15)\n", - " dark.redshift_object = light_result.instance.galaxies.lens.redshift\n", - " dark.redshift_source = light_result.instance.galaxies.source.redshift\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=light_result.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=lens_disk,\n", - " dark=dark,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_light_dark[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=250,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"mass_stellar_dark\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/mass_stellar_dark/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step.\n", - "\n", - "Between SOURCE PIX PIPELINE 1 and 2, adaptive over-sampling is applied to the dataset using the pixelized source\n", - "reconstruction from search 1. This improves the pixelization accuracy in search 2 and all subsequent pipelines." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", - " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", - " noise_map=dataset.noise_map,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=over_sampling)\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")\n", - "\n", - "mass_result = mass_light_dark(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM (Source, Light and Mass): Mass Light Dark\n", + "===============================================\n", + "\n", + "This example shows how to use the SLaM pipelines to end with a mass model which decomposes the lens into its\n", + "stars and dark matter, using a light plus dark matter mass model.\n", + "\n", + "Unlike other example SLaM pipelines, which end with the MASS TOTAL PIPELINE, this script ends with the\n", + "MASS LIGHT DARK PIPELINE.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light.\n", + "- **MASS LIGHT DARK PIPELINE:** The MASS LIGHT DARK PIPELINE fits a mass model where the stellar mass is tied to the lens light.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Model__\n", + "\n", + "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and a MASS LIGHT DARK PIPELINE this SLaM script\n", + "fits `Imaging` dataset of a strong lens system, where in the final model:\n", + "\n", + " - The lens galaxy's light is a `Sersic` linear light profile.\n", + " - The lens galaxy's stellar mass distribution is a `Sersic` tied to the light model above.\n", + " - The lens galaxy's dark matter mass distribution is modeled as a `NFWMCRLudlow`.\n", + " - The source galaxy's light is a `Pixelization`.\n", + "\n", + "Each SLaM pipeline is implemented as a Python function below, with a documentation string above each function\n", + "describing the pipeline in detail. The full pipeline is run at the bottom of the script.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `slam_start_here` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except:\n", + "\n", + " - The lens galaxy's bulge uses a linear `Sersic` light profile (`al.lp_linear.Sersic`) instead of an MGE.\n", + " - The source galaxy's MGE uses `gaussian_per_basis=1`.\n", + "\n", + "The linear `Sersic` profile is used here because the MASS LIGHT DARK PIPELINE requires a `LightMassProfile`\n", + "(`al.lmp.Sersic`) for the lens stellar mass, which shares the same profile type as the light model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = af.Model(al.lp_linear.Sersic)\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`.\n", + "\n", + "Note that between SOURCE PIX PIPELINE 1 and this search, the calling section applies adaptive over-sampling to\n", + "the dataset using the pixelized source reconstruction from search 1. This improves the accuracy of the\n", + "pixelization by ensuring the over-sampling is adapted to the source morphology." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light profile\n", + "(`al.lp_linear.Sersic`) instead of an MGE. This ensures the light model is consistent with the\n", + "MASS LIGHT DARK PIPELINE, which links stellar mass to a `Sersic` light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " )\n", + "\n", + " lens_bulge = af.Model(al.lp_linear.Sersic)\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS LIGHT DARK PIPELINE__\n", + "\n", + "The MASS LIGHT DARK PIPELINE fits a mass model where the stellar mass is tied to the lens light model and a\n", + "separate dark matter halo is included.\n", + "\n", + "The lens bulge is modeled as a `LightMassProfile` (`al.lmp.Sersic`) whose parameters are initialized from the\n", + "LIGHT LP PIPELINE result via `al.util.chaining.mass_light_dark_from`. The dark matter halo (`NFWMCRLudlow`) centre\n", + "is aligned with the stellar bulge centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_light_dark(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Whether to use the gradient of the mass-to-light ratio profile.\n", + " use_gradient = False\n", + " # Whether to link the mass-to-light ratios of the bulge and disk to the same value.\n", + " link_mass_to_light_ratios = True\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " lp_chain_tracer = al.util.chaining.lp_chain_tracer_from(\n", + " light_result=light_result, settings_search=settings_search\n", + " )\n", + "\n", + " lens_bulge = al.util.chaining.mass_light_dark_from(\n", + " light_result=light_result,\n", + " lp_chain_tracer=lp_chain_tracer,\n", + " name=\"bulge\",\n", + " use_gradient=use_gradient,\n", + " )\n", + " lens_disk = al.util.chaining.mass_light_dark_from(\n", + " light_result=light_result,\n", + " lp_chain_tracer=lp_chain_tracer,\n", + " name=\"disk\",\n", + " use_gradient=use_gradient,\n", + " )\n", + "\n", + " lens_bulge, lens_disk = al.util.chaining.link_ratios(\n", + " link_mass_to_light_ratios=link_mass_to_light_ratios,\n", + " light_result=light_result,\n", + " bulge=lens_bulge,\n", + " disk=lens_disk,\n", + " )\n", + "\n", + " dark = af.Model(al.mp.NFWMCRLudlow)\n", + "\n", + " try:\n", + " dark.centre = lens_bulge.centre\n", + " except AttributeError:\n", + " dark.centre = lens_bulge.profile_list[0].centre\n", + "\n", + " dark.mass_at_200 = af.LogUniformPrior(lower_limit=1e10, upper_limit=1e15)\n", + " dark.redshift_object = light_result.instance.galaxies.lens.redshift\n", + " dark.redshift_source = light_result.instance.galaxies.source.redshift\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=light_result.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=lens_disk,\n", + " dark=dark,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_light_dark[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=250,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"mass_stellar_dark\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/mass_stellar_dark/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step.\n", + "\n", + "Between SOURCE PIX PIPELINE 1 and 2, adaptive over-sampling is applied to the dataset using the pixelized source\n", + "reconstruction from search 1. This improves the pixelization accuracy in search 2 and all subsequent pipelines." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", + " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", + " noise_map=dataset.noise_map,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=over_sampling)\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")\n", + "\n", + "mass_result = mass_light_dark(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/operated_light_profile/modeling.ipynb b/notebooks/imaging/features/advanced/operated_light_profile/modeling.ipynb index d5166454d..5c6b79659 100644 --- a/notebooks/imaging/features/advanced/operated_light_profile/modeling.ipynb +++ b/notebooks/imaging/features/advanced/operated_light_profile/modeling.ipynb @@ -1,477 +1,514 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Operated Light Profiles\n", - "==========================================\n", - "\n", - "It is common for galaxies to have point-source emission, for example bright emission right at their centre due to\n", - "an active galactic nuclei or a compact knot of star formation.\n", - "\n", - "This point-source emission is subject to blurring during data acquiziton due to the telescope optics, and therefore\n", - "is not seen as a single pixel of light but spread over multiple pixels as a convolution with the telescope\n", - "Point Spread Function (PSF).\n", - "\n", - "It is difficult to model this compact point source emission using a point-source light profile (or an extremely\n", - "compact Gaussian / Sersic profile). This is because when the model-image of a compact point source of light is\n", - "convolved with the PSF, the solution to this convolution is extremely sensitive to which pixel (and sub-pixel) the\n", - "compact model emission lands in.\n", - "\n", - "Operated light profiles offer an alternative approach, whereby the light profile is assumed to have already been\n", - "convolved with the PSF. This operated light profile is then fitted directly to the point-source emission, which as\n", - "discussed above shows the PSF features.\n", - "\n", - "Operated light profiles bypass the convolution step entirely, and therefore if you had a use-case which\n", - "required fitting other components of a galaxy without convolution they could be used for this purpose too.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a linear `Sersic` bulge.\n", - " - The lens galaxy includes a linear `Gaussian` psf.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a linear `SersicCore`.\n", - "\n", - "__Fit__\n", - "\n", - "For operated light profiles, there is no `fit.py` example found for standard light profiles, linear light profiles\n", - "and other examples.\n", - "\n", - "This is done purely to keep the number of examples in the workspace manageable. to perform a fit with operated light\n", - "profiles, simply follow one of the other `modeling/imaging/fit.py` examples and replace the light profiles\n", - "with operated light profiles using the API described below.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"light_operated\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/operated_light_profile/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's light is a linear `Sersic` bulge [6 parameters].\n", - "\n", - " - The lens galaxy's point source emission is a linear operated `Gaussian` centred on the bulge [3 parameters].\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is a linear `SersicCore` [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=25.\n", - "\n", - "This lens galaxy includes a `Gaussian` operated light profile (`lp_operated`) PSF, which accounts for the PSF\n", - "convolution of the point source emission.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = af.Model(al.lp_linear.Sersic)\n", - "psf = af.Model(al.lp_operated.Gaussian)\n", - "bulge.centre = psf.centre\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " psf=psf,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - ")\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There is also a linear variant of every operated light profile (see `linear_light_profiles.py`).\n", - "\n", - "We will use this, as it simplifies parameter space, which is particularly important for operated light profiles \n", - "which can prove quite difficult to sample robustly.\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space for this model is N=22." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "psf = af.Model(al.lp_linear_operated.Gaussian)\n", - "bulge.centre = psf.centre\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, psf=psf, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the lens galaxy's light has a `Gaussian` PSF component." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"operated_light_profiles\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to\n", - "print the estimated VRAM required by a model.\n", - "\n", - "For each operated light profile in the model extra is used VRAM. For 3-10 linear Sersic light profiles this is a tiny \n", - "amount of VRAM (e.g. < 10MB per batched likelihood). Even for large batch sizes (e.g. over 100) you probably \n", - "will not use enough VRAM to require monitoring.\n", - "\n", - "__Run Time__\n", - "\n", - "The likelihood evaluation time for operated light profiles are faster than all other light profiles. This is because\n", - "the computationally expensive PSF convolution step is omitted.\n", - "\n", - "The overall run-time may be a little slower than other models though, because the `psf` component adds a few\n", - "extra parameters.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this does \n", - "not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the lens galaxy's light has a `Gaussian` PSF component." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "__Wrap Up__\n", - "\n", - "This script shows how fit a lens model where bright point source emission is modeled using a `Gaussian` PSF.\n", - "\n", - "It is uncommon for this feature to be necessary, as there are very few lens galaxies which have such bright point\n", - "source emission in their centre. Nevertheless, it has been observed in a select few systems, and is therefore \n", - "something you may need to model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Operated Light Profiles\n", + "==========================================\n", + "\n", + "It is common for galaxies to have point-source emission, for example bright emission right at their centre due to\n", + "an active galactic nuclei or a compact knot of star formation.\n", + "\n", + "This point-source emission is subject to blurring during data acquiziton due to the telescope optics, and therefore\n", + "is not seen as a single pixel of light but spread over multiple pixels as a convolution with the telescope\n", + "Point Spread Function (PSF).\n", + "\n", + "It is difficult to model this compact point source emission using a point-source light profile (or an extremely\n", + "compact Gaussian / Sersic profile). This is because when the model-image of a compact point source of light is\n", + "convolved with the PSF, the solution to this convolution is extremely sensitive to which pixel (and sub-pixel) the\n", + "compact model emission lands in.\n", + "\n", + "Operated light profiles offer an alternative approach, whereby the light profile is assumed to have already been\n", + "convolved with the PSF. This operated light profile is then fitted directly to the point-source emission, which as\n", + "discussed above shows the PSF features.\n", + "\n", + "Operated light profiles bypass the convolution step entirely, and therefore if you had a use-case which\n", + "required fitting other components of a galaxy without convolution they could be used for this purpose too.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a linear `Sersic` bulge.\n", + " - The lens galaxy includes a linear `Gaussian` psf.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a linear `SersicCore`.\n", + "\n", + "__Fit__\n", + "\n", + "For operated light profiles, there is no `fit.py` example found for standard light profiles, linear light profiles\n", + "and other examples.\n", + "\n", + "This is done purely to keep the number of examples in the workspace manageable. to perform a fit with operated light\n", + "profiles, simply follow one of the other `modeling/imaging/fit.py` examples and replace the light profiles\n", + "with operated light profiles using the API described below.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"light_operated\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/operated_light_profile/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's light is a linear `Sersic` bulge [6 parameters].\n", + "\n", + " - The lens galaxy's point source emission is a linear operated `Gaussian` centred on the bulge [3 parameters].\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is a linear `SersicCore` [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=25.\n", + "\n", + "This lens galaxy includes a `Gaussian` operated light profile (`lp_operated`) PSF, which accounts for the PSF\n", + "convolution of the point source emission.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = af.Model(al.lp_linear.Sersic)\n", + "psf = af.Model(al.lp_operated.Gaussian)\n", + "bulge.centre = psf.centre\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " psf=psf,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + ")\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There is also a linear variant of every operated light profile (see `linear_light_profiles.py`).\n", + "\n", + "We will use this, as it simplifies parameter space, which is particularly important for operated light profiles \n", + "which can prove quite difficult to sample robustly.\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space for this model is N=22." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "psf = af.Model(al.lp_linear_operated.Gaussian)\n", + "bulge.centre = psf.centre\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, psf=psf, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the lens galaxy's light has a `Gaussian` PSF component." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"operated_light_profiles\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to\n", + "print the estimated VRAM required by a model.\n", + "\n", + "For each operated light profile in the model extra is used VRAM. For 3-10 linear Sersic light profiles this is a tiny \n", + "amount of VRAM (e.g. < 10MB per batched likelihood). Even for large batch sizes (e.g. over 100) you probably \n", + "will not use enough VRAM to require monitoring.\n", + "\n", + "__Run Time__\n", + "\n", + "The likelihood evaluation time for operated light profiles are faster than all other light profiles. This is because\n", + "the computationally expensive PSF convolution step is omitted.\n", + "\n", + "The overall run-time may be a little slower than other models though, because the `psf` component adds a few\n", + "extra parameters.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this does \n", + "not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the lens galaxy's light has a `Gaussian` PSF component." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "__Wrap Up__\n", + "\n", + "This script shows how fit a lens model where bright point source emission is modeled using a `Gaussian` PSF.\n", + "\n", + "It is uncommon for this feature to be necessary, as there are very few lens galaxies which have such bright point\n", + "source emission in their centre. Nevertheless, it has been observed in a select few systems, and is therefore \n", + "something you may need to model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/operated_light_profile/simulator.ipynb b/notebooks/imaging/features/advanced/operated_light_profile/simulator.ipynb index a0e528131..1cafcb490 100644 --- a/notebooks/imaging/features/advanced/operated_light_profile/simulator.ipynb +++ b/notebooks/imaging/features/advanced/operated_light_profile/simulator.ipynb @@ -1,377 +1,414 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Light Operated\n", - "=========================\n", - "\n", - "It is common for galaxies to have point-source emission, for example bright emission right at their centre due to\n", - "an active galactic nuclei or a compact knot of star formation.\n", - "\n", - "This point-source emission is subject to blurring during data acquisition due to the telescope optics, and therefore\n", - "is not seen as a single pixel of light but spread over multiple pixels as a convolution with the telescope\n", - "Point Spread Function (PSF).\n", - "\n", - "This script simulate an `Imaging` dataset of a 'galaxy-scale' strong lens which has this point-source emission in\n", - "its centre. This emission uses an operated `Gaussian` light profile, which therefore is already convolved with the\n", - "PSF.\n", - "\n", - "This dataset is used in `modeling/features/operated_light_profiles.py` to demonstrate how to fit this point-source\n", - "emission using an operated light profile.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's light profile is an `Sersic` bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The lens galaxy has a point source of emission at its centre which is modeled as a operated `Gaussian`.\n", - " - The source galaxy's light is an `Sersic`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"light_operated\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "\n", - "This includes an operated `Gaussian` component which represents the PSF convolved emission of a point-source \n", - "emission at the galaxy's centre. \n", - "\n", - "It uses the same parameters as the `psf` defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " psf=al.lp_operated.Gaussian(\n", - " centre=(0.0, 0.0), ell_comps=(0.0, 0.0), intensity=100.0, sigma=0.1\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/light_operated`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Light Operated\n", + "=========================\n", + "\n", + "It is common for galaxies to have point-source emission, for example bright emission right at their centre due to\n", + "an active galactic nuclei or a compact knot of star formation.\n", + "\n", + "This point-source emission is subject to blurring during data acquisition due to the telescope optics, and therefore\n", + "is not seen as a single pixel of light but spread over multiple pixels as a convolution with the telescope\n", + "Point Spread Function (PSF).\n", + "\n", + "This script simulate an `Imaging` dataset of a 'galaxy-scale' strong lens which has this point-source emission in\n", + "its centre. This emission uses an operated `Gaussian` light profile, which therefore is already convolved with the\n", + "PSF.\n", + "\n", + "This dataset is used in `modeling/features/operated_light_profiles.py` to demonstrate how to fit this point-source\n", + "emission using an operated light profile.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's light profile is an `Sersic` bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The lens galaxy has a point source of emission at its centre which is modeled as a operated `Gaussian`.\n", + " - The source galaxy's light is an `Sersic`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"light_operated\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "\n", + "This includes an operated `Gaussian` component which represents the PSF convolved emission of a point-source \n", + "emission at the galaxy's centre. \n", + "\n", + "It uses the same parameters as the `psf` defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " psf=al.lp_operated.Gaussian(\n", + " centre=(0.0, 0.0), ell_comps=(0.0, 0.0), intensity=100.0, sigma=0.1\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/light_operated`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/shapelets/fit.ipynb b/notebooks/imaging/features/advanced/shapelets/fit.ipynb index 99caa822d..5d91c7f83 100644 --- a/notebooks/imaging/features/advanced/shapelets/fit.ipynb +++ b/notebooks/imaging/features/advanced/shapelets/fit.ipynb @@ -1,683 +1,720 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Shapelets\n", - "============================\n", - "\n", - "A shapelet is a basis function that is appropriate for capturing the exponential / disk-like features of a galaxy. It\n", - "has been employed in many strong lensing studies to model the light of the lensed source galaxy, because it can\n", - "represent features of disky star forming galaxies that a single Sersic function cannot.\n", - "\n", - "- https://ui.adsabs.harvard.edu/abs/2016MNRAS.457.3066T\n", - "- https://iopscience.iop.org/article/10.1088/0004-637X/813/2/102\n", - "\n", - "Shapelets are described in full in the following paper:\n", - "\n", - " https://arxiv.org/abs/astro-ph/0105178\n", - "\n", - "This script performs a model-fit using shapelets, where it decomposes the galaxy light into ~20\n", - "Shapelets. The `intensity` of every Shapelet is solved for via linear algebra (see the `linear_light_profiles.py`\n", - "feature).\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Symmetric light profiles (e.g.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Basis:** We first build a `Basis`, which is built from multiple light profiles (in this case, shapelets).\n", - "- **Coefficients:** The `Basis` is composed of many shapelets, each with different coefficients (n and m) values and a.\n", - "- **Linear Light Profiles:** We now show Composing a basis of multiple shapelets and use them to fit the source galaxy's light.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Positive Negative Solver:** In other examples which use linear algebra to fit the data, for example linear light profiles, the.\n", - "- **Intensities:** The fit contains the solved for intensity values.\n", - "- **Shapelet Cartesian:** The shapelets above were defined on a polar grid, which is suitable for modeling radially symmetric.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals, because they fail to capture\n", - "irregular and asymmetric morphological of galaxies (e.g. isophotal twists, an ellipticity which varies radially).\n", - "Shapelets can capture some of these features and can therefore better represent the emission of complex source galaxies.\n", - "\n", - "The shapelet model can be composed in a way that has fewer non-linear parameters than an elliptical Sersic. In this\n", - "example, the ~20 shapelets which represent the `bulge` of the source are composed in a model corresponding to just\n", - "N=3 non-linear parameters (a `bulge` comprising a linear Sersic would give N=6).\n", - "\n", - "Therefore, shapelets fit more complex source galaxy morphologies using fewer non-linear parameters than the standard\n", - "light profile models.\n", - "\n", - "__Disadvantages__\n", - "\n", - "- There are many types of galaxy structure which shapelets may struggle to represent, such as a bar or asymmetric\n", - "knots of star formation. They also rely on the galaxy having a distinct centre over which the shapelets can be\n", - "centered, which is not the case if the galaxy is multiple merging systems or has bright companion galaxies.\n", - "\n", - "- The linear algebra used to solve for the `intensity` of each shapelet has to allow for negative values of intensity\n", - "in order for shapelets to work. Negative surface brightnesses are unphysical, and are often inferred in a shapelet\n", - "decomposition, for example if the true galaxy has structure that cannot be captured by the shapelet basis. Other\n", - "approaches can force positive-only intensities on the solution, such as the Multi-Gaussian Expansion (MGE) or a pixelization.\n", - "\n", - "- Computationally slower than standard light profiles like the Sersic.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's bulge is a superposition of `ShapeletPolar` profiles.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the galaxy dataset `light_basis` via .fits files, which we will fit with \n", - "the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular_annular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " inner_radius=0.4,\n", - " outer_radius=3.0,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Basis__\n", - "\n", - "We first build a `Basis`, which is built from multiple light profiles (in this case, shapelets). \n", - "\n", - "Below, we make a `Basis` out of 20 elliptical polar shapelet light profiles which: \n", - "\n", - " - All share the same centre and elliptical components.\n", - " - The size of the Shapelet basis is controlled by a `beta` parameter, which is the same for all shapelet basis \n", - " functions.\n", - "\n", - "Note that any light profile can be used to compose a Basis. This includes Gaussians, which are often used to \n", - "represent the light of elliptical galaxies (see `modeling/features/multi_gaussian_expansion.py`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_n = 5\n", - "total_m = sum(range(2, total_n + 1)) + 1\n", - "\n", - "n_count = 1\n", - "m_count = -1\n", - "\n", - "shapelets_bulge_list = []\n", - "\n", - "shapelet_0 = al.lp.ShapeletPolar(\n", - " n=0,\n", - " m=0,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=1.0,\n", - " beta=1.0,\n", - ")\n", - "\n", - "shapelets_bulge_list.append(shapelet_0)\n", - "\n", - "for i in range(total_n + total_m):\n", - " shapelet = al.lp.ShapeletPolarSph(\n", - " n=n_count, m=m_count, centre=(0.0, 0.0), intensity=1.0, beta=1.0\n", - " )\n", - "\n", - " shapelets_bulge_list.append(shapelet)\n", - "\n", - " m_count += 2\n", - "\n", - " if m_count > n_count:\n", - " n_count += 1\n", - " m_count = -n_count\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Coefficients__\n", - "\n", - "The `Basis` is composed of many shapelets, each with different coefficients (n and m) values and a size parameter \n", - "`beta`.\n", - "\n", - "Each combination of coefficients creates shapelets with different radial and azimuthal features. They capture \n", - "emission on different scales, with low coefficients corresponding to smooth features and high coefficients \n", - "corresponding to more variable wave-like features. The size of the coefficients is determined by the input \n", - "parameter `beta`, where larger values correspond to larger coefficients and therefore larger shapelets.\n", - "\n", - "These coefficients are visualized below using `subplot_basis_image`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "subplot_basis_image(basis=bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Light Profiles__\n", - "\n", - "We now show Composing a basis of multiple shapelets and use them to fit the source galaxy's light in data.\n", - "\n", - "This does not perform a model-fit via a non-linear search, and therefore requires us to manually specify and guess\n", - "suitable parameter values for the shapelets (e.g. the `centre`, `ell_comps`, `beta`). However, shapelets are\n", - "very flexible and will give us a decent looking source reconstruction even if we just guess sensible values\n", - "for each parameter. \n", - "\n", - "The one parameter that is tricky to guess is the `intensity` of each shapelet. A wide range of positive\n", - "and negative `intensity` values are required to decompose the source galaxy's light accurately. We certainly\n", - "cannot obtain a good solution by guessing the `intensity` values by eye.\n", - "\n", - "We therefore use linear light profile shapelets, which determine the optimal value for each shapelet's `intensity` \n", - "via linear algebra. Linear light profiles are described in the `linear_light_profiles.py` example and you should\n", - "familiarize yourself with this example before using shapelets.\n", - "\n", - "We therefore again setup a `Basis` in an analogous fashion to the previous example, but this time we use linear\n", - "shapelets (via the `lp_linear.linear` module)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_n = 5\n", - "total_m = sum(range(2, total_n + 1)) + 1\n", - "\n", - "n_count = 1\n", - "m_count = -1\n", - "\n", - "shapelets_bulge_list = []\n", - "\n", - "shapelet_0 = al.lp_linear.ShapeletPolar(\n", - " n=0,\n", - " m=0,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " beta=1.0,\n", - ")\n", - "\n", - "shapelets_bulge_list.append(shapelet_0)\n", - "\n", - "for i in range(total_n + total_m):\n", - " shapelet = al.lp_linear.ShapeletPolar(\n", - " n=n_count,\n", - " m=m_count,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " beta=1.0,\n", - " )\n", - "\n", - " shapelets_bulge_list.append(shapelet)\n", - "\n", - " m_count += 2\n", - "\n", - " if m_count > n_count:\n", - " n_count += 1\n", - " m_count = -n_count\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We now illustrate the API for fitting shapelets using standard autolens objects like the `Galaxy`, `Tracer` \n", - "and `FitImaging`.\n", - "\n", - "Once we have a `Basis`, we can treat it like any other light profile in order to create a `Galaxy` and `Tracer` and \n", - "use it to fit data.\n", - "\n", - "We are applying shapelets to reconstruct the source galaxy's light, which means we need an accurate mass model of the\n", - "lens galaxy. We use the true lens mass model from the simulator script to do this, noting that later in the example\n", - "we will infer the lens mass model using a non-linear search.\n", - "\n", - "__Positive Negative Solver__\n", - "\n", - "In other examples which use linear algebra to fit the data, for example linear light profiles, the Multi Gaussian\n", - "Expansion (MGE) and pixelization, we use a `positive_only` solver, which forces all solved for intensities to be\n", - "positive. This is a physical and sensible approach, because the surface brightnesses of a galaxy cannot be negative.\n", - "\n", - "Shapelets cannot be solved for using a `positive_only` solver, because the shapelets ability to decompose the\n", - "light of a galaxy relies on the ability to use negative intensities. This is because the shapelets are not\n", - "physically motivated light profiles, but instead a mathematical basis that can represent any light profile.\n", - "\n", - "This means shapelets may include negative flux in the reconstructed source galaxy, which is unphysical, and\n", - "a disadvantage of using shapelets.\n", - "\n", - "The `Settings` object below uses a `use_positive_only_solver=False` to allow for negative intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=bulge,\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_positive_only_solver=False),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the shapelets do a reasonable job at capturing the appearance of the source galaxy,\n", - "with only faint residuals visible where the lensed source is located.\n", - "\n", - "This is despite the beta parameter of the shapelets being a complete guess and not the optimal value for fitting the\n", - "source galaxy's light. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can use `subplot_basis_image` to plot each individual shapelet in the reconstructed basis.\n", - "\n", - "This plot shows each shapelet has a unique `intensity` that was solved for via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer.galaxies[1].bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The fit contains the solved for intensity values.\n", - "\n", - "These are computed using a fit's `linear_light_profile_intensity_dict`, which maps each linear light profile \n", - "in the model parameterization above to its `intensity`.\n", - "\n", - "The code below shows how to use this dictionary, as an alternative to using the max_log_likelihood quantities above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_bulge = fit.tracer.galaxies[1].bulge\n", - "\n", - "print(\n", - " f\"\\n Intensity of source galaxy's first shapelet in bulge = {fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the solved \n", - "for `intensity` values is also accessible from a fit.\n", - "\n", - "For example, the first linear light profile of the shapelet `bulge` component above printed it solved for intensity \n", - "value, but it was still represented as a linear light profile. \n", - "\n", - "The `tracer` created below instead has a standard light profile with an `intensity` actually set.\n", - "\n", - "The benefit of using a tracer with standard light profiles is it can be visualized, as performed above (linear \n", - "light profiles cannot by default because they do not have `intensity` values)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "print(tracer.galaxies[1].bulge.profile_list[0].intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Shapelet Cartesian__\n", - "\n", - "The shapelets above were defined on a polar grid, which is suitable for modeling radially symmetric sources like\n", - "most galaxies.\n", - "\n", - "An alternative approach is to define the shapelets on a Cartesian grid, which we plot the basis of below\n", - "and show an example fit.\n", - "\n", - "These are generally not recommended for modeling galaxies, but may be better in certain situations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_xy = 5\n", - "\n", - "shapelets_bulge_list = []\n", - "\n", - "for x in range(total_xy):\n", - " for y in range(total_xy):\n", - " shapelet = al.lp.ShapeletCartesian(\n", - " n_y=y,\n", - " n_x=x,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=1.0,\n", - " beta=1.0,\n", - " )\n", - "\n", - " shapelets_bulge_list.append(shapelet)\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "subplot_basis_image(basis=bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For fitting, we again use the linear light profile version of the Cartesian shapelets, which solves for the\n", - "optimal intensity of each shapelet via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_xy = 5\n", - "\n", - "shapelets_bulge_list = []\n", - "\n", - "for x in range(total_xy):\n", - " for y in range(total_xy):\n", - " shapelet = al.lp_linear.ShapeletCartesian(\n", - " n_y=y, n_x=x, centre=(0.0, 0.0), ell_comps=(0.0, 0.0), beta=1.0\n", - " )\n", - "\n", - " shapelets_bulge_list.append(shapelet)\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=bulge,\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_positive_only_solver=False),\n", - ")\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer.galaxies[1].bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script has illustrated how to use shapelets to model the light of galaxies.\n", - "\n", - "Shapelets are a powerful basis function for capturing complex morphological features of galaxies that standard\n", - "light profiles struggle to represent. However, they do have drawbacks, such as the need to allow for negative\n", - "intensities in the solution, which is unphysical. \n", - "\n", - "As a rule of thumb, modeling is generally better if a pixelization is used to reconstruct the source galaxy's light,\n", - "but shapelets can be a useful middle-ground between standard light profiles and a pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Shapelets\n", + "============================\n", + "\n", + "A shapelet is a basis function that is appropriate for capturing the exponential / disk-like features of a galaxy. It\n", + "has been employed in many strong lensing studies to model the light of the lensed source galaxy, because it can\n", + "represent features of disky star forming galaxies that a single Sersic function cannot.\n", + "\n", + "- https://ui.adsabs.harvard.edu/abs/2016MNRAS.457.3066T\n", + "- https://iopscience.iop.org/article/10.1088/0004-637X/813/2/102\n", + "\n", + "Shapelets are described in full in the following paper:\n", + "\n", + " https://arxiv.org/abs/astro-ph/0105178\n", + "\n", + "This script performs a model-fit using shapelets, where it decomposes the galaxy light into ~20\n", + "Shapelets. The `intensity` of every Shapelet is solved for via linear algebra (see the `linear_light_profiles.py`\n", + "feature).\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Symmetric light profiles (e.g.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Basis:** We first build a `Basis`, which is built from multiple light profiles (in this case, shapelets).\n", + "- **Coefficients:** The `Basis` is composed of many shapelets, each with different coefficients (n and m) values and a.\n", + "- **Linear Light Profiles:** We now show Composing a basis of multiple shapelets and use them to fit the source galaxy's light.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Positive Negative Solver:** In other examples which use linear algebra to fit the data, for example linear light profiles, the.\n", + "- **Intensities:** The fit contains the solved for intensity values.\n", + "- **Shapelet Cartesian:** The shapelets above were defined on a polar grid, which is suitable for modeling radially symmetric.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals, because they fail to capture\n", + "irregular and asymmetric morphological of galaxies (e.g. isophotal twists, an ellipticity which varies radially).\n", + "Shapelets can capture some of these features and can therefore better represent the emission of complex source galaxies.\n", + "\n", + "The shapelet model can be composed in a way that has fewer non-linear parameters than an elliptical Sersic. In this\n", + "example, the ~20 shapelets which represent the `bulge` of the source are composed in a model corresponding to just\n", + "N=3 non-linear parameters (a `bulge` comprising a linear Sersic would give N=6).\n", + "\n", + "Therefore, shapelets fit more complex source galaxy morphologies using fewer non-linear parameters than the standard\n", + "light profile models.\n", + "\n", + "__Disadvantages__\n", + "\n", + "- There are many types of galaxy structure which shapelets may struggle to represent, such as a bar or asymmetric\n", + "knots of star formation. They also rely on the galaxy having a distinct centre over which the shapelets can be\n", + "centered, which is not the case if the galaxy is multiple merging systems or has bright companion galaxies.\n", + "\n", + "- The linear algebra used to solve for the `intensity` of each shapelet has to allow for negative values of intensity\n", + "in order for shapelets to work. Negative surface brightnesses are unphysical, and are often inferred in a shapelet\n", + "decomposition, for example if the true galaxy has structure that cannot be captured by the shapelet basis. Other\n", + "approaches can force positive-only intensities on the solution, such as the Multi-Gaussian Expansion (MGE) or a pixelization.\n", + "\n", + "- Computationally slower than standard light profiles like the Sersic.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's bulge is a superposition of `ShapeletPolar` profiles.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the galaxy dataset `light_basis` via .fits files, which we will fit with \n", + "the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular_annular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " inner_radius=0.4,\n", + " outer_radius=3.0,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Basis__\n", + "\n", + "We first build a `Basis`, which is built from multiple light profiles (in this case, shapelets). \n", + "\n", + "Below, we make a `Basis` out of 20 elliptical polar shapelet light profiles which: \n", + "\n", + " - All share the same centre and elliptical components.\n", + " - The size of the Shapelet basis is controlled by a `beta` parameter, which is the same for all shapelet basis \n", + " functions.\n", + "\n", + "Note that any light profile can be used to compose a Basis. This includes Gaussians, which are often used to \n", + "represent the light of elliptical galaxies (see `modeling/features/multi_gaussian_expansion.py`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_n = 5\n", + "total_m = sum(range(2, total_n + 1)) + 1\n", + "\n", + "n_count = 1\n", + "m_count = -1\n", + "\n", + "shapelets_bulge_list = []\n", + "\n", + "shapelet_0 = al.lp.ShapeletPolar(\n", + " n=0,\n", + " m=0,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=1.0,\n", + " beta=1.0,\n", + ")\n", + "\n", + "shapelets_bulge_list.append(shapelet_0)\n", + "\n", + "for i in range(total_n + total_m):\n", + " shapelet = al.lp.ShapeletPolarSph(\n", + " n=n_count, m=m_count, centre=(0.0, 0.0), intensity=1.0, beta=1.0\n", + " )\n", + "\n", + " shapelets_bulge_list.append(shapelet)\n", + "\n", + " m_count += 2\n", + "\n", + " if m_count > n_count:\n", + " n_count += 1\n", + " m_count = -n_count\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Coefficients__\n", + "\n", + "The `Basis` is composed of many shapelets, each with different coefficients (n and m) values and a size parameter \n", + "`beta`.\n", + "\n", + "Each combination of coefficients creates shapelets with different radial and azimuthal features. They capture \n", + "emission on different scales, with low coefficients corresponding to smooth features and high coefficients \n", + "corresponding to more variable wave-like features. The size of the coefficients is determined by the input \n", + "parameter `beta`, where larger values correspond to larger coefficients and therefore larger shapelets.\n", + "\n", + "These coefficients are visualized below using `subplot_basis_image`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "subplot_basis_image(basis=bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Light Profiles__\n", + "\n", + "We now show Composing a basis of multiple shapelets and use them to fit the source galaxy's light in data.\n", + "\n", + "This does not perform a model-fit via a non-linear search, and therefore requires us to manually specify and guess\n", + "suitable parameter values for the shapelets (e.g. the `centre`, `ell_comps`, `beta`). However, shapelets are\n", + "very flexible and will give us a decent looking source reconstruction even if we just guess sensible values\n", + "for each parameter. \n", + "\n", + "The one parameter that is tricky to guess is the `intensity` of each shapelet. A wide range of positive\n", + "and negative `intensity` values are required to decompose the source galaxy's light accurately. We certainly\n", + "cannot obtain a good solution by guessing the `intensity` values by eye.\n", + "\n", + "We therefore use linear light profile shapelets, which determine the optimal value for each shapelet's `intensity` \n", + "via linear algebra. Linear light profiles are described in the `linear_light_profiles.py` example and you should\n", + "familiarize yourself with this example before using shapelets.\n", + "\n", + "We therefore again setup a `Basis` in an analogous fashion to the previous example, but this time we use linear\n", + "shapelets (via the `lp_linear.linear` module)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_n = 5\n", + "total_m = sum(range(2, total_n + 1)) + 1\n", + "\n", + "n_count = 1\n", + "m_count = -1\n", + "\n", + "shapelets_bulge_list = []\n", + "\n", + "shapelet_0 = al.lp_linear.ShapeletPolar(\n", + " n=0,\n", + " m=0,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " beta=1.0,\n", + ")\n", + "\n", + "shapelets_bulge_list.append(shapelet_0)\n", + "\n", + "for i in range(total_n + total_m):\n", + " shapelet = al.lp_linear.ShapeletPolar(\n", + " n=n_count,\n", + " m=m_count,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " beta=1.0,\n", + " )\n", + "\n", + " shapelets_bulge_list.append(shapelet)\n", + "\n", + " m_count += 2\n", + "\n", + " if m_count > n_count:\n", + " n_count += 1\n", + " m_count = -n_count\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We now illustrate the API for fitting shapelets using standard autolens objects like the `Galaxy`, `Tracer` \n", + "and `FitImaging`.\n", + "\n", + "Once we have a `Basis`, we can treat it like any other light profile in order to create a `Galaxy` and `Tracer` and \n", + "use it to fit data.\n", + "\n", + "We are applying shapelets to reconstruct the source galaxy's light, which means we need an accurate mass model of the\n", + "lens galaxy. We use the true lens mass model from the simulator script to do this, noting that later in the example\n", + "we will infer the lens mass model using a non-linear search.\n", + "\n", + "__Positive Negative Solver__\n", + "\n", + "In other examples which use linear algebra to fit the data, for example linear light profiles, the Multi Gaussian\n", + "Expansion (MGE) and pixelization, we use a `positive_only` solver, which forces all solved for intensities to be\n", + "positive. This is a physical and sensible approach, because the surface brightnesses of a galaxy cannot be negative.\n", + "\n", + "Shapelets cannot be solved for using a `positive_only` solver, because the shapelets ability to decompose the\n", + "light of a galaxy relies on the ability to use negative intensities. This is because the shapelets are not\n", + "physically motivated light profiles, but instead a mathematical basis that can represent any light profile.\n", + "\n", + "This means shapelets may include negative flux in the reconstructed source galaxy, which is unphysical, and\n", + "a disadvantage of using shapelets.\n", + "\n", + "The `Settings` object below uses a `use_positive_only_solver=False` to allow for negative intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=bulge,\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_positive_only_solver=False),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the shapelets do a reasonable job at capturing the appearance of the source galaxy,\n", + "with only faint residuals visible where the lensed source is located.\n", + "\n", + "This is despite the beta parameter of the shapelets being a complete guess and not the optimal value for fitting the\n", + "source galaxy's light. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can use `subplot_basis_image` to plot each individual shapelet in the reconstructed basis.\n", + "\n", + "This plot shows each shapelet has a unique `intensity` that was solved for via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer.galaxies[1].bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The fit contains the solved for intensity values.\n", + "\n", + "These are computed using a fit's `linear_light_profile_intensity_dict`, which maps each linear light profile \n", + "in the model parameterization above to its `intensity`.\n", + "\n", + "The code below shows how to use this dictionary, as an alternative to using the max_log_likelihood quantities above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_bulge = fit.tracer.galaxies[1].bulge\n", + "\n", + "print(\n", + " f\"\\n Intensity of source galaxy's first shapelet in bulge = {fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the solved \n", + "for `intensity` values is also accessible from a fit.\n", + "\n", + "For example, the first linear light profile of the shapelet `bulge` component above printed it solved for intensity \n", + "value, but it was still represented as a linear light profile. \n", + "\n", + "The `tracer` created below instead has a standard light profile with an `intensity` actually set.\n", + "\n", + "The benefit of using a tracer with standard light profiles is it can be visualized, as performed above (linear \n", + "light profiles cannot by default because they do not have `intensity` values)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "print(tracer.galaxies[1].bulge.profile_list[0].intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Shapelet Cartesian__\n", + "\n", + "The shapelets above were defined on a polar grid, which is suitable for modeling radially symmetric sources like\n", + "most galaxies.\n", + "\n", + "An alternative approach is to define the shapelets on a Cartesian grid, which we plot the basis of below\n", + "and show an example fit.\n", + "\n", + "These are generally not recommended for modeling galaxies, but may be better in certain situations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_xy = 5\n", + "\n", + "shapelets_bulge_list = []\n", + "\n", + "for x in range(total_xy):\n", + " for y in range(total_xy):\n", + " shapelet = al.lp.ShapeletCartesian(\n", + " n_y=y,\n", + " n_x=x,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=1.0,\n", + " beta=1.0,\n", + " )\n", + "\n", + " shapelets_bulge_list.append(shapelet)\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "subplot_basis_image(basis=bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For fitting, we again use the linear light profile version of the Cartesian shapelets, which solves for the\n", + "optimal intensity of each shapelet via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_xy = 5\n", + "\n", + "shapelets_bulge_list = []\n", + "\n", + "for x in range(total_xy):\n", + " for y in range(total_xy):\n", + " shapelet = al.lp_linear.ShapeletCartesian(\n", + " n_y=y, n_x=x, centre=(0.0, 0.0), ell_comps=(0.0, 0.0), beta=1.0\n", + " )\n", + "\n", + " shapelets_bulge_list.append(shapelet)\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=bulge,\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_positive_only_solver=False),\n", + ")\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer.galaxies[1].bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script has illustrated how to use shapelets to model the light of galaxies.\n", + "\n", + "Shapelets are a powerful basis function for capturing complex morphological features of galaxies that standard\n", + "light profiles struggle to represent. However, they do have drawbacks, such as the need to allow for negative\n", + "intensities in the solution, which is unphysical. \n", + "\n", + "As a rule of thumb, modeling is generally better if a pixelization is used to reconstruct the source galaxy's light,\n", + "but shapelets can be a useful middle-ground between standard light profiles and a pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/shapelets/modeling.ipynb b/notebooks/imaging/features/advanced/shapelets/modeling.ipynb index 92cbb8658..10e4cd019 100644 --- a/notebooks/imaging/features/advanced/shapelets/modeling.ipynb +++ b/notebooks/imaging/features/advanced/shapelets/modeling.ipynb @@ -1,837 +1,874 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Shapelets\n", - "============================\n", - "\n", - "A shapelet is a basis function that is appropriate for capturing the exponential / disk-like features of a galaxy. It\n", - "has been employed in many strong lensing studies to model the light of the lensed source galaxy, because it can\n", - "represent features of disky star forming galaxies that a single Sersic function cannot.\n", - "\n", - "- https://ui.adsabs.harvard.edu/abs/2016MNRAS.457.3066T\n", - "- https://iopscience.iop.org/article/10.1088/0004-637X/813/2/102\n", - "\n", - "Shapelets are described in full in the following paper:\n", - "\n", - " https://arxiv.org/abs/astro-ph/0105178\n", - "\n", - "This script performs a model-fit using shapelets, where it decomposes the galaxy light into ~20\n", - "Shapelets. The `intensity` of every Shapelet is solved for via linear algebra (see the `linear_light_profiles.py`\n", - "feature).\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Symmetric light profiles (e.g.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Positive Negative Solver:** In other examples which use linear algebra to fit the data, for example linear light profiles, the.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Position Likelihood:** We add a penalty term ot the likelihood function, which penalizes models where the brightest.\n", - "- **Brief Description:** Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH`.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Shapelet Cartesian:** The shapelets above were defined on a polar grid, which is suitable for modeling radially symmetric.\n", - "- **Cartesian Shapelets:** Overview of cartesian shapelets for this example.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Lens Shapelets:** A model where the lens is modeled as shapelets can be composed and fitted as shown below.\n", - "- **Regularization:** There is one downside to `Basis` functions, we may compose a model with too much freedom.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals, because they fail to capture\n", - "irregular and asymmetric morphological of galaxies (e.g. isophotal twists, an ellipticity which varies radially).\n", - "Shapelets can capture some of these features and can therefore better represent the emission of complex source galaxies.\n", - "\n", - "The shapelet model can be composed in a way that has fewer non-linear parameters than an elliptical Sersic. In this\n", - "example, the ~20 shapelets which represent the `bulge` of the source are composed in a model corresponding to just\n", - "N=3 non-linear parameters (a `bulge` comprising a linear Sersic would give N=6).\n", - "\n", - "Therefore, shapelets fit more complex source galaxy morphologies using fewer non-linear parameters than the standard\n", - "light profile models.\n", - "\n", - "__Disadvantages__\n", - "\n", - "- There are many types of galaxy structure which shapelets may struggle to represent, such as a bar or asymmetric\n", - "knots of star formation. They also rely on the galaxy having a distinct centre over which the shapelets can be\n", - "centered, which is not the case if the galaxy is multiple merging systems or has bright companion galaxies.\n", - "\n", - "- The linear algebra used to solve for the `intensity` of each shapelet has to allow for negative values of intensity\n", - "in order for shapelets to work. Negative surface brightnesses are unphysical, and are often inferred in a shapelet\n", - "decomposition, for example if the true galaxy has structure that cannot be captured by the shapelet basis. Other\n", - "approaches can force positive-only intensities on the solution, such as the Multi-Gaussian Expansion (MGE) or a pixelization.\n", - "\n", - "- Computationally slower than standard light profiles like the Sersic.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's bulge is a superposition of `ShapeletPolar` profiles.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the galaxy dataset `light_basis` via .fits files, which we will fit with \n", - "the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "if not (dataset_path / \"positions.json\").exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular_annular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " inner_radius=0.4,\n", - " outer_radius=3.0,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positive Negative Solver__\n", - "\n", - "In other examples which use linear algebra to fit the data, for example linear light profiles, the Multi Gaussian\n", - "Expansion (MGE) and pixelization, we use a `positive_only` solver, which forces all solved for intensities to be\n", - "positive. This is a physical and sensible approach, because the surface brightnesses of a galaxy cannot be negative.\n", - "\n", - "Shapelets cannot be solved for using a `positive_only` solver, because the shapelets ability to decompose the\n", - "light of a galaxy relies on the ability to use negative intensities. This is because the shapelets are not\n", - "physically motivated light profiles, but instead a mathematical basis that can represent any light profile.\n", - "\n", - "This means shapelets may include negative flux in the reconstructed source galaxy, which is unphysical, and\n", - "a disadvantage of using shapelets.\n", - "\n", - "The `Settings` object below uses a `use_positive_only_solver=False` to allow for negative intensities.\n", - "\n", - "__Model__\n", - "\n", - "The shapelet decomposition above produced residuals, which we now rectify by fitting the shapelets in a non-linear \n", - "search, simultaneously fitting the lens's mass and source galaxies.\n", - "\n", - "We compose our model using `Model` objects, which represent the galaxies we fit to our data. In this \n", - "example we fit a model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - " - The source galaxy's bulge is a superposition of 10 linear `ShapeletPolar` profiles [3 parameters].\n", - " - The centres of the Shapelets are all linked together.\n", - " - The size of the Shapelet basis is controlled by a `beta` parameter, which is the same for all Shapelet basis \n", - " functions.\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=10.\n", - "\n", - "Note how this Shapelet model can capture features more complex than a Sersic, but has fewer non-linear parameters\n", - "(N=3 compared to N=7 for a `Sersic`).\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "total_n = 10\n", - "total_m = sum(range(2, total_n + 1)) + 1\n", - "\n", - "shapelets_bulge_list = af.Collection(\n", - " af.Model(al.lp_linear.ShapeletPolar) for _ in range(total_n + total_m + 1)\n", - ")\n", - "\n", - "n_count = 1\n", - "m_count = -1\n", - "\n", - "for i, shapelet in enumerate(shapelets_bulge_list):\n", - " if i == 0:\n", - " shapelet.n = 0\n", - " shapelet.m = 0\n", - "\n", - " else:\n", - " shapelet.n = n_count\n", - " shapelet.m = m_count\n", - "\n", - " m_count += 2\n", - "\n", - " if m_count > n_count:\n", - " n_count += 1\n", - " m_count = -n_count\n", - "\n", - " shapelet.centre = shapelets_bulge_list[0].centre\n", - " shapelet.ell_comps = shapelets_bulge_list[0].ell_comps\n", - " shapelet.beta = shapelets_bulge_list[0].beta\n", - "\n", - "bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=shapelets_bulge_list,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the source galaxy is made of many `ShapeletPolar` profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"shapelets\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Position Likelihood__\n", - "\n", - "We add a penalty term ot the likelihood function, which penalizes models where the brightest multiple images of\n", - "the lensed source galaxy do not trace close to one another in the source plane. This removes \"demagnified source\n", - "solutions\" from the source model, which one is likely to infer without this penalty.\n", - "\n", - "A comprehensive description of why we do this is given at the following readthedocs page. I strongly recommend you \n", - "read this page in full if you are not familiar with the positions likelihood penalty and demagnified source \n", - "reconstructions:\n", - "\n", - "NOTE: The description at the URL is for pixelized source reconstructions (e.g using a retangular or Delaunay mesh \n", - "to reconstruct the source), as opposed to shapelets. However, the same issues with demagnified solutions exist for\n", - "shapelets, whereby they go to large size (beta) values and reconstruct a demagnified image of the source for\n", - "mass models which do not accurately trace the multiple images to one another in the source plane.\n", - "\n", - " https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", - "\n", - "__Brief Description__\n", - "\n", - "Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH` object, which\n", - "includes the positions we loaded above, alongside a `threshold`.\n", - "\n", - "This is because `Inversion`'s suffer a bias whereby they fit unphysical lens models where the source galaxy is \n", - "reconstructed as a demagnified version of the lensed source. \n", - "\n", - "To prevent these solutions biasing the model-fit we specify a `position_threshold` of 0.5\", which requires that a \n", - "mass model traces the four (y,x) coordinates specified by our positions (that correspond to the brightest regions of the \n", - "lensed source) within 0.5\" of one another in the source-plane. If this criteria is not met, a large penalty term is\n", - "added to likelihood that massively reduces the overall likelihood. This penalty is larger if the ``positions``\n", - "trace further from one another.\n", - "\n", - "This ensures the unphysical solutions that bias a pixelization have a lower likelihood that the physical solutions\n", - "we desire. Furthermore, the penalty term reduces as the image-plane multiple image positions trace closer in the \n", - "source-plane, ensuring Nautilus converges towards an accurate mass model. It does this very fast, as \n", - "ray-tracing just a few multiple image positions is computationally cheap. \n", - "\n", - "The threshold of 0.3\" is large. For an accurate lens model we would anticipate the positions trace within < 0.01\" of\n", - "one another. The high threshold ensures only the initial mass models at the start of the fit are penalized.\n", - "\n", - "Position thresholding is described in more detail in the \n", - "script `autolens_workspace/*/guides/modeling/customize`\n", - "\n", - "The arc-second positions of the multiply imaged lensed source galaxy were drawn onto the\n", - "image via the GUI described in the file `autolens_workspace/*/imaging/data_preparation/gui/positions.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n", - "\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " settings=al.Settings(use_positive_only_solver=False),\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", - "required by a model.\n", - "\n", - "For each shapelet, extra VRAM is used. For around 60 shapelets this typically requires a modest amount of \n", - "VRAM (e.g. 10\u201350 MB per batched likelihood). Models that use hundreds of shapelets, especially in combination with a \n", - "large batch size, may therefore exceed GBs of VRAM and require you to adjust the batch size to fit within your GPU's VRAM.\n", - "\n", - "__Run Time__\n", - "\n", - "The likelihood evaluation time for a shapelets is significantly slower than standard light profiles.\n", - "This is because the image of every shapelets must be computed and evaluated, and each must be blurred with the PSF.\n", - "In this example, the evaluation time is ~0.37s, compared to ~0.01 seconds for standard light profiles.\n", - "\n", - "Gains in the overall run-time however are made thanks to the models reduced complexity and lower\n", - "number of free parameters. The source is modeled with 3 free parameters, compared to 6+ for a linear light profile \n", - "Sersic.\n", - "\n", - "However, the multi-gaussian expansion (MGE) approach is even faster than shapelets. It uses fewer Gaussian basis\n", - "functions (speed up the likelihood evaluation) and has fewer free parameters (speeding up the non-linear search).\n", - "Furthermore, none of the free parameters scale the size of the source galaxy, which means the non-linear search\n", - "can converge faster.\n", - "\n", - "I recommend you try using an MGE approach alongside shapelets. For many science cases, the MGE approach will be\n", - "faster and give higher quality results. Shapelets may perform better for irregular sources, but this is not\n", - "guaranteed.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_galaxies_images(tracer=result.max_log_likelihood_tracer, grid=dataset.grid)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Shapelet Cartesian__\n", - "\n", - "The shapelets above were defined on a polar grid, which is suitable for modeling radially symmetric sources like\n", - "most galaxies.\n", - "\n", - "An alternative approach is to define the shapelets on a Cartesian grid, which we plot the basis of below\n", - "and show an example fit.\n", - "\n", - "These are generally not recommended for modeling galaxies, but may be better in certain situations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_xy = 5\n", - "\n", - "shapelets_bulge_list = []\n", - "\n", - "for x in range(total_xy):\n", - " for y in range(total_xy):\n", - " shapelet = al.lp.ShapeletCartesian(\n", - " n_y=y,\n", - " n_x=x,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=1.0,\n", - " beta=1.0,\n", - " )\n", - "\n", - " shapelets_bulge_list.append(shapelet)\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "subplot_basis_image(basis=bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cartesian Shapelets__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_xy = 5\n", - "\n", - "shapelets_bulge_list = []\n", - "\n", - "for x in range(total_xy):\n", - " for y in range(total_xy):\n", - " shapelet = al.lp_linear.ShapeletCartesian(\n", - " n_y=y, n_x=x, centre=(0.0, 0.0), ell_comps=(0.0, 0.0), beta=1.0\n", - " )\n", - "\n", - " shapelets_bulge_list.append(shapelet)\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=bulge,\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_positive_only_solver=False),\n", - ")\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer.galaxies[1].bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "Here is how we compose a model using Cartesian shapelets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "total_xy = 5\n", - "\n", - "shapelets_bulge_list = af.Collection(\n", - " af.Model(al.lp_linear.ShapeletCartesian) for _ in range(total_xy**2)\n", - ")\n", - "\n", - "for x in range(total_xy):\n", - " for y in range(total_xy):\n", - " shapelet.n_y = y\n", - " shapelet.n_x = x\n", - "\n", - " shapelet.centre = shapelets_bulge_list[0].centre\n", - " shapelet.ell_comps = shapelets_bulge_list[0].ell_comps\n", - " shapelet.beta = shapelets_bulge_list[0].beta\n", - "\n", - "bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=shapelets_bulge_list,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "print(model.info)\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"shapelets_cartesian\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Shapelets__\n", - "\n", - "A model where the lens is modeled as shapelets can be composed and fitted as shown below.\n", - "\n", - "I have not seen this model used in the literature, and am not clear on its advantages over a standard light profile\n", - "model. However, it is worth trying if you are fitting a lens galaxy with a complex morphology.\n", - "\n", - "For most massive early-type galaxies, an MGE model will be faster and give higher quality results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=shapelets_bulge_list,\n", - ")\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=shapelets_bulge_list,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results in **Pyautolens**, which\n", - "includes a dedicated tutorial for linear objects like basis functions.\n", - "\n", - "__Basis Regularization (Advanced / Unused)__\n", - "\n", - "A shapelet `Basis` can additionally carry a regularization term (e.g. `al.reg.Constant`) that penalises non-smooth\n", - "intensity solutions. This is a research-only feature and is not used by any production scientific analysis.\n", - "\n", - "The code and rationale for the regularization branch have been moved out of this user-facing script. If you want to\n", - "experiment with adding a regularization to a shapelet `Basis`, see:\n", - "\n", - " autolens_workspace_developer/basis_regularization/shapelets_lens.py\n", - "\n", - "__Wrap Up__\n", - "\n", - "This script has illustrated how to use shapelets to model the light of galaxies.\n", - "\n", - "Shapelets are a powerful basis function for capturing complex morphological features of galaxies that standard\n", - "light profiles struggle to represent. However, they do have drawbacks, such as the need to allow for negative\n", - "intensities in the solution, which is unphysical. \n", - "\n", - "As a rule of thumb, modeling is generally better if a pixelization is used to reconstruct the source galaxy's light,\n", - "but shapelets can be a useful middle-ground between standard light profiles and a pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Shapelets\n", + "============================\n", + "\n", + "A shapelet is a basis function that is appropriate for capturing the exponential / disk-like features of a galaxy. It\n", + "has been employed in many strong lensing studies to model the light of the lensed source galaxy, because it can\n", + "represent features of disky star forming galaxies that a single Sersic function cannot.\n", + "\n", + "- https://ui.adsabs.harvard.edu/abs/2016MNRAS.457.3066T\n", + "- https://iopscience.iop.org/article/10.1088/0004-637X/813/2/102\n", + "\n", + "Shapelets are described in full in the following paper:\n", + "\n", + " https://arxiv.org/abs/astro-ph/0105178\n", + "\n", + "This script performs a model-fit using shapelets, where it decomposes the galaxy light into ~20\n", + "Shapelets. The `intensity` of every Shapelet is solved for via linear algebra (see the `linear_light_profiles.py`\n", + "feature).\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Symmetric light profiles (e.g.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Positive Negative Solver:** In other examples which use linear algebra to fit the data, for example linear light profiles, the.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Position Likelihood:** We add a penalty term ot the likelihood function, which penalizes models where the brightest.\n", + "- **Brief Description:** Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH`.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Shapelet Cartesian:** The shapelets above were defined on a polar grid, which is suitable for modeling radially symmetric.\n", + "- **Cartesian Shapelets:** Overview of cartesian shapelets for this example.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Lens Shapelets:** A model where the lens is modeled as shapelets can be composed and fitted as shown below.\n", + "- **Regularization:** There is one downside to `Basis` functions, we may compose a model with too much freedom.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals, because they fail to capture\n", + "irregular and asymmetric morphological of galaxies (e.g. isophotal twists, an ellipticity which varies radially).\n", + "Shapelets can capture some of these features and can therefore better represent the emission of complex source galaxies.\n", + "\n", + "The shapelet model can be composed in a way that has fewer non-linear parameters than an elliptical Sersic. In this\n", + "example, the ~20 shapelets which represent the `bulge` of the source are composed in a model corresponding to just\n", + "N=3 non-linear parameters (a `bulge` comprising a linear Sersic would give N=6).\n", + "\n", + "Therefore, shapelets fit more complex source galaxy morphologies using fewer non-linear parameters than the standard\n", + "light profile models.\n", + "\n", + "__Disadvantages__\n", + "\n", + "- There are many types of galaxy structure which shapelets may struggle to represent, such as a bar or asymmetric\n", + "knots of star formation. They also rely on the galaxy having a distinct centre over which the shapelets can be\n", + "centered, which is not the case if the galaxy is multiple merging systems or has bright companion galaxies.\n", + "\n", + "- The linear algebra used to solve for the `intensity` of each shapelet has to allow for negative values of intensity\n", + "in order for shapelets to work. Negative surface brightnesses are unphysical, and are often inferred in a shapelet\n", + "decomposition, for example if the true galaxy has structure that cannot be captured by the shapelet basis. Other\n", + "approaches can force positive-only intensities on the solution, such as the Multi-Gaussian Expansion (MGE) or a pixelization.\n", + "\n", + "- Computationally slower than standard light profiles like the Sersic.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's bulge is a superposition of `ShapeletPolar` profiles.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the galaxy dataset `light_basis` via .fits files, which we will fit with \n", + "the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "if not (dataset_path / \"positions.json\").exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular_annular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " inner_radius=0.4,\n", + " outer_radius=3.0,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positive Negative Solver__\n", + "\n", + "In other examples which use linear algebra to fit the data, for example linear light profiles, the Multi Gaussian\n", + "Expansion (MGE) and pixelization, we use a `positive_only` solver, which forces all solved for intensities to be\n", + "positive. This is a physical and sensible approach, because the surface brightnesses of a galaxy cannot be negative.\n", + "\n", + "Shapelets cannot be solved for using a `positive_only` solver, because the shapelets ability to decompose the\n", + "light of a galaxy relies on the ability to use negative intensities. This is because the shapelets are not\n", + "physically motivated light profiles, but instead a mathematical basis that can represent any light profile.\n", + "\n", + "This means shapelets may include negative flux in the reconstructed source galaxy, which is unphysical, and\n", + "a disadvantage of using shapelets.\n", + "\n", + "The `Settings` object below uses a `use_positive_only_solver=False` to allow for negative intensities.\n", + "\n", + "__Model__\n", + "\n", + "The shapelet decomposition above produced residuals, which we now rectify by fitting the shapelets in a non-linear \n", + "search, simultaneously fitting the lens's mass and source galaxies.\n", + "\n", + "We compose our model using `Model` objects, which represent the galaxies we fit to our data. In this \n", + "example we fit a model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + " - The source galaxy's bulge is a superposition of 10 linear `ShapeletPolar` profiles [3 parameters].\n", + " - The centres of the Shapelets are all linked together.\n", + " - The size of the Shapelet basis is controlled by a `beta` parameter, which is the same for all Shapelet basis \n", + " functions.\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=10.\n", + "\n", + "Note how this Shapelet model can capture features more complex than a Sersic, but has fewer non-linear parameters\n", + "(N=3 compared to N=7 for a `Sersic`).\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "total_n = 10\n", + "total_m = sum(range(2, total_n + 1)) + 1\n", + "\n", + "shapelets_bulge_list = af.Collection(\n", + " af.Model(al.lp_linear.ShapeletPolar) for _ in range(total_n + total_m + 1)\n", + ")\n", + "\n", + "n_count = 1\n", + "m_count = -1\n", + "\n", + "for i, shapelet in enumerate(shapelets_bulge_list):\n", + " if i == 0:\n", + " shapelet.n = 0\n", + " shapelet.m = 0\n", + "\n", + " else:\n", + " shapelet.n = n_count\n", + " shapelet.m = m_count\n", + "\n", + " m_count += 2\n", + "\n", + " if m_count > n_count:\n", + " n_count += 1\n", + " m_count = -n_count\n", + "\n", + " shapelet.centre = shapelets_bulge_list[0].centre\n", + " shapelet.ell_comps = shapelets_bulge_list[0].ell_comps\n", + " shapelet.beta = shapelets_bulge_list[0].beta\n", + "\n", + "bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=shapelets_bulge_list,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the source galaxy is made of many `ShapeletPolar` profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"shapelets\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Position Likelihood__\n", + "\n", + "We add a penalty term ot the likelihood function, which penalizes models where the brightest multiple images of\n", + "the lensed source galaxy do not trace close to one another in the source plane. This removes \"demagnified source\n", + "solutions\" from the source model, which one is likely to infer without this penalty.\n", + "\n", + "A comprehensive description of why we do this is given at the following readthedocs page. I strongly recommend you \n", + "read this page in full if you are not familiar with the positions likelihood penalty and demagnified source \n", + "reconstructions:\n", + "\n", + "NOTE: The description at the URL is for pixelized source reconstructions (e.g using a retangular or Delaunay mesh \n", + "to reconstruct the source), as opposed to shapelets. However, the same issues with demagnified solutions exist for\n", + "shapelets, whereby they go to large size (beta) values and reconstruct a demagnified image of the source for\n", + "mass models which do not accurately trace the multiple images to one another in the source plane.\n", + "\n", + " https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", + "\n", + "__Brief Description__\n", + "\n", + "Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH` object, which\n", + "includes the positions we loaded above, alongside a `threshold`.\n", + "\n", + "This is because `Inversion`'s suffer a bias whereby they fit unphysical lens models where the source galaxy is \n", + "reconstructed as a demagnified version of the lensed source. \n", + "\n", + "To prevent these solutions biasing the model-fit we specify a `position_threshold` of 0.5\", which requires that a \n", + "mass model traces the four (y,x) coordinates specified by our positions (that correspond to the brightest regions of the \n", + "lensed source) within 0.5\" of one another in the source-plane. If this criteria is not met, a large penalty term is\n", + "added to likelihood that massively reduces the overall likelihood. This penalty is larger if the ``positions``\n", + "trace further from one another.\n", + "\n", + "This ensures the unphysical solutions that bias a pixelization have a lower likelihood that the physical solutions\n", + "we desire. Furthermore, the penalty term reduces as the image-plane multiple image positions trace closer in the \n", + "source-plane, ensuring Nautilus converges towards an accurate mass model. It does this very fast, as \n", + "ray-tracing just a few multiple image positions is computationally cheap. \n", + "\n", + "The threshold of 0.3\" is large. For an accurate lens model we would anticipate the positions trace within < 0.01\" of\n", + "one another. The high threshold ensures only the initial mass models at the start of the fit are penalized.\n", + "\n", + "Position thresholding is described in more detail in the \n", + "script `autolens_workspace/*/guides/modeling/customize`\n", + "\n", + "The arc-second positions of the multiply imaged lensed source galaxy were drawn onto the\n", + "image via the GUI described in the file `autolens_workspace/*/imaging/data_preparation/gui/positions.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n", + "\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " settings=al.Settings(use_positive_only_solver=False),\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", + "required by a model.\n", + "\n", + "For each shapelet, extra VRAM is used. For around 60 shapelets this typically requires a modest amount of \n", + "VRAM (e.g. 10\u201350 MB per batched likelihood). Models that use hundreds of shapelets, especially in combination with a \n", + "large batch size, may therefore exceed GBs of VRAM and require you to adjust the batch size to fit within your GPU's VRAM.\n", + "\n", + "__Run Time__\n", + "\n", + "The likelihood evaluation time for a shapelets is significantly slower than standard light profiles.\n", + "This is because the image of every shapelets must be computed and evaluated, and each must be blurred with the PSF.\n", + "In this example, the evaluation time is ~0.37s, compared to ~0.01 seconds for standard light profiles.\n", + "\n", + "Gains in the overall run-time however are made thanks to the models reduced complexity and lower\n", + "number of free parameters. The source is modeled with 3 free parameters, compared to 6+ for a linear light profile \n", + "Sersic.\n", + "\n", + "However, the multi-gaussian expansion (MGE) approach is even faster than shapelets. It uses fewer Gaussian basis\n", + "functions (speed up the likelihood evaluation) and has fewer free parameters (speeding up the non-linear search).\n", + "Furthermore, none of the free parameters scale the size of the source galaxy, which means the non-linear search\n", + "can converge faster.\n", + "\n", + "I recommend you try using an MGE approach alongside shapelets. For many science cases, the MGE approach will be\n", + "faster and give higher quality results. Shapelets may perform better for irregular sources, but this is not\n", + "guaranteed.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_galaxies_images(tracer=result.max_log_likelihood_tracer, grid=dataset.grid)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Shapelet Cartesian__\n", + "\n", + "The shapelets above were defined on a polar grid, which is suitable for modeling radially symmetric sources like\n", + "most galaxies.\n", + "\n", + "An alternative approach is to define the shapelets on a Cartesian grid, which we plot the basis of below\n", + "and show an example fit.\n", + "\n", + "These are generally not recommended for modeling galaxies, but may be better in certain situations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_xy = 5\n", + "\n", + "shapelets_bulge_list = []\n", + "\n", + "for x in range(total_xy):\n", + " for y in range(total_xy):\n", + " shapelet = al.lp.ShapeletCartesian(\n", + " n_y=y,\n", + " n_x=x,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=1.0,\n", + " beta=1.0,\n", + " )\n", + "\n", + " shapelets_bulge_list.append(shapelet)\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "subplot_basis_image(basis=bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cartesian Shapelets__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_xy = 5\n", + "\n", + "shapelets_bulge_list = []\n", + "\n", + "for x in range(total_xy):\n", + " for y in range(total_xy):\n", + " shapelet = al.lp_linear.ShapeletCartesian(\n", + " n_y=y, n_x=x, centre=(0.0, 0.0), ell_comps=(0.0, 0.0), beta=1.0\n", + " )\n", + "\n", + " shapelets_bulge_list.append(shapelet)\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=bulge,\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_positive_only_solver=False),\n", + ")\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer.galaxies[1].bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "Here is how we compose a model using Cartesian shapelets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "total_xy = 5\n", + "\n", + "shapelets_bulge_list = af.Collection(\n", + " af.Model(al.lp_linear.ShapeletCartesian) for _ in range(total_xy**2)\n", + ")\n", + "\n", + "for x in range(total_xy):\n", + " for y in range(total_xy):\n", + " shapelet.n_y = y\n", + " shapelet.n_x = x\n", + "\n", + " shapelet.centre = shapelets_bulge_list[0].centre\n", + " shapelet.ell_comps = shapelets_bulge_list[0].ell_comps\n", + " shapelet.beta = shapelets_bulge_list[0].beta\n", + "\n", + "bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=shapelets_bulge_list,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"shapelets_cartesian\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Shapelets__\n", + "\n", + "A model where the lens is modeled as shapelets can be composed and fitted as shown below.\n", + "\n", + "I have not seen this model used in the literature, and am not clear on its advantages over a standard light profile\n", + "model. However, it is worth trying if you are fitting a lens galaxy with a complex morphology.\n", + "\n", + "For most massive early-type galaxies, an MGE model will be faster and give higher quality results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=shapelets_bulge_list,\n", + ")\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=shapelets_bulge_list,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results in **Pyautolens**, which\n", + "includes a dedicated tutorial for linear objects like basis functions.\n", + "\n", + "__Basis Regularization (Advanced / Unused)__\n", + "\n", + "A shapelet `Basis` can additionally carry a regularization term (e.g. `al.reg.Constant`) that penalises non-smooth\n", + "intensity solutions. This is a research-only feature and is not used by any production scientific analysis.\n", + "\n", + "The code and rationale for the regularization branch have been moved out of this user-facing script. If you want to\n", + "experiment with adding a regularization to a shapelet `Basis`, see:\n", + "\n", + " autolens_workspace_developer/basis_regularization/shapelets_lens.py\n", + "\n", + "__Wrap Up__\n", + "\n", + "This script has illustrated how to use shapelets to model the light of galaxies.\n", + "\n", + "Shapelets are a powerful basis function for capturing complex morphological features of galaxies that standard\n", + "light profiles struggle to represent. However, they do have drawbacks, such as the need to allow for negative\n", + "intensities in the solution, which is unphysical. \n", + "\n", + "As a rule of thumb, modeling is generally better if a pixelization is used to reconstruct the source galaxy's light,\n", + "but shapelets can be a useful middle-ground between standard light profiles and a pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/sky_background/fit.ipynb b/notebooks/imaging/features/advanced/sky_background/fit.ipynb index 90f87e936..466070c9b 100644 --- a/notebooks/imaging/features/advanced/sky_background/fit.ipynb +++ b/notebooks/imaging/features/advanced/sky_background/fit.ipynb @@ -1,301 +1,338 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Sky Background\n", - "=================================\n", - "\n", - "The background of an image is the light that is not associated with the strong lens we are interested in. This is due to\n", - "light from the sky, zodiacal light, and light from other galaxies in the field of view.\n", - "\n", - "The background sky is often subtracted from image data during the data reduction procedure. If this subtraction is\n", - "perfect, there is then no need to include the sky in the model-fitting. However, it is difficult to achieve a perfect\n", - "subtraction and there is some uncertainty in the procedure.\n", - "\n", - "The residuals of an imperfect back sky subtraction can leave a signal in the image which is degenerate with the\n", - "light profile of the lens galaxy. This is especially true for low surface brightness features, such as the faint\n", - "outskirts of the galaxy.\n", - "\n", - "Fitting the sky can therefore ensure errors on light profile parameters which fit the low surface brightness features\n", - "further out, like the effective radius and Sersic index, fully account for the uncertainties in the sky background.\n", - "\n", - "This example script illustrates how to include the sky background in the model-fitting of an `Imaging` dataset as\n", - "a non-linear free parameter (e.g. an extra dimension in the non-linear parameter space).\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a galaxy with a model where:\n", - "\n", - " - The sky background is included as part of a `DatasetModel`.\n", - " - The lens galaxy's light is a linear `Sersic` bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a linear `SersicCore`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `fit` examples." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the galaxy dataset `sky_background` via .fits files, which we will fit with the model.\n", - "\n", - "This dataset has not had the sky background subtracted from it, therefore the sky background is included in the\n", - "image data when we fit it. \n", - "\n", - "This is seen clearly in the plot, where the outskirts of the image do not go to values near 0.0 electrons per second\n", - "like other datasets but instead have values of 5.0 electrons per second, the sky background level used to simulate\n", - "the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"sky_background\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/sky_background/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We first show how to use a `DatasetModel` object to fit the sky background in the data.\n", - "\n", - "This illustrates the API for performing a sky background fit using standard objects like a `Galaxy` and `FitImaging` .\n", - "\n", - "This does not perform a model-fit via a non-linear search, and therefore requires us to manually specify and guess\n", - "suitable parameter values for the sky. We will use the true value of 5.0 electrons per second.\n", - "\n", - "For the galaxies, we will use the true parameters used to simulate the data, for illustrative purposes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "dataset_model = al.DatasetModel(background_sky_level=5.0)\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer, dataset_model=dataset_model)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the sky is subtracted from the data such that the outskirts are zero.\n", - "\n", - "There are few residuals, except for perhaps some central regions where the light profile is not perfectly fitted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This example shows how to include the sky background as part of a fit using a `DatasetModel` object.\n", - "\n", - "It is useful for ensuring uncertainties on lens galaxy light profile parameters fully account for uncertainties\n", - "in the sky background subtraction, especially for low surface brightness features in the outskirts of the galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Sky Background\n", + "=================================\n", + "\n", + "The background of an image is the light that is not associated with the strong lens we are interested in. This is due to\n", + "light from the sky, zodiacal light, and light from other galaxies in the field of view.\n", + "\n", + "The background sky is often subtracted from image data during the data reduction procedure. If this subtraction is\n", + "perfect, there is then no need to include the sky in the model-fitting. However, it is difficult to achieve a perfect\n", + "subtraction and there is some uncertainty in the procedure.\n", + "\n", + "The residuals of an imperfect back sky subtraction can leave a signal in the image which is degenerate with the\n", + "light profile of the lens galaxy. This is especially true for low surface brightness features, such as the faint\n", + "outskirts of the galaxy.\n", + "\n", + "Fitting the sky can therefore ensure errors on light profile parameters which fit the low surface brightness features\n", + "further out, like the effective radius and Sersic index, fully account for the uncertainties in the sky background.\n", + "\n", + "This example script illustrates how to include the sky background in the model-fitting of an `Imaging` dataset as\n", + "a non-linear free parameter (e.g. an extra dimension in the non-linear parameter space).\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a galaxy with a model where:\n", + "\n", + " - The sky background is included as part of a `DatasetModel`.\n", + " - The lens galaxy's light is a linear `Sersic` bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a linear `SersicCore`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `fit` examples." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the galaxy dataset `sky_background` via .fits files, which we will fit with the model.\n", + "\n", + "This dataset has not had the sky background subtracted from it, therefore the sky background is included in the\n", + "image data when we fit it. \n", + "\n", + "This is seen clearly in the plot, where the outskirts of the image do not go to values near 0.0 electrons per second\n", + "like other datasets but instead have values of 5.0 electrons per second, the sky background level used to simulate\n", + "the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"sky_background\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/sky_background/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We first show how to use a `DatasetModel` object to fit the sky background in the data.\n", + "\n", + "This illustrates the API for performing a sky background fit using standard objects like a `Galaxy` and `FitImaging` .\n", + "\n", + "This does not perform a model-fit via a non-linear search, and therefore requires us to manually specify and guess\n", + "suitable parameter values for the sky. We will use the true value of 5.0 electrons per second.\n", + "\n", + "For the galaxies, we will use the true parameters used to simulate the data, for illustrative purposes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "dataset_model = al.DatasetModel(background_sky_level=5.0)\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer, dataset_model=dataset_model)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the sky is subtracted from the data such that the outskirts are zero.\n", + "\n", + "There are few residuals, except for perhaps some central regions where the light profile is not perfectly fitted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This example shows how to include the sky background as part of a fit using a `DatasetModel` object.\n", + "\n", + "It is useful for ensuring uncertainties on lens galaxy light profile parameters fully account for uncertainties\n", + "in the sky background subtraction, especially for low surface brightness features in the outskirts of the galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/sky_background/modeling.ipynb b/notebooks/imaging/features/advanced/sky_background/modeling.ipynb index 96c60dfe7..a972d43d4 100644 --- a/notebooks/imaging/features/advanced/sky_background/modeling.ipynb +++ b/notebooks/imaging/features/advanced/sky_background/modeling.ipynb @@ -1,434 +1,471 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Sky Background\n", - "=================================\n", - "\n", - "The background of an image is the light that is not associated with the strong lens we are interested in. This is due to\n", - "light from the sky, zodiacal light, and light from other galaxies in the field of view.\n", - "\n", - "The background sky is often subtracted from image data during the data reduction procedure. If this subtraction is\n", - "perfect, there is then no need to include the sky in the model-fitting. However, it is difficult to achieve a perfect\n", - "subtraction and there is some uncertainty in the procedure.\n", - "\n", - "The residuals of an imperfect back sky subtraction can leave a signal in the image which is degenerate with the\n", - "light profile of the lens galaxy. This is especially true for low surface brightness features, such as the faint\n", - "outskirts of the galaxy.\n", - "\n", - "Fitting the sky can therefore ensure errors on light profile parameters which fit the low surface brightness features\n", - "further out, like the effective radius and Sersic index, fully account for the uncertainties in the sky background.\n", - "\n", - "This example script illustrates how to include the sky background in the model-fitting of an `Imaging` dataset as\n", - "a non-linear free parameter (e.g. an extra dimension in the non-linear parameter space).\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a galaxy with a model where:\n", - "\n", - " - The sky background is included as part of a `DatasetModel`.\n", - " - The lens galaxy's light is a linear `Sersic` bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a linear `SersicCore`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the galaxy dataset `sky_background` via .fits files, which we will fit with the model.\n", - "\n", - "This dataset has not had the sky background subtracted from it, therefore the sky background is included in the\n", - "image data when we fit it. \n", - "\n", - "This is seen clearly in the plot, where the outskirts of the image do not go to values near 0.0 electrons per second\n", - "like other datasets but instead have values of 5.0 electrons per second, the sky background level used to simulate\n", - "the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"sky_background\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/advanced/sky_background/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "In this example we compose a lens model where:\n", - "\n", - " - The sky background is included as a `DatasetModel` [1 parameter].\n", - "\n", - " - The lens galaxy's light is a linear `Sersic` bulge [6 parameters].\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is a linear `SersicCore` [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=22.\n", - "\n", - "The sky is not included in the `galaxies` collection, but is its own separate component in the overall model.\n", - "\n", - "We update the prior on the `background_sky_level` manually, such that it surrounds the true value of 5.0 electrons\n", - "per second. \n", - "\n", - "You must always update the prior on the sky's intensity manually (unlike light profile priors), because the appropriate\n", - "prior depends on the dataset being fitted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - ")\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "dataset_model = af.Model(al.DatasetModel)\n", - "dataset_model.background_sky_level = af.UniformPrior(lower_limit=0.0, upper_limit=5.0)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " dataset_model=dataset_model, galaxies=af.Collection(lens=lens, source=source)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the sky is a model component that is not part of the `galaxies` collection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"sky_background\",\n", - " unique_tag=dataset_name,\n", - " n_live=125,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", - "required by a model.\n", - "\n", - "The sky background model has a negligible impact on the VRAM required for the model-fit.\n", - "\n", - "__Run Time__\n", - "\n", - "For standard light profiles, the log likelihood evaluation time is of order ~0.01 seconds for this dataset.\n", - "\n", - "Adding the background sky model to the analysis has a negligible impact on the run time, as it requires simply adding\n", - "a constant value to the data. The run time is therefore still of order ~0.01 seconds.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that a `background_sky_level` of approximately 5.0 electrons per second was inferred, as expected." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To print the exact value, the `background_sky_level` attribute of the result contains the `intensity` of the sky." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.instance.dataset_model.background_sky_level)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This example shows how to include the sky background as part of a fit using a `DatasetModel` object.\n", - "\n", - "It is useful for ensuring uncertainties on lens galaxy light profile parameters fully account for uncertainties\n", - "in the sky background subtraction, especially for low surface brightness features in the outskirts of the galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Sky Background\n", + "=================================\n", + "\n", + "The background of an image is the light that is not associated with the strong lens we are interested in. This is due to\n", + "light from the sky, zodiacal light, and light from other galaxies in the field of view.\n", + "\n", + "The background sky is often subtracted from image data during the data reduction procedure. If this subtraction is\n", + "perfect, there is then no need to include the sky in the model-fitting. However, it is difficult to achieve a perfect\n", + "subtraction and there is some uncertainty in the procedure.\n", + "\n", + "The residuals of an imperfect back sky subtraction can leave a signal in the image which is degenerate with the\n", + "light profile of the lens galaxy. This is especially true for low surface brightness features, such as the faint\n", + "outskirts of the galaxy.\n", + "\n", + "Fitting the sky can therefore ensure errors on light profile parameters which fit the low surface brightness features\n", + "further out, like the effective radius and Sersic index, fully account for the uncertainties in the sky background.\n", + "\n", + "This example script illustrates how to include the sky background in the model-fitting of an `Imaging` dataset as\n", + "a non-linear free parameter (e.g. an extra dimension in the non-linear parameter space).\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a galaxy with a model where:\n", + "\n", + " - The sky background is included as part of a `DatasetModel`.\n", + " - The lens galaxy's light is a linear `Sersic` bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a linear `SersicCore`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the galaxy dataset `sky_background` via .fits files, which we will fit with the model.\n", + "\n", + "This dataset has not had the sky background subtracted from it, therefore the sky background is included in the\n", + "image data when we fit it. \n", + "\n", + "This is seen clearly in the plot, where the outskirts of the image do not go to values near 0.0 electrons per second\n", + "like other datasets but instead have values of 5.0 electrons per second, the sky background level used to simulate\n", + "the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"sky_background\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/advanced/sky_background/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "In this example we compose a lens model where:\n", + "\n", + " - The sky background is included as a `DatasetModel` [1 parameter].\n", + "\n", + " - The lens galaxy's light is a linear `Sersic` bulge [6 parameters].\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is a linear `SersicCore` [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=22.\n", + "\n", + "The sky is not included in the `galaxies` collection, but is its own separate component in the overall model.\n", + "\n", + "We update the prior on the `background_sky_level` manually, such that it surrounds the true value of 5.0 electrons\n", + "per second. \n", + "\n", + "You must always update the prior on the sky's intensity manually (unlike light profile priors), because the appropriate\n", + "prior depends on the dataset being fitted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + ")\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "dataset_model = af.Model(al.DatasetModel)\n", + "dataset_model.background_sky_level = af.UniformPrior(lower_limit=0.0, upper_limit=5.0)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " dataset_model=dataset_model, galaxies=af.Collection(lens=lens, source=source)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the sky is a model component that is not part of the `galaxies` collection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"sky_background\",\n", + " unique_tag=dataset_name,\n", + " n_live=125,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", + "required by a model.\n", + "\n", + "The sky background model has a negligible impact on the VRAM required for the model-fit.\n", + "\n", + "__Run Time__\n", + "\n", + "For standard light profiles, the log likelihood evaluation time is of order ~0.01 seconds for this dataset.\n", + "\n", + "Adding the background sky model to the analysis has a negligible impact on the run time, as it requires simply adding\n", + "a constant value to the data. The run time is therefore still of order ~0.01 seconds.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that a `background_sky_level` of approximately 5.0 electrons per second was inferred, as expected." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To print the exact value, the `background_sky_level` attribute of the result contains the `intensity` of the sky." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.instance.dataset_model.background_sky_level)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This example shows how to include the sky background as part of a fit using a `DatasetModel` object.\n", + "\n", + "It is useful for ensuring uncertainties on lens galaxy light profile parameters fully account for uncertainties\n", + "in the sky background subtraction, especially for low surface brightness features in the outskirts of the galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/sky_background/simulator.ipynb b/notebooks/imaging/features/advanced/sky_background/simulator.ipynb index 2be6990f6..2c9384ae5 100644 --- a/notebooks/imaging/features/advanced/sky_background/simulator.ipynb +++ b/notebooks/imaging/features/advanced/sky_background/simulator.ipynb @@ -1,360 +1,397 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Start Here\n", - "=====================\n", - "\n", - "This script simulates `Imaging` of a strong lens where the sky background is not subtracted from the image and therefore\n", - "appears in the dataset.\n", - "\n", - "It is used to demonstrate sky background modeling in\n", - "the `autolens_workspace/*/modeling/features/sky_background.py` example.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's light profile is a `Sersic`.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a `Sersic`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"sky_background\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf.\n", - "\n", - "The input `subtract_background_sky=False` ensures the background sky is not subtracted from the image, but is \n", - "included in the data and used to estimate its noise-map.\n", - "\n", - "The `background_sky_level` is also increased to 5.0 electrons per second, which is much higer than the value of \n", - "0.1 electrons per second used in previous examples. This is to better illustrate the sky background in the image\n", - "and how it can be modeled." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=5.0,\n", - " add_poisson_noise_to_data=True,\n", - " subtract_background_sky=False,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now plot the simulated `Imaging` dataset before outputting it to fits.\n", - "\n", - "Note how unlike the `Tracer` image above, the simulated `Imaging` dataset includes the blurring effects of the \n", - "telescope's PSF and also has noise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/sky_background`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Start Here\n", + "=====================\n", + "\n", + "This script simulates `Imaging` of a strong lens where the sky background is not subtracted from the image and therefore\n", + "appears in the dataset.\n", + "\n", + "It is used to demonstrate sky background modeling in\n", + "the `autolens_workspace/*/modeling/features/sky_background.py` example.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's light profile is a `Sersic`.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a `Sersic`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"sky_background\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf.\n", + "\n", + "The input `subtract_background_sky=False` ensures the background sky is not subtracted from the image, but is \n", + "included in the data and used to estimate its noise-map.\n", + "\n", + "The `background_sky_level` is also increased to 5.0 electrons per second, which is much higer than the value of \n", + "0.1 electrons per second used in previous examples. This is to better illustrate the sky background in the image\n", + "and how it can be modeled." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=5.0,\n", + " add_poisson_noise_to_data=True,\n", + " subtract_background_sky=False,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now plot the simulated `Imaging` dataset before outputting it to fits.\n", + "\n", + "Note how unlike the `Tracer` image above, the simulated `Imaging` dataset includes the blurring effects of the \n", + "telescope's PSF and also has noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/sky_background`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/subhalo/detect/database.ipynb b/notebooks/imaging/features/advanced/subhalo/detect/database.ipynb index 3f353915d..ebcf07b9e 100644 --- a/notebooks/imaging/features/advanced/subhalo/detect/database.ipynb +++ b/notebooks/imaging/features/advanced/subhalo/detect/database.ipynb @@ -1,384 +1,421 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Subhalo Detection: Database\n", - "===========================\n", - "\n", - "The example `subhalo/detect/start_here.ipynb` shows how to perform dark matter subhalo detection in strong lens\n", - "with **PyAutoLens**, including using results to inspect and visualize the fit.\n", - "\n", - "This example shows how to load the results of subhalo detection analysis into a `.sqlite` database, which can be\n", - "manipulated stand-alone in this Python script or in a Jupyter notebook. This is useful when fits are performed on a\n", - "super computer and results are downloaded separately for inspection.\n", - "\n", - "The database in this example is built by scraping the results of the `subhalo/detect/start_here.ipynb` example. You\n", - "can also write results directly to the database during the fit by using a session.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Start Here Notebooks:** You should be familiar with dark matter subhalo detection, by reading the example.\n", - "- **Grid Searches:** If the results of the database include a grid search of non-linear searches, the aggregator has a.\n", - "- **Grid Search Visualization:** The grid search visualization tools can also be used to plot the results of the grid search.\n", - "- **Best Fit:** We can retrieve a new aggregator containing only the maximum log evidence results of the grid.\n", - "\n", - "__Model__\n", - "\n", - "This script uses the results of the `subhalo/detect/start_here.ipynb` example. You must run this script to completion\n", - "first to ensure the results the database uses are available.\n", - "\n", - "__Start Here Notebooks__\n", - "\n", - "You should be familiar with dark matter subhalo detection, by reading the example `subhalo/detect/start_here.ipynb`.\n", - "\n", - "You should also be familiar with the database, by reading the example `imaging/advanced/database/start_here.ipynb`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import json\n", - "import os\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "___Database__\n", - "\n", - "The name of the database, which corresponds to the output results folder." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "database_name = \"subhalo_detect\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If the `.sqlite` file of the database is already in the output folder we delete it and create a new database immediately\n", - "afterwards.\n", - "\n", - "This ensures we don't double up on results if we run the script multiple times, and if new results are added to the\n", - "output folder (e.g. download from a super computer) they are added to the database." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " os.remove(Path(\"output\", f\"{database_name}.sqlite\"))\n", - "except FileNotFoundError:\n", - " pass" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the database file `subhalo_detect.sqlite` in the output folder." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg = af.Aggregator.from_database(\n", - " filename=f\"{database_name}.sqlite\", completed_only=False\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Add all results in the directory \"output\" to the database, which we manipulate below via the aggregator." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg.add_directory(directory=Path(\"output\", database_name))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Agg No / With Subhalo__\n", - "\n", - "Standard aggregator querying can be used to get aggregates of results for lens models with and without a subhalo.\n", - "\n", - "The easiest query uses the name of the subhalo searches in the SLaM subhalo pipeline." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_no_subhalo = agg.query(agg.search.name == \"subhalo[1]\")\n", - "agg_with_subhalo = agg.query(agg.search.name == \"subhalo[3]_[single_plane_refine]\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can extract the `log_evidence` values of the results with and without and DM subhalo via the aggregators.\n", - "\n", - "We create a dictionary of these values where the keys are the `unique_tag` of each search, which is the name of the\n", - "dataset fitted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_evidence_no_subhalo_dict = {}\n", - "\n", - "for search, samples in zip(\n", - " agg_no_subhalo.values(\"search\"), agg_no_subhalo.values(\"samples\")\n", - "):\n", - " log_evidence_no_subhalo_dict[search.unique_tag] = samples.log_evidence\n", - "\n", - "print(\"\\nLog Evidence No Subhalo\")\n", - "print(log_evidence_no_subhalo_dict)\n", - "\n", - "log_evidence_with_subhalo_dict = {}\n", - "\n", - "for search, samples in zip(\n", - " agg_with_subhalo.values(\"search\"), agg_with_subhalo.values(\"samples\")\n", - "):\n", - " log_evidence_with_subhalo_dict[search.unique_tag] = samples.log_evidence\n", - "\n", - "print(\"\\nLog Evidence With Subhalo\")\n", - "print(log_evidence_with_subhalo_dict)\n", - "\n", - "log_evidence_difference_dict = {}\n", - "\n", - "# for key in log_evidence_no_subhalo_dict.keys():\n", - "\n", - "# log_evidence_difference_dict[key] = log_evidence_with_subhalo_dict[key] - log_evidence_no_subhalo_dict[key]\n", - "\n", - "print(\"\\nLog Evidence Difference\")\n", - "print(log_evidence_difference_dict)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From these, we can create the maximum likelihood instances of the lens model and corresponding `FitImaging` objects.\n", - "\n", - "These can then be passed to the `aplt.subplot_detection_imaging` to visualize the results of the subhalo detection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit_agg_no_subhalo = al.agg.FitImagingAgg(aggregator=agg_no_subhalo)\n", - "fit_no_subhalo_gen = fit_agg_no_subhalo.max_log_likelihood_gen_from()\n", - "fit_no_subhalo = list(fit_no_subhalo_gen)[0]\n", - "\n", - "fit_agg_with_subhalo = al.agg.FitImagingAgg(aggregator=agg_with_subhalo)\n", - "fit_with_subhalo_gen = fit_agg_with_subhalo.max_log_likelihood_gen_from()\n", - "fit_with_subhalo = list(fit_with_subhalo_gen)[0]\n", - "\n", - "\n", - "aplt.subplot_detection_fits(\n", - " fit_imaging_no_subhalo=fit_no_subhalo,\n", - " fit_imaging_with_subhalo=fit_imaging_with_subhalo,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid Searches__\n", - "\n", - "If the results of the database include a grid search of non-linear searches, the aggregator has a dedicated method\n", - "to return the grid of results.\n", - "\n", - "We iterate over these results using a for loop below, where each iteration will correspond to a different lens in \n", - "our analysis (e.g. if there are multiple lenses in the dataset that are fitted). In the `start_here.ipynb` example,\n", - "only one lens is fitted, so this for loop is only iterated over once." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for agg_grid, search in zip(\n", - " agg.grid_searches(), agg.grid_searches().best_fits().values(\"search\")\n", - "):\n", - " # Extract the `GridSearchResult` which the `start_here.ipynb` example uses\n", - " # for result inspection and visualization.\n", - "\n", - " result_subhalo_grid_search = agg_grid[\"result\"]\n", - "\n", - " # This can be manipulated in the ways shown in `start_here.ipynb`, for example\n", - " # to plot the log evidence of each cell.\n", - "\n", - " result_subhalo_grid_search = al.subhalo.SubhaloGridSearchResult(\n", - " result=result_subhalo_grid_search\n", - " )\n", - "\n", - " log_evidence_array = result_subhalo_grid_search.figure_of_merit_array(\n", - " use_log_evidences=True,\n", - " relative_to_value=log_evidence_no_subhalo_dict[search.unique_tag],\n", - " )\n", - "\n", - " print(log_evidence_array)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid Search Visualization__\n", - "\n", - "The grid search visualization tools can also be used to plot the results of the grid search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "samples_no_subhalo_gen = agg_no_subhalo.values(\"samples\")\n", - "\n", - "fit_agg_no_subhalo = al.agg.FitImagingAgg(aggregator=agg_no_subhalo)\n", - "fit_no_subhalo_gen = fit_agg_no_subhalo.max_log_likelihood_gen_from()\n", - "\n", - "fit_agg_with_subhalo = al.agg.FitImagingAgg(aggregator=agg_with_subhalo)\n", - "fit_with_subhalo_gen = fit_agg_with_subhalo.max_log_likelihood_gen_from()\n", - "\n", - "for agg_grid, fit_no_subhalo, fit_with_subhalo, samples_no_subhalo in zip(\n", - " agg.grid_searches(),\n", - " fit_no_subhalo_gen,\n", - " fit_with_subhalo_gen,\n", - " samples_no_subhalo_gen,\n", - "):\n", - " # Extract the `GridSearchResult` which the `start_here.ipynb` example uses\n", - " # for result inspection and visualization.\n", - "\n", - " result_subhalo_grid_search = agg_grid[\"result\"]\n", - "\n", - " # This can be manipulated in the ways shown in `start_here.ipynb`, for example\n", - " # to plot the log evidence of each cell.\n", - "\n", - " result_subhalo_grid_search = al.subhalo.SubhaloGridSearchResult(\n", - " result=result_subhalo_grid_search\n", - " )\n", - "\n", - " aplt.subplot_detection_imaging(\n", - " result=result_subhalo_grid_search,\n", - " fit_imaging_with_subhalo=fit_imaging_with_subhalo,\n", - " )\n", - " aplt.subplot_detection_fits(\n", - " fit_imaging_no_subhalo=fit_no_subhalo,\n", - " fit_imaging_with_subhalo=fit_imaging_with_subhalo,\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Best Fit__\n", - "\n", - "We can retrieve a new aggregator containing only the maximum log evidence results of the grid search. \n", - "\n", - "This can then be used as a normal aggregator to inspect the `Samples` of the fit or plot the best-fit `FitImaging`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "agg_best_fit = agg.grid_searches().best_fits()\n", - "\n", - "samples_gen = agg_best_fit.values(\"samples\")\n", - "\n", - "for samples in samples_gen:\n", - " print(samples.log_evidence)\n", - "\n", - "fit_agg = al.agg.FitImagingAgg(\n", - " aggregator=agg_best_fit,\n", - ")\n", - "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", - "\n", - "for fit_list in fit_gen:\n", - " # Only one `Analysis` so take first and only dataset.\n", - " fit = fit_list[0]\n", - "\n", - " aplt.subplot_fit_imaging(fit=fit)\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Subhalo Detection: Database\n", + "===========================\n", + "\n", + "The example `subhalo/detect/start_here.ipynb` shows how to perform dark matter subhalo detection in strong lens\n", + "with **PyAutoLens**, including using results to inspect and visualize the fit.\n", + "\n", + "This example shows how to load the results of subhalo detection analysis into a `.sqlite` database, which can be\n", + "manipulated stand-alone in this Python script or in a Jupyter notebook. This is useful when fits are performed on a\n", + "super computer and results are downloaded separately for inspection.\n", + "\n", + "The database in this example is built by scraping the results of the `subhalo/detect/start_here.ipynb` example. You\n", + "can also write results directly to the database during the fit by using a session.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Start Here Notebooks:** You should be familiar with dark matter subhalo detection, by reading the example.\n", + "- **Grid Searches:** If the results of the database include a grid search of non-linear searches, the aggregator has a.\n", + "- **Grid Search Visualization:** The grid search visualization tools can also be used to plot the results of the grid search.\n", + "- **Best Fit:** We can retrieve a new aggregator containing only the maximum log evidence results of the grid.\n", + "\n", + "__Model__\n", + "\n", + "This script uses the results of the `subhalo/detect/start_here.ipynb` example. You must run this script to completion\n", + "first to ensure the results the database uses are available.\n", + "\n", + "__Start Here Notebooks__\n", + "\n", + "You should be familiar with dark matter subhalo detection, by reading the example `subhalo/detect/start_here.ipynb`.\n", + "\n", + "You should also be familiar with the database, by reading the example `imaging/advanced/database/start_here.ipynb`." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import json\n", + "import os\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "___Database__\n", + "\n", + "The name of the database, which corresponds to the output results folder." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "database_name = \"subhalo_detect\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If the `.sqlite` file of the database is already in the output folder we delete it and create a new database immediately\n", + "afterwards.\n", + "\n", + "This ensures we don't double up on results if we run the script multiple times, and if new results are added to the\n", + "output folder (e.g. download from a super computer) they are added to the database." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " os.remove(Path(\"output\", f\"{database_name}.sqlite\"))\n", + "except FileNotFoundError:\n", + " pass" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the database file `subhalo_detect.sqlite` in the output folder." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg = af.Aggregator.from_database(\n", + " filename=f\"{database_name}.sqlite\", completed_only=False\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Add all results in the directory \"output\" to the database, which we manipulate below via the aggregator." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg.add_directory(directory=Path(\"output\", database_name))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Agg No / With Subhalo__\n", + "\n", + "Standard aggregator querying can be used to get aggregates of results for lens models with and without a subhalo.\n", + "\n", + "The easiest query uses the name of the subhalo searches in the SLaM subhalo pipeline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_no_subhalo = agg.query(agg.search.name == \"subhalo[1]\")\n", + "agg_with_subhalo = agg.query(agg.search.name == \"subhalo[3]_[single_plane_refine]\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can extract the `log_evidence` values of the results with and without and DM subhalo via the aggregators.\n", + "\n", + "We create a dictionary of these values where the keys are the `unique_tag` of each search, which is the name of the\n", + "dataset fitted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_evidence_no_subhalo_dict = {}\n", + "\n", + "for search, samples in zip(\n", + " agg_no_subhalo.values(\"search\"), agg_no_subhalo.values(\"samples\")\n", + "):\n", + " log_evidence_no_subhalo_dict[search.unique_tag] = samples.log_evidence\n", + "\n", + "print(\"\\nLog Evidence No Subhalo\")\n", + "print(log_evidence_no_subhalo_dict)\n", + "\n", + "log_evidence_with_subhalo_dict = {}\n", + "\n", + "for search, samples in zip(\n", + " agg_with_subhalo.values(\"search\"), agg_with_subhalo.values(\"samples\")\n", + "):\n", + " log_evidence_with_subhalo_dict[search.unique_tag] = samples.log_evidence\n", + "\n", + "print(\"\\nLog Evidence With Subhalo\")\n", + "print(log_evidence_with_subhalo_dict)\n", + "\n", + "log_evidence_difference_dict = {}\n", + "\n", + "# for key in log_evidence_no_subhalo_dict.keys():\n", + "\n", + "# log_evidence_difference_dict[key] = log_evidence_with_subhalo_dict[key] - log_evidence_no_subhalo_dict[key]\n", + "\n", + "print(\"\\nLog Evidence Difference\")\n", + "print(log_evidence_difference_dict)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From these, we can create the maximum likelihood instances of the lens model and corresponding `FitImaging` objects.\n", + "\n", + "These can then be passed to the `aplt.subplot_detection_imaging` to visualize the results of the subhalo detection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_agg_no_subhalo = al.agg.FitImagingAgg(aggregator=agg_no_subhalo)\n", + "fit_no_subhalo_gen = fit_agg_no_subhalo.max_log_likelihood_gen_from()\n", + "fit_no_subhalo = list(fit_no_subhalo_gen)[0]\n", + "\n", + "fit_agg_with_subhalo = al.agg.FitImagingAgg(aggregator=agg_with_subhalo)\n", + "fit_with_subhalo_gen = fit_agg_with_subhalo.max_log_likelihood_gen_from()\n", + "fit_with_subhalo = list(fit_with_subhalo_gen)[0]\n", + "\n", + "\n", + "aplt.subplot_detection_fits(\n", + " fit_imaging_no_subhalo=fit_no_subhalo,\n", + " fit_imaging_with_subhalo=fit_imaging_with_subhalo,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid Searches__\n", + "\n", + "If the results of the database include a grid search of non-linear searches, the aggregator has a dedicated method\n", + "to return the grid of results.\n", + "\n", + "We iterate over these results using a for loop below, where each iteration will correspond to a different lens in \n", + "our analysis (e.g. if there are multiple lenses in the dataset that are fitted). In the `start_here.ipynb` example,\n", + "only one lens is fitted, so this for loop is only iterated over once." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for agg_grid, search in zip(\n", + " agg.grid_searches(), agg.grid_searches().best_fits().values(\"search\")\n", + "):\n", + " # Extract the `GridSearchResult` which the `start_here.ipynb` example uses\n", + " # for result inspection and visualization.\n", + "\n", + " result_subhalo_grid_search = agg_grid[\"result\"]\n", + "\n", + " # This can be manipulated in the ways shown in `start_here.ipynb`, for example\n", + " # to plot the log evidence of each cell.\n", + "\n", + " result_subhalo_grid_search = al.subhalo.SubhaloGridSearchResult(\n", + " result=result_subhalo_grid_search\n", + " )\n", + "\n", + " log_evidence_array = result_subhalo_grid_search.figure_of_merit_array(\n", + " use_log_evidences=True,\n", + " relative_to_value=log_evidence_no_subhalo_dict[search.unique_tag],\n", + " )\n", + "\n", + " print(log_evidence_array)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid Search Visualization__\n", + "\n", + "The grid search visualization tools can also be used to plot the results of the grid search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "samples_no_subhalo_gen = agg_no_subhalo.values(\"samples\")\n", + "\n", + "fit_agg_no_subhalo = al.agg.FitImagingAgg(aggregator=agg_no_subhalo)\n", + "fit_no_subhalo_gen = fit_agg_no_subhalo.max_log_likelihood_gen_from()\n", + "\n", + "fit_agg_with_subhalo = al.agg.FitImagingAgg(aggregator=agg_with_subhalo)\n", + "fit_with_subhalo_gen = fit_agg_with_subhalo.max_log_likelihood_gen_from()\n", + "\n", + "for agg_grid, fit_no_subhalo, fit_with_subhalo, samples_no_subhalo in zip(\n", + " agg.grid_searches(),\n", + " fit_no_subhalo_gen,\n", + " fit_with_subhalo_gen,\n", + " samples_no_subhalo_gen,\n", + "):\n", + " # Extract the `GridSearchResult` which the `start_here.ipynb` example uses\n", + " # for result inspection and visualization.\n", + "\n", + " result_subhalo_grid_search = agg_grid[\"result\"]\n", + "\n", + " # This can be manipulated in the ways shown in `start_here.ipynb`, for example\n", + " # to plot the log evidence of each cell.\n", + "\n", + " result_subhalo_grid_search = al.subhalo.SubhaloGridSearchResult(\n", + " result=result_subhalo_grid_search\n", + " )\n", + "\n", + " aplt.subplot_detection_imaging(\n", + " result=result_subhalo_grid_search,\n", + " fit_imaging_with_subhalo=fit_imaging_with_subhalo,\n", + " )\n", + " aplt.subplot_detection_fits(\n", + " fit_imaging_no_subhalo=fit_no_subhalo,\n", + " fit_imaging_with_subhalo=fit_imaging_with_subhalo,\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Best Fit__\n", + "\n", + "We can retrieve a new aggregator containing only the maximum log evidence results of the grid search. \n", + "\n", + "This can then be used as a normal aggregator to inspect the `Samples` of the fit or plot the best-fit `FitImaging`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "agg_best_fit = agg.grid_searches().best_fits()\n", + "\n", + "samples_gen = agg_best_fit.values(\"samples\")\n", + "\n", + "for samples in samples_gen:\n", + " print(samples.log_evidence)\n", + "\n", + "fit_agg = al.agg.FitImagingAgg(\n", + " aggregator=agg_best_fit,\n", + ")\n", + "fit_gen = fit_agg.max_log_likelihood_gen_from()\n", + "\n", + "for fit_list in fit_gen:\n", + " # Only one `Analysis` so take first and only dataset.\n", + " fit = fit_list[0]\n", + "\n", + " aplt.subplot_fit_imaging(fit=fit)\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/subhalo/detect/start_here.ipynb b/notebooks/imaging/features/advanced/subhalo/detect/start_here.ipynb index 5e61c41a4..edf17c695 100644 --- a/notebooks/imaging/features/advanced/subhalo/detect/start_here.ipynb +++ b/notebooks/imaging/features/advanced/subhalo/detect/start_here.ipynb @@ -1,1044 +1,1081 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Subhalo Detection: Start Here\n", - "=============================\n", - "\n", - "Strong gravitational lenses can be used to detect the presence of small-scale dark matter (DM) subhalos. This occurs\n", - "when the DM subhalo overlaps the lensed source emission, and therefore gravitationally perturbs the observed image of\n", - "the lensed source galaxy.\n", - "\n", - "When a DM subhalo is not included in the lens model, residuals will be present in the fit to the data in the lensed\n", - "source regions near the subhalo. By adding a DM subhalo to the lens model, these residuals can be reduced. Bayesian\n", - "model comparison can then be used to quantify whether or not the improvement to the fit is significant enough to\n", - "claim the detection of a DM subhalo.\n", - "\n", - "The example illustrates DM subhalo detection with **PyAutoLens**.\n", - "\n", - "__Contents__\n", - "\n", - "- **SLaM Pipelines:** The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which.\n", - "- **Grid Search:** The second stage of the SUBHALO PIPELINE uses a grid-search of non-linear searches to determine the.\n", - "- **Pixelized Source:** Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "- **Bayesian Evidence:** To determine if a DM subhalo was detected by the pipeline, we can compare the log of the Bayesian.\n", - "- **Log Likelihood:** Different metrics can be used to inspect whether a DM subhalo was detected.\n", - "- **Grid Search Result:** The grid search results have attributes which can be used to inspect the results of the DM subhalo.\n", - "\n", - "__SLaM Pipelines__\n", - "\n", - "The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which automate the fitting\n", - "of complex lens models. The SLaM pipelines are used for all DM subhalo detection analyses. Therefore\n", - "you should be familiar with the SLaM pipelines before performing DM subhalo detection yourself. If you are unfamiliar\n", - "with the SLaM pipelines, checkout the\n", - "example `autolens_workspace/notebooks/guides/modeling/slam_start_here`.\n", - "\n", - "Dark matter subhalo detection runs the standard SLaM pipelines, and then extends them with a SUBHALO PIPELINE which\n", - "performs the following three chained non-linear searches:\n", - "\n", - " 1) Fits the lens model fitted in the MASS PIPELINE again, without a DM subhalo, to estimate the Bayesian evidence\n", - " of the model without a DM subhalo.\n", - "\n", - " 2) Performs a grid-search of non-linear searches, where each grid cell includes a DM subhalo whose (y,x) centre is\n", - " confined to a small 2D section of the image plane via uniform priors (we explain this in more detail below).\n", - "\n", - " 3) Fit the lens model again, including a DM subhalo whose (y,x) centre is initialized from the highest log evidence\n", - " grid cell of the grid-search. The Bayesian evidence estimated in this model-fit is compared to the model-fit\n", - " which did not include a DM subhalo, to determine whether or not a DM subhalo was detected.\n", - "\n", - "__Grid Search__\n", - "\n", - "The second stage of the SUBHALO PIPELINE uses a grid-search of non-linear searches to determine the highest log\n", - "evidence model with a DM subhalo. This grid search confines each DM subhalo in the lens model to a small 2D section\n", - "of the image plane via priors on its (y,x) centre. The reasons for this are as follows:\n", - "\n", - " - Lens models including a DM subhalo often have a multi-model parameter space. This means there are multiple lens\n", - " models with high likelihood solutions, each of which place the DM subhalo in different (y,x) image-plane location.\n", - " Multi-modal parameter spaces are synonomously difficult for non-linear searches to fit, and often produce\n", - " incorrect or inefficient fitting. The grid search breaks the multi-modal parameter space into many single-peaked\n", - " parameter spaces, making the model-fitting faster and more reliable.\n", - "\n", - " - By inferring how placing a DM subhalo at different locations in the image-plane changes the Bayesian evidence, we\n", - " map out spatial information on where a DM subhalo is detected. This can help our interpretation of the DM subhalo\n", - " detection.\n", - "\n", - "__Pixelized Source__\n", - "\n", - "Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of the source's light\n", - "are at a level where the subhalo's perturbing lensing effects can be detected.\n", - "\n", - "This requires the source reconstruction to be performed using a pixelized source, as this provides a more detailed\n", - "reconstruction of the source's light than fits using light profiles.\n", - "\n", - "This example therefore using a pixelized source and the corresponding SLaM pipelines.\n", - "\n", - "__Model__\n", - "\n", - "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE, MASS TOTAL PIPELINE and SUBHALO PIPELINE this\n", - "SLaM script fits `Imaging` of a strong lens system, where in the final model:\n", - "\n", - " - The lens galaxy's light is an MGE bulge.\n", - " - The lens galaxy's total mass distribution is an `PowerLaw`.\n", - " - A dark matter subhalo near the lens galaxy mass is included as a `NFWMCRLudlowSph`.\n", - " - The source galaxy is an `Inversion`.\n", - "\n", - "Each SLaM pipeline is implemented as a Python function below, with a documentation string above each function\n", - "describing the pipeline in detail. The full pipeline is run at the bottom of the script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except:\n", - "\n", - " - The lens galaxy's MGE uses 30 Gaussians (instead of 20) to better capture complex lens light morphology.\n", - " - The source galaxy's MGE uses `gaussian_per_basis=1` (instead of 2) for a simpler source model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`.\n", - "\n", - "Note that between SOURCE PIX PIPELINE 1 and this search, the calling section applies adaptive over-sampling to\n", - "the dataset using the pixelized source reconstruction from search 1. This improves the accuracy of the\n", - "pixelization by ensuring the over-sampling is adapted to the source morphology." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except:\n", - "\n", - " - The lens galaxy's MGE uses 30 Gaussians (consistent with SOURCE LP PIPELINE).\n", - " - `use_jax=True` is passed to the analysis for faster computation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " bulge = light_result.instance.galaxies.lens.bulge\n", - " disk = light_result.instance.galaxies.lens.disk\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=bulge,\n", - " disk=disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (grid search)__\n", - "\n", - "The second search of the SUBHALO PIPELINE performs a [number_of_steps x number_of_steps] grid search of\n", - "non-linear searches. Each grid cell includes a DM subhalo whose (y,x) centre is confined to a small 2D section\n", - "of the image plane via uniform priors.\n", - "\n", - "This grid search maps out where in the image plane including a DM subhalo provides a better fit to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def subhalo_grid_search(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " mass_result: af.Result,\n", - " subhalo_mass: af.Model,\n", - " grid_dimension_arcsec: float = 3.0,\n", - " number_of_steps: int = 2,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " )\n", - "\n", - " subhalo = af.Model(al.Galaxy, mass=subhalo_mass)\n", - "\n", - " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", - " subhalo.mass.centre_0 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - " )\n", - " subhalo.mass.centre_1 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - " )\n", - "\n", - " subhalo.redshift = mass_result.instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_object = mass_result.instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_source = mass_result.instance.galaxies.source.redshift\n", - "\n", - " lens = mass_result.model.galaxies.lens\n", - " source = al.util.chaining.source_from(result=mass_result)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, subhalo=subhalo, source=source),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[1]_[search_lens_plane]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " subhalo_grid_search = af.SearchGridSearch(\n", - " search=search,\n", - " number_of_steps=number_of_steps,\n", - " )\n", - "\n", - " return subhalo_grid_search.fit(\n", - " model=model,\n", - " analysis=analysis,\n", - " grid_priors=[\n", - " model.galaxies.subhalo.mass.centre_1,\n", - " model.galaxies.subhalo.mass.centre_0,\n", - " ],\n", - " info=settings_search.info,\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (refine)__\n", - "\n", - "The third search of the SUBHALO PIPELINE refits the lens model including a DM subhalo, initializing the\n", - "subhalo centre from the highest log evidence grid cell of the grid search.\n", - "\n", - "The Bayesian evidence from this fit is compared to the no-subhalo fit to determine whether a DM subhalo\n", - "was detected." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def subhalo_refine(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " mass_result: af.Result,\n", - " subhalo_grid_search_result: af.Result,\n", - " subhalo_mass: af.Model,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " )\n", - "\n", - " subhalo = af.Model(\n", - " al.Galaxy,\n", - " redshift=mass_result.instance.galaxies.lens.redshift,\n", - " mass=subhalo_mass,\n", - " )\n", - "\n", - " subhalo.redshift = mass_result.instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_object = mass_result.instance.galaxies.lens.redshift\n", - " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", - " subhalo.mass.centre = subhalo_grid_search_result.model_centred_absolute(\n", - " a=1.0\n", - " ).galaxies.subhalo.mass.centre\n", - "\n", - " subhalo.redshift = subhalo_grid_search_result.model.galaxies.subhalo.redshift\n", - " subhalo.mass.redshift_object = subhalo.redshift\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=subhalo_grid_search_result.model.galaxies.lens,\n", - " subhalo=subhalo,\n", - " source=subhalo_grid_search_result.model.galaxies.source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[2]_[single_plane_refine]\",\n", - " **settings_search.search_dict,\n", - " n_live=600,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"dark_matter_subhalo\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/advanced/subhalo/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.05,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"subhalo_detect\"),\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step.\n", - "\n", - "Between SOURCE PIX PIPELINE 1 and 2, adaptive over-sampling is applied to the dataset using the pixelized source\n", - "reconstruction from search 1. This improves the pixelization accuracy in search 2 and all subsequent pipelines." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", - " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", - " noise_map=dataset.noise_map,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=over_sampling)\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")\n", - "\n", - "result_subhalo_grid_search = subhalo_grid_search(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mass_result=mass_result,\n", - " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", - " grid_dimension_arcsec=3.0,\n", - " number_of_steps=2,\n", - ")\n", - "\n", - "result_with_subhalo = subhalo_refine(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mass_result=mass_result,\n", - " subhalo_grid_search_result=result_subhalo_grid_search,\n", - " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Bayesian Evidence__\n", - "\n", - "To determine if a DM subhalo was detected by the pipeline, we can compare the log of the Bayesian evidences of the\n", - "model-fits performed with and without a subhalo.\n", - "\n", - "The following scale describes how different log evidence increases correspond to difference detection significances:\n", - "\n", - " - Negative log evidence increase: No detection.\n", - " - Log evidence increase between 0 and 3: No detection.\n", - " - Log evidence increase between 3 and 5: Weak evidence, should consider it a non-detection.\n", - " - Log evidence increase between 5 and 10: Medium evidence, but still inconclusive.\n", - " - Log evidence increase between 10 and 20: Strong evidence, consider it a detection.\n", - " - Log evidence increase > 20: Very strong evidence, definitive detection.\n", - "\n", - "Lets inspect the log evidence increase for the model-fit performed in this example:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "evidence_no_subhalo = mass_result.samples.log_evidence\n", - "evidence_with_subhalo = result_with_subhalo.samples.log_evidence\n", - "\n", - "log_evidence_increase = evidence_with_subhalo - evidence_no_subhalo\n", - "\n", - "print(\"Evidence Increase: \", log_evidence_increase)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood__\n", - "\n", - "Different metrics can be used to inspect whether a DM subhalo was detected.\n", - "\n", - "The Bayesian evidence is the most rigorous because it penalizes models based on their complexity. An alternative\n", - "goodness of fit is the `log_likelihood`, which is directly related to the residuals of the model or the chi-squared\n", - "value.\n", - "\n", - "The benefit of the log likelihood is it is a straight forward value indicating how well a model fitted the data.\n", - "The `log_likelihood` of the lens model without a subhalo must always be less than the model with a subhalo. If\n", - "this is not the case, something must have gone wrong with one of the model-fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_likelihood_no_subhalo = mass_result.samples.log_likelihood\n", - "log_likelihood_with_subhalo = result_with_subhalo.samples.log_likelihood\n", - "\n", - "log_likelihood_increase = log_likelihood_with_subhalo - log_likelihood_no_subhalo\n", - "\n", - "print(\"Log Likelihood Increase: \", log_likelihood_increase)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid Search Result__\n", - "\n", - "The grid search results have attributes which can be used to inspect the results of the DM subhalo grid-search.\n", - "\n", - "For example, we can produce a 2D array of the log evidence values computed for each grid cell of the grid-search,\n", - "computed relative to the `log_evidence` of the model-fit which did not include a subhalo." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "subhalo_grid_search_result = al.subhalo.SubhaloGridSearchResult(\n", - " result=result_subhalo_grid_search\n", - ")\n", - "\n", - "log_evidence_array = subhalo_grid_search_result.figure_of_merit_array(\n", - " use_log_evidences=True,\n", - " relative_to_value=mass_result.samples.log_evidence,\n", - ")\n", - "\n", - "print(\"Log Evidence Array: \\n\")\n", - "print(log_evidence_array)\n", - "\n", - "aplt.plot_array(array=log_evidence_array, title=\"\")\n", - "\n", - "mass_array = subhalo_grid_search_result.subhalo_mass_array\n", - "\n", - "print(\"Mass Array: \\n\")\n", - "print(mass_array)\n", - "\n", - "subhalo_centres_grid = subhalo_grid_search_result.subhalo_centres_grid\n", - "\n", - "print(\"Subhalo Centres Grid: \\n\")\n", - "print(subhalo_centres_grid)\n", - "\n", - "einstein_radius_array = subhalo_grid_search_result.attribute_grid(\n", - " \"galaxies.lens.mass.einstein_radius\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Subhalo Detection: Start Here\n", + "=============================\n", + "\n", + "Strong gravitational lenses can be used to detect the presence of small-scale dark matter (DM) subhalos. This occurs\n", + "when the DM subhalo overlaps the lensed source emission, and therefore gravitationally perturbs the observed image of\n", + "the lensed source galaxy.\n", + "\n", + "When a DM subhalo is not included in the lens model, residuals will be present in the fit to the data in the lensed\n", + "source regions near the subhalo. By adding a DM subhalo to the lens model, these residuals can be reduced. Bayesian\n", + "model comparison can then be used to quantify whether or not the improvement to the fit is significant enough to\n", + "claim the detection of a DM subhalo.\n", + "\n", + "The example illustrates DM subhalo detection with **PyAutoLens**.\n", + "\n", + "__Contents__\n", + "\n", + "- **SLaM Pipelines:** The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which.\n", + "- **Grid Search:** The second stage of the SUBHALO PIPELINE uses a grid-search of non-linear searches to determine the.\n", + "- **Pixelized Source:** Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "- **Bayesian Evidence:** To determine if a DM subhalo was detected by the pipeline, we can compare the log of the Bayesian.\n", + "- **Log Likelihood:** Different metrics can be used to inspect whether a DM subhalo was detected.\n", + "- **Grid Search Result:** The grid search results have attributes which can be used to inspect the results of the DM subhalo.\n", + "\n", + "__SLaM Pipelines__\n", + "\n", + "The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which automate the fitting\n", + "of complex lens models. The SLaM pipelines are used for all DM subhalo detection analyses. Therefore\n", + "you should be familiar with the SLaM pipelines before performing DM subhalo detection yourself. If you are unfamiliar\n", + "with the SLaM pipelines, checkout the\n", + "example `autolens_workspace/notebooks/guides/modeling/slam_start_here`.\n", + "\n", + "Dark matter subhalo detection runs the standard SLaM pipelines, and then extends them with a SUBHALO PIPELINE which\n", + "performs the following three chained non-linear searches:\n", + "\n", + " 1) Fits the lens model fitted in the MASS PIPELINE again, without a DM subhalo, to estimate the Bayesian evidence\n", + " of the model without a DM subhalo.\n", + "\n", + " 2) Performs a grid-search of non-linear searches, where each grid cell includes a DM subhalo whose (y,x) centre is\n", + " confined to a small 2D section of the image plane via uniform priors (we explain this in more detail below).\n", + "\n", + " 3) Fit the lens model again, including a DM subhalo whose (y,x) centre is initialized from the highest log evidence\n", + " grid cell of the grid-search. The Bayesian evidence estimated in this model-fit is compared to the model-fit\n", + " which did not include a DM subhalo, to determine whether or not a DM subhalo was detected.\n", + "\n", + "__Grid Search__\n", + "\n", + "The second stage of the SUBHALO PIPELINE uses a grid-search of non-linear searches to determine the highest log\n", + "evidence model with a DM subhalo. This grid search confines each DM subhalo in the lens model to a small 2D section\n", + "of the image plane via priors on its (y,x) centre. The reasons for this are as follows:\n", + "\n", + " - Lens models including a DM subhalo often have a multi-model parameter space. This means there are multiple lens\n", + " models with high likelihood solutions, each of which place the DM subhalo in different (y,x) image-plane location.\n", + " Multi-modal parameter spaces are synonomously difficult for non-linear searches to fit, and often produce\n", + " incorrect or inefficient fitting. The grid search breaks the multi-modal parameter space into many single-peaked\n", + " parameter spaces, making the model-fitting faster and more reliable.\n", + "\n", + " - By inferring how placing a DM subhalo at different locations in the image-plane changes the Bayesian evidence, we\n", + " map out spatial information on where a DM subhalo is detected. This can help our interpretation of the DM subhalo\n", + " detection.\n", + "\n", + "__Pixelized Source__\n", + "\n", + "Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of the source's light\n", + "are at a level where the subhalo's perturbing lensing effects can be detected.\n", + "\n", + "This requires the source reconstruction to be performed using a pixelized source, as this provides a more detailed\n", + "reconstruction of the source's light than fits using light profiles.\n", + "\n", + "This example therefore using a pixelized source and the corresponding SLaM pipelines.\n", + "\n", + "__Model__\n", + "\n", + "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE, MASS TOTAL PIPELINE and SUBHALO PIPELINE this\n", + "SLaM script fits `Imaging` of a strong lens system, where in the final model:\n", + "\n", + " - The lens galaxy's light is an MGE bulge.\n", + " - The lens galaxy's total mass distribution is an `PowerLaw`.\n", + " - A dark matter subhalo near the lens galaxy mass is included as a `NFWMCRLudlowSph`.\n", + " - The source galaxy is an `Inversion`.\n", + "\n", + "Each SLaM pipeline is implemented as a Python function below, with a documentation string above each function\n", + "describing the pipeline in detail. The full pipeline is run at the bottom of the script." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except:\n", + "\n", + " - The lens galaxy's MGE uses 30 Gaussians (instead of 20) to better capture complex lens light morphology.\n", + " - The source galaxy's MGE uses `gaussian_per_basis=1` (instead of 2) for a simpler source model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`.\n", + "\n", + "Note that between SOURCE PIX PIPELINE 1 and this search, the calling section applies adaptive over-sampling to\n", + "the dataset using the pixelized source reconstruction from search 1. This improves the accuracy of the\n", + "pixelization by ensuring the over-sampling is adapted to the source morphology." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except:\n", + "\n", + " - The lens galaxy's MGE uses 30 Gaussians (consistent with SOURCE LP PIPELINE).\n", + " - `use_jax=True` is passed to the analysis for faster computation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " bulge = light_result.instance.galaxies.lens.bulge\n", + " disk = light_result.instance.galaxies.lens.disk\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=bulge,\n", + " disk=disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (grid search)__\n", + "\n", + "The second search of the SUBHALO PIPELINE performs a [number_of_steps x number_of_steps] grid search of\n", + "non-linear searches. Each grid cell includes a DM subhalo whose (y,x) centre is confined to a small 2D section\n", + "of the image plane via uniform priors.\n", + "\n", + "This grid search maps out where in the image plane including a DM subhalo provides a better fit to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def subhalo_grid_search(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " mass_result: af.Result,\n", + " subhalo_mass: af.Model,\n", + " grid_dimension_arcsec: float = 3.0,\n", + " number_of_steps: int = 2,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " )\n", + "\n", + " subhalo = af.Model(al.Galaxy, mass=subhalo_mass)\n", + "\n", + " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", + " subhalo.mass.centre_0 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + " )\n", + " subhalo.mass.centre_1 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + " )\n", + "\n", + " subhalo.redshift = mass_result.instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_object = mass_result.instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_source = mass_result.instance.galaxies.source.redshift\n", + "\n", + " lens = mass_result.model.galaxies.lens\n", + " source = al.util.chaining.source_from(result=mass_result)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, subhalo=subhalo, source=source),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[1]_[search_lens_plane]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " subhalo_grid_search = af.SearchGridSearch(\n", + " search=search,\n", + " number_of_steps=number_of_steps,\n", + " )\n", + "\n", + " return subhalo_grid_search.fit(\n", + " model=model,\n", + " analysis=analysis,\n", + " grid_priors=[\n", + " model.galaxies.subhalo.mass.centre_1,\n", + " model.galaxies.subhalo.mass.centre_0,\n", + " ],\n", + " info=settings_search.info,\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (refine)__\n", + "\n", + "The third search of the SUBHALO PIPELINE refits the lens model including a DM subhalo, initializing the\n", + "subhalo centre from the highest log evidence grid cell of the grid search.\n", + "\n", + "The Bayesian evidence from this fit is compared to the no-subhalo fit to determine whether a DM subhalo\n", + "was detected." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def subhalo_refine(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " mass_result: af.Result,\n", + " subhalo_grid_search_result: af.Result,\n", + " subhalo_mass: af.Model,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " )\n", + "\n", + " subhalo = af.Model(\n", + " al.Galaxy,\n", + " redshift=mass_result.instance.galaxies.lens.redshift,\n", + " mass=subhalo_mass,\n", + " )\n", + "\n", + " subhalo.redshift = mass_result.instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_object = mass_result.instance.galaxies.lens.redshift\n", + " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", + " subhalo.mass.centre = subhalo_grid_search_result.model_centred_absolute(\n", + " a=1.0\n", + " ).galaxies.subhalo.mass.centre\n", + "\n", + " subhalo.redshift = subhalo_grid_search_result.model.galaxies.subhalo.redshift\n", + " subhalo.mass.redshift_object = subhalo.redshift\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=subhalo_grid_search_result.model.galaxies.lens,\n", + " subhalo=subhalo,\n", + " source=subhalo_grid_search_result.model.galaxies.source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[2]_[single_plane_refine]\",\n", + " **settings_search.search_dict,\n", + " n_live=600,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"dark_matter_subhalo\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/advanced/subhalo/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.05,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"subhalo_detect\"),\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step.\n", + "\n", + "Between SOURCE PIX PIPELINE 1 and 2, adaptive over-sampling is applied to the dataset using the pixelized source\n", + "reconstruction from search 1. This improves the pixelization accuracy in search 2 and all subsequent pipelines." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", + " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", + " noise_map=dataset.noise_map,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_pixelization=over_sampling)\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")\n", + "\n", + "result_subhalo_grid_search = subhalo_grid_search(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mass_result=mass_result,\n", + " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", + " grid_dimension_arcsec=3.0,\n", + " number_of_steps=2,\n", + ")\n", + "\n", + "result_with_subhalo = subhalo_refine(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mass_result=mass_result,\n", + " subhalo_grid_search_result=result_subhalo_grid_search,\n", + " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Bayesian Evidence__\n", + "\n", + "To determine if a DM subhalo was detected by the pipeline, we can compare the log of the Bayesian evidences of the\n", + "model-fits performed with and without a subhalo.\n", + "\n", + "The following scale describes how different log evidence increases correspond to difference detection significances:\n", + "\n", + " - Negative log evidence increase: No detection.\n", + " - Log evidence increase between 0 and 3: No detection.\n", + " - Log evidence increase between 3 and 5: Weak evidence, should consider it a non-detection.\n", + " - Log evidence increase between 5 and 10: Medium evidence, but still inconclusive.\n", + " - Log evidence increase between 10 and 20: Strong evidence, consider it a detection.\n", + " - Log evidence increase > 20: Very strong evidence, definitive detection.\n", + "\n", + "Lets inspect the log evidence increase for the model-fit performed in this example:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "evidence_no_subhalo = mass_result.samples.log_evidence\n", + "evidence_with_subhalo = result_with_subhalo.samples.log_evidence\n", + "\n", + "log_evidence_increase = evidence_with_subhalo - evidence_no_subhalo\n", + "\n", + "print(\"Evidence Increase: \", log_evidence_increase)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood__\n", + "\n", + "Different metrics can be used to inspect whether a DM subhalo was detected.\n", + "\n", + "The Bayesian evidence is the most rigorous because it penalizes models based on their complexity. An alternative\n", + "goodness of fit is the `log_likelihood`, which is directly related to the residuals of the model or the chi-squared\n", + "value.\n", + "\n", + "The benefit of the log likelihood is it is a straight forward value indicating how well a model fitted the data.\n", + "The `log_likelihood` of the lens model without a subhalo must always be less than the model with a subhalo. If\n", + "this is not the case, something must have gone wrong with one of the model-fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_likelihood_no_subhalo = mass_result.samples.log_likelihood\n", + "log_likelihood_with_subhalo = result_with_subhalo.samples.log_likelihood\n", + "\n", + "log_likelihood_increase = log_likelihood_with_subhalo - log_likelihood_no_subhalo\n", + "\n", + "print(\"Log Likelihood Increase: \", log_likelihood_increase)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid Search Result__\n", + "\n", + "The grid search results have attributes which can be used to inspect the results of the DM subhalo grid-search.\n", + "\n", + "For example, we can produce a 2D array of the log evidence values computed for each grid cell of the grid-search,\n", + "computed relative to the `log_evidence` of the model-fit which did not include a subhalo." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "subhalo_grid_search_result = al.subhalo.SubhaloGridSearchResult(\n", + " result=result_subhalo_grid_search\n", + ")\n", + "\n", + "log_evidence_array = subhalo_grid_search_result.figure_of_merit_array(\n", + " use_log_evidences=True,\n", + " relative_to_value=mass_result.samples.log_evidence,\n", + ")\n", + "\n", + "print(\"Log Evidence Array: \\n\")\n", + "print(log_evidence_array)\n", + "\n", + "aplt.plot_array(array=log_evidence_array, title=\"\")\n", + "\n", + "mass_array = subhalo_grid_search_result.subhalo_mass_array\n", + "\n", + "print(\"Mass Array: \\n\")\n", + "print(mass_array)\n", + "\n", + "subhalo_centres_grid = subhalo_grid_search_result.subhalo_centres_grid\n", + "\n", + "print(\"Subhalo Centres Grid: \\n\")\n", + "print(subhalo_centres_grid)\n", + "\n", + "einstein_radius_array = subhalo_grid_search_result.attribute_grid(\n", + " \"galaxies.lens.mass.einstein_radius\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/subhalo/sensitivity/slam_source_parametric.ipynb b/notebooks/imaging/features/advanced/subhalo/sensitivity/slam_source_parametric.ipynb index 00c70c2b0..e2daf8e29 100644 --- a/notebooks/imaging/features/advanced/subhalo/sensitivity/slam_source_parametric.ipynb +++ b/notebooks/imaging/features/advanced/subhalo/sensitivity/slam_source_parametric.ipynb @@ -1,1224 +1,1261 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM (Source, Light and Mass): Subhalo Source Parametric Sensitivity Mapping\n", - "============================================================================\n", - "\n", - "This example illustrates how to perform DM subhalo sensitivity mapping using a SLaM pipeline for a source modeling\n", - "using light profiles.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except the lens mass uses an `Isothermal` with its centre fixed.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **Simulate Function Class:** We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above).\n", - "- **Base Fit:** We have defined a `Simulate` class that will be used to simulate every dataset simulated by the.\n", - "- **Perturb Fit:** We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", - "\n", - "__Model__\n", - "\n", - "Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE, MASS TOTAL PIPELINE and SUBHALO PIPELINE this SLaM script\n", - "fits `Imaging` of a strong lens system, where in the final model:\n", - "\n", - " - The lens galaxy's light is an MGE bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - A dark matter subhalo near The lens galaxy mass is included as a`NFWMCRLudlowSph`.\n", - " - The source galaxy is an `Inversion`.\n", - "\n", - "It ends by performing sensitivity mapping of the data using the above model, so as to determine where in the data\n", - "subhalos of a given mass could have been detected if present.\n", - "\n", - "This uses the SLaM pipelines:\n", - "\n", - " `source_lp`\n", - " `source_pix`\n", - " `light_lp`\n", - " `mass_total`\n", - " `subhalo/detection`\n", - "\n", - "Check them out for a full description of the analysis!\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `subhalo/detect/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import os\n", - "from pathlib import Path\n", - "from typing import List, Optional, Tuple, Union\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except the lens mass uses an `Isothermal` with its centre fixed to (0.0, 0.0)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search,\n", - " dataset,\n", - " mask_radius,\n", - " redshift_lens,\n", - " redshift_source,\n", - "):\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = (0.0, 0.0)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " disk=None,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search,\n", - " dataset,\n", - " mask_radius,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - "):\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " )\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search,\n", - " dataset,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " light_result,\n", - "):\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=light_result.instance.galaxies.lens.bulge,\n", - " disk=light_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate Function Class__\n", - "\n", - "We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above) and uses it to\n", - "simulate a strong lens dataset, which include a dark matter subhalo, which is subsequently fitted.\n", - "\n", - "Additional attributes required to simulate the data (mask, PSF) can be passed to the `__init__` method, and the\n", - "simulation is performed in the `__call__` method.\n", - "\n", - "When this dataset is simulated, the quantity `instance.perturb` is used in `__call__`. This is an instance\n", - "of the `NFWMCRLudlowSph`, and it is different every time the `simulate_cls` is called based on the value of sensitivity\n", - "being computed.\n", - "\n", - "In this example, this `instance.perturb` corresponds to two different subhalos with values of `mass_at_200` of\n", - "1e6 MSun and 1e13 MSun." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class SimulateImaging:\n", - " def __init__(self, mask, psf):\n", - " \"\"\"\n", - " Class used to simulate the strong lens imaging used for sensitivity mapping.\n", - "\n", - " Parameters\n", - " ----------\n", - " mask\n", - " The mask applied to the real image data, which is applied to every simulated imaging.\n", - " psf\n", - " The PSF of the real image data, which is applied to every simulated imaging and used for each fit.\n", - " \"\"\"\n", - " self.mask = mask\n", - " self.psf = psf\n", - "\n", - " def __call__(self, instance: af.ModelInstance, simulate_path: str):\n", - " \"\"\"\n", - " The `simulate_function` called by the `Sensitivity` class which simulates each strong lens image fitted\n", - " by the sensitivity mapper.\n", - "\n", - " The simulation procedure is as follows:\n", - "\n", - " 1) Use the input galaxies of the sensitivity `instance` to set up a tracer, which generates the image-plane\n", - " image of the strong lens system.\n", - "\n", - " 2) Simulate this image using the input dataset noise (Poisson) and PSF.\n", - "\n", - " 3) Apply the mask used in the analysis of the real image to the simulated image.\n", - "\n", - " 4) Output information about the simulation to hard-disk.\n", - "\n", - " The `subhalo` in the sensitivity `instance` changes for every iteration of the sensitivity mapping, ensuring\n", - " that we map out the sensitivity of the analysis to the subhalo properties (centre, mass, etc.).\n", - "\n", - " Parameters\n", - " ----------\n", - " instance\n", - " The sensitivity instance, which includes the galaxies whose parameters are varied to perform sensitivity.\n", - " The subhalo in this instance changes for every iteration of the sensitivity mapping.\n", - " simulate_path\n", - " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", - " mapping.\n", - "\n", - " Returns\n", - " -------\n", - " A simulated image of a strong lens, which id input into the fits of the sensitivity mapper.\n", - " \"\"\"\n", - "\n", - " \"\"\"\n", - " __Resume Fit__\n", - "\n", - " If sensitivity mapping already began on this grid cell, the dataset will have been simulated already and we\n", - " do not want to resimulate it and change its noise properties.\n", - "\n", - " We therefore load it from the `simulate_path` instead.\n", - " \"\"\"\n", - " try:\n", - " dataset = al.Imaging.from_fits(\n", - " data_path=f\"{simulate_path}/data.fits\",\n", - " psf_path=f\"{simulate_path}/psf.fits\",\n", - " noise_map_path=f\"{simulate_path}/noise_map.fits\",\n", - " pixel_scales=self.mask.pixel_scales,\n", - " check_noise_map=False,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=self.mask)\n", - "\n", - " tracer = al.Tracer(\n", - " galaxies=[\n", - " instance.galaxies.lens,\n", - " instance.perturb,\n", - " instance.galaxies.source,\n", - " ]\n", - " )\n", - "\n", - " traced_grid = tracer.traced_grid_2d_list_from(\n", - " grid=dataset.grid,\n", - " )[-1]\n", - "\n", - " source_centre = instance.galaxies.source.bulge.centre\n", - "\n", - " over_sample_size = (\n", - " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=traced_grid,\n", - " sub_size_list=[16, 8, 2],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=[source_centre],\n", - " )\n", - " )\n", - "\n", - " over_sample_size_lens = (\n", - " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[8, 4, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - " )\n", - "\n", - " over_sample_size = np.where(\n", - " over_sample_size > over_sample_size_lens,\n", - " over_sample_size,\n", - " over_sample_size_lens,\n", - " )\n", - " over_sample_size = al.Array2D(values=over_sample_size, mask=self.mask)\n", - "\n", - " return dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " except FileNotFoundError:\n", - " pass\n", - "\n", - " \"\"\"\n", - " __Ray Tracing__\n", - "\n", - " Set up the `Tracer` which is used to simulate the strong lens imaging, which may include the subhalo in\n", - " addition to the lens and source galaxy.\n", - " \"\"\"\n", - " tracer = al.Tracer(\n", - " galaxies=[\n", - " instance.galaxies.lens,\n", - " instance.perturb,\n", - " instance.galaxies.source,\n", - " ]\n", - " )\n", - "\n", - " \"\"\"\n", - " __Simulate__\n", - "\n", - " Set up the grid, PSF and simulator settings used to simulate imaging of the strong lens. These should be\n", - " tuned to match the S/N and noise properties of the observed data you are performing sensitivity mapping on.\n", - " \"\"\"\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=self.mask.shape_native,\n", - " pixel_scales=self.mask.pixel_scales,\n", - " )\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " simulator = al.SimulatorImaging(\n", - " exposure_time=1000.0,\n", - " psf=self.psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - " noise_seed=1,\n", - " )\n", - "\n", - " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - " \"\"\"\n", - " __Masking__\n", - "\n", - " The data generated by the simulate function is what is ultimately fitted.\n", - "\n", - " Therefore, we also apply the mask for the analysis before we return the simulated data.\n", - " \"\"\"\n", - " dataset = dataset.apply_mask(mask=self.mask)\n", - "\n", - " traced_grid = tracer.traced_grid_2d_list_from(\n", - " grid=dataset.grid,\n", - " )[-1]\n", - "\n", - " source_centre = instance.galaxies.source.bulge.centre\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=traced_grid,\n", - " sub_size_list=[16, 8, 2],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=[source_centre],\n", - " )\n", - "\n", - " over_sample_size_lens = (\n", - " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[8, 4, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - " )\n", - "\n", - " over_sample_size = np.where(\n", - " over_sample_size > over_sample_size_lens,\n", - " over_sample_size,\n", - " over_sample_size_lens,\n", - " )\n", - " over_sample_size = al.Array2D(values=over_sample_size, mask=self.mask)\n", - "\n", - " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " \"\"\"\n", - " Outputs info about the `Tracer` to the fit, so we know exactly how we simulated the image.\n", - " \"\"\"\n", - " self.output_info(simulate_path=simulate_path, dataset=dataset, tracer=tracer)\n", - "\n", - " return dataset\n", - "\n", - " def output_info(self, simulate_path: str, dataset: al.Imaging, tracer: al.Imaging):\n", - " \"\"\"\n", - " Output information about the data simulated for this iteration of sensitivity mapping.\n", - "\n", - " This information output is as follows:\n", - "\n", - " - A subplot of the simulated imaging dataset.\n", - " - A subplot of the tracer used to simulate this imaging dataset.\n", - " - A .json file containing the tracer galaxies.\n", - " - Output the simulated dataset to .fits files which are used to load the data if a run is resumed.\n", - "\n", - " Parameters\n", - " ----------\n", - " simulate_path\n", - " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", - " mapping.\n", - " dataset\n", - " The simulated imaging dataset which is visualized.\n", - " tracer\n", - " The tracer used to simulate the imaging dataset, which is visualized and output to a .json file.\n", - " \"\"\"\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " aplt.subplot_lensed_images(\n", - " tracer=tracer,\n", - " grid=dataset.grid,\n", - " output_path=simulate_path,\n", - " output_format=\"png\",\n", - " )\n", - "\n", - " al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(simulate_path) / \"tracer.json\",\n", - " )\n", - "\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=Path(simulate_path) / \"data.fits\",\n", - " psf_path=Path(simulate_path) / \"psf.fits\",\n", - " noise_map_path=Path(simulate_path) / \"noise_map.fits\",\n", - " overwrite=True,\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Base Fit__\n", - "\n", - "We have defined a `Simulate` class that will be used to simulate every dataset simulated by the sensitivity mapper.\n", - "Each simulated dataset will have a unique set of parameters for the `subhalo` (e.g. due to different values of\n", - "`perturb_model`.\n", - "\n", - "We will fit each simulated dataset using the `base_model`, which quantifies whether not including the dark matter\n", - "subhalo in the model changess the goodness-of-fit and therefore indicates if we are sensitive to the subhalo.\n", - "\n", - "We now write a `BaseFit` class, defining how the `base_model` is fitted to each simulated dataset and the\n", - "goodness-of-fit used to quantify whether the model fits the data well. As above, the `__init__` method can be\n", - "extended with new inputs to control how the model is fitted and the `__call__` method performs the fit.\n", - "\n", - "In this example, we use a full non-linear search to fit the `base_model` to the simulated data and return\n", - "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", - "more computationally efficient, for example performing a single log likelihood evaluation of the `base_model` fit\n", - "to the simulated data.\n", - "\n", - "Fucntionality which adapts the mesh and regularization of a pixelized source reconstruction to the unlensed source's\n", - "morphology require an `adapt_images`. This is an input of the __init__ constructor which is passed to the `Analysis`\n", - "for every simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class BaseFit:\n", - " def __init__(self, adapt_images, number_of_cores: int = 1):\n", - " \"\"\"\n", - " Class used to fit every dataset used for sensitivity mapping with the base model (the model without the\n", - " perturbed feature sensitivity mapping maps out).\n", - "\n", - " In this example, the base model therefore does not include the dark matter subhalo, but the simulated\n", - " dataset includes one.\n", - "\n", - " The base fit is repeated for every parameter on the sensitivity grid and compared to the perturbed fit. This\n", - " maps out the sensitivity of every parameter is (e.g. the sensitivity of the mass of the subhalo).\n", - "\n", - " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", - " fitted, below we include an input `analysis_cls` which is the `AnalysisImaging` class used to fit the model\n", - " to the dataset.\n", - "\n", - " Parameters\n", - " ----------\n", - " adapt_images\n", - " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", - " reconstructed galaxy's morphology.\n", - " number_of_cores\n", - " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", - " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", - " \"\"\"\n", - " self.adapt_images = adapt_images\n", - " self.number_of_cores = number_of_cores\n", - "\n", - " def __call__(self, dataset, model, paths, instance):\n", - " \"\"\"\n", - " The base fitting function which fits every dataset used for sensitivity mapping with the base model.\n", - "\n", - " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", - " quantify how sensitive the model is to the perturbed feature.\n", - "\n", - " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", - " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", - "\n", - " Parameters\n", - " ----------\n", - " dataset\n", - " The dataset which is simulated with the perturbed model and which is fitted.\n", - " model\n", - " The model instance which is fitted to the dataset, which does not include the perturbed feature.\n", - " paths\n", - " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", - " instance\n", - " The simulation instance, which includes the perturbed feature that is used to simulate the dataset.\n", - " This is often not used, but may be useful for certain sensitivity mapping tasks, for example using\n", - " true values of the simulated instance to set up aspects of the model-fit (e.g. the priors).\n", - " \"\"\"\n", - "\n", - " search = af.Nautilus(\n", - " paths=paths,\n", - " n_live=50,\n", - " number_of_cores=self.number_of_cores,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - " analysis._adapt_images = self.adapt_images\n", - "\n", - " return search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Perturb Fit__\n", - "\n", - "We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each simulated dataset. This\n", - "behaves analogously to the `BaseFit` class above, but now fits the `perturb_model` to the simulated data (as\n", - "opposed to the `base_model`).\n", - "\n", - "Again, in this example we use a full non-linear search to fit the `perturb_model` to the simulated data and return\n", - "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", - "more computationally efficient, for example performing a single log likelihood evaluation of the `perturb_model` fit\n", - "to the simulated data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class PerturbFit:\n", - " def __init__(self, adapt_images, number_of_cores: int = 1):\n", - " \"\"\"\n", - " Class used to fit every dataset used for sensitivity mapping with the perturbed model (the model with the\n", - " perturbed feature sensitivity mapping maps out).\n", - "\n", - " In this example, the perturbed model therefore includes the dark matter subhalo, which is also in the\n", - " simulated dataset.\n", - "\n", - " The perturbed fit is repeated for every parameter on the sensitivity grid and compared to the base fit. This\n", - " maps out the sensitivity of every parameter is (e.g. the sensitivity of mass of the subhalo).\n", - "\n", - " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", - " fitted, below we include an input `analysis_cls` which is the `Analysis` class used to fit the model to the\n", - " dataset.\n", - "\n", - " Parameters\n", - " ----------\n", - " adapt_images\n", - " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", - " reconstructed galaxy's morphology.\n", - " number_of_cores\n", - " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", - " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", - " \"\"\"\n", - " self.adapt_images = adapt_images\n", - " self.number_of_cores = number_of_cores\n", - "\n", - " def __call__(self, dataset, model, paths, instance):\n", - " \"\"\"\n", - " The perturbed fitting function which fits every dataset used for sensitivity mapping with the perturbed model.\n", - "\n", - " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", - " quantify how sensitive the model is to the perturbed feature.\n", - "\n", - " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", - " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", - "\n", - " Parameters\n", - " ----------\n", - " dataset\n", - " The dataset which is simulated with the perturbed model and which is fitted.\n", - " model\n", - " The model instance which is fitted to the dataset, which includes the perturbed feature.\n", - " paths\n", - " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", - " instance\n", - " The simulation instance, which includes the perturbed feature that is used to simulate the dataset.\n", - " This is often not used, but may be useful for certain sensitivity mapping tasks, for example using\n", - " true values of the simulated instance to set up aspects of the model-fit (e.g. the priors).\n", - " \"\"\"\n", - "\n", - " search = af.Nautilus(\n", - " paths=paths,\n", - " n_live=50,\n", - " number_of_cores=self.number_of_cores,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - " analysis._adapt_images = self.adapt_images\n", - "\n", - " return search.fit(model=model, analysis=analysis)\n", - "\n", - "\n", - "def base_model_narrow_priors_from(base_model, result, stretch: float = 1.0):\n", - " \"\"\"\n", - " Returns a base model where priors are updated to `UniformPriors` with a `lower_limit` and `upper_limit` which\n", - " are a narrow range over the `simulation_instance` parameter values.\n", - "\n", - " Using this updated model can significiantly speed up sensitivity mapping, as it dramatically reduces\n", - " the volume of parameter space that needs to be sampled.\n", - "\n", - " The downside of this approach is that these narrow priors may remove viable solutions that alter the sensitivity\n", - " or lead to an inaccurate estimate of the Bayesian evidence.\n", - "\n", - " The size of each parameter bounds have been chosen based on previous lens modeling intuition. For example, I have\n", - " never seen a DM subhalo change the centre of a lens mass model by more than 0.01\", therefore this is the value\n", - " used for bounding that parameter.\n", - "\n", - " Parameters\n", - " ----------\n", - " base_model\n", - " The base model which will be used for sensitivity mapping, which this function updates to have narrower priors.\n", - " result\n", - " The result used to set up the base model and which is used to set these updated priors.\n", - " stretch\n", - " A multiplicative factor which can be used to shrink or broaden the priors more.\n", - "\n", - " Returns\n", - " -------\n", - " A base model with priors updated to narrow uniform priors.\n", - " \"\"\"\n", - "\n", - " if hasattr(base_model.galaxies.lens, \"mass\"):\n", - " base_model.galaxies.lens.mass.centre.centre_0 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.01 * stretch\n", - " ).galaxies.lens.mass.centre.centre_0\n", - " )\n", - " base_model.galaxies.lens.mass.centre.centre_1 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.01 * stretch\n", - " ).galaxies.lens.mass.centre.centre_1\n", - " )\n", - " base_model.galaxies.lens.mass.ell_comps.ell_comps_0 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.05 * stretch\n", - " ).galaxies.lens.mass.ell_comps.ell_comps_0\n", - " )\n", - " base_model.galaxies.lens.mass.ell_comps.ell_comps_1 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.05 * stretch\n", - " ).galaxies.lens.mass.ell_comps.ell_comps_1\n", - " )\n", - " base_model.galaxies.lens.mass.einstein_radius = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.1 * stretch\n", - " ).galaxies.lens.mass.einstein_radius\n", - " )\n", - " base_model.galaxies.lens.mass.slope = result.model_centred_max_lh_bounded(\n", - " b=0.1 * stretch\n", - " ).galaxies.lens.mass.slope\n", - "\n", - " if hasattr(base_model.galaxies.lens, \"shear\"):\n", - " base_model.galaxies.lens.shear.gamma_1 = result.model_centred_max_lh_bounded(\n", - " b=0.05 * stretch\n", - " ).galaxies.lens.shear.gamma_1\n", - " base_model.galaxies.lens.shear.gamma_2 = result.model_centred_max_lh_bounded(\n", - " b=0.05 * stretch\n", - " ).galaxies.lens.shear.gamma_2\n", - "\n", - " if hasattr(base_model.galaxies.source, \"bulge\"):\n", - " base_model.galaxies.source.bulge.centre.centre_0 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.1 * stretch\n", - " ).galaxies.source.bulge.centre.centre_0\n", - " )\n", - " base_model.galaxies.source.bulge.centre.centre_1 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.1 * stretch\n", - " ).galaxies.source.bulge.centre.centre_1\n", - " )\n", - " base_model.galaxies.source.bulge.ell_comps.ell_comps_0 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.1 * stretch\n", - " ).galaxies.source.bulge.ell_comps.ell_comps_0\n", - " )\n", - " base_model.galaxies.source.bulge.ell_comps.ell_comps_1 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.1 * stretch\n", - " ).galaxies.source.bulge.ell_comps.ell_comps_1\n", - " )\n", - " base_model.galaxies.source.bulge.effective_radius = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.2 * stretch\n", - " ).galaxies.source.bulge.effective_radius\n", - " )\n", - " base_model.galaxies.source.bulge.sersic_index = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.2 * stretch\n", - " ).galaxies.source.bulge.sersic_index\n", - " )\n", - "\n", - " if base_model.galaxies.source.bulge.effective_radius.lower_limit < 0.0:\n", - " base_model.galaxies.source.bulge.effective_radius = af.UniformPrior(\n", - " lower_limit=0.0,\n", - " upper_limit=base_model.galaxies.source.bulge.effective_radius.upper_limit,\n", - " )\n", - "\n", - " if base_model.galaxies.source.bulge.sersic_index.lower_limit < 0.0:\n", - " base_model.galaxies.source.bulge.sersic_index = af.UniformPrior(\n", - " lower_limit=0.0,\n", - " upper_limit=base_model.galaxies.source.bulge.sersic_index.upper_limit,\n", - " )\n", - "\n", - " return base_model\n", - "\n", - "\n", - "def visualize_sensitivity(\n", - " result,\n", - " paths: af.DirectoryPaths,\n", - " mass_result: af.Result,\n", - " mask: al.Mask2D,\n", - "):\n", - " \"\"\"\n", - " Visualize the results of strong lens sensitivity mapping via the SLaM pipeline.\n", - "\n", - " Parameters\n", - " ----------\n", - " result\n", - " The result of the sensitivity mapping, which contains grids of the log evidence and log likelihood differences.\n", - " paths\n", - " The paths object which defines the output path for the results of the sensitivity mapping.\n", - " mass_result\n", - " The result of the mass pipeline, which is used to subtract the lens light from the dataset.\n", - " mask\n", - " The mask used to mask the dataset, which is plotted over the lens subtracted image.\n", - " \"\"\"\n", - "\n", - " result = al.SubhaloSensitivityResult(\n", - " result=result,\n", - " )\n", - "\n", - " data_subtracted = (\n", - " mass_result.max_log_likelihood_fit.subtracted_images_of_planes_list[-1]\n", - " )\n", - "\n", - " data_subtracted = data_subtracted.apply_mask(mask=mask)\n", - "\n", - " aplt.subplot_sensitivity(result=result, data_subtracted=data_subtracted)\n", - "\n", - "\n", - "class Visualizer:\n", - " def __init__(self, mass_result: af.Result, mask: al.Mask2D):\n", - " \"\"\"\n", - " Performs on-the-fly visualization of the sensitivity mapping, outputting the results of the sensitivity\n", - " mapping so far to hard disk after each sensitivity cell fit is complete.\n", - "\n", - " Parameters\n", - " ----------\n", - " mass_result\n", - " The result of the SLaM MASS PIPELINE which ran before this pipeline.\n", - " mask\n", - " The Mask2D that is applied to the imaging data for model-fitting.\n", - " \"\"\"\n", - "\n", - " self.mass_result = mass_result\n", - " self.mask = mask\n", - "\n", - " def __call__(self, sensitivity_result, paths: af.DirectoryPaths):\n", - " \"\"\"\n", - " Called by the `Sensitivity` class after every sensitivity cell has been fitted, to visualize results so far.\n", - "\n", - " Parameters\n", - " ----------\n", - " sensitivity_result\n", - " The result of the sensitivity mapping search so far.\n", - " paths\n", - " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", - " \"\"\"\n", - " visualize_sensitivity(\n", - " result=sensitivity_result,\n", - " paths=paths,\n", - " mass_result=self.mass_result,\n", - " mask=self.mask,\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"dark_matter_subhalo\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_Path().exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/advanced/subhalo/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.05,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_result_for_lens=source_lp_result,\n", - " source_result_for_source=source_lp_result,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_lp_result,\n", - " source_result_for_source=source_lp_result,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (sensitivity mapping)__\n", - "\n", - "The SUBHALO PIPELINE (sensitivity mapping) performs sensitivity mapping of the data using the lens model\n", - "fitted above, so as to determine where subhalos of what mass could be detected in the data. A full description of\n", - "Sensitivity mapping if given in the SLaM pipeline script `slam/subhalo/sensitivity_imaging.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "subhalo_mass = af.Model(al.mp.NFWMCRLudlowSph)\n", - "grid_dimension_arcsec = 3.0\n", - "number_of_steps = 2\n", - "sensitivity_mask = None\n", - "\n", - "base_model = mass_result.model\n", - "\n", - "base_model = base_model_narrow_priors_from(base_model=base_model, result=mass_result)\n", - "\n", - "perturb_model = af.Model(al.Galaxy, redshift=0.5, mass=subhalo_mass)\n", - "\n", - "perturb_model.mass.log10m_vir = 9.0\n", - "perturb_model.mass.c_gNFW = 12.0\n", - "perturb_model.mass.overdens = 200.0\n", - "perturb_model.mass.inner_slope = 2.2\n", - "\n", - "perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - ")\n", - "perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - ")\n", - "perturb_model.mass.redshift_object = mass_result.model.galaxies.lens.redshift\n", - "perturb_model.mass.redshift_source = mass_result.model.galaxies.source.redshift\n", - "\n", - "\n", - "def perturb_model_prior_func(perturb_instance, perturb_model):\n", - " b = 0.05\n", - "\n", - " perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", - " lower_limit=perturb_instance.mass.centre[0] - b,\n", - " upper_limit=perturb_instance.mass.centre[0] + b,\n", - " )\n", - "\n", - " perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", - " lower_limit=perturb_instance.mass.centre[1] - b,\n", - " upper_limit=perturb_instance.mass.centre[1] + b,\n", - " )\n", - "\n", - " perturb_model.mass.log10m_vir = af.UniformPrior(lower_limit=6, upper_limit=12)\n", - "\n", - " return perturb_model\n", - "\n", - "\n", - "simulation_instance = mass_result.instance\n", - "\n", - "fit = mass_result.max_log_likelihood_fit\n", - "\n", - "simulation_instance.galaxies.lens = (\n", - " fit.model_obj_linear_light_profiles_to_light_profiles.galaxies[0]\n", - ")\n", - "simulation_instance.galaxies.source = (\n", - " fit.model_obj_linear_light_profiles_to_light_profiles.galaxies[-1]\n", - ")\n", - "\n", - "paths = af.DirectoryPaths(\n", - " name=f\"subhalo__sensitivity\",\n", - " path_prefix=settings_search.path_prefix,\n", - " unique_tag=settings_search.unique_tag,\n", - ")\n", - "\n", - "sensitivity = af.Sensitivity(\n", - " paths=paths,\n", - " simulation_instance=simulation_instance,\n", - " base_model=base_model,\n", - " perturb_model=perturb_model,\n", - " simulate_cls=SimulateImaging(mask=mask, psf=dataset.psf),\n", - " base_fit_cls=BaseFit(\n", - " adapt_images=None, number_of_cores=settings_search.number_of_cores\n", - " ),\n", - " perturb_fit_cls=PerturbFit(\n", - " adapt_images=None, number_of_cores=settings_search.number_of_cores\n", - " ),\n", - " perturb_model_prior_func=perturb_model_prior_func,\n", - " visualizer_cls=Visualizer(mass_result=mass_result, mask=mask),\n", - " number_of_steps=number_of_steps,\n", - " batch_range=None,\n", - " mask=sensitivity_mask,\n", - ")\n", - "\n", - "subhalo_results = sensitivity.run()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM (Source, Light and Mass): Subhalo Source Parametric Sensitivity Mapping\n", + "============================================================================\n", + "\n", + "This example illustrates how to perform DM subhalo sensitivity mapping using a SLaM pipeline for a source modeling\n", + "using light profiles.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except the lens mass uses an `Isothermal` with its centre fixed.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **Simulate Function Class:** We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above).\n", + "- **Base Fit:** We have defined a `Simulate` class that will be used to simulate every dataset simulated by the.\n", + "- **Perturb Fit:** We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", + "\n", + "__Model__\n", + "\n", + "Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE, MASS TOTAL PIPELINE and SUBHALO PIPELINE this SLaM script\n", + "fits `Imaging` of a strong lens system, where in the final model:\n", + "\n", + " - The lens galaxy's light is an MGE bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - A dark matter subhalo near The lens galaxy mass is included as a`NFWMCRLudlowSph`.\n", + " - The source galaxy is an `Inversion`.\n", + "\n", + "It ends by performing sensitivity mapping of the data using the above model, so as to determine where in the data\n", + "subhalos of a given mass could have been detected if present.\n", + "\n", + "This uses the SLaM pipelines:\n", + "\n", + " `source_lp`\n", + " `source_pix`\n", + " `light_lp`\n", + " `mass_total`\n", + " `subhalo/detection`\n", + "\n", + "Check them out for a full description of the analysis!\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `subhalo/detect/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import os\n", + "from pathlib import Path\n", + "from typing import List, Optional, Tuple, Union\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except the lens mass uses an `Isothermal` with its centre fixed to (0.0, 0.0)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search,\n", + " dataset,\n", + " mask_radius,\n", + " redshift_lens,\n", + " redshift_source,\n", + "):\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = (0.0, 0.0)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " disk=None,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search,\n", + " dataset,\n", + " mask_radius,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + "):\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " )\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search,\n", + " dataset,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " light_result,\n", + "):\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=light_result.instance.galaxies.lens.bulge,\n", + " disk=light_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate Function Class__\n", + "\n", + "We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above) and uses it to\n", + "simulate a strong lens dataset, which include a dark matter subhalo, which is subsequently fitted.\n", + "\n", + "Additional attributes required to simulate the data (mask, PSF) can be passed to the `__init__` method, and the\n", + "simulation is performed in the `__call__` method.\n", + "\n", + "When this dataset is simulated, the quantity `instance.perturb` is used in `__call__`. This is an instance\n", + "of the `NFWMCRLudlowSph`, and it is different every time the `simulate_cls` is called based on the value of sensitivity\n", + "being computed.\n", + "\n", + "In this example, this `instance.perturb` corresponds to two different subhalos with values of `mass_at_200` of\n", + "1e6 MSun and 1e13 MSun." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class SimulateImaging:\n", + " def __init__(self, mask, psf):\n", + " \"\"\"\n", + " Class used to simulate the strong lens imaging used for sensitivity mapping.\n", + "\n", + " Parameters\n", + " ----------\n", + " mask\n", + " The mask applied to the real image data, which is applied to every simulated imaging.\n", + " psf\n", + " The PSF of the real image data, which is applied to every simulated imaging and used for each fit.\n", + " \"\"\"\n", + " self.mask = mask\n", + " self.psf = psf\n", + "\n", + " def __call__(self, instance: af.ModelInstance, simulate_path: str):\n", + " \"\"\"\n", + " The `simulate_function` called by the `Sensitivity` class which simulates each strong lens image fitted\n", + " by the sensitivity mapper.\n", + "\n", + " The simulation procedure is as follows:\n", + "\n", + " 1) Use the input galaxies of the sensitivity `instance` to set up a tracer, which generates the image-plane\n", + " image of the strong lens system.\n", + "\n", + " 2) Simulate this image using the input dataset noise (Poisson) and PSF.\n", + "\n", + " 3) Apply the mask used in the analysis of the real image to the simulated image.\n", + "\n", + " 4) Output information about the simulation to hard-disk.\n", + "\n", + " The `subhalo` in the sensitivity `instance` changes for every iteration of the sensitivity mapping, ensuring\n", + " that we map out the sensitivity of the analysis to the subhalo properties (centre, mass, etc.).\n", + "\n", + " Parameters\n", + " ----------\n", + " instance\n", + " The sensitivity instance, which includes the galaxies whose parameters are varied to perform sensitivity.\n", + " The subhalo in this instance changes for every iteration of the sensitivity mapping.\n", + " simulate_path\n", + " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", + " mapping.\n", + "\n", + " Returns\n", + " -------\n", + " A simulated image of a strong lens, which id input into the fits of the sensitivity mapper.\n", + " \"\"\"\n", + "\n", + " \"\"\"\n", + " __Resume Fit__\n", + "\n", + " If sensitivity mapping already began on this grid cell, the dataset will have been simulated already and we\n", + " do not want to resimulate it and change its noise properties.\n", + "\n", + " We therefore load it from the `simulate_path` instead.\n", + " \"\"\"\n", + " try:\n", + " dataset = al.Imaging.from_fits(\n", + " data_path=f\"{simulate_path}/data.fits\",\n", + " psf_path=f\"{simulate_path}/psf.fits\",\n", + " noise_map_path=f\"{simulate_path}/noise_map.fits\",\n", + " pixel_scales=self.mask.pixel_scales,\n", + " check_noise_map=False,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=self.mask)\n", + "\n", + " tracer = al.Tracer(\n", + " galaxies=[\n", + " instance.galaxies.lens,\n", + " instance.perturb,\n", + " instance.galaxies.source,\n", + " ]\n", + " )\n", + "\n", + " traced_grid = tracer.traced_grid_2d_list_from(\n", + " grid=dataset.grid,\n", + " )[-1]\n", + "\n", + " source_centre = instance.galaxies.source.bulge.centre\n", + "\n", + " over_sample_size = (\n", + " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=traced_grid,\n", + " sub_size_list=[16, 8, 2],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=[source_centre],\n", + " )\n", + " )\n", + "\n", + " over_sample_size_lens = (\n", + " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[8, 4, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + " )\n", + "\n", + " over_sample_size = np.where(\n", + " over_sample_size > over_sample_size_lens,\n", + " over_sample_size,\n", + " over_sample_size_lens,\n", + " )\n", + " over_sample_size = al.Array2D(values=over_sample_size, mask=self.mask)\n", + "\n", + " return dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " except FileNotFoundError:\n", + " pass\n", + "\n", + " \"\"\"\n", + " __Ray Tracing__\n", + "\n", + " Set up the `Tracer` which is used to simulate the strong lens imaging, which may include the subhalo in\n", + " addition to the lens and source galaxy.\n", + " \"\"\"\n", + " tracer = al.Tracer(\n", + " galaxies=[\n", + " instance.galaxies.lens,\n", + " instance.perturb,\n", + " instance.galaxies.source,\n", + " ]\n", + " )\n", + "\n", + " \"\"\"\n", + " __Simulate__\n", + "\n", + " Set up the grid, PSF and simulator settings used to simulate imaging of the strong lens. These should be\n", + " tuned to match the S/N and noise properties of the observed data you are performing sensitivity mapping on.\n", + " \"\"\"\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=self.mask.shape_native,\n", + " pixel_scales=self.mask.pixel_scales,\n", + " )\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " simulator = al.SimulatorImaging(\n", + " exposure_time=1000.0,\n", + " psf=self.psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + " noise_seed=1,\n", + " )\n", + "\n", + " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + " \"\"\"\n", + " __Masking__\n", + "\n", + " The data generated by the simulate function is what is ultimately fitted.\n", + "\n", + " Therefore, we also apply the mask for the analysis before we return the simulated data.\n", + " \"\"\"\n", + " dataset = dataset.apply_mask(mask=self.mask)\n", + "\n", + " traced_grid = tracer.traced_grid_2d_list_from(\n", + " grid=dataset.grid,\n", + " )[-1]\n", + "\n", + " source_centre = instance.galaxies.source.bulge.centre\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=traced_grid,\n", + " sub_size_list=[16, 8, 2],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=[source_centre],\n", + " )\n", + "\n", + " over_sample_size_lens = (\n", + " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[8, 4, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + " )\n", + "\n", + " over_sample_size = np.where(\n", + " over_sample_size > over_sample_size_lens,\n", + " over_sample_size,\n", + " over_sample_size_lens,\n", + " )\n", + " over_sample_size = al.Array2D(values=over_sample_size, mask=self.mask)\n", + "\n", + " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " \"\"\"\n", + " Outputs info about the `Tracer` to the fit, so we know exactly how we simulated the image.\n", + " \"\"\"\n", + " self.output_info(simulate_path=simulate_path, dataset=dataset, tracer=tracer)\n", + "\n", + " return dataset\n", + "\n", + " def output_info(self, simulate_path: str, dataset: al.Imaging, tracer: al.Imaging):\n", + " \"\"\"\n", + " Output information about the data simulated for this iteration of sensitivity mapping.\n", + "\n", + " This information output is as follows:\n", + "\n", + " - A subplot of the simulated imaging dataset.\n", + " - A subplot of the tracer used to simulate this imaging dataset.\n", + " - A .json file containing the tracer galaxies.\n", + " - Output the simulated dataset to .fits files which are used to load the data if a run is resumed.\n", + "\n", + " Parameters\n", + " ----------\n", + " simulate_path\n", + " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", + " mapping.\n", + " dataset\n", + " The simulated imaging dataset which is visualized.\n", + " tracer\n", + " The tracer used to simulate the imaging dataset, which is visualized and output to a .json file.\n", + " \"\"\"\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " aplt.subplot_lensed_images(\n", + " tracer=tracer,\n", + " grid=dataset.grid,\n", + " output_path=simulate_path,\n", + " output_format=\"png\",\n", + " )\n", + "\n", + " al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(simulate_path) / \"tracer.json\",\n", + " )\n", + "\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=Path(simulate_path) / \"data.fits\",\n", + " psf_path=Path(simulate_path) / \"psf.fits\",\n", + " noise_map_path=Path(simulate_path) / \"noise_map.fits\",\n", + " overwrite=True,\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Base Fit__\n", + "\n", + "We have defined a `Simulate` class that will be used to simulate every dataset simulated by the sensitivity mapper.\n", + "Each simulated dataset will have a unique set of parameters for the `subhalo` (e.g. due to different values of\n", + "`perturb_model`.\n", + "\n", + "We will fit each simulated dataset using the `base_model`, which quantifies whether not including the dark matter\n", + "subhalo in the model changess the goodness-of-fit and therefore indicates if we are sensitive to the subhalo.\n", + "\n", + "We now write a `BaseFit` class, defining how the `base_model` is fitted to each simulated dataset and the\n", + "goodness-of-fit used to quantify whether the model fits the data well. As above, the `__init__` method can be\n", + "extended with new inputs to control how the model is fitted and the `__call__` method performs the fit.\n", + "\n", + "In this example, we use a full non-linear search to fit the `base_model` to the simulated data and return\n", + "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", + "more computationally efficient, for example performing a single log likelihood evaluation of the `base_model` fit\n", + "to the simulated data.\n", + "\n", + "Fucntionality which adapts the mesh and regularization of a pixelized source reconstruction to the unlensed source's\n", + "morphology require an `adapt_images`. This is an input of the __init__ constructor which is passed to the `Analysis`\n", + "for every simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class BaseFit:\n", + " def __init__(self, adapt_images, number_of_cores: int = 1):\n", + " \"\"\"\n", + " Class used to fit every dataset used for sensitivity mapping with the base model (the model without the\n", + " perturbed feature sensitivity mapping maps out).\n", + "\n", + " In this example, the base model therefore does not include the dark matter subhalo, but the simulated\n", + " dataset includes one.\n", + "\n", + " The base fit is repeated for every parameter on the sensitivity grid and compared to the perturbed fit. This\n", + " maps out the sensitivity of every parameter is (e.g. the sensitivity of the mass of the subhalo).\n", + "\n", + " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", + " fitted, below we include an input `analysis_cls` which is the `AnalysisImaging` class used to fit the model\n", + " to the dataset.\n", + "\n", + " Parameters\n", + " ----------\n", + " adapt_images\n", + " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", + " reconstructed galaxy's morphology.\n", + " number_of_cores\n", + " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", + " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", + " \"\"\"\n", + " self.adapt_images = adapt_images\n", + " self.number_of_cores = number_of_cores\n", + "\n", + " def __call__(self, dataset, model, paths, instance):\n", + " \"\"\"\n", + " The base fitting function which fits every dataset used for sensitivity mapping with the base model.\n", + "\n", + " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", + " quantify how sensitive the model is to the perturbed feature.\n", + "\n", + " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", + " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", + "\n", + " Parameters\n", + " ----------\n", + " dataset\n", + " The dataset which is simulated with the perturbed model and which is fitted.\n", + " model\n", + " The model instance which is fitted to the dataset, which does not include the perturbed feature.\n", + " paths\n", + " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", + " instance\n", + " The simulation instance, which includes the perturbed feature that is used to simulate the dataset.\n", + " This is often not used, but may be useful for certain sensitivity mapping tasks, for example using\n", + " true values of the simulated instance to set up aspects of the model-fit (e.g. the priors).\n", + " \"\"\"\n", + "\n", + " search = af.Nautilus(\n", + " paths=paths,\n", + " n_live=50,\n", + " number_of_cores=self.number_of_cores,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + " analysis._adapt_images = self.adapt_images\n", + "\n", + " return search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Perturb Fit__\n", + "\n", + "We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each simulated dataset. This\n", + "behaves analogously to the `BaseFit` class above, but now fits the `perturb_model` to the simulated data (as\n", + "opposed to the `base_model`).\n", + "\n", + "Again, in this example we use a full non-linear search to fit the `perturb_model` to the simulated data and return\n", + "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", + "more computationally efficient, for example performing a single log likelihood evaluation of the `perturb_model` fit\n", + "to the simulated data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class PerturbFit:\n", + " def __init__(self, adapt_images, number_of_cores: int = 1):\n", + " \"\"\"\n", + " Class used to fit every dataset used for sensitivity mapping with the perturbed model (the model with the\n", + " perturbed feature sensitivity mapping maps out).\n", + "\n", + " In this example, the perturbed model therefore includes the dark matter subhalo, which is also in the\n", + " simulated dataset.\n", + "\n", + " The perturbed fit is repeated for every parameter on the sensitivity grid and compared to the base fit. This\n", + " maps out the sensitivity of every parameter is (e.g. the sensitivity of mass of the subhalo).\n", + "\n", + " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", + " fitted, below we include an input `analysis_cls` which is the `Analysis` class used to fit the model to the\n", + " dataset.\n", + "\n", + " Parameters\n", + " ----------\n", + " adapt_images\n", + " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", + " reconstructed galaxy's morphology.\n", + " number_of_cores\n", + " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", + " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", + " \"\"\"\n", + " self.adapt_images = adapt_images\n", + " self.number_of_cores = number_of_cores\n", + "\n", + " def __call__(self, dataset, model, paths, instance):\n", + " \"\"\"\n", + " The perturbed fitting function which fits every dataset used for sensitivity mapping with the perturbed model.\n", + "\n", + " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", + " quantify how sensitive the model is to the perturbed feature.\n", + "\n", + " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", + " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", + "\n", + " Parameters\n", + " ----------\n", + " dataset\n", + " The dataset which is simulated with the perturbed model and which is fitted.\n", + " model\n", + " The model instance which is fitted to the dataset, which includes the perturbed feature.\n", + " paths\n", + " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", + " instance\n", + " The simulation instance, which includes the perturbed feature that is used to simulate the dataset.\n", + " This is often not used, but may be useful for certain sensitivity mapping tasks, for example using\n", + " true values of the simulated instance to set up aspects of the model-fit (e.g. the priors).\n", + " \"\"\"\n", + "\n", + " search = af.Nautilus(\n", + " paths=paths,\n", + " n_live=50,\n", + " number_of_cores=self.number_of_cores,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + " analysis._adapt_images = self.adapt_images\n", + "\n", + " return search.fit(model=model, analysis=analysis)\n", + "\n", + "\n", + "def base_model_narrow_priors_from(base_model, result, stretch: float = 1.0):\n", + " \"\"\"\n", + " Returns a base model where priors are updated to `UniformPriors` with a `lower_limit` and `upper_limit` which\n", + " are a narrow range over the `simulation_instance` parameter values.\n", + "\n", + " Using this updated model can significiantly speed up sensitivity mapping, as it dramatically reduces\n", + " the volume of parameter space that needs to be sampled.\n", + "\n", + " The downside of this approach is that these narrow priors may remove viable solutions that alter the sensitivity\n", + " or lead to an inaccurate estimate of the Bayesian evidence.\n", + "\n", + " The size of each parameter bounds have been chosen based on previous lens modeling intuition. For example, I have\n", + " never seen a DM subhalo change the centre of a lens mass model by more than 0.01\", therefore this is the value\n", + " used for bounding that parameter.\n", + "\n", + " Parameters\n", + " ----------\n", + " base_model\n", + " The base model which will be used for sensitivity mapping, which this function updates to have narrower priors.\n", + " result\n", + " The result used to set up the base model and which is used to set these updated priors.\n", + " stretch\n", + " A multiplicative factor which can be used to shrink or broaden the priors more.\n", + "\n", + " Returns\n", + " -------\n", + " A base model with priors updated to narrow uniform priors.\n", + " \"\"\"\n", + "\n", + " if hasattr(base_model.galaxies.lens, \"mass\"):\n", + " base_model.galaxies.lens.mass.centre.centre_0 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.01 * stretch\n", + " ).galaxies.lens.mass.centre.centre_0\n", + " )\n", + " base_model.galaxies.lens.mass.centre.centre_1 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.01 * stretch\n", + " ).galaxies.lens.mass.centre.centre_1\n", + " )\n", + " base_model.galaxies.lens.mass.ell_comps.ell_comps_0 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.05 * stretch\n", + " ).galaxies.lens.mass.ell_comps.ell_comps_0\n", + " )\n", + " base_model.galaxies.lens.mass.ell_comps.ell_comps_1 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.05 * stretch\n", + " ).galaxies.lens.mass.ell_comps.ell_comps_1\n", + " )\n", + " base_model.galaxies.lens.mass.einstein_radius = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.1 * stretch\n", + " ).galaxies.lens.mass.einstein_radius\n", + " )\n", + " base_model.galaxies.lens.mass.slope = result.model_centred_max_lh_bounded(\n", + " b=0.1 * stretch\n", + " ).galaxies.lens.mass.slope\n", + "\n", + " if hasattr(base_model.galaxies.lens, \"shear\"):\n", + " base_model.galaxies.lens.shear.gamma_1 = result.model_centred_max_lh_bounded(\n", + " b=0.05 * stretch\n", + " ).galaxies.lens.shear.gamma_1\n", + " base_model.galaxies.lens.shear.gamma_2 = result.model_centred_max_lh_bounded(\n", + " b=0.05 * stretch\n", + " ).galaxies.lens.shear.gamma_2\n", + "\n", + " if hasattr(base_model.galaxies.source, \"bulge\"):\n", + " base_model.galaxies.source.bulge.centre.centre_0 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.1 * stretch\n", + " ).galaxies.source.bulge.centre.centre_0\n", + " )\n", + " base_model.galaxies.source.bulge.centre.centre_1 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.1 * stretch\n", + " ).galaxies.source.bulge.centre.centre_1\n", + " )\n", + " base_model.galaxies.source.bulge.ell_comps.ell_comps_0 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.1 * stretch\n", + " ).galaxies.source.bulge.ell_comps.ell_comps_0\n", + " )\n", + " base_model.galaxies.source.bulge.ell_comps.ell_comps_1 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.1 * stretch\n", + " ).galaxies.source.bulge.ell_comps.ell_comps_1\n", + " )\n", + " base_model.galaxies.source.bulge.effective_radius = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.2 * stretch\n", + " ).galaxies.source.bulge.effective_radius\n", + " )\n", + " base_model.galaxies.source.bulge.sersic_index = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.2 * stretch\n", + " ).galaxies.source.bulge.sersic_index\n", + " )\n", + "\n", + " if base_model.galaxies.source.bulge.effective_radius.lower_limit < 0.0:\n", + " base_model.galaxies.source.bulge.effective_radius = af.UniformPrior(\n", + " lower_limit=0.0,\n", + " upper_limit=base_model.galaxies.source.bulge.effective_radius.upper_limit,\n", + " )\n", + "\n", + " if base_model.galaxies.source.bulge.sersic_index.lower_limit < 0.0:\n", + " base_model.galaxies.source.bulge.sersic_index = af.UniformPrior(\n", + " lower_limit=0.0,\n", + " upper_limit=base_model.galaxies.source.bulge.sersic_index.upper_limit,\n", + " )\n", + "\n", + " return base_model\n", + "\n", + "\n", + "def visualize_sensitivity(\n", + " result,\n", + " paths: af.DirectoryPaths,\n", + " mass_result: af.Result,\n", + " mask: al.Mask2D,\n", + "):\n", + " \"\"\"\n", + " Visualize the results of strong lens sensitivity mapping via the SLaM pipeline.\n", + "\n", + " Parameters\n", + " ----------\n", + " result\n", + " The result of the sensitivity mapping, which contains grids of the log evidence and log likelihood differences.\n", + " paths\n", + " The paths object which defines the output path for the results of the sensitivity mapping.\n", + " mass_result\n", + " The result of the mass pipeline, which is used to subtract the lens light from the dataset.\n", + " mask\n", + " The mask used to mask the dataset, which is plotted over the lens subtracted image.\n", + " \"\"\"\n", + "\n", + " result = al.SubhaloSensitivityResult(\n", + " result=result,\n", + " )\n", + "\n", + " data_subtracted = (\n", + " mass_result.max_log_likelihood_fit.subtracted_images_of_planes_list[-1]\n", + " )\n", + "\n", + " data_subtracted = data_subtracted.apply_mask(mask=mask)\n", + "\n", + " aplt.subplot_sensitivity(result=result, data_subtracted=data_subtracted)\n", + "\n", + "\n", + "class Visualizer:\n", + " def __init__(self, mass_result: af.Result, mask: al.Mask2D):\n", + " \"\"\"\n", + " Performs on-the-fly visualization of the sensitivity mapping, outputting the results of the sensitivity\n", + " mapping so far to hard disk after each sensitivity cell fit is complete.\n", + "\n", + " Parameters\n", + " ----------\n", + " mass_result\n", + " The result of the SLaM MASS PIPELINE which ran before this pipeline.\n", + " mask\n", + " The Mask2D that is applied to the imaging data for model-fitting.\n", + " \"\"\"\n", + "\n", + " self.mass_result = mass_result\n", + " self.mask = mask\n", + "\n", + " def __call__(self, sensitivity_result, paths: af.DirectoryPaths):\n", + " \"\"\"\n", + " Called by the `Sensitivity` class after every sensitivity cell has been fitted, to visualize results so far.\n", + "\n", + " Parameters\n", + " ----------\n", + " sensitivity_result\n", + " The result of the sensitivity mapping search so far.\n", + " paths\n", + " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", + " \"\"\"\n", + " visualize_sensitivity(\n", + " result=sensitivity_result,\n", + " paths=paths,\n", + " mass_result=self.mass_result,\n", + " mask=self.mask,\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"dark_matter_subhalo\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_Path().exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/advanced/subhalo/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.05,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_result_for_lens=source_lp_result,\n", + " source_result_for_source=source_lp_result,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_lp_result,\n", + " source_result_for_source=source_lp_result,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (sensitivity mapping)__\n", + "\n", + "The SUBHALO PIPELINE (sensitivity mapping) performs sensitivity mapping of the data using the lens model\n", + "fitted above, so as to determine where subhalos of what mass could be detected in the data. A full description of\n", + "Sensitivity mapping if given in the SLaM pipeline script `slam/subhalo/sensitivity_imaging.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "subhalo_mass = af.Model(al.mp.NFWMCRLudlowSph)\n", + "grid_dimension_arcsec = 3.0\n", + "number_of_steps = 2\n", + "sensitivity_mask = None\n", + "\n", + "base_model = mass_result.model\n", + "\n", + "base_model = base_model_narrow_priors_from(base_model=base_model, result=mass_result)\n", + "\n", + "perturb_model = af.Model(al.Galaxy, redshift=0.5, mass=subhalo_mass)\n", + "\n", + "perturb_model.mass.log10m_vir = 9.0\n", + "perturb_model.mass.c_gNFW = 12.0\n", + "perturb_model.mass.overdens = 200.0\n", + "perturb_model.mass.inner_slope = 2.2\n", + "\n", + "perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + ")\n", + "perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + ")\n", + "perturb_model.mass.redshift_object = mass_result.model.galaxies.lens.redshift\n", + "perturb_model.mass.redshift_source = mass_result.model.galaxies.source.redshift\n", + "\n", + "\n", + "def perturb_model_prior_func(perturb_instance, perturb_model):\n", + " b = 0.05\n", + "\n", + " perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", + " lower_limit=perturb_instance.mass.centre[0] - b,\n", + " upper_limit=perturb_instance.mass.centre[0] + b,\n", + " )\n", + "\n", + " perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", + " lower_limit=perturb_instance.mass.centre[1] - b,\n", + " upper_limit=perturb_instance.mass.centre[1] + b,\n", + " )\n", + "\n", + " perturb_model.mass.log10m_vir = af.UniformPrior(lower_limit=6, upper_limit=12)\n", + "\n", + " return perturb_model\n", + "\n", + "\n", + "simulation_instance = mass_result.instance\n", + "\n", + "fit = mass_result.max_log_likelihood_fit\n", + "\n", + "simulation_instance.galaxies.lens = (\n", + " fit.model_obj_linear_light_profiles_to_light_profiles.galaxies[0]\n", + ")\n", + "simulation_instance.galaxies.source = (\n", + " fit.model_obj_linear_light_profiles_to_light_profiles.galaxies[-1]\n", + ")\n", + "\n", + "paths = af.DirectoryPaths(\n", + " name=f\"subhalo__sensitivity\",\n", + " path_prefix=settings_search.path_prefix,\n", + " unique_tag=settings_search.unique_tag,\n", + ")\n", + "\n", + "sensitivity = af.Sensitivity(\n", + " paths=paths,\n", + " simulation_instance=simulation_instance,\n", + " base_model=base_model,\n", + " perturb_model=perturb_model,\n", + " simulate_cls=SimulateImaging(mask=mask, psf=dataset.psf),\n", + " base_fit_cls=BaseFit(\n", + " adapt_images=None, number_of_cores=settings_search.number_of_cores\n", + " ),\n", + " perturb_fit_cls=PerturbFit(\n", + " adapt_images=None, number_of_cores=settings_search.number_of_cores\n", + " ),\n", + " perturb_model_prior_func=perturb_model_prior_func,\n", + " visualizer_cls=Visualizer(mass_result=mass_result, mask=mask),\n", + " number_of_steps=number_of_steps,\n", + " batch_range=None,\n", + " mask=sensitivity_mask,\n", + ")\n", + "\n", + "subhalo_results = sensitivity.run()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/subhalo/sensitivity/slam_source_pixelized.ipynb b/notebooks/imaging/features/advanced/subhalo/sensitivity/slam_source_pixelized.ipynb index e13179b8e..1ca3e1c42 100644 --- a/notebooks/imaging/features/advanced/subhalo/sensitivity/slam_source_pixelized.ipynb +++ b/notebooks/imaging/features/advanced/subhalo/sensitivity/slam_source_pixelized.ipynb @@ -1,1414 +1,1451 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM (Source, Light and Mass): Subhalo Source Pixelized Sensitivity Mapping\n", - "===========================================================================\n", - "\n", - "This example illustrates how to perform DM subhalo sensitivity mapping using a SLaM pipeline for a dataset where the\n", - "source is modeled using a pixelization.\n", - "\n", - "The sensitivity mapping simulation procedure for a pixelized source is different light profile sources. When pixelized\n", - "sources are used, the source reconstruction on the mesh is used, such that the simulations capture the irregular\n", - "morphologies of real source galaxies.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except the lens mass uses an `Isothermal` with its centre fixed.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **Simulate Function Class:** We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above).\n", - "- **Base Fit:** We have defined a `Simulate` class that will be used to simulate every dataset simulated by the.\n", - "- **Perturb Fit:** We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", - "\n", - "__Model__\n", - "\n", - "Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE, MASS TOTAL PIPELINE and SUBHALO PIPELINE this SLaM script\n", - "fits `Imaging` of a strong lens system, where in the final model:\n", - "\n", - " - The lens galaxy's light is an MGE bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - A dark matter subhalo near The lens galaxy mass is included as a`NFWMCRLudlowSph`.\n", - " - The source galaxy is an `Inversion`.\n", - "\n", - "This uses the SLaM pipelines:\n", - "\n", - " `source_lp`\n", - " `source_pix`\n", - " `light_lp`\n", - " `mass_total`\n", - " `subhalo/detection`\n", - "\n", - "Check them out for a full description of the analysis!\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `subhalo/detect/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import os\n", - "from pathlib import Path\n", - "from typing import List, Optional, Tuple, Union\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except the lens mass uses an `Isothermal` with its centre fixed to (0.0, 0.0)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search,\n", - " dataset,\n", - " mask_radius,\n", - " redshift_lens,\n", - " redshift_source,\n", - "):\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = (0.0, 0.0)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search,\n", - " dataset,\n", - " source_lp_result,\n", - " mesh_shape,\n", - "):\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_lp_result.model.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search,\n", - " dataset,\n", - " source_lp_result,\n", - " source_pix_result_1,\n", - " mesh_shape,\n", - "):\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", - " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", - " noise_map=dataset.noise_map,\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_pixelization=over_sampling,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search,\n", - " dataset,\n", - " mask_radius,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - "):\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " )\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search,\n", - " dataset,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " light_result,\n", - "):\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=light_result.instance.galaxies.lens.bulge,\n", - " disk=light_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate Function Class__\n", - "\n", - "We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above) and uses it to\n", - "simulate a strong lens dataset, which include a dark matter subhalo, which is subsequently fitted.\n", - "\n", - "Additional attributes required to simulate the data (mask, PSF) can be passed to the `__init__` method, and the\n", - "simulation is performed in the `__call__` method.\n", - "\n", - "When this dataset is simulated, the quantity `instance.perturb` is used in `__call__`. This is an instance\n", - "of the `NFWMCRLudlowSph`, and it is different every time the `simulate_cls` is called based on the value of sensitivity\n", - "being computed.\n", - "\n", - "In this example, this `instance.perturb` corresponds to two different subhalos with values of `mass_at_200` of\n", - "1e6 MSun and 1e13 MSun." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class SimulateImagingPixelized:\n", - " def __init__(\n", - " self,\n", - " mask,\n", - " psf,\n", - " inversion,\n", - " interpolated_pixelized_shape: Union[tuple, tuple] = (1001, 1001),\n", - " image_plane_subgrid_size=8,\n", - " ):\n", - " \"\"\"\n", - " Class used to simulate the strong lens dataset used for sensitivity mapping.\n", - "\n", - " Parameters\n", - " ----------\n", - " mask\n", - " The mask applied to the real image data, which is applied to every simulated dataset.\n", - " psf\n", - " The PSF of the real image data, which is applied to every simulated dataset and used for each fit.\n", - " inversion\n", - " The `Inversion` used to reconstruct the source of the real image, included the pixelized source\n", - " reconstruction on a mesh (e.g. Delaunay / Voronoi).\n", - " interpolated_pixelized_shape\n", - " The pixelized source reconstruction is interpolated from an irregular mesh to a rectangular uniform array\n", - " and grid of this shape.\n", - " image_plane_subgrid_size\n", - " The size of the subgrid used to create the image-plane grid, whereby multiple image pixels\n", - " are traced to the source-plane image and evaluated to compute the flux of the simulated image.\n", - " \"\"\"\n", - " self.mask = mask\n", - " self.psf = psf\n", - " self.inversion = inversion\n", - " self.interpolated_pixelized_shape = interpolated_pixelized_shape\n", - " self.image_plane_subgrid_size = image_plane_subgrid_size\n", - "\n", - " def __call__(self, instance: af.ModelInstance, simulate_path: str):\n", - " \"\"\"\n", - " The `simulate_function` called by the `Sensitivity` class which simulates each strong lens image fitted\n", - " by the sensitivity mapper.\n", - "\n", - " The simulation procedure is as follows:\n", - "\n", - " 1) Extract the pixelized reconstructed source of a previous fit, which is likely on an irregular mesh\n", - " (e.g. Delaunay / Voronoi), and interpolate the source emission onto a rectangular uniform array and grid.\n", - "\n", - " 1) Use the input galaxies of the sensitivity `instance` to set up a tracer, which generates the image-plane\n", - " image of the strong lens system.\n", - "\n", - " 2) Simulate this image using the input dataset noise (Poisson) and PSF.\n", - "\n", - " 3) Apply the mask used in the analysis of the real image to the simulated image.\n", - "\n", - " 4) Output information about the simulation to hard-disk.\n", - "\n", - " The `subhalo` in the sensitivity `instance` changes for every iteration of the sensitivity mapping, ensuring\n", - " that we map out the sensitivity of the analysis to the subhalo properties (centre, mass, etc.).\n", - "\n", - " Parameters\n", - " ----------\n", - " instance\n", - " The sensitivity instance, which includes the galaxies whose parameters are varied to perform sensitivity.\n", - " The subhalo in this instance changes for every iteration of the sensitivity mapping.\n", - " simulate_path\n", - " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", - " mapping.\n", - "\n", - " Returns\n", - " -------\n", - " A simulated image of a strong lens, which id input into the fits of the sensitivity mapper.\n", - " \"\"\"\n", - "\n", - " \"\"\"\n", - " __Resume Fit__\n", - "\n", - " If sensitivity mapping already began on this grid cell, the dataset will have been simulated already and we\n", - " do not want to resimulate it and change its noise properties.\n", - "\n", - " We therefore load it from the `simulate_path` instead.\n", - " \"\"\"\n", - " try:\n", - " dataset = al.Imaging.from_fits(\n", - " data_path=f\"{simulate_path}/data.fits\",\n", - " psf_path=f\"{simulate_path}/psf.fits\",\n", - " noise_map_path=f\"{simulate_path}/noise_map.fits\",\n", - " pixel_scales=self.mask.pixel_scales,\n", - " check_noise_map=False,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=self.mask)\n", - "\n", - " over_sample_size = (\n", - " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[8, 4, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - " )\n", - "\n", - " return dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size, over_sample_size_pixelization=4\n", - " )\n", - "\n", - " except FileNotFoundError:\n", - " pass\n", - "\n", - " \"\"\"\n", - " __Source Galaxy Image__\n", - "\n", - " We now load the source galaxy image from the pixelized inversion of a previous fit, which could be on an irregular\n", - " Delaunay or Voronoi mesh.\n", - "\n", - " Irregular meshes cannot be used to simulate lensed images of a source. Therefore, we interpolate the mesh to\n", - " a uniform grid of shape `interpolated_pixelized_shape`. This should be high resolution (e.g. 1000 x 1000)\n", - " to ensure the interpolated source array captures all structure resolved on the Delaunay / Voronoi mesh.\n", - "\n", - " Loads source array from previous reconstruction, maps to square and wraps in AutoLens Array.\n", - " Loads lens galaxy and perturb from provided instance\n", - " Loads source galaxy redshift and sets up a `galaxy` class object at that redshift.\n", - " \"\"\"\n", - "\n", - " mapper = self.inversion.cls_list_from(cls=al.Mapper)[0]\n", - "\n", - " mapper_valued = al.MapperValued(\n", - " mapper=mapper,\n", - " values=self.inversion.reconstruction_dict[mapper],\n", - " )\n", - "\n", - " source_image = mapper_valued.interpolated_array_from(\n", - " shape_native=self.interpolated_pixelized_shape,\n", - " extent=(-2.0, 2.0, -2.0, 2.0),\n", - " )\n", - "\n", - " \"\"\"\n", - " __Create Grids__\n", - "\n", - " To create the lensed image, we will ray-trace image pixels to the source-plane and interpolate them onto the\n", - " source galaxy image.\n", - "\n", - " We therefore need the image-plane grid of (y,x) coordinates.\n", - " \"\"\"\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=self.mask.shape_native,\n", - " pixel_scales=self.mask.pixel_scales,\n", - " over_sample_size=self.image_plane_subgrid_size,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Ray Tracing__\n", - "\n", - " We create a tracer, which will create the lensed grid we overlay the interpolated source galaxy image above\n", - " in order to create the lensed source galaxy image.\n", - "\n", - " This creates the grid we will overlay the source image, in order to created the lensed source image.\n", - "\n", - " The source-plane requires a source-galaxy with a `redshift` in order for the tracer to trace it. We therefore\n", - " make one, noting it has no light profiles because its emission is entirely defined by the source galaxy image.\n", - " \"\"\"\n", - " tracer = al.Tracer(\n", - " galaxies=[\n", - " instance.galaxies.lens,\n", - " instance.perturb,\n", - " al.Galaxy(redshift=instance.galaxies.source.redshift),\n", - " ]\n", - " )\n", - "\n", - " \"\"\"\n", - " __Simulate__\n", - "\n", - " Using the tracer above, we create the image of the lensed source galaxy on the image-plane grid. This\n", - " uses the `source_image` and therefore capture the source's irregular and asymmetric morphological features\n", - " which the source reconstruction procedure fitted.\n", - "\n", - " Set up the grid, PSF and simulator settings used to simulate imaging of the strong lens. These should be\n", - " tuned to match the S/N and noise properties of the observed data you are performing sensitivity mapping on.\n", - "\n", - " The `SimulatorImaging` will be passed directly the image of the strong lens we created above, which\n", - " will be convolved with the psf before noise is added.\n", - "\n", - " To ensure the PSF convolution extends over the whole image, the image is padded before convolution to mitigate\n", - " edge effects and trimmed after the simulation so it retains the original `shape_native`.\n", - " \"\"\"\n", - " simulator = al.SimulatorImaging(\n", - " exposure_time=1000.0,\n", - " psf=self.psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - " noise_seed=1,\n", - " )\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " dataset = simulator.via_source_image_from(\n", - " tracer=tracer, grid=grid, source_image=source_image\n", - " )\n", - "\n", - " \"\"\"\n", - " __Masking__\n", - "\n", - " The data generated by the simulate function is what is ultimately fitted.\n", - "\n", - " Therefore, we also apply the mask for the analysis before we return the simulated data.\n", - " \"\"\"\n", - " dataset = dataset.apply_mask(mask=self.mask)\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[8, 4, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(\n", - " over_sample_size_lp=over_sample_size, over_sample_size_pixelization=4\n", - " )\n", - "\n", - " \"\"\"\n", - " Outputs info about the `Tracer` to the fit, so we know exactly how we simulated the image.\n", - " \"\"\"\n", - " tracer_no_perturb = al.Tracer(\n", - " galaxies=[\n", - " instance.galaxies.lens,\n", - " al.Galaxy(redshift=instance.galaxies.source.redshift),\n", - " ]\n", - " )\n", - "\n", - " self.output_info(\n", - " simulate_path=simulate_path,\n", - " grid=grid,\n", - " dataset=dataset,\n", - " source_image=source_image,\n", - " tracer=tracer,\n", - " tracer_no_perturb=tracer_no_perturb,\n", - " )\n", - "\n", - " return dataset\n", - "\n", - " def output_info(\n", - " self,\n", - " simulate_path: str,\n", - " grid: al.Grid2D,\n", - " dataset: al.Imaging,\n", - " source_image: al.Array2D,\n", - " tracer: al.Tracer,\n", - " tracer_no_perturb: al.Tracer,\n", - " ):\n", - " \"\"\"\n", - " Output information about the data simulated for this iteration of sensitivity mapping.\n", - "\n", - " This information output is as follows:\n", - "\n", - " - A subplot of the simulated imaging dataset.\n", - " - A subplot of the tracer used to simulate this imaging dataset.\n", - " - A .json file containing the tracer galaxies.\n", - " - Output the simulated dataset to .fits files which are used to load the data if a run is resumed.\n", - "\n", - " Parameters\n", - " ----------\n", - " simulate_path\n", - " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", - " mapping.\n", - " dataset\n", - " The simulated dataset dataset which is visualized.\n", - " tracer\n", - " The tracer used to simulate the dataset dataset, which is visualized and output to a .json file.\n", - " \"\"\"\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(simulate_path) / \"tracer.json\",\n", - " )\n", - "\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=Path(simulate_path) / \"data.fits\",\n", - " psf_path=Path(simulate_path) / \"psf.fits\",\n", - " noise_map_path=Path(simulate_path) / \"noise_map.fits\",\n", - " overwrite=True,\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Base Fit__\n", - "\n", - "We have defined a `Simulate` class that will be used to simulate every dataset simulated by the sensitivity mapper.\n", - "Each simulated dataset will have a unique set of parameters for the `subhalo` (e.g. due to different values of\n", - "`perturb_model`.\n", - "\n", - "We will fit each simulated dataset using the `base_model`, which quantifies whether not including the dark matter\n", - "subhalo in the model changess the goodness-of-fit and therefore indicates if we are sensitive to the subhalo.\n", - "\n", - "We now write a `BaseFit` class, defining how the `base_model` is fitted to each simulated dataset and the\n", - "goodness-of-fit used to quantify whether the model fits the data well. As above, the `__init__` method can be\n", - "extended with new inputs to control how the model is fitted and the `__call__` method performs the fit.\n", - "\n", - "In this example, we use a full non-linear search to fit the `base_model` to the simulated data and return\n", - "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", - "more computationally efficient, for example performing a single log likelihood evaluation of the `base_model` fit\n", - "to the simulated data.\n", - "\n", - "Fucntionality which adapts the mesh and regularization of a pixelized source reconstruction to the unlensed source's\n", - "morphology require an `adapt_images`. This is an input of the __init__ constructor which is passed to the `Analysis`\n", - "for every simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class BaseFit:\n", - " def __init__(self, adapt_images, number_of_cores: int = 1):\n", - " \"\"\"\n", - " Class used to fit every dataset used for sensitivity mapping with the base model (the model without the\n", - " perturbed feature sensitivity mapping maps out).\n", - "\n", - " In this example, the base model therefore does not include the dark matter subhalo, but the simulated\n", - " dataset includes one.\n", - "\n", - " The base fit is repeated for every parameter on the sensitivity grid and compared to the perturbed fit. This\n", - " maps out the sensitivity of every parameter is (e.g. the sensitivity of the mass of the subhalo).\n", - "\n", - " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", - " fitted, below we include an input `analysis_cls` which is the `AnalysisImaging` class used to fit the model\n", - " to the dataset.\n", - "\n", - " Parameters\n", - " ----------\n", - " adapt_images\n", - " The result of the previous search containing adapt images used to adapt certain pixelized source meshs's\n", - " and regularizations to the unlensed source morphology.\n", - " number_of_cores\n", - " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", - " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", - " \"\"\"\n", - " self.adapt_images = adapt_images\n", - " self.number_of_cores = number_of_cores\n", - "\n", - " def __call__(self, dataset, model, paths, instance):\n", - " \"\"\"\n", - " The base fitting function which fits every dataset used for sensitivity mapping with the base model.\n", - "\n", - " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", - " quantify how sensitive the model is to the perturbed feature.\n", - "\n", - " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", - " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", - "\n", - " Parameters\n", - " ----------\n", - " dataset\n", - " The dataset which is simulated with the perturbed model and which is fitted.\n", - " model\n", - " The model instance which is fitted to the dataset, which does not include the perturbed feature.\n", - " paths\n", - " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", - " \"\"\"\n", - "\n", - " search = af.Nautilus(\n", - " paths=paths,\n", - " n_live=50,\n", - " number_of_cores=self.number_of_cores,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - " analysis._adapt_images = self.adapt_images\n", - "\n", - " return search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Perturb Fit__\n", - "\n", - "We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each simulated dataset. This\n", - "behaves analogously to the `BaseFit` class above, but now fits the `perturb_model` to the simulated data (as\n", - "opposed to the `base_model`).\n", - "\n", - "Again, in this example we use a full non-linear search to fit the `perturb_model` to the simulated data and return\n", - "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", - "more computationally efficient, for example performing a single log likelihood evaluation of the `perturb_model` fit\n", - "to the simulated data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class PerturbFit:\n", - " def __init__(self, adapt_images, number_of_cores: int = 1):\n", - " \"\"\"\n", - " Class used to fit every dataset used for sensitivity mapping with the perturbed model (the model with the\n", - " perturbed feature sensitivity mapping maps out).\n", - "\n", - " In this example, the perturbed model therefore includes the dark matter subhalo, which is also in the\n", - " simulated dataset.\n", - "\n", - " The perturbed fit is repeated for every parameter on the sensitivity grid and compared to the base fit. This\n", - " maps out the sensitivity of every parameter is (e.g. the sensitivity of mass of the subhalo).\n", - "\n", - " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", - " fitted, below we include an input `analysis_cls` which is the `Analysis` class used to fit the model to the\n", - " dataset.\n", - "\n", - " Parameters\n", - " ----------\n", - " adapt_images\n", - " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", - " reconstructed galaxy's morphology.\n", - " number_of_cores\n", - " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", - " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", - " \"\"\"\n", - " self.adapt_images = adapt_images\n", - " self.number_of_cores = number_of_cores\n", - "\n", - " def __call__(self, dataset, model, paths, instance):\n", - " \"\"\"\n", - " The perturbed fitting function which fits every dataset used for sensitivity mapping with the perturbed model.\n", - "\n", - " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", - " quantify how sensitive the model is to the perturbed feature.\n", - "\n", - " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", - " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", - "\n", - " Parameters\n", - " ----------\n", - " dataset\n", - " The dataset which is simulated with the perturbed model and which is fitted.\n", - " model\n", - " The model instance which is fitted to the dataset, which includes the perturbed feature.\n", - " paths\n", - " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", - " \"\"\"\n", - "\n", - " search = af.Nautilus(\n", - " paths=paths,\n", - " n_live=50,\n", - " number_of_cores=self.number_of_cores,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - " analysis._adapt_images = self.adapt_images\n", - "\n", - " return search.fit(model=model, analysis=analysis)\n", - "\n", - "\n", - "def base_model_narrow_priors_from(base_model, result, stretch: float = 1.0):\n", - " \"\"\"\n", - " Returns a base model where priors are updated to `UniformPriors` with a `lower_limit` and `upper_limit` which\n", - " are a narrow range over the `simulation_instance` parameter values.\n", - "\n", - " Using this updated model can significiantly speed up sensitivity mapping, as it dramatically reduces\n", - " the volume of parameter space that needs to be sampled.\n", - "\n", - " The downside of this approach is that these narrow priors may remove viable solutions that alter the sensitivity\n", - " or lead to an inaccurate estimate of the Bayesian evidence.\n", - "\n", - " The size of each parameter bounds have been chosen based on previous lens modeling intuition. For example, I have\n", - " never seen a DM subhalo change the centre of a lens mass model by more than 0.01\", therefore this is the value\n", - " used for bounding that parameter.\n", - "\n", - " Parameters\n", - " ----------\n", - " base_model\n", - " The base model which will be used for sensitivity mapping, which this function updates to have narrower priors.\n", - " result\n", - " The result used to set up the base model and which is used to set these updated priors.\n", - " stretch\n", - " A multiplicative factor which can be used to shrink or broaden the priors more.\n", - "\n", - " Returns\n", - " -------\n", - " A base model with priors updated to narrow uniform priors.\n", - " \"\"\"\n", - "\n", - " if hasattr(base_model.galaxies.lens, \"mass\"):\n", - " base_model.galaxies.lens.mass.centre.centre_0 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.01 * stretch\n", - " ).galaxies.lens.mass.centre.centre_0\n", - " )\n", - " base_model.galaxies.lens.mass.centre.centre_1 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.01 * stretch\n", - " ).galaxies.lens.mass.centre.centre_1\n", - " )\n", - " base_model.galaxies.lens.mass.ell_comps.ell_comps_0 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.05 * stretch\n", - " ).galaxies.lens.mass.ell_comps.ell_comps_0\n", - " )\n", - " base_model.galaxies.lens.mass.ell_comps.ell_comps_1 = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.05 * stretch\n", - " ).galaxies.lens.mass.ell_comps.ell_comps_1\n", - " )\n", - " base_model.galaxies.lens.mass.einstein_radius = (\n", - " result.model_centred_max_lh_bounded(\n", - " b=0.1 * stretch\n", - " ).galaxies.lens.mass.einstein_radius\n", - " )\n", - " base_model.galaxies.lens.mass.slope = result.model_centred_max_lh_bounded(\n", - " b=0.1 * stretch\n", - " ).galaxies.lens.mass.slope\n", - "\n", - " if hasattr(base_model.galaxies.lens, \"shear\"):\n", - " base_model.galaxies.lens.shear.gamma_1 = result.model_centred_max_lh_bounded(\n", - " b=0.05 * stretch\n", - " ).galaxies.lens.shear.gamma_1\n", - " base_model.galaxies.lens.shear.gamma_2 = result.model_centred_max_lh_bounded(\n", - " b=0.05 * stretch\n", - " ).galaxies.lens.shear.gamma_2\n", - "\n", - " return base_model\n", - "\n", - "\n", - "def visualize_sensitivity(\n", - " result,\n", - " paths: af.DirectoryPaths,\n", - " mass_result: af.Result,\n", - " mask: al.Mask2D,\n", - "):\n", - " \"\"\"\n", - " Visualize the results of strong lens sensitivity mapping via the SLaM pipeline.\n", - "\n", - " Parameters\n", - " ----------\n", - " result\n", - " The result of the sensitivity mapping, which contains grids of the log evidence and log likelihood differences.\n", - " paths\n", - " The paths object which defines the output path for the results of the sensitivity mapping.\n", - " mass_result\n", - " The result of the mass pipeline, which is used to subtract the lens light from the dataset.\n", - " mask\n", - " The mask used to mask the dataset, which is plotted over the lens subtracted image.\n", - " \"\"\"\n", - "\n", - " result = al.SubhaloSensitivityResult(\n", - " result=result,\n", - " )\n", - "\n", - " data_subtracted = (\n", - " mass_result.max_log_likelihood_fit.subtracted_images_of_planes_list[-1]\n", - " )\n", - "\n", - " data_subtracted = data_subtracted.apply_mask(mask=mask)\n", - "\n", - " aplt.subplot_sensitivity(result=result, data_subtracted=data_subtracted)\n", - "\n", - "\n", - "class Visualizer:\n", - " def __init__(self, mass_result: af.Result, mask: al.Mask2D):\n", - " \"\"\"\n", - " Performs on-the-fly visualization of the sensitivity mapping, outputting the results of the sensitivity\n", - " mapping so far to hard disk after each sensitivity cell fit is complete.\n", - "\n", - " Parameters\n", - " ----------\n", - " mass_result\n", - " The result of the SLaM MASS PIPELINE which ran before this pipeline.\n", - " mask\n", - " The Mask2D that is applied to the imaging data for model-fitting.\n", - " \"\"\"\n", - "\n", - " self.mass_result = mass_result\n", - " self.mask = mask\n", - "\n", - " def __call__(self, sensitivity_result, paths: af.DirectoryPaths):\n", - " \"\"\"\n", - " Called by the `Sensitivity` class after every sensitivity cell has been fitted, to visualize results so far.\n", - "\n", - " Parameters\n", - " ----------\n", - " sensitivity_result\n", - " The result of the sensitivity mapping search so far.\n", - " paths\n", - " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", - " \"\"\"\n", - " visualize_sensitivity(\n", - " result=sensitivity_result,\n", - " paths=paths,\n", - " mass_result=self.mass_result,\n", - " mask=self.mask,\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"dark_matter_subhalo\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_Path().exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/advanced/subhalo/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.05,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_shape=mesh_shape,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh_shape=mesh_shape,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (sensitivity mapping)__\n", - "\n", - "The SUBHALO PIPELINE (sensitivity mapping) performs sensitivity mapping of the data using the lens model\n", - "fitted above, so as to determine where subhalos of what mass could be detected in the data. A full description of\n", - "Sensitivity mapping if given in the SLaM pipeline script `slam/subhalo/sensitivity_imaging.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "subhalo_mass = af.Model(al.mp.NFWMCRLudlowSph)\n", - "grid_dimension_arcsec = 3.0\n", - "number_of_steps = 2\n", - "sensitivity_mask = None\n", - "\n", - "base_model = mass_result.model\n", - "\n", - "base_model = base_model_narrow_priors_from(base_model=base_model, result=mass_result)\n", - "\n", - "perturb_model = af.Model(al.Galaxy, redshift=0.5, mass=subhalo_mass)\n", - "\n", - "perturb_model.mass.log10m_vir = 9.0\n", - "perturb_model.mass.c_gNFW = 12.0\n", - "perturb_model.mass.overdens = 200.0\n", - "perturb_model.mass.inner_slope = 2.2\n", - "perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - ")\n", - "perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - ")\n", - "perturb_model.mass.redshift_object = mass_result.model.galaxies.lens.redshift\n", - "perturb_model.mass.redshift_source = mass_result.model.galaxies.source.redshift\n", - "\n", - "\n", - "def perturb_model_prior_func(perturb_instance, perturb_model):\n", - " b = 0.05\n", - "\n", - " perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", - " lower_limit=perturb_instance.mass.centre[0] - b,\n", - " upper_limit=perturb_instance.mass.centre[0] + b,\n", - " )\n", - "\n", - " perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", - " lower_limit=perturb_instance.mass.centre[1] - b,\n", - " upper_limit=perturb_instance.mass.centre[1] + b,\n", - " )\n", - "\n", - " perturb_model.mass.log10m_vir = af.UniformPrior(lower_limit=6, upper_limit=12)\n", - "\n", - " return perturb_model\n", - "\n", - "\n", - "simulation_instance = mass_result.instance\n", - "\n", - "fit = mass_result.max_log_likelihood_fit\n", - "\n", - "simulation_instance.galaxies.lens = (\n", - " fit.model_obj_linear_light_profiles_to_light_profiles.galaxies[0]\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - ")\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "paths = af.DirectoryPaths(\n", - " name=f\"subhalo__sensitivity\",\n", - " path_prefix=settings_search.path_prefix,\n", - " unique_tag=settings_search.unique_tag,\n", - ")\n", - "\n", - "simulate_cls = SimulateImagingPixelized(\n", - " mask=mask, psf=dataset.psf, inversion=mass_result.max_log_likelihood_fit.inversion\n", - ")\n", - "\n", - "sensitivity = af.Sensitivity(\n", - " paths=paths,\n", - " simulation_instance=simulation_instance,\n", - " base_model=base_model,\n", - " perturb_model=perturb_model,\n", - " simulate_cls=simulate_cls,\n", - " base_fit_cls=BaseFit(\n", - " adapt_images=adapt_images, number_of_cores=settings_search.number_of_cores\n", - " ),\n", - " perturb_fit_cls=PerturbFit(\n", - " adapt_images=adapt_images, number_of_cores=settings_search.number_of_cores\n", - " ),\n", - " perturb_model_prior_func=perturb_model_prior_func,\n", - " visualizer_cls=Visualizer(mass_result=mass_result, mask=mask),\n", - " number_of_steps=number_of_steps,\n", - " batch_range=None,\n", - " mask=sensitivity_mask,\n", - ")\n", - "\n", - "subhalo_results = sensitivity.run()\n", - "\n", - "visualize_sensitivity(\n", - " result=subhalo_results, paths=paths, mass_result=mass_result, mask=mask\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM (Source, Light and Mass): Subhalo Source Pixelized Sensitivity Mapping\n", + "===========================================================================\n", + "\n", + "This example illustrates how to perform DM subhalo sensitivity mapping using a SLaM pipeline for a dataset where the\n", + "source is modeled using a pixelization.\n", + "\n", + "The sensitivity mapping simulation procedure for a pixelized source is different light profile sources. When pixelized\n", + "sources are used, the source reconstruction on the mesh is used, such that the simulations capture the irregular\n", + "morphologies of real source galaxies.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except the lens mass uses an `Isothermal` with its centre fixed.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **Simulate Function Class:** We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above).\n", + "- **Base Fit:** We have defined a `Simulate` class that will be used to simulate every dataset simulated by the.\n", + "- **Perturb Fit:** We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", + "\n", + "__Model__\n", + "\n", + "Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE, MASS TOTAL PIPELINE and SUBHALO PIPELINE this SLaM script\n", + "fits `Imaging` of a strong lens system, where in the final model:\n", + "\n", + " - The lens galaxy's light is an MGE bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - A dark matter subhalo near The lens galaxy mass is included as a`NFWMCRLudlowSph`.\n", + " - The source galaxy is an `Inversion`.\n", + "\n", + "This uses the SLaM pipelines:\n", + "\n", + " `source_lp`\n", + " `source_pix`\n", + " `light_lp`\n", + " `mass_total`\n", + " `subhalo/detection`\n", + "\n", + "Check them out for a full description of the analysis!\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `subhalo/detect/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import os\n", + "from pathlib import Path\n", + "from typing import List, Optional, Tuple, Union\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except the lens mass uses an `Isothermal` with its centre fixed to (0.0, 0.0)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search,\n", + " dataset,\n", + " mask_radius,\n", + " redshift_lens,\n", + " redshift_source,\n", + "):\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = (0.0, 0.0)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search,\n", + " dataset,\n", + " source_lp_result,\n", + " mesh_shape,\n", + "):\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_lp_result.model.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search,\n", + " dataset,\n", + " source_lp_result,\n", + " source_pix_result_1,\n", + " mesh_shape,\n", + "):\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " over_sampling = al.util.over_sample.over_sample_size_via_adapt_from(\n", + " data=adapt_images.galaxy_name_image_dict[\"('galaxies', 'source')\"],\n", + " noise_map=dataset.noise_map,\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_pixelization=over_sampling,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search,\n", + " dataset,\n", + " mask_radius,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + "):\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " )\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search,\n", + " dataset,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " light_result,\n", + "):\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=light_result.instance.galaxies.lens.bulge,\n", + " disk=light_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate Function Class__\n", + "\n", + "We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above) and uses it to\n", + "simulate a strong lens dataset, which include a dark matter subhalo, which is subsequently fitted.\n", + "\n", + "Additional attributes required to simulate the data (mask, PSF) can be passed to the `__init__` method, and the\n", + "simulation is performed in the `__call__` method.\n", + "\n", + "When this dataset is simulated, the quantity `instance.perturb` is used in `__call__`. This is an instance\n", + "of the `NFWMCRLudlowSph`, and it is different every time the `simulate_cls` is called based on the value of sensitivity\n", + "being computed.\n", + "\n", + "In this example, this `instance.perturb` corresponds to two different subhalos with values of `mass_at_200` of\n", + "1e6 MSun and 1e13 MSun." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class SimulateImagingPixelized:\n", + " def __init__(\n", + " self,\n", + " mask,\n", + " psf,\n", + " inversion,\n", + " interpolated_pixelized_shape: Union[tuple, tuple] = (1001, 1001),\n", + " image_plane_subgrid_size=8,\n", + " ):\n", + " \"\"\"\n", + " Class used to simulate the strong lens dataset used for sensitivity mapping.\n", + "\n", + " Parameters\n", + " ----------\n", + " mask\n", + " The mask applied to the real image data, which is applied to every simulated dataset.\n", + " psf\n", + " The PSF of the real image data, which is applied to every simulated dataset and used for each fit.\n", + " inversion\n", + " The `Inversion` used to reconstruct the source of the real image, included the pixelized source\n", + " reconstruction on a mesh (e.g. Delaunay / Voronoi).\n", + " interpolated_pixelized_shape\n", + " The pixelized source reconstruction is interpolated from an irregular mesh to a rectangular uniform array\n", + " and grid of this shape.\n", + " image_plane_subgrid_size\n", + " The size of the subgrid used to create the image-plane grid, whereby multiple image pixels\n", + " are traced to the source-plane image and evaluated to compute the flux of the simulated image.\n", + " \"\"\"\n", + " self.mask = mask\n", + " self.psf = psf\n", + " self.inversion = inversion\n", + " self.interpolated_pixelized_shape = interpolated_pixelized_shape\n", + " self.image_plane_subgrid_size = image_plane_subgrid_size\n", + "\n", + " def __call__(self, instance: af.ModelInstance, simulate_path: str):\n", + " \"\"\"\n", + " The `simulate_function` called by the `Sensitivity` class which simulates each strong lens image fitted\n", + " by the sensitivity mapper.\n", + "\n", + " The simulation procedure is as follows:\n", + "\n", + " 1) Extract the pixelized reconstructed source of a previous fit, which is likely on an irregular mesh\n", + " (e.g. Delaunay / Voronoi), and interpolate the source emission onto a rectangular uniform array and grid.\n", + "\n", + " 1) Use the input galaxies of the sensitivity `instance` to set up a tracer, which generates the image-plane\n", + " image of the strong lens system.\n", + "\n", + " 2) Simulate this image using the input dataset noise (Poisson) and PSF.\n", + "\n", + " 3) Apply the mask used in the analysis of the real image to the simulated image.\n", + "\n", + " 4) Output information about the simulation to hard-disk.\n", + "\n", + " The `subhalo` in the sensitivity `instance` changes for every iteration of the sensitivity mapping, ensuring\n", + " that we map out the sensitivity of the analysis to the subhalo properties (centre, mass, etc.).\n", + "\n", + " Parameters\n", + " ----------\n", + " instance\n", + " The sensitivity instance, which includes the galaxies whose parameters are varied to perform sensitivity.\n", + " The subhalo in this instance changes for every iteration of the sensitivity mapping.\n", + " simulate_path\n", + " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", + " mapping.\n", + "\n", + " Returns\n", + " -------\n", + " A simulated image of a strong lens, which id input into the fits of the sensitivity mapper.\n", + " \"\"\"\n", + "\n", + " \"\"\"\n", + " __Resume Fit__\n", + "\n", + " If sensitivity mapping already began on this grid cell, the dataset will have been simulated already and we\n", + " do not want to resimulate it and change its noise properties.\n", + "\n", + " We therefore load it from the `simulate_path` instead.\n", + " \"\"\"\n", + " try:\n", + " dataset = al.Imaging.from_fits(\n", + " data_path=f\"{simulate_path}/data.fits\",\n", + " psf_path=f\"{simulate_path}/psf.fits\",\n", + " noise_map_path=f\"{simulate_path}/noise_map.fits\",\n", + " pixel_scales=self.mask.pixel_scales,\n", + " check_noise_map=False,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=self.mask)\n", + "\n", + " over_sample_size = (\n", + " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[8, 4, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + " )\n", + "\n", + " return dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size, over_sample_size_pixelization=4\n", + " )\n", + "\n", + " except FileNotFoundError:\n", + " pass\n", + "\n", + " \"\"\"\n", + " __Source Galaxy Image__\n", + "\n", + " We now load the source galaxy image from the pixelized inversion of a previous fit, which could be on an irregular\n", + " Delaunay or Voronoi mesh.\n", + "\n", + " Irregular meshes cannot be used to simulate lensed images of a source. Therefore, we interpolate the mesh to\n", + " a uniform grid of shape `interpolated_pixelized_shape`. This should be high resolution (e.g. 1000 x 1000)\n", + " to ensure the interpolated source array captures all structure resolved on the Delaunay / Voronoi mesh.\n", + "\n", + " Loads source array from previous reconstruction, maps to square and wraps in AutoLens Array.\n", + " Loads lens galaxy and perturb from provided instance\n", + " Loads source galaxy redshift and sets up a `galaxy` class object at that redshift.\n", + " \"\"\"\n", + "\n", + " mapper = self.inversion.cls_list_from(cls=al.Mapper)[0]\n", + "\n", + " mapper_valued = al.MapperValued(\n", + " mapper=mapper,\n", + " values=self.inversion.reconstruction_dict[mapper],\n", + " )\n", + "\n", + " source_image = mapper_valued.interpolated_array_from(\n", + " shape_native=self.interpolated_pixelized_shape,\n", + " extent=(-2.0, 2.0, -2.0, 2.0),\n", + " )\n", + "\n", + " \"\"\"\n", + " __Create Grids__\n", + "\n", + " To create the lensed image, we will ray-trace image pixels to the source-plane and interpolate them onto the\n", + " source galaxy image.\n", + "\n", + " We therefore need the image-plane grid of (y,x) coordinates.\n", + " \"\"\"\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=self.mask.shape_native,\n", + " pixel_scales=self.mask.pixel_scales,\n", + " over_sample_size=self.image_plane_subgrid_size,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Ray Tracing__\n", + "\n", + " We create a tracer, which will create the lensed grid we overlay the interpolated source galaxy image above\n", + " in order to create the lensed source galaxy image.\n", + "\n", + " This creates the grid we will overlay the source image, in order to created the lensed source image.\n", + "\n", + " The source-plane requires a source-galaxy with a `redshift` in order for the tracer to trace it. We therefore\n", + " make one, noting it has no light profiles because its emission is entirely defined by the source galaxy image.\n", + " \"\"\"\n", + " tracer = al.Tracer(\n", + " galaxies=[\n", + " instance.galaxies.lens,\n", + " instance.perturb,\n", + " al.Galaxy(redshift=instance.galaxies.source.redshift),\n", + " ]\n", + " )\n", + "\n", + " \"\"\"\n", + " __Simulate__\n", + "\n", + " Using the tracer above, we create the image of the lensed source galaxy on the image-plane grid. This\n", + " uses the `source_image` and therefore capture the source's irregular and asymmetric morphological features\n", + " which the source reconstruction procedure fitted.\n", + "\n", + " Set up the grid, PSF and simulator settings used to simulate imaging of the strong lens. These should be\n", + " tuned to match the S/N and noise properties of the observed data you are performing sensitivity mapping on.\n", + "\n", + " The `SimulatorImaging` will be passed directly the image of the strong lens we created above, which\n", + " will be convolved with the psf before noise is added.\n", + "\n", + " To ensure the PSF convolution extends over the whole image, the image is padded before convolution to mitigate\n", + " edge effects and trimmed after the simulation so it retains the original `shape_native`.\n", + " \"\"\"\n", + " simulator = al.SimulatorImaging(\n", + " exposure_time=1000.0,\n", + " psf=self.psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + " noise_seed=1,\n", + " )\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " dataset = simulator.via_source_image_from(\n", + " tracer=tracer, grid=grid, source_image=source_image\n", + " )\n", + "\n", + " \"\"\"\n", + " __Masking__\n", + "\n", + " The data generated by the simulate function is what is ultimately fitted.\n", + "\n", + " Therefore, we also apply the mask for the analysis before we return the simulated data.\n", + " \"\"\"\n", + " dataset = dataset.apply_mask(mask=self.mask)\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[8, 4, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(\n", + " over_sample_size_lp=over_sample_size, over_sample_size_pixelization=4\n", + " )\n", + "\n", + " \"\"\"\n", + " Outputs info about the `Tracer` to the fit, so we know exactly how we simulated the image.\n", + " \"\"\"\n", + " tracer_no_perturb = al.Tracer(\n", + " galaxies=[\n", + " instance.galaxies.lens,\n", + " al.Galaxy(redshift=instance.galaxies.source.redshift),\n", + " ]\n", + " )\n", + "\n", + " self.output_info(\n", + " simulate_path=simulate_path,\n", + " grid=grid,\n", + " dataset=dataset,\n", + " source_image=source_image,\n", + " tracer=tracer,\n", + " tracer_no_perturb=tracer_no_perturb,\n", + " )\n", + "\n", + " return dataset\n", + "\n", + " def output_info(\n", + " self,\n", + " simulate_path: str,\n", + " grid: al.Grid2D,\n", + " dataset: al.Imaging,\n", + " source_image: al.Array2D,\n", + " tracer: al.Tracer,\n", + " tracer_no_perturb: al.Tracer,\n", + " ):\n", + " \"\"\"\n", + " Output information about the data simulated for this iteration of sensitivity mapping.\n", + "\n", + " This information output is as follows:\n", + "\n", + " - A subplot of the simulated imaging dataset.\n", + " - A subplot of the tracer used to simulate this imaging dataset.\n", + " - A .json file containing the tracer galaxies.\n", + " - Output the simulated dataset to .fits files which are used to load the data if a run is resumed.\n", + "\n", + " Parameters\n", + " ----------\n", + " simulate_path\n", + " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", + " mapping.\n", + " dataset\n", + " The simulated dataset dataset which is visualized.\n", + " tracer\n", + " The tracer used to simulate the dataset dataset, which is visualized and output to a .json file.\n", + " \"\"\"\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(simulate_path) / \"tracer.json\",\n", + " )\n", + "\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=Path(simulate_path) / \"data.fits\",\n", + " psf_path=Path(simulate_path) / \"psf.fits\",\n", + " noise_map_path=Path(simulate_path) / \"noise_map.fits\",\n", + " overwrite=True,\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Base Fit__\n", + "\n", + "We have defined a `Simulate` class that will be used to simulate every dataset simulated by the sensitivity mapper.\n", + "Each simulated dataset will have a unique set of parameters for the `subhalo` (e.g. due to different values of\n", + "`perturb_model`.\n", + "\n", + "We will fit each simulated dataset using the `base_model`, which quantifies whether not including the dark matter\n", + "subhalo in the model changess the goodness-of-fit and therefore indicates if we are sensitive to the subhalo.\n", + "\n", + "We now write a `BaseFit` class, defining how the `base_model` is fitted to each simulated dataset and the\n", + "goodness-of-fit used to quantify whether the model fits the data well. As above, the `__init__` method can be\n", + "extended with new inputs to control how the model is fitted and the `__call__` method performs the fit.\n", + "\n", + "In this example, we use a full non-linear search to fit the `base_model` to the simulated data and return\n", + "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", + "more computationally efficient, for example performing a single log likelihood evaluation of the `base_model` fit\n", + "to the simulated data.\n", + "\n", + "Fucntionality which adapts the mesh and regularization of a pixelized source reconstruction to the unlensed source's\n", + "morphology require an `adapt_images`. This is an input of the __init__ constructor which is passed to the `Analysis`\n", + "for every simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class BaseFit:\n", + " def __init__(self, adapt_images, number_of_cores: int = 1):\n", + " \"\"\"\n", + " Class used to fit every dataset used for sensitivity mapping with the base model (the model without the\n", + " perturbed feature sensitivity mapping maps out).\n", + "\n", + " In this example, the base model therefore does not include the dark matter subhalo, but the simulated\n", + " dataset includes one.\n", + "\n", + " The base fit is repeated for every parameter on the sensitivity grid and compared to the perturbed fit. This\n", + " maps out the sensitivity of every parameter is (e.g. the sensitivity of the mass of the subhalo).\n", + "\n", + " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", + " fitted, below we include an input `analysis_cls` which is the `AnalysisImaging` class used to fit the model\n", + " to the dataset.\n", + "\n", + " Parameters\n", + " ----------\n", + " adapt_images\n", + " The result of the previous search containing adapt images used to adapt certain pixelized source meshs's\n", + " and regularizations to the unlensed source morphology.\n", + " number_of_cores\n", + " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", + " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", + " \"\"\"\n", + " self.adapt_images = adapt_images\n", + " self.number_of_cores = number_of_cores\n", + "\n", + " def __call__(self, dataset, model, paths, instance):\n", + " \"\"\"\n", + " The base fitting function which fits every dataset used for sensitivity mapping with the base model.\n", + "\n", + " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", + " quantify how sensitive the model is to the perturbed feature.\n", + "\n", + " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", + " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", + "\n", + " Parameters\n", + " ----------\n", + " dataset\n", + " The dataset which is simulated with the perturbed model and which is fitted.\n", + " model\n", + " The model instance which is fitted to the dataset, which does not include the perturbed feature.\n", + " paths\n", + " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", + " \"\"\"\n", + "\n", + " search = af.Nautilus(\n", + " paths=paths,\n", + " n_live=50,\n", + " number_of_cores=self.number_of_cores,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + " analysis._adapt_images = self.adapt_images\n", + "\n", + " return search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Perturb Fit__\n", + "\n", + "We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each simulated dataset. This\n", + "behaves analogously to the `BaseFit` class above, but now fits the `perturb_model` to the simulated data (as\n", + "opposed to the `base_model`).\n", + "\n", + "Again, in this example we use a full non-linear search to fit the `perturb_model` to the simulated data and return\n", + "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", + "more computationally efficient, for example performing a single log likelihood evaluation of the `perturb_model` fit\n", + "to the simulated data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class PerturbFit:\n", + " def __init__(self, adapt_images, number_of_cores: int = 1):\n", + " \"\"\"\n", + " Class used to fit every dataset used for sensitivity mapping with the perturbed model (the model with the\n", + " perturbed feature sensitivity mapping maps out).\n", + "\n", + " In this example, the perturbed model therefore includes the dark matter subhalo, which is also in the\n", + " simulated dataset.\n", + "\n", + " The perturbed fit is repeated for every parameter on the sensitivity grid and compared to the base fit. This\n", + " maps out the sensitivity of every parameter is (e.g. the sensitivity of mass of the subhalo).\n", + "\n", + " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", + " fitted, below we include an input `analysis_cls` which is the `Analysis` class used to fit the model to the\n", + " dataset.\n", + "\n", + " Parameters\n", + " ----------\n", + " adapt_images\n", + " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", + " reconstructed galaxy's morphology.\n", + " number_of_cores\n", + " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", + " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", + " \"\"\"\n", + " self.adapt_images = adapt_images\n", + " self.number_of_cores = number_of_cores\n", + "\n", + " def __call__(self, dataset, model, paths, instance):\n", + " \"\"\"\n", + " The perturbed fitting function which fits every dataset used for sensitivity mapping with the perturbed model.\n", + "\n", + " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", + " quantify how sensitive the model is to the perturbed feature.\n", + "\n", + " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", + " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", + "\n", + " Parameters\n", + " ----------\n", + " dataset\n", + " The dataset which is simulated with the perturbed model and which is fitted.\n", + " model\n", + " The model instance which is fitted to the dataset, which includes the perturbed feature.\n", + " paths\n", + " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", + " \"\"\"\n", + "\n", + " search = af.Nautilus(\n", + " paths=paths,\n", + " n_live=50,\n", + " number_of_cores=self.number_of_cores,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + " analysis._adapt_images = self.adapt_images\n", + "\n", + " return search.fit(model=model, analysis=analysis)\n", + "\n", + "\n", + "def base_model_narrow_priors_from(base_model, result, stretch: float = 1.0):\n", + " \"\"\"\n", + " Returns a base model where priors are updated to `UniformPriors` with a `lower_limit` and `upper_limit` which\n", + " are a narrow range over the `simulation_instance` parameter values.\n", + "\n", + " Using this updated model can significiantly speed up sensitivity mapping, as it dramatically reduces\n", + " the volume of parameter space that needs to be sampled.\n", + "\n", + " The downside of this approach is that these narrow priors may remove viable solutions that alter the sensitivity\n", + " or lead to an inaccurate estimate of the Bayesian evidence.\n", + "\n", + " The size of each parameter bounds have been chosen based on previous lens modeling intuition. For example, I have\n", + " never seen a DM subhalo change the centre of a lens mass model by more than 0.01\", therefore this is the value\n", + " used for bounding that parameter.\n", + "\n", + " Parameters\n", + " ----------\n", + " base_model\n", + " The base model which will be used for sensitivity mapping, which this function updates to have narrower priors.\n", + " result\n", + " The result used to set up the base model and which is used to set these updated priors.\n", + " stretch\n", + " A multiplicative factor which can be used to shrink or broaden the priors more.\n", + "\n", + " Returns\n", + " -------\n", + " A base model with priors updated to narrow uniform priors.\n", + " \"\"\"\n", + "\n", + " if hasattr(base_model.galaxies.lens, \"mass\"):\n", + " base_model.galaxies.lens.mass.centre.centre_0 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.01 * stretch\n", + " ).galaxies.lens.mass.centre.centre_0\n", + " )\n", + " base_model.galaxies.lens.mass.centre.centre_1 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.01 * stretch\n", + " ).galaxies.lens.mass.centre.centre_1\n", + " )\n", + " base_model.galaxies.lens.mass.ell_comps.ell_comps_0 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.05 * stretch\n", + " ).galaxies.lens.mass.ell_comps.ell_comps_0\n", + " )\n", + " base_model.galaxies.lens.mass.ell_comps.ell_comps_1 = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.05 * stretch\n", + " ).galaxies.lens.mass.ell_comps.ell_comps_1\n", + " )\n", + " base_model.galaxies.lens.mass.einstein_radius = (\n", + " result.model_centred_max_lh_bounded(\n", + " b=0.1 * stretch\n", + " ).galaxies.lens.mass.einstein_radius\n", + " )\n", + " base_model.galaxies.lens.mass.slope = result.model_centred_max_lh_bounded(\n", + " b=0.1 * stretch\n", + " ).galaxies.lens.mass.slope\n", + "\n", + " if hasattr(base_model.galaxies.lens, \"shear\"):\n", + " base_model.galaxies.lens.shear.gamma_1 = result.model_centred_max_lh_bounded(\n", + " b=0.05 * stretch\n", + " ).galaxies.lens.shear.gamma_1\n", + " base_model.galaxies.lens.shear.gamma_2 = result.model_centred_max_lh_bounded(\n", + " b=0.05 * stretch\n", + " ).galaxies.lens.shear.gamma_2\n", + "\n", + " return base_model\n", + "\n", + "\n", + "def visualize_sensitivity(\n", + " result,\n", + " paths: af.DirectoryPaths,\n", + " mass_result: af.Result,\n", + " mask: al.Mask2D,\n", + "):\n", + " \"\"\"\n", + " Visualize the results of strong lens sensitivity mapping via the SLaM pipeline.\n", + "\n", + " Parameters\n", + " ----------\n", + " result\n", + " The result of the sensitivity mapping, which contains grids of the log evidence and log likelihood differences.\n", + " paths\n", + " The paths object which defines the output path for the results of the sensitivity mapping.\n", + " mass_result\n", + " The result of the mass pipeline, which is used to subtract the lens light from the dataset.\n", + " mask\n", + " The mask used to mask the dataset, which is plotted over the lens subtracted image.\n", + " \"\"\"\n", + "\n", + " result = al.SubhaloSensitivityResult(\n", + " result=result,\n", + " )\n", + "\n", + " data_subtracted = (\n", + " mass_result.max_log_likelihood_fit.subtracted_images_of_planes_list[-1]\n", + " )\n", + "\n", + " data_subtracted = data_subtracted.apply_mask(mask=mask)\n", + "\n", + " aplt.subplot_sensitivity(result=result, data_subtracted=data_subtracted)\n", + "\n", + "\n", + "class Visualizer:\n", + " def __init__(self, mass_result: af.Result, mask: al.Mask2D):\n", + " \"\"\"\n", + " Performs on-the-fly visualization of the sensitivity mapping, outputting the results of the sensitivity\n", + " mapping so far to hard disk after each sensitivity cell fit is complete.\n", + "\n", + " Parameters\n", + " ----------\n", + " mass_result\n", + " The result of the SLaM MASS PIPELINE which ran before this pipeline.\n", + " mask\n", + " The Mask2D that is applied to the imaging data for model-fitting.\n", + " \"\"\"\n", + "\n", + " self.mass_result = mass_result\n", + " self.mask = mask\n", + "\n", + " def __call__(self, sensitivity_result, paths: af.DirectoryPaths):\n", + " \"\"\"\n", + " Called by the `Sensitivity` class after every sensitivity cell has been fitted, to visualize results so far.\n", + "\n", + " Parameters\n", + " ----------\n", + " sensitivity_result\n", + " The result of the sensitivity mapping search so far.\n", + " paths\n", + " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", + " \"\"\"\n", + " visualize_sensitivity(\n", + " result=sensitivity_result,\n", + " paths=paths,\n", + " mass_result=self.mass_result,\n", + " mask=self.mask,\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"dark_matter_subhalo\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_Path().exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/advanced/subhalo/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.05,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_shape=mesh_shape,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh_shape=mesh_shape,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (sensitivity mapping)__\n", + "\n", + "The SUBHALO PIPELINE (sensitivity mapping) performs sensitivity mapping of the data using the lens model\n", + "fitted above, so as to determine where subhalos of what mass could be detected in the data. A full description of\n", + "Sensitivity mapping if given in the SLaM pipeline script `slam/subhalo/sensitivity_imaging.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "subhalo_mass = af.Model(al.mp.NFWMCRLudlowSph)\n", + "grid_dimension_arcsec = 3.0\n", + "number_of_steps = 2\n", + "sensitivity_mask = None\n", + "\n", + "base_model = mass_result.model\n", + "\n", + "base_model = base_model_narrow_priors_from(base_model=base_model, result=mass_result)\n", + "\n", + "perturb_model = af.Model(al.Galaxy, redshift=0.5, mass=subhalo_mass)\n", + "\n", + "perturb_model.mass.log10m_vir = 9.0\n", + "perturb_model.mass.c_gNFW = 12.0\n", + "perturb_model.mass.overdens = 200.0\n", + "perturb_model.mass.inner_slope = 2.2\n", + "perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + ")\n", + "perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + ")\n", + "perturb_model.mass.redshift_object = mass_result.model.galaxies.lens.redshift\n", + "perturb_model.mass.redshift_source = mass_result.model.galaxies.source.redshift\n", + "\n", + "\n", + "def perturb_model_prior_func(perturb_instance, perturb_model):\n", + " b = 0.05\n", + "\n", + " perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", + " lower_limit=perturb_instance.mass.centre[0] - b,\n", + " upper_limit=perturb_instance.mass.centre[0] + b,\n", + " )\n", + "\n", + " perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", + " lower_limit=perturb_instance.mass.centre[1] - b,\n", + " upper_limit=perturb_instance.mass.centre[1] + b,\n", + " )\n", + "\n", + " perturb_model.mass.log10m_vir = af.UniformPrior(lower_limit=6, upper_limit=12)\n", + "\n", + " return perturb_model\n", + "\n", + "\n", + "simulation_instance = mass_result.instance\n", + "\n", + "fit = mass_result.max_log_likelihood_fit\n", + "\n", + "simulation_instance.galaxies.lens = (\n", + " fit.model_obj_linear_light_profiles_to_light_profiles.galaxies[0]\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + ")\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "paths = af.DirectoryPaths(\n", + " name=f\"subhalo__sensitivity\",\n", + " path_prefix=settings_search.path_prefix,\n", + " unique_tag=settings_search.unique_tag,\n", + ")\n", + "\n", + "simulate_cls = SimulateImagingPixelized(\n", + " mask=mask, psf=dataset.psf, inversion=mass_result.max_log_likelihood_fit.inversion\n", + ")\n", + "\n", + "sensitivity = af.Sensitivity(\n", + " paths=paths,\n", + " simulation_instance=simulation_instance,\n", + " base_model=base_model,\n", + " perturb_model=perturb_model,\n", + " simulate_cls=simulate_cls,\n", + " base_fit_cls=BaseFit(\n", + " adapt_images=adapt_images, number_of_cores=settings_search.number_of_cores\n", + " ),\n", + " perturb_fit_cls=PerturbFit(\n", + " adapt_images=adapt_images, number_of_cores=settings_search.number_of_cores\n", + " ),\n", + " perturb_model_prior_func=perturb_model_prior_func,\n", + " visualizer_cls=Visualizer(mass_result=mass_result, mask=mask),\n", + " number_of_steps=number_of_steps,\n", + " batch_range=None,\n", + " mask=sensitivity_mask,\n", + ")\n", + "\n", + "subhalo_results = sensitivity.run()\n", + "\n", + "visualize_sensitivity(\n", + " result=subhalo_results, paths=paths, mass_result=mass_result, mask=mask\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/subhalo/sensitivity/start_here.ipynb b/notebooks/imaging/features/advanced/subhalo/sensitivity/start_here.ipynb index 5aa79cd5c..94176568b 100644 --- a/notebooks/imaging/features/advanced/subhalo/sensitivity/start_here.ipynb +++ b/notebooks/imaging/features/advanced/subhalo/sensitivity/start_here.ipynb @@ -1,999 +1,1036 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Sensitivity Mapping: Start Here\n", - "===============================\n", - "\n", - "Bayesian model comparison allows us to take a dataset, fit it with multiple models and use the Bayesian evidence to\n", - "quantify which model objectively gives the best-fit following the principles of Occam's Razor.\n", - "\n", - "However, a complex model may not be favoured by model comparison not because it is the 'wrong' model, but simply\n", - "because the dataset being fitted is not of a sufficient quality for the more complex model to be favoured. Sensitivity\n", - "mapping addresses what quality of data would be needed for the more complex model to be favoured.\n", - "\n", - "In order to do this, sensitivity mapping involves us writing a function that uses the model(s) to simulate a dataset.\n", - "We then use this function to simulate many datasets, for many different models, and fit each dataset to quantify\n", - "how much the change in the model led to a measurable change in the data. This is called computing the sensitivity.\n", - "\n", - "How we compute the sensitivity is chosen by us, the user. In this example, we will perform multiple model-fits\n", - "with a nested sampling search, and therefore perform Bayesian model comparison to compute the sensitivity. This allows\n", - "us to infer how much of a Bayesian evidence increase we should expect for datasets of varying quality and / or models\n", - "with different parameters.\n", - "\n", - "__Contents__\n", - "\n", - "- **Subhalo Detection Discussion:** For strong lensing, this process is crucial for dark matter substructure detection, as discussed in.\n", - "- **Subhalo Sensitivity Mapping:** Sensitivity mapping is a process where we simulate many thousands of strong lens images.\n", - "- **SLaM Pipelines:** The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which.\n", - "- **Pixelized Source:** Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of.\n", - "- **Base Model:** We now define the `base_model` that we use to perform sensitivity mapping.\n", - "- **Perturb Model:** We now define the `perturb_model`, which is the model component whose parameters we iterate over to.\n", - "- **Mapping Grid:** Sensitivity mapping is typically performed over a large range of parameters.\n", - "- **Perturb Model Prior Func:** The default priors on the `perturb_model` are `UniformPrior`'s bounded around each sensitivity grid.\n", - "- **Simulation Instance:** We are performing sensitivity mapping to determine where a subhalo is detectable.\n", - "- **Simulate Function Class:** We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above).\n", - "- **Base Fit:** We have defined a `Simulate` class that will be used to simulate every dataset simulated by the.\n", - "- **Perturb Fit:** We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each.\n", - "- **Results:** You should now look at the results of the sensitivity mapping in the folder.\n", - "\n", - "__Subhalo Detection Discussion__\n", - "\n", - "For strong lensing, this process is crucial for dark matter substructure detection, as discussed in the following paper:\n", - "\n", - "https://arxiv.org/abs/0903.4752\n", - "\n", - "In subhalo detection, our strong lens modeling informs us of whether there is a dark matter subhalo at a given (y,x)\n", - "image-plane location of the strong lens. We determine this by fitting a lens models which includes a subhalo. However,\n", - "we are only able to detect dark matter subhalos with (y,x) locations near the lensed source light, and when the\n", - "subhalo is massive enough to perturb its light in an observable way.\n", - "\n", - "Subhalo detection analysis therefore does not tell us where we could detect subhalos and of what mass. To know this,\n", - "we must perform sensitivity mapping.\n", - "\n", - "__Subhalo Sensitivity Mapping__\n", - "\n", - "Sensitivity mapping is a process where we simulate many thousands of strong lens images. Each simulated image includes a\n", - "dark matter subhalo at a given (y,x) coordinate and at a given mass. We fit each simulated dataset twice,\n", - "with a lens model which does not include a subhalo and with a lens model that does.\n", - "\n", - "If the Bayesian evidence of the lens model including a subhalo is higher than the model which does not, a subhalo at t\n", - "hat (y,x) location and mass is therefore detectable.\n", - "\n", - "For many simulated datasets, we will find the evidence does not increase when we include a subhalo in the model-fit,\n", - "informing us that regions of the image-plane away from the lensed source are not sensitive to subhalos.\n", - "\n", - "The sensitivity map is performed over a three dimensional (or higher) grid of subhalo (y,x) and mass. Thus, once\n", - "sensitivity mapping is complete, we have a complete map of where in the image-plane subhalos of what mass are\n", - "detectable. We can plot 2D plots of this grid to visualize where we are sensitive to dark matter subhalos.\n", - "\n", - "The information provided by a sensitivity map is ultimately required to turn dark matter subhalo detections into\n", - "constraints on the dark matter subhalo mass function, which is the primary goal of subhalo detection. Thus, it is\n", - "necessary to make statements about the nature of dark matter: cold, warm, fuzzy, or something else entirely?\n", - "\n", - "__SLaM Pipelines__\n", - "\n", - "The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which automate the fitting\n", - "of complex lens models. The SLaM pipelines are used for all DM subhalo detection analyses.\n", - "\n", - "This example script does not use a SLaM pipeline, to keep the sensitivity mapping self contained. However, it is\n", - "anticipated that any user performing sensitivity mapping on real data will use the SLaM pipelines, which in the\n", - "`subhalo` package have dedicated extensions for performing sensitivity mapping to both imaging and interferometer data.\n", - "\n", - "Therefore you should be familiar with the SLaM pipelines before performing DM subhalo sensitivity mapping on real\n", - "data. If you are unfamiliar with the SLaM pipelines, checkout the\n", - "example `autolens_workspace/notebooks/guides/modeling/slam_start_here`.\n", - "\n", - "__Pixelized Source__\n", - "\n", - "Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of the source's light\n", - "are at a level where the subhalo's perturbing lensing effects can be detected.\n", - "\n", - "This requires the source reconstruction to be performed using a pixelized source, as this provides a more detailed\n", - "reconstruction of the source's light than fits using light profiles.\n", - "\n", - "Therefore, the corresponding sensitivity mapping should also be performed using pixelized sources. This example\n", - "sticks to light profile sources, to provide faster run times illustrative purposes.\n", - "\n", - "The example `subhalo/sensitivity/examples/source_pixelized.ipynb` extends the SLaM pipelines with pixelized sources\n", - "and therefore shows how to perform sensitivity mapping using pixelized sources.\n", - "\n", - "Note that the simulation procedure for a pixelized source is different to the one shown here. In this example, the\n", - "light profile source parameters are used to simulate each sensitivity mapping dataset. When pixelized sources are\n", - "used, the source reconstruction on the mesh is used, such that the simulations capture the irregular morphologies\n", - "of real source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import os\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__ \n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"dark_matter_subhalo\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/advanced/subhalo/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.05,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model + Search + Analysis + Model-Fit (Base Search)__\n", - "\n", - "We are performing sensitivity mapping to determine where a subhalo is detectable. This will require us to simulate \n", - "many realizations of our dataset with a lens model, called the `simulation_instance`. To get this model, we therefore \n", - "fit the data before performing sensitivity mapping so that we can set the `simulation_instance` as the maximum \n", - "likelihood model.\n", - "\n", - "We perform this fit using the lens model we will use to perform sensitivity mapping, which we call the `base_model`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "base_model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", - " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge),\n", - " ),\n", - ")\n", - "\n", - "search_base = af.Nautilus(\n", - " path_prefix=Path(\"imaging\", \"advanced\", \"subhalo\", \"sensitivity\"),\n", - " name=\"sensitivity_mapping_base\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result = search_base.fit(model=base_model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Base Model__\n", - "\n", - "We now define the `base_model` that we use to perform sensitivity mapping. This is the lens model that is fitted to \n", - "every simulated strong lens without a subhalo, giving us the Bayesian evidence which we compare to the model which \n", - "includes one!). \n", - "\n", - "For this model, we can use the `base_model` above, however we will use the result of fitting this model to the dataset\n", - "before sensitivity mapping. This ensures the priors associated with each parameter are initialized so as to speed up\n", - "each non-linear search performed during sensitivity mapping." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "base_model = result.model" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Perturb Model__\n", - "\n", - "We now define the `perturb_model`, which is the model component whose parameters we iterate over to perform \n", - "sensitivity mapping. In this case, this model is a `NFWMCRLudlowSph` model and we will iterate over its\n", - "`centre` and `mass_at_200`. We set it up as a `Model` so it has an associated redshift and can be directly\n", - "passed to the tracer in the simulate function below.\n", - "\n", - "Many instances of the `perturb_model` are created and used to simulate the many strong lens datasets that we fit. \n", - "However, it is only included in half of the model-fits; corresponding to the lens models which include a dark matter \n", - "subhalo and whose Bayesian evidence we compare to the simpler model-fits consisting of just the `base_model` to \n", - "determine if the subhalo was detectable.\n", - "\n", - "By fitting both models to every simulated lens, we therefore infer the Bayesian evidence of every model to every \n", - "dataset. Sensitivity mapping therefore maps out for what values of `centre` and `mass_at_200` in the dark matter \n", - "subhalo the model-fit including a subhalo provide higher values of Bayesian evidence than the simpler model-fit (and\n", - "therefore when it is detectable!)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "perturb_model = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.NFWMCRLudlowSph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapping Grid__\n", - "\n", - "Sensitivity mapping is typically performed over a large range of parameters. However, to make this demonstration quick\n", - "and clear we are going to fix the `centre` of the subhalo to a value near the Einstein ring of (1.6, 0.0). We will \n", - "iterate over just two `mass_at_200` values corresponding to subhalos of mass 1e6 and 1e13, of which only the latter\n", - "will be shown to be detectable." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_dimension_arcsec = 3.0\n", - "\n", - "perturb_model.mass.mass_at_200 = 1e10\n", - "perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - ")\n", - "perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - ")\n", - "perturb_model.mass.redshift_object = 0.5\n", - "perturb_model.mass.redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Perturb Model Prior Func__\n", - "\n", - "The default priors on the `perturb_model` are `UniformPrior`'s bounded around each sensitivity grid cell.\n", - "\n", - "For example, the first simulated dark matter subhalo is at location (1.5, -1.5) and its priors are: \n", - "\n", - "- y is `UniformPrior(lower_limit=0.0, upper_limit=3.0)`.\n", - "- x is `UniformPrior(lower_limit=-3.0, upper_limit=0.0)`.\n", - "\n", - "The `mass_at_200` is fixed to a value of 1e10 in the `perturb_model` above, which is the fixed value used by\n", - "the model fit.\n", - "\n", - "By passing a `perturb_model_prior_func` to the sensitivity mapper, we can manually overwrite the priors on \n", - "the `perturb_model` which are used instead for the fit.\n", - "\n", - "Below, we update the priors as follows:\n", - "\n", - "- The y and x priors are trimmed to much narrower bounded priors, confined to regions 0.05\" each side of the \n", - " true DM subhalo.\n", - "\n", - "- The `mass_at_200` is made a free parameter with `LogUniformPrior(lower_limit=1e6, 1e12)`. This is a large range, \n", - " but ensures there are solutions where the DM subhalo can go to lower masses. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def perturb_model_prior_func(perturb_instance, perturb_model):\n", - " b = 0.05\n", - "\n", - " perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", - " lower_limit=perturb_instance.mass.centre[0] - b,\n", - " upper_limit=perturb_instance.mass.centre[0] + b,\n", - " )\n", - "\n", - " perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", - " lower_limit=perturb_instance.mass.centre[1] - b,\n", - " upper_limit=perturb_instance.mass.centre[1] + b,\n", - " )\n", - "\n", - " perturb_model.mass.mass_at_200 = af.LogUniformPrior(\n", - " lower_limit=1e6, upper_limit=1e12\n", - " )\n", - "\n", - " return perturb_model\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulation Instance__\n", - "\n", - "We are performing sensitivity mapping to determine where a subhalo is detectable. This will require us to \n", - "simulate many realizations of our dataset with a lens model, called the `simulation_instance`. This model uses the\n", - "result of the fit above.\n", - "\n", - "The code below ensures that the lens light, mass and source parameters of the strong lens are used when simulating\n", - "each dataset with a dark matter subhalo." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulation_instance = result.instance" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate Function Class__\n", - "\n", - "We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above) and uses it to \n", - "simulate a strong lens dataset, which include a dark matter subhalo, which is subsequently fitted.\n", - "\n", - "Additional attributes required to simulate the data (mask, PSF) can be passed to the `__init__` method, and the \n", - "simulation is performed in the `__call__` method.\n", - "\n", - "When this dataset is simulated, the quantity `instance.perturb` is used in `__call__`. This is an instance \n", - "of the `NFWMCRLudlowSph`, and it is different every time the `simulate_cls` is called based on the value of sensitivity \n", - "being computed. \n", - "\n", - "In this example, this `instance.perturb` corresponds to two different subhalos with values of `mass_at_200` of \n", - "1e6 MSun and 1e13 MSun." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class SimulateImaging:\n", - " def __init__(self, mask, psf):\n", - " \"\"\"\n", - " Class used to simulate the strong lens imaging used for sensitivity mapping.\n", - "\n", - " Parameters\n", - " ----------\n", - " mask\n", - " The mask applied to the real image data, which is applied to every simulated imaging.\n", - " psf\n", - " The PSF of the real image data, which is applied to every simulated imaging and used for each fit.\n", - " \"\"\"\n", - " self.mask = mask\n", - " self.psf = psf\n", - "\n", - " def __call__(self, instance: af.ModelInstance, simulate_path: str):\n", - " \"\"\"\n", - " The `simulate_function` called by the `Sensitivity` class which simulates each strong lens image fitted\n", - " by the sensitivity mapper.\n", - "\n", - " The simulation procedure is as follows:\n", - "\n", - " 1) Use the input galaxies of the sensitivity `instance` to set up a tracer, which generates the image-plane\n", - " image of the strong lens system.\n", - "\n", - " 2) Simulate this image using the input dataset noise (Poisson) and PSF.\n", - "\n", - " 3) Apply the mask used in the analysis of the real image to the simulated image.\n", - "\n", - " 4) Output information about the simulation to hard-disk.\n", - "\n", - " The `subhalo` in the sensitivity `instance` changes for every iteration of the sensitivity mapping, ensuring\n", - " that we map out the sensitivity of the analysis to the subhalo properties (centre, mass, etc.).\n", - "\n", - " Parameters\n", - " ----------\n", - " instance\n", - " The sensitivity instance, which includes the galaxies whose parameters are varied to perform sensitivity.\n", - " The subhalo in this instance changes for every iteration of the sensitivity mapping.\n", - " simulate_path\n", - " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", - " mapping.\n", - "\n", - " Returns\n", - " -------\n", - " A simulated image of a strong lens, which id input into the fits of the sensitivity mapper.\n", - " \"\"\"\n", - "\n", - " \"\"\"\n", - " __Resume Fit__\n", - "\n", - " If sensitivity mapping already began on this grid cell, the dataset will have been simulated already and we\n", - " do not want to resimulate it and change its noise properties. \n", - "\n", - " We therefore load it from the `simulate_path` instead.\n", - " \"\"\"\n", - " try:\n", - " dataset = al.Imaging.from_fits(\n", - " data_path=f\"{simulate_path}/data.fits\",\n", - " psf_path=f\"{simulate_path}/psf.fits\",\n", - " noise_map_path=f\"{simulate_path}/noise_map.fits\",\n", - " pixel_scales=self.mask.pixel_scales,\n", - " check_noise_map=False,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=self.mask)\n", - "\n", - " tracer = al.Tracer(\n", - " galaxies=[\n", - " instance.galaxies.lens,\n", - " instance.perturb,\n", - " instance.galaxies.source,\n", - " ]\n", - " )\n", - "\n", - " traced_grid = tracer.traced_grid_2d_list_from(\n", - " grid=dataset.grid,\n", - " )[-1]\n", - "\n", - " source_centre = instance.galaxies.source.bulge.centre\n", - "\n", - " over_sample_size = (\n", - " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=traced_grid,\n", - " sub_size_list=[16, 8, 2],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=[source_centre],\n", - " )\n", - " )\n", - "\n", - " over_sample_size_lens = (\n", - " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - " )\n", - "\n", - " over_sample_size = np.where(\n", - " over_sample_size > over_sample_size_lens,\n", - " over_sample_size,\n", - " over_sample_size_lens,\n", - " )\n", - " over_sample_size = al.Array2D(values=over_sample_size, mask=self.mask)\n", - "\n", - " return dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " except FileNotFoundError:\n", - " pass\n", - "\n", - " \"\"\"\n", - " __Ray Tracing__\n", - "\n", - " Set up the `Tracer` which is used to simulate the strong lens imaging, which may include the subhalo in\n", - " addition to the lens and source galaxy.\n", - " \"\"\"\n", - " tracer = al.Tracer(\n", - " galaxies=[\n", - " instance.galaxies.lens,\n", - " instance.perturb,\n", - " instance.galaxies.source,\n", - " ]\n", - " )\n", - "\n", - " \"\"\"\n", - " __Simulate__\n", - "\n", - " Set up the grid, PSF and simulator settings used to simulate imaging of the strong lens. These should be \n", - " tuned to match the S/N and noise properties of the observed data you are performing sensitivity mapping on.\n", - " \"\"\"\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=self.mask.shape_native,\n", - " pixel_scales=self.mask.pixel_scales,\n", - " )\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " simulator = al.SimulatorImaging(\n", - " exposure_time=1000.0,\n", - " psf=self.psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - " noise_seed=1,\n", - " )\n", - "\n", - " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - " \"\"\"\n", - " __Masking__\n", - "\n", - " The data generated by the simulate function is what is ultimately fitted.\n", - "\n", - " Therefore, we also apply the mask for the analysis before we return the simulated data.\n", - " \"\"\"\n", - " dataset = dataset.apply_mask(mask=self.mask)\n", - "\n", - " traced_grid = tracer.traced_grid_2d_list_from(\n", - " grid=dataset.grid,\n", - " )[-1]\n", - "\n", - " source_centre = instance.galaxies.source.bulge.centre\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=traced_grid,\n", - " sub_size_list=[16, 8, 2],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=[source_centre],\n", - " )\n", - "\n", - " over_sample_size_lens = (\n", - " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - " )\n", - "\n", - " over_sample_size = np.where(\n", - " over_sample_size > over_sample_size_lens,\n", - " over_sample_size,\n", - " over_sample_size_lens,\n", - " )\n", - " over_sample_size = al.Array2D(values=over_sample_size, mask=self.mask)\n", - "\n", - " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " \"\"\"\n", - " Outputs info about the `Tracer` to the fit, so we know exactly how we simulated the image.\n", - " \"\"\"\n", - " self.output_info(simulate_path=simulate_path, dataset=dataset, tracer=tracer)\n", - "\n", - " return dataset\n", - "\n", - " def output_info(self, simulate_path: str, dataset: al.Imaging, tracer: al.Imaging):\n", - " \"\"\"\n", - " Output information about the data simulated for this iteration of sensitivity mapping.\n", - "\n", - " This information output is as follows:\n", - "\n", - " - A subplot of the simulated imaging dataset.\n", - " - A subplot of the tracer used to simulate this imaging dataset.\n", - " - A .json file containing the tracer galaxies.\n", - "\n", - " Parameters\n", - " ----------\n", - " simulate_path\n", - " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", - " mapping.\n", - " dataset\n", - " The simulated imaging dataset which is visualized.\n", - " tracer\n", - " The tracer used to simulate the imaging dataset, which is visualized and output to a .json file.\n", - " \"\"\"\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " aplt.subplot_lensed_images(\n", - " tracer=tracer,\n", - " grid=dataset.grid,\n", - " output_path=simulate_path,\n", - " output_format=\"png\",\n", - " )\n", - "\n", - " al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(simulate_path) / \"tracer.json\",\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Base Fit__\n", - "\n", - "We have defined a `Simulate` class that will be used to simulate every dataset simulated by the sensitivity mapper.\n", - "Each simulated dataset will have a unique set of parameters for the `subhalo` (e.g. due to different values of\n", - "`perturb_model`.\n", - "\n", - "We will fit each simulated dataset using the `base_model`, which quantifies whether not including the dark matter\n", - "subhalo in the model changess the goodness-of-fit and therefore indicates if we are sensitive to the subhalo.\n", - "\n", - "We now write a `BaseFit` class, defining how the `base_model` is fitted to each simulated dataset and the \n", - "goodness-of-fit used to quantify whether the model fits the data well. As above, the `__init__` method can be\n", - "extended with new inputs to control how the model is fitted and the `__call__` method performs the fit.\n", - "\n", - "In this example, we use a full non-linear search to fit the `base_model` to the simulated data and return\n", - "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", - "more computationally efficient, for example performing a single log likelihood evaluation of the `base_model` fit\n", - "to the simulated data.\n", - "\n", - "Fucntionality which adapts the mesh and regularization of a pixelized source reconstruction to the unlensed source's \n", - "morphology require an `adapt_images`. This is an input of the __init__ constructor which is passed to the `Analysis` \n", - "for every simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class BaseFit:\n", - " def __init__(self, adapt_images, number_of_cores: int = 1):\n", - " \"\"\"\n", - " Class used to fit every dataset used for sensitivity mapping with the base model (the model without the\n", - " perturbed feature sensitivity mapping maps out).\n", - "\n", - " In this example, the base model therefore does not include the dark matter subhalo, but the simulated\n", - " dataset includes one.\n", - "\n", - " The base fit is repeated for every parameter on the sensitivity grid and compared to the perturbed fit. This\n", - " maps out the sensitivity of every parameter is (e.g. the sensitivity of the mass of the subhalo).\n", - "\n", - " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", - " fitted, below we include an input `analysis_cls` which is the `AnalysisImaging` class used to fit the model\n", - " to the dataset.\n", - "\n", - " Parameters\n", - " ----------\n", - " adapt_images\n", - " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", - " reconstructed galaxy's morphology.\n", - " number_of_cores\n", - " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", - " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", - " \"\"\"\n", - " self.adapt_images = adapt_images\n", - " self.number_of_cores = number_of_cores\n", - "\n", - " def __call__(self, dataset, model, paths, instance):\n", - " \"\"\"\n", - " The base fitting function which fits every dataset used for sensitivity mapping with the base model.\n", - "\n", - " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", - " quantify how sensitive the model is to the perturbed feature.\n", - "\n", - " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", - " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", - "\n", - " Parameters\n", - " ----------\n", - " dataset\n", - " The dataset which is simulated with the perturbed model and which is fitted.\n", - " model\n", - " The model instance which is fitted to the dataset, which does not include the perturbed feature.\n", - " paths\n", - " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", - " instance\n", - " The simulation instance, which includes the perturbed feature that is used to simulate the dataset.\n", - " This is often not used, but may be useful for certain sensitivity mapping tasks, for example using\n", - " true values of the simulated instance to set up aspects of the model-fit (e.g. the priors).\n", - " \"\"\"\n", - "\n", - " search = af.Nautilus(\n", - " paths=paths,\n", - " n_live=50,\n", - " n_batch=50,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - " analysis._adapt_images = self.adapt_images\n", - "\n", - " return search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Perturb Fit__\n", - "\n", - "We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each simulated dataset. This\n", - "behaves analogously to the `BaseFit` class above, but now fits the `perturb_model` to the simulated data (as\n", - "opposed to the `base_model`).\n", - "\n", - "Again, in this example we use a full non-linear search to fit the `perturb_model` to the simulated data and return\n", - "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", - "more computationally efficient, for example performing a single log likelihood evaluation of the `perturb_model` fit\n", - "to the simulated data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "class PerturbFit:\n", - " def __init__(self, adapt_images, number_of_cores: int = 1):\n", - " \"\"\"\n", - " Class used to fit every dataset used for sensitivity mapping with the perturbed model (the model with the\n", - " perturbed feature sensitivity mapping maps out).\n", - "\n", - " In this example, the perturbed model therefore includes the dark matter subhalo, which is also in the\n", - " simulated dataset.\n", - "\n", - " The perturbed fit is repeated for every parameter on the sensitivity grid and compared to the base fit. This\n", - " maps out the sensitivity of every parameter is (e.g. the sensitivity of mass of the subhalo).\n", - "\n", - " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", - " fitted, below we include an input `analysis_cls` which is the `Analysis` class used to fit the model to the\n", - " dataset.\n", - "\n", - " Parameters\n", - " ----------\n", - " adapt_images\n", - " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", - " reconstructed galaxy's morphology.\n", - " number_of_cores\n", - " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", - " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", - " \"\"\"\n", - " self.adapt_images = adapt_images\n", - " self.number_of_cores = number_of_cores\n", - "\n", - " def __call__(self, dataset, model, paths, instance):\n", - " \"\"\"\n", - " The perturbed fitting function which fits every dataset used for sensitivity mapping with the perturbed model.\n", - "\n", - " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", - " quantify how sensitive the model is to the perturbed feature.\n", - "\n", - " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", - " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", - "\n", - " Parameters\n", - " ----------\n", - " dataset\n", - " The dataset which is simulated with the perturbed model and which is fitted.\n", - " model\n", - " The model instance which is fitted to the dataset, which includes the perturbed feature.\n", - " paths\n", - " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", - " instance\n", - " The simulation instance, which includes the perturbed feature that is used to simulate the dataset.\n", - " This is often not used, but may be useful for certain sensitivity mapping tasks, for example using\n", - " true values of the simulated instance to set up aspects of the model-fit (e.g. the priors).\n", - " \"\"\"\n", - "\n", - " search = af.Nautilus(\n", - " paths=paths,\n", - " n_live=50,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - " analysis._adapt_images = self.adapt_images\n", - "\n", - " return search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now combine all of the objects created above and perform sensitivity mapping. The inputs to the `Sensitivity`\n", - "object below are:\n", - "\n", - "- `simulation_instance`: This is an instance of the model used to simulate every dataset that is fitted. In this example \n", - "it is a lens model that does not include a subhalo, which was inferred by fitting the dataset we perform sensitivity \n", - "mapping on.\n", - "\n", - "- `base_model`: This is the lens model that is fitted to every simulated dataset, which does not include a subhalo. In \n", - "this example is composed of an `Isothermal` lens and MGE source.\n", - "\n", - "- `perturb_model`: This is the extra model component that alongside the `base_model` is fitted to every simulated \n", - "dataset. In this example it is a `NFWMCRLudlowSph` dark matter subhalo.\n", - "\n", - "- `simulate_cls`: This is the function that uses the `simulation_instance` and many instances of the `perturb_model` \n", - "to simulate many datasets that are fitted with the `base_model` and `base_model` + `perturb_model`.\n", - "\n", - "- `base_fit_cls`: This is the function that fits the `base_model` to every simulated dataset and returns the\n", - "goodness-of-fit of the model to the data.\n", - "\n", - "- `perturb_fit_cls`: This is the function that fits the `base_model` + `perturb_model` to every simulated dataset and\n", - "returns the goodness-of-fit of the model to the data.\n", - "\n", - "- `number_of_steps`: The number of steps over which the parameters in the `perturb_model` are iterated. In this \n", - "example, each subhalo ``centre` has a `UniformPrior` with lower limit -3.0 and upper limit 3.0, therefore \n", - "the `number_of_steps=2` will simulate and fit 4 datasets where the `centre` values \n", - "are [(-1.5, -1.5), (-1.5, 1.5), (1.5, -1.5), (1.5, 1.5)].\n", - "\n", - "- `number_of_cores`: The number of cores over which the sensitivity mapping is performed, enabling parallel processing\n", - "if set above 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "paths = af.DirectoryPaths(\n", - " path_prefix=Path(\"features\"),\n", - " name=\"sensitivity_mapping\",\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result)\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "sensitivity = af.Sensitivity(\n", - " paths=paths,\n", - " simulation_instance=simulation_instance,\n", - " base_model=base_model,\n", - " perturb_model=perturb_model,\n", - " simulate_cls=SimulateImaging(mask=mask, psf=dataset.psf),\n", - " base_fit_cls=BaseFit(adapt_images=adapt_images),\n", - " perturb_fit_cls=PerturbFit(adapt_images=adapt_images),\n", - " perturb_model_prior_func=perturb_model_prior_func,\n", - " number_of_steps=2,\n", - " # number_of_steps=(4, 2),\n", - " number_of_cores=2,\n", - ")\n", - "\n", - "sensitivity_result = sensitivity.run()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Results__\n", - "\n", - "You should now look at the results of the sensitivity mapping in the folder `output/features/sensitivity_mapping`. \n", - "\n", - "You will note the following 4 sets of x2 model-fits have been performed:\n", - "\n", - " - The `base_model` is fitted to a simulated dataset where a subhalo is included at the (y,x) \n", - " coorindates [(-1.5, -1.5), (-1.5, 1.5), (1.5, -1.5), (1.5, 1.5)].\n", - "\n", - " - The `base_model` + `perturb_model` is fitted to a simulated dataset where a subhalo is included at the (y,x) \n", - " coorindates [(-1.5, -1.5), (-1.5, 1.5), (1.5, -1.5), (1.5, 1.5)].\n", - "\n", - "The fit produces a `sensitivity_result`. \n", - "\n", - "We can print the `log_evidence_differences` of every cell of the sensitivity map." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(sensitivity_result.log_evidence_differences.native)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Sensitivity Mapping: Start Here\n", + "===============================\n", + "\n", + "Bayesian model comparison allows us to take a dataset, fit it with multiple models and use the Bayesian evidence to\n", + "quantify which model objectively gives the best-fit following the principles of Occam's Razor.\n", + "\n", + "However, a complex model may not be favoured by model comparison not because it is the 'wrong' model, but simply\n", + "because the dataset being fitted is not of a sufficient quality for the more complex model to be favoured. Sensitivity\n", + "mapping addresses what quality of data would be needed for the more complex model to be favoured.\n", + "\n", + "In order to do this, sensitivity mapping involves us writing a function that uses the model(s) to simulate a dataset.\n", + "We then use this function to simulate many datasets, for many different models, and fit each dataset to quantify\n", + "how much the change in the model led to a measurable change in the data. This is called computing the sensitivity.\n", + "\n", + "How we compute the sensitivity is chosen by us, the user. In this example, we will perform multiple model-fits\n", + "with a nested sampling search, and therefore perform Bayesian model comparison to compute the sensitivity. This allows\n", + "us to infer how much of a Bayesian evidence increase we should expect for datasets of varying quality and / or models\n", + "with different parameters.\n", + "\n", + "__Contents__\n", + "\n", + "- **Subhalo Detection Discussion:** For strong lensing, this process is crucial for dark matter substructure detection, as discussed in.\n", + "- **Subhalo Sensitivity Mapping:** Sensitivity mapping is a process where we simulate many thousands of strong lens images.\n", + "- **SLaM Pipelines:** The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which.\n", + "- **Pixelized Source:** Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of.\n", + "- **Base Model:** We now define the `base_model` that we use to perform sensitivity mapping.\n", + "- **Perturb Model:** We now define the `perturb_model`, which is the model component whose parameters we iterate over to.\n", + "- **Mapping Grid:** Sensitivity mapping is typically performed over a large range of parameters.\n", + "- **Perturb Model Prior Func:** The default priors on the `perturb_model` are `UniformPrior`'s bounded around each sensitivity grid.\n", + "- **Simulation Instance:** We are performing sensitivity mapping to determine where a subhalo is detectable.\n", + "- **Simulate Function Class:** We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above).\n", + "- **Base Fit:** We have defined a `Simulate` class that will be used to simulate every dataset simulated by the.\n", + "- **Perturb Fit:** We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each.\n", + "- **Results:** You should now look at the results of the sensitivity mapping in the folder.\n", + "\n", + "__Subhalo Detection Discussion__\n", + "\n", + "For strong lensing, this process is crucial for dark matter substructure detection, as discussed in the following paper:\n", + "\n", + "https://arxiv.org/abs/0903.4752\n", + "\n", + "In subhalo detection, our strong lens modeling informs us of whether there is a dark matter subhalo at a given (y,x)\n", + "image-plane location of the strong lens. We determine this by fitting a lens models which includes a subhalo. However,\n", + "we are only able to detect dark matter subhalos with (y,x) locations near the lensed source light, and when the\n", + "subhalo is massive enough to perturb its light in an observable way.\n", + "\n", + "Subhalo detection analysis therefore does not tell us where we could detect subhalos and of what mass. To know this,\n", + "we must perform sensitivity mapping.\n", + "\n", + "__Subhalo Sensitivity Mapping__\n", + "\n", + "Sensitivity mapping is a process where we simulate many thousands of strong lens images. Each simulated image includes a\n", + "dark matter subhalo at a given (y,x) coordinate and at a given mass. We fit each simulated dataset twice,\n", + "with a lens model which does not include a subhalo and with a lens model that does.\n", + "\n", + "If the Bayesian evidence of the lens model including a subhalo is higher than the model which does not, a subhalo at t\n", + "hat (y,x) location and mass is therefore detectable.\n", + "\n", + "For many simulated datasets, we will find the evidence does not increase when we include a subhalo in the model-fit,\n", + "informing us that regions of the image-plane away from the lensed source are not sensitive to subhalos.\n", + "\n", + "The sensitivity map is performed over a three dimensional (or higher) grid of subhalo (y,x) and mass. Thus, once\n", + "sensitivity mapping is complete, we have a complete map of where in the image-plane subhalos of what mass are\n", + "detectable. We can plot 2D plots of this grid to visualize where we are sensitive to dark matter subhalos.\n", + "\n", + "The information provided by a sensitivity map is ultimately required to turn dark matter subhalo detections into\n", + "constraints on the dark matter subhalo mass function, which is the primary goal of subhalo detection. Thus, it is\n", + "necessary to make statements about the nature of dark matter: cold, warm, fuzzy, or something else entirely?\n", + "\n", + "__SLaM Pipelines__\n", + "\n", + "The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which automate the fitting\n", + "of complex lens models. The SLaM pipelines are used for all DM subhalo detection analyses.\n", + "\n", + "This example script does not use a SLaM pipeline, to keep the sensitivity mapping self contained. However, it is\n", + "anticipated that any user performing sensitivity mapping on real data will use the SLaM pipelines, which in the\n", + "`subhalo` package have dedicated extensions for performing sensitivity mapping to both imaging and interferometer data.\n", + "\n", + "Therefore you should be familiar with the SLaM pipelines before performing DM subhalo sensitivity mapping on real\n", + "data. If you are unfamiliar with the SLaM pipelines, checkout the\n", + "example `autolens_workspace/notebooks/guides/modeling/slam_start_here`.\n", + "\n", + "__Pixelized Source__\n", + "\n", + "Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of the source's light\n", + "are at a level where the subhalo's perturbing lensing effects can be detected.\n", + "\n", + "This requires the source reconstruction to be performed using a pixelized source, as this provides a more detailed\n", + "reconstruction of the source's light than fits using light profiles.\n", + "\n", + "Therefore, the corresponding sensitivity mapping should also be performed using pixelized sources. This example\n", + "sticks to light profile sources, to provide faster run times illustrative purposes.\n", + "\n", + "The example `subhalo/sensitivity/examples/source_pixelized.ipynb` extends the SLaM pipelines with pixelized sources\n", + "and therefore shows how to perform sensitivity mapping using pixelized sources.\n", + "\n", + "Note that the simulation procedure for a pixelized source is different to the one shown here. In this example, the\n", + "light profile source parameters are used to simulate each sensitivity mapping dataset. When pixelized sources are\n", + "used, the source reconstruction on the mesh is used, such that the simulations capture the irregular morphologies\n", + "of real source galaxies." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import os\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__ \n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"dark_matter_subhalo\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/advanced/subhalo/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.05,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model + Search + Analysis + Model-Fit (Base Search)__\n", + "\n", + "We are performing sensitivity mapping to determine where a subhalo is detectable. This will require us to simulate \n", + "many realizations of our dataset with a lens model, called the `simulation_instance`. To get this model, we therefore \n", + "fit the data before performing sensitivity mapping so that we can set the `simulation_instance` as the maximum \n", + "likelihood model.\n", + "\n", + "We perform this fit using the lens model we will use to perform sensitivity mapping, which we call the `base_model`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "base_model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", + " source=af.Model(al.Galaxy, redshift=1.0, bulge=bulge),\n", + " ),\n", + ")\n", + "\n", + "search_base = af.Nautilus(\n", + " path_prefix=Path(\"imaging\", \"advanced\", \"subhalo\", \"sensitivity\"),\n", + " name=\"sensitivity_mapping_base\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result = search_base.fit(model=base_model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Base Model__\n", + "\n", + "We now define the `base_model` that we use to perform sensitivity mapping. This is the lens model that is fitted to \n", + "every simulated strong lens without a subhalo, giving us the Bayesian evidence which we compare to the model which \n", + "includes one!). \n", + "\n", + "For this model, we can use the `base_model` above, however we will use the result of fitting this model to the dataset\n", + "before sensitivity mapping. This ensures the priors associated with each parameter are initialized so as to speed up\n", + "each non-linear search performed during sensitivity mapping." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "base_model = result.model" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Perturb Model__\n", + "\n", + "We now define the `perturb_model`, which is the model component whose parameters we iterate over to perform \n", + "sensitivity mapping. In this case, this model is a `NFWMCRLudlowSph` model and we will iterate over its\n", + "`centre` and `mass_at_200`. We set it up as a `Model` so it has an associated redshift and can be directly\n", + "passed to the tracer in the simulate function below.\n", + "\n", + "Many instances of the `perturb_model` are created and used to simulate the many strong lens datasets that we fit. \n", + "However, it is only included in half of the model-fits; corresponding to the lens models which include a dark matter \n", + "subhalo and whose Bayesian evidence we compare to the simpler model-fits consisting of just the `base_model` to \n", + "determine if the subhalo was detectable.\n", + "\n", + "By fitting both models to every simulated lens, we therefore infer the Bayesian evidence of every model to every \n", + "dataset. Sensitivity mapping therefore maps out for what values of `centre` and `mass_at_200` in the dark matter \n", + "subhalo the model-fit including a subhalo provide higher values of Bayesian evidence than the simpler model-fit (and\n", + "therefore when it is detectable!)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "perturb_model = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.NFWMCRLudlowSph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapping Grid__\n", + "\n", + "Sensitivity mapping is typically performed over a large range of parameters. However, to make this demonstration quick\n", + "and clear we are going to fix the `centre` of the subhalo to a value near the Einstein ring of (1.6, 0.0). We will \n", + "iterate over just two `mass_at_200` values corresponding to subhalos of mass 1e6 and 1e13, of which only the latter\n", + "will be shown to be detectable." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_dimension_arcsec = 3.0\n", + "\n", + "perturb_model.mass.mass_at_200 = 1e10\n", + "perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + ")\n", + "perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + ")\n", + "perturb_model.mass.redshift_object = 0.5\n", + "perturb_model.mass.redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Perturb Model Prior Func__\n", + "\n", + "The default priors on the `perturb_model` are `UniformPrior`'s bounded around each sensitivity grid cell.\n", + "\n", + "For example, the first simulated dark matter subhalo is at location (1.5, -1.5) and its priors are: \n", + "\n", + "- y is `UniformPrior(lower_limit=0.0, upper_limit=3.0)`.\n", + "- x is `UniformPrior(lower_limit=-3.0, upper_limit=0.0)`.\n", + "\n", + "The `mass_at_200` is fixed to a value of 1e10 in the `perturb_model` above, which is the fixed value used by\n", + "the model fit.\n", + "\n", + "By passing a `perturb_model_prior_func` to the sensitivity mapper, we can manually overwrite the priors on \n", + "the `perturb_model` which are used instead for the fit.\n", + "\n", + "Below, we update the priors as follows:\n", + "\n", + "- The y and x priors are trimmed to much narrower bounded priors, confined to regions 0.05\" each side of the \n", + " true DM subhalo.\n", + "\n", + "- The `mass_at_200` is made a free parameter with `LogUniformPrior(lower_limit=1e6, 1e12)`. This is a large range, \n", + " but ensures there are solutions where the DM subhalo can go to lower masses. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def perturb_model_prior_func(perturb_instance, perturb_model):\n", + " b = 0.05\n", + "\n", + " perturb_model.mass.centre.centre_0 = af.UniformPrior(\n", + " lower_limit=perturb_instance.mass.centre[0] - b,\n", + " upper_limit=perturb_instance.mass.centre[0] + b,\n", + " )\n", + "\n", + " perturb_model.mass.centre.centre_1 = af.UniformPrior(\n", + " lower_limit=perturb_instance.mass.centre[1] - b,\n", + " upper_limit=perturb_instance.mass.centre[1] + b,\n", + " )\n", + "\n", + " perturb_model.mass.mass_at_200 = af.LogUniformPrior(\n", + " lower_limit=1e6, upper_limit=1e12\n", + " )\n", + "\n", + " return perturb_model\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulation Instance__\n", + "\n", + "We are performing sensitivity mapping to determine where a subhalo is detectable. This will require us to \n", + "simulate many realizations of our dataset with a lens model, called the `simulation_instance`. This model uses the\n", + "result of the fit above.\n", + "\n", + "The code below ensures that the lens light, mass and source parameters of the strong lens are used when simulating\n", + "each dataset with a dark matter subhalo." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulation_instance = result.instance" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate Function Class__\n", + "\n", + "We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above) and uses it to \n", + "simulate a strong lens dataset, which include a dark matter subhalo, which is subsequently fitted.\n", + "\n", + "Additional attributes required to simulate the data (mask, PSF) can be passed to the `__init__` method, and the \n", + "simulation is performed in the `__call__` method.\n", + "\n", + "When this dataset is simulated, the quantity `instance.perturb` is used in `__call__`. This is an instance \n", + "of the `NFWMCRLudlowSph`, and it is different every time the `simulate_cls` is called based on the value of sensitivity \n", + "being computed. \n", + "\n", + "In this example, this `instance.perturb` corresponds to two different subhalos with values of `mass_at_200` of \n", + "1e6 MSun and 1e13 MSun." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class SimulateImaging:\n", + " def __init__(self, mask, psf):\n", + " \"\"\"\n", + " Class used to simulate the strong lens imaging used for sensitivity mapping.\n", + "\n", + " Parameters\n", + " ----------\n", + " mask\n", + " The mask applied to the real image data, which is applied to every simulated imaging.\n", + " psf\n", + " The PSF of the real image data, which is applied to every simulated imaging and used for each fit.\n", + " \"\"\"\n", + " self.mask = mask\n", + " self.psf = psf\n", + "\n", + " def __call__(self, instance: af.ModelInstance, simulate_path: str):\n", + " \"\"\"\n", + " The `simulate_function` called by the `Sensitivity` class which simulates each strong lens image fitted\n", + " by the sensitivity mapper.\n", + "\n", + " The simulation procedure is as follows:\n", + "\n", + " 1) Use the input galaxies of the sensitivity `instance` to set up a tracer, which generates the image-plane\n", + " image of the strong lens system.\n", + "\n", + " 2) Simulate this image using the input dataset noise (Poisson) and PSF.\n", + "\n", + " 3) Apply the mask used in the analysis of the real image to the simulated image.\n", + "\n", + " 4) Output information about the simulation to hard-disk.\n", + "\n", + " The `subhalo` in the sensitivity `instance` changes for every iteration of the sensitivity mapping, ensuring\n", + " that we map out the sensitivity of the analysis to the subhalo properties (centre, mass, etc.).\n", + "\n", + " Parameters\n", + " ----------\n", + " instance\n", + " The sensitivity instance, which includes the galaxies whose parameters are varied to perform sensitivity.\n", + " The subhalo in this instance changes for every iteration of the sensitivity mapping.\n", + " simulate_path\n", + " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", + " mapping.\n", + "\n", + " Returns\n", + " -------\n", + " A simulated image of a strong lens, which id input into the fits of the sensitivity mapper.\n", + " \"\"\"\n", + "\n", + " \"\"\"\n", + " __Resume Fit__\n", + "\n", + " If sensitivity mapping already began on this grid cell, the dataset will have been simulated already and we\n", + " do not want to resimulate it and change its noise properties. \n", + "\n", + " We therefore load it from the `simulate_path` instead.\n", + " \"\"\"\n", + " try:\n", + " dataset = al.Imaging.from_fits(\n", + " data_path=f\"{simulate_path}/data.fits\",\n", + " psf_path=f\"{simulate_path}/psf.fits\",\n", + " noise_map_path=f\"{simulate_path}/noise_map.fits\",\n", + " pixel_scales=self.mask.pixel_scales,\n", + " check_noise_map=False,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=self.mask)\n", + "\n", + " tracer = al.Tracer(\n", + " galaxies=[\n", + " instance.galaxies.lens,\n", + " instance.perturb,\n", + " instance.galaxies.source,\n", + " ]\n", + " )\n", + "\n", + " traced_grid = tracer.traced_grid_2d_list_from(\n", + " grid=dataset.grid,\n", + " )[-1]\n", + "\n", + " source_centre = instance.galaxies.source.bulge.centre\n", + "\n", + " over_sample_size = (\n", + " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=traced_grid,\n", + " sub_size_list=[16, 8, 2],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=[source_centre],\n", + " )\n", + " )\n", + "\n", + " over_sample_size_lens = (\n", + " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + " )\n", + "\n", + " over_sample_size = np.where(\n", + " over_sample_size > over_sample_size_lens,\n", + " over_sample_size,\n", + " over_sample_size_lens,\n", + " )\n", + " over_sample_size = al.Array2D(values=over_sample_size, mask=self.mask)\n", + "\n", + " return dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " except FileNotFoundError:\n", + " pass\n", + "\n", + " \"\"\"\n", + " __Ray Tracing__\n", + "\n", + " Set up the `Tracer` which is used to simulate the strong lens imaging, which may include the subhalo in\n", + " addition to the lens and source galaxy.\n", + " \"\"\"\n", + " tracer = al.Tracer(\n", + " galaxies=[\n", + " instance.galaxies.lens,\n", + " instance.perturb,\n", + " instance.galaxies.source,\n", + " ]\n", + " )\n", + "\n", + " \"\"\"\n", + " __Simulate__\n", + "\n", + " Set up the grid, PSF and simulator settings used to simulate imaging of the strong lens. These should be \n", + " tuned to match the S/N and noise properties of the observed data you are performing sensitivity mapping on.\n", + " \"\"\"\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=self.mask.shape_native,\n", + " pixel_scales=self.mask.pixel_scales,\n", + " )\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " simulator = al.SimulatorImaging(\n", + " exposure_time=1000.0,\n", + " psf=self.psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + " noise_seed=1,\n", + " )\n", + "\n", + " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + " \"\"\"\n", + " __Masking__\n", + "\n", + " The data generated by the simulate function is what is ultimately fitted.\n", + "\n", + " Therefore, we also apply the mask for the analysis before we return the simulated data.\n", + " \"\"\"\n", + " dataset = dataset.apply_mask(mask=self.mask)\n", + "\n", + " traced_grid = tracer.traced_grid_2d_list_from(\n", + " grid=dataset.grid,\n", + " )[-1]\n", + "\n", + " source_centre = instance.galaxies.source.bulge.centre\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=traced_grid,\n", + " sub_size_list=[16, 8, 2],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=[source_centre],\n", + " )\n", + "\n", + " over_sample_size_lens = (\n", + " al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + " )\n", + "\n", + " over_sample_size = np.where(\n", + " over_sample_size > over_sample_size_lens,\n", + " over_sample_size,\n", + " over_sample_size_lens,\n", + " )\n", + " over_sample_size = al.Array2D(values=over_sample_size, mask=self.mask)\n", + "\n", + " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " \"\"\"\n", + " Outputs info about the `Tracer` to the fit, so we know exactly how we simulated the image.\n", + " \"\"\"\n", + " self.output_info(simulate_path=simulate_path, dataset=dataset, tracer=tracer)\n", + "\n", + " return dataset\n", + "\n", + " def output_info(self, simulate_path: str, dataset: al.Imaging, tracer: al.Imaging):\n", + " \"\"\"\n", + " Output information about the data simulated for this iteration of sensitivity mapping.\n", + "\n", + " This information output is as follows:\n", + "\n", + " - A subplot of the simulated imaging dataset.\n", + " - A subplot of the tracer used to simulate this imaging dataset.\n", + " - A .json file containing the tracer galaxies.\n", + "\n", + " Parameters\n", + " ----------\n", + " simulate_path\n", + " The path where the simulated dataset is output, contained within each sub-folder of the sensitivity\n", + " mapping.\n", + " dataset\n", + " The simulated imaging dataset which is visualized.\n", + " tracer\n", + " The tracer used to simulate the imaging dataset, which is visualized and output to a .json file.\n", + " \"\"\"\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " aplt.subplot_lensed_images(\n", + " tracer=tracer,\n", + " grid=dataset.grid,\n", + " output_path=simulate_path,\n", + " output_format=\"png\",\n", + " )\n", + "\n", + " al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(simulate_path) / \"tracer.json\",\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Base Fit__\n", + "\n", + "We have defined a `Simulate` class that will be used to simulate every dataset simulated by the sensitivity mapper.\n", + "Each simulated dataset will have a unique set of parameters for the `subhalo` (e.g. due to different values of\n", + "`perturb_model`.\n", + "\n", + "We will fit each simulated dataset using the `base_model`, which quantifies whether not including the dark matter\n", + "subhalo in the model changess the goodness-of-fit and therefore indicates if we are sensitive to the subhalo.\n", + "\n", + "We now write a `BaseFit` class, defining how the `base_model` is fitted to each simulated dataset and the \n", + "goodness-of-fit used to quantify whether the model fits the data well. As above, the `__init__` method can be\n", + "extended with new inputs to control how the model is fitted and the `__call__` method performs the fit.\n", + "\n", + "In this example, we use a full non-linear search to fit the `base_model` to the simulated data and return\n", + "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", + "more computationally efficient, for example performing a single log likelihood evaluation of the `base_model` fit\n", + "to the simulated data.\n", + "\n", + "Fucntionality which adapts the mesh and regularization of a pixelized source reconstruction to the unlensed source's \n", + "morphology require an `adapt_images`. This is an input of the __init__ constructor which is passed to the `Analysis` \n", + "for every simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class BaseFit:\n", + " def __init__(self, adapt_images, number_of_cores: int = 1):\n", + " \"\"\"\n", + " Class used to fit every dataset used for sensitivity mapping with the base model (the model without the\n", + " perturbed feature sensitivity mapping maps out).\n", + "\n", + " In this example, the base model therefore does not include the dark matter subhalo, but the simulated\n", + " dataset includes one.\n", + "\n", + " The base fit is repeated for every parameter on the sensitivity grid and compared to the perturbed fit. This\n", + " maps out the sensitivity of every parameter is (e.g. the sensitivity of the mass of the subhalo).\n", + "\n", + " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", + " fitted, below we include an input `analysis_cls` which is the `AnalysisImaging` class used to fit the model\n", + " to the dataset.\n", + "\n", + " Parameters\n", + " ----------\n", + " adapt_images\n", + " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", + " reconstructed galaxy's morphology.\n", + " number_of_cores\n", + " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", + " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", + " \"\"\"\n", + " self.adapt_images = adapt_images\n", + " self.number_of_cores = number_of_cores\n", + "\n", + " def __call__(self, dataset, model, paths, instance):\n", + " \"\"\"\n", + " The base fitting function which fits every dataset used for sensitivity mapping with the base model.\n", + "\n", + " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", + " quantify how sensitive the model is to the perturbed feature.\n", + "\n", + " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", + " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", + "\n", + " Parameters\n", + " ----------\n", + " dataset\n", + " The dataset which is simulated with the perturbed model and which is fitted.\n", + " model\n", + " The model instance which is fitted to the dataset, which does not include the perturbed feature.\n", + " paths\n", + " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", + " instance\n", + " The simulation instance, which includes the perturbed feature that is used to simulate the dataset.\n", + " This is often not used, but may be useful for certain sensitivity mapping tasks, for example using\n", + " true values of the simulated instance to set up aspects of the model-fit (e.g. the priors).\n", + " \"\"\"\n", + "\n", + " search = af.Nautilus(\n", + " paths=paths,\n", + " n_live=50,\n", + " n_batch=50,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + " analysis._adapt_images = self.adapt_images\n", + "\n", + " return search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Perturb Fit__\n", + "\n", + "We now define a `PerturbFit` class, which defines how the `perturb_model` is fitted to each simulated dataset. This\n", + "behaves analogously to the `BaseFit` class above, but now fits the `perturb_model` to the simulated data (as\n", + "opposed to the `base_model`).\n", + "\n", + "Again, in this example we use a full non-linear search to fit the `perturb_model` to the simulated data and return\n", + "the `log_evidence` of the model fit as the goodness-of-fit. This fit could easily be something much simpler and\n", + "more computationally efficient, for example performing a single log likelihood evaluation of the `perturb_model` fit\n", + "to the simulated data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "class PerturbFit:\n", + " def __init__(self, adapt_images, number_of_cores: int = 1):\n", + " \"\"\"\n", + " Class used to fit every dataset used for sensitivity mapping with the perturbed model (the model with the\n", + " perturbed feature sensitivity mapping maps out).\n", + "\n", + " In this example, the perturbed model therefore includes the dark matter subhalo, which is also in the\n", + " simulated dataset.\n", + "\n", + " The perturbed fit is repeated for every parameter on the sensitivity grid and compared to the base fit. This\n", + " maps out the sensitivity of every parameter is (e.g. the sensitivity of mass of the subhalo).\n", + "\n", + " The `__init__` constructor can be extended with new inputs which can be used to control how the dataset is\n", + " fitted, below we include an input `analysis_cls` which is the `Analysis` class used to fit the model to the\n", + " dataset.\n", + "\n", + " Parameters\n", + " ----------\n", + " adapt_images\n", + " Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the\n", + " reconstructed galaxy's morphology.\n", + " number_of_cores\n", + " The number of cores used to perform the non-linear search. If 1, each model-fit on the grid is performed\n", + " in serial, if > 1 fits are distributed in parallel using the Python multiprocessing module.\n", + " \"\"\"\n", + " self.adapt_images = adapt_images\n", + " self.number_of_cores = number_of_cores\n", + "\n", + " def __call__(self, dataset, model, paths, instance):\n", + " \"\"\"\n", + " The perturbed fitting function which fits every dataset used for sensitivity mapping with the perturbed model.\n", + "\n", + " This function receives as input each simulated dataset of the sensitivity map and fits it, in order to\n", + " quantify how sensitive the model is to the perturbed feature.\n", + "\n", + " In this example, a full non-linear search is performed to determine how well the model fits the dataset.\n", + " The `log_evidence` of the fit is returned which acts as the sensitivity map figure of merit.\n", + "\n", + " Parameters\n", + " ----------\n", + " dataset\n", + " The dataset which is simulated with the perturbed model and which is fitted.\n", + " model\n", + " The model instance which is fitted to the dataset, which includes the perturbed feature.\n", + " paths\n", + " The `Paths` instance which contains the path to the folder where the results of the fit are written to.\n", + " instance\n", + " The simulation instance, which includes the perturbed feature that is used to simulate the dataset.\n", + " This is often not used, but may be useful for certain sensitivity mapping tasks, for example using\n", + " true values of the simulated instance to set up aspects of the model-fit (e.g. the priors).\n", + " \"\"\"\n", + "\n", + " search = af.Nautilus(\n", + " paths=paths,\n", + " n_live=50,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + " analysis._adapt_images = self.adapt_images\n", + "\n", + " return search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now combine all of the objects created above and perform sensitivity mapping. The inputs to the `Sensitivity`\n", + "object below are:\n", + "\n", + "- `simulation_instance`: This is an instance of the model used to simulate every dataset that is fitted. In this example \n", + "it is a lens model that does not include a subhalo, which was inferred by fitting the dataset we perform sensitivity \n", + "mapping on.\n", + "\n", + "- `base_model`: This is the lens model that is fitted to every simulated dataset, which does not include a subhalo. In \n", + "this example is composed of an `Isothermal` lens and MGE source.\n", + "\n", + "- `perturb_model`: This is the extra model component that alongside the `base_model` is fitted to every simulated \n", + "dataset. In this example it is a `NFWMCRLudlowSph` dark matter subhalo.\n", + "\n", + "- `simulate_cls`: This is the function that uses the `simulation_instance` and many instances of the `perturb_model` \n", + "to simulate many datasets that are fitted with the `base_model` and `base_model` + `perturb_model`.\n", + "\n", + "- `base_fit_cls`: This is the function that fits the `base_model` to every simulated dataset and returns the\n", + "goodness-of-fit of the model to the data.\n", + "\n", + "- `perturb_fit_cls`: This is the function that fits the `base_model` + `perturb_model` to every simulated dataset and\n", + "returns the goodness-of-fit of the model to the data.\n", + "\n", + "- `number_of_steps`: The number of steps over which the parameters in the `perturb_model` are iterated. In this \n", + "example, each subhalo ``centre` has a `UniformPrior` with lower limit -3.0 and upper limit 3.0, therefore \n", + "the `number_of_steps=2` will simulate and fit 4 datasets where the `centre` values \n", + "are [(-1.5, -1.5), (-1.5, 1.5), (1.5, -1.5), (1.5, 1.5)].\n", + "\n", + "- `number_of_cores`: The number of cores over which the sensitivity mapping is performed, enabling parallel processing\n", + "if set above 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "paths = af.DirectoryPaths(\n", + " path_prefix=Path(\"features\"),\n", + " name=\"sensitivity_mapping\",\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result)\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "sensitivity = af.Sensitivity(\n", + " paths=paths,\n", + " simulation_instance=simulation_instance,\n", + " base_model=base_model,\n", + " perturb_model=perturb_model,\n", + " simulate_cls=SimulateImaging(mask=mask, psf=dataset.psf),\n", + " base_fit_cls=BaseFit(adapt_images=adapt_images),\n", + " perturb_fit_cls=PerturbFit(adapt_images=adapt_images),\n", + " perturb_model_prior_func=perturb_model_prior_func,\n", + " number_of_steps=2,\n", + " # number_of_steps=(4, 2),\n", + " number_of_cores=2,\n", + ")\n", + "\n", + "sensitivity_result = sensitivity.run()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Results__\n", + "\n", + "You should now look at the results of the sensitivity mapping in the folder `output/features/sensitivity_mapping`. \n", + "\n", + "You will note the following 4 sets of x2 model-fits have been performed:\n", + "\n", + " - The `base_model` is fitted to a simulated dataset where a subhalo is included at the (y,x) \n", + " coorindates [(-1.5, -1.5), (-1.5, 1.5), (1.5, -1.5), (1.5, 1.5)].\n", + "\n", + " - The `base_model` + `perturb_model` is fitted to a simulated dataset where a subhalo is included at the (y,x) \n", + " coorindates [(-1.5, -1.5), (-1.5, 1.5), (1.5, -1.5), (1.5, 1.5)].\n", + "\n", + "The fit produces a `sensitivity_result`. \n", + "\n", + "We can print the `log_evidence_differences` of every cell of the sensitivity map." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(sensitivity_result.log_evidence_differences.native)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/advanced/subhalo/simulator.ipynb b/notebooks/imaging/features/advanced/subhalo/simulator.ipynb index f73079920..90a1d0566 100644 --- a/notebooks/imaging/features/advanced/subhalo/simulator.ipynb +++ b/notebooks/imaging/features/advanced/subhalo/simulator.ipynb @@ -1,482 +1,519 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Dark Matter Subhalo\n", - "==============================\n", - "\n", - "If a low mass dark matter halo overlaps the lensed source emission, it perturbs it in a unique and observable way.\n", - "\n", - "This script simulates an imaging dataset which includes a dark matter subhalo, which is high enough mass to\n", - "detect for Hubble Space Telescope imaging.\n", - "\n", - "This is used in `advanced/subhalo` to illustrate how to fit a lens model which includes a dark matter subhalo.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "- **Subhalo Difference Image:** An informative way to visualize the effect of a subhalo on a strong lens is to subtract the.\n", - "- **No Lens Light:** The code below simulates the same lens, but without a lens light component.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's light profiles are an `Sersic` and `Exponential`.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The subhalo`s `MassProfile` is a `NFWSph`.\n", - " - The source galaxy's light is an `Sersic`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"dark_matter_subhalo\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "\n", - "We also over high levels of adaptive over sampling around the centre of the dark matter subhalo." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(150, 150),\n", - " pixel_scales=0.5,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0), (1.601, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "\n", - "the `lens_galaxy` below includes a dark matter `subhalo` mass component." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " subhalo=al.mp.NFWTruncatedMCRLudlowSph(centre=(1.601, 0.0), mass_at_200=1.0e10),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Subhalo Difference Image__\n", - "\n", - "An informative way to visualize the effect of a subhalo on a strong lens is to subtract the image-plane image of the \n", - "tracer with and without the subhalo included. \n", - "\n", - "This effectively creates a subhalo residual-map, showing the regions of the image-plane where the subhalo's\n", - "effects are located (e.g. near the location of the subhalo).\n", - "\n", - "If this image creates very small residuals (e.g. below the noise level), it means that the subhalo is not detectable \n", - "in the image. Inspecting this image will therefore save you a lot of time, as you will avoid searching for\n", - "subhalos that do not produce strong enough effects to be visible in the image!\n", - "\n", - "On the other hand, if the resduals are large, it does not necessarily confirm that the subhalo is detectable. This is\n", - "because the subhalo effect may be degenerate with the lens model, whereby the lens mass model or source parameters\n", - "can change their parameters to account for the subhalo's effect. Only full end-to-end lens modeling can robustly\n", - "confirm whether a subhalo is detectable." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy_no_subhalo = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "tracer_no_subhalo = al.Tracer(galaxies=[lens_galaxy_no_subhalo, source_galaxy])\n", - "\n", - "image = tracer.image_2d_from(grid=grid)\n", - "image_no_subhalo = tracer_no_subhalo.image_2d_from(grid=grid)\n", - "\n", - "subhalo_residual_image = image - image_no_subhalo\n", - "\n", - "aplt.plot_array(\n", - " array=subhalo_residual_image,\n", - " title=\"\",\n", - " output_path=dataset_path,\n", - " output_format=\"png\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__No Lens Light__\n", - "\n", - "The code below simulates the same lens, but without a lens light component." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"dark_matter_subhalo_no_lens_light\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " subhalo=al.mp.NFWTruncatedMCRLudlowSph(centre=(1.601, 0.0), mass_at_200=1.0e10),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")\n", - "\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/dark_matter_subhalo`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Dark Matter Subhalo\n", + "==============================\n", + "\n", + "If a low mass dark matter halo overlaps the lensed source emission, it perturbs it in a unique and observable way.\n", + "\n", + "This script simulates an imaging dataset which includes a dark matter subhalo, which is high enough mass to\n", + "detect for Hubble Space Telescope imaging.\n", + "\n", + "This is used in `advanced/subhalo` to illustrate how to fit a lens model which includes a dark matter subhalo.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "- **Subhalo Difference Image:** An informative way to visualize the effect of a subhalo on a strong lens is to subtract the.\n", + "- **No Lens Light:** The code below simulates the same lens, but without a lens light component.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's light profiles are an `Sersic` and `Exponential`.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The subhalo`s `MassProfile` is a `NFWSph`.\n", + " - The source galaxy's light is an `Sersic`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"dark_matter_subhalo\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "\n", + "We also over high levels of adaptive over sampling around the centre of the dark matter subhalo." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(150, 150),\n", + " pixel_scales=0.5,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0), (1.601, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "\n", + "the `lens_galaxy` below includes a dark matter `subhalo` mass component." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " subhalo=al.mp.NFWTruncatedMCRLudlowSph(centre=(1.601, 0.0), mass_at_200=1.0e10),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Subhalo Difference Image__\n", + "\n", + "An informative way to visualize the effect of a subhalo on a strong lens is to subtract the image-plane image of the \n", + "tracer with and without the subhalo included. \n", + "\n", + "This effectively creates a subhalo residual-map, showing the regions of the image-plane where the subhalo's\n", + "effects are located (e.g. near the location of the subhalo).\n", + "\n", + "If this image creates very small residuals (e.g. below the noise level), it means that the subhalo is not detectable \n", + "in the image. Inspecting this image will therefore save you a lot of time, as you will avoid searching for\n", + "subhalos that do not produce strong enough effects to be visible in the image!\n", + "\n", + "On the other hand, if the resduals are large, it does not necessarily confirm that the subhalo is detectable. This is\n", + "because the subhalo effect may be degenerate with the lens model, whereby the lens mass model or source parameters\n", + "can change their parameters to account for the subhalo's effect. Only full end-to-end lens modeling can robustly\n", + "confirm whether a subhalo is detectable." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy_no_subhalo = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "tracer_no_subhalo = al.Tracer(galaxies=[lens_galaxy_no_subhalo, source_galaxy])\n", + "\n", + "image = tracer.image_2d_from(grid=grid)\n", + "image_no_subhalo = tracer_no_subhalo.image_2d_from(grid=grid)\n", + "\n", + "subhalo_residual_image = image - image_no_subhalo\n", + "\n", + "aplt.plot_array(\n", + " array=subhalo_residual_image,\n", + " title=\"\",\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__No Lens Light__\n", + "\n", + "The code below simulates the same lens, but without a lens light component." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"dark_matter_subhalo_no_lens_light\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " subhalo=al.mp.NFWTruncatedMCRLudlowSph(centre=(1.601, 0.0), mass_at_200=1.0e10),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")\n", + "\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/dark_matter_subhalo`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/extra_galaxies/modeling.ipynb b/notebooks/imaging/features/extra_galaxies/modeling.ipynb index 082e9aaeb..d6aed3d3f 100644 --- a/notebooks/imaging/features/extra_galaxies/modeling.ipynb +++ b/notebooks/imaging/features/extra_galaxies/modeling.ipynb @@ -1,745 +1,782 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Extra Galaxies\n", - "=================================\n", - "\n", - "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source\n", - "and whose mass may contribute to the ray-tracing and lens model.\n", - "\n", - "If their emission is significant, and close enough to the lens and source, we may simply mask it from the data\n", - "to ensure it does not impact the model-fit. In this script, we first illustrate how to do this, and outline two\n", - "different approaches to masking the emission of these extra galaxies which is appropriate for different lens models.\n", - "\n", - "Next, we consider a different approach which extends the modeling API to include these extra galaxies in the model-fit.\n", - "This includes both light profiles which fit and subtract their emission and mass profiles which account for their\n", - "contribution to the lensing of the source galaxy. The centres of each galaxy (e.g. their brightest pixels in the data)\n", - "are used as the centre of the light and mass profiles of these galaxies, in order to reduce model complexity.\n", - "\n", - "The second approach is more complex and computationally expensive, but if the emission of the extra galaxies blends\n", - "significantly with the lensed source emission, or if their mass is anticipated to contributed signficiantly, it is the\n", - "best approach to take.\n", - "\n", - "__Contents__\n", - "\n", - "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Extra Galaxies Over Sampling:** Over sampling is a numerical technique where the images of light profiles and galaxies are.\n", - "- **Extra Galaxies Noise Scaling:** To prevent extra galaxies from impacting the model-fit, we do not mask them entirely from the fit.\n", - "- **Extra Galaxies Dataset:** We are now going to model the dataset with extra galaxies included in the model, where these.\n", - "- **Extra Galaxies Centres:** To set up a lens model including each extra galaxy with light and / or mass profile, we input.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Extra Galaxies Model:** We now use the modeling API to create the model for the extra galaxies.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Approaches to Extra Galaxies:** We illustrated two extremes of how to prevent the emission of extra galaxies impacting the.\n", - "- **Scaling Relations:** The modeling API has full support for composing the extra galaxies such that their light and or.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Data Preparation__\n", - "\n", - "To perform modeling which accounts for extra galaxies, a mask of their emission or list of the centre of each extra\n", - "galaxy are used to set up the model-fit. For the example dataset used here, these tasks have already been performed and\n", - "the metadata (`mask_extra_galaxies.fits` and `extra_galaxies_centres.json` are already included in results folder.\n", - "\n", - "The tutorial `autolens_workspace/*/imaging/data_preparation/optional/extra_galaxies_centres.py`\n", - "describes how to create these centres and output them to a `.json` file.\n", - "\n", - "To mask the emission of extra galaxies and omit them from the fit, a `mask_extra_galaxies.fits` file is required.\n", - "The `data_preparation` tutorial `autolens_workspace/*/imaging/data_preparation/optional/mask_extra_galaxies.py`\n", - "describes how to create this mask.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `extra_galaxies` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Visualization of this dataset shows two galaxies either side of the lensed source. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define a bigger circular mask of 6.0\" than the 3.0\" masks used in other tutorials, to ensure the extra galaxy's \n", - "emission is included." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 6.0\n", - "\n", - "mask_main = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask_main)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets plot the masked imaging to make sure the extra galaxies are included." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Over Sampling__\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated \n", - "on a higher resolution grid than the image data to ensure the calculation is accurate. \n", - "\n", - "For a new user, the details of over-sampling are not important, therefore just be aware that below we make it so that \n", - "all calculations use an adaptive over sampling scheme which ensures high accuracy and precision.\n", - "\n", - "Crucially, this over sampling is applied at the centre of both extra galaxies, ensuring the light of both are over \n", - "sampled correctly.\n", - "\n", - "Once you are more experienced, you should read up on over-sampling in more detail via \n", - "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0), (1.0, 3.5), (-2.0, -3.5)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Noise Scaling__\n", - "\n", - "To prevent extra galaxies from impacting the model-fit, we do not mask them entirely from the fit, which\n", - "would be analogous to making the circular mask smaller or using a more refined mask. When pixels are masked and\n", - "removed entirely from the fit, their coordinates are not used when performing ray-tracing and the light of the\n", - "lens and source galaxies in these pixels not evaluated.\n", - "\n", - "Instead, the pixels are kept in the fit, but their data values are scaled to zero and their noise-map values\n", - "are increased to very large values. This means that during the model-fit, these pixels contribute negligibly to\n", - "the likelihood of the fit, and therefore do not impact the lens model.\n", - "\n", - "This approach is used because for certain types of modeling approaches, like a pixelized source reconstruction, \n", - "masking regions of the image in a way that removes their image pixels entirely from the fit can produce \n", - "discontinuities in the pixelixation. This can lead to unexpected systematics and unsatisfactory results\n", - "\n", - "In this case, applying the mask in a way where the image pixels are not removed from the fit, but their data and \n", - "noise-map values are scaled such that they contribute negligibly to the fit, is a better approach. \n", - "\n", - "We illustrate the API for doing this below, and show the subplot imaging where the extra galaxies mask has scaled\n", - "the data values to zeros, increasing the noise-map values to large values and in turn made the signal to noise\n", - "of its pixels effectively zero.\n", - "\n", - "We reload the dataset to ensure that the trimming performed for the FFT, based on the previous mask,\n", - "does not impact the noise-scaling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_extra_galaxies = al.Mask2D.from_fits(\n", - " file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"),\n", - " pixel_scales=0.1,\n", - " invert=True, # Note that we invert the mask here as `True` means a pixel is scaled.\n", - ")\n", - "\n", - "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native, pixel_scales=0.1, centre=(0.0, 0.0), radius=6.0\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now perform a model-fit using the standard API, where the extra galaxies are not included in the model.\n", - "\n", - "The mask we have applied ensures the extra galaxies do not impact the fit, and the model-fit returns a good fit to the\n", - "lensed source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - ")\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"extra_galaxies_noise_scaling\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " iterations_per_quick_update=20000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the `features/pixelization` example we perform a fit using this noise scaling scheme and a pixelization,\n", - "so check this out if you are interested in how to do this.\n", - "\n", - "__Extra Galaxies Dataset__\n", - "\n", - "We are now going to model the dataset with extra galaxies included in the model, where these galaxies include\n", - "both the light and mass profiles of the extra galaxies.\n", - "\n", - "We therefore reload the dataset and apply the 6.0\" circular mask to it, but do not use the extra galaxies mask\n", - "as the emission of the extra galaxies is included in the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_main = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=6.0\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask_main)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Centres__\n", - "\n", - "To set up a lens model including each extra galaxy with light and / or mass profile, we input manually the centres of\n", - "the extra galaxies.\n", - "\n", - "In principle, a lens model including the extra galaxies could be composed without these centres. For example, if \n", - "there were two extra galaxies in the data, we could simply add two additional light and mass profiles into the model. \n", - "The modeling API does support this, but we will not use it in this example.\n", - "\n", - "This is because models where the extra galaxies have free centres are often too complex to fit. It is likely the fit \n", - "will infer an inaccurate lens model and local maxima, because the parameter space is too complex.\n", - "\n", - "For example, a common problem is that one of the extra galaxy light profiles intended to model a nearby galaxy instead \n", - "fit one of the lensed source's multiple images. Alternatively, an extra galaxy's mass profile may recenter itself and \n", - "act as part of the main lens galaxy's mass distribution.\n", - "\n", - "Therefore, when modeling extra galaxies we input the centre of each, in order to fix their light and mass profile \n", - "centres or set up priors centre around these values.\n", - "\n", - "The `data_preparation` tutorial `autolens_workspace/*/imaging/data_preparation/examples/optional/extra_galaxies_centres.py` \n", - "describes how to create these centres. Using this script they have been output to the `.json` file we load below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_centres = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"extra_galaxies_centres.json\"))\n", - ")\n", - "\n", - "print(extra_galaxies_centres)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Over Sampling__\n", - "\n", - "Over sampling was discussed above, below we show how to apply it using the loaded centres of the extra galaxies.\n", - "\n", - "There is still a galaxy at the centre of the image so we include this in the `centre_list` with a centre \n", - "of (0.0\", 0.0\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)] + extra_galaxies_centres.in_list,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__ \n", - "\n", - "Perform the normal steps to set up the main model of the lens galaxy and source.\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - ")\n", - "mass = af.Model(al.mp.Isothermal)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Model__ \n", - "\n", - "We now use the modeling API to create the model for the extra galaxies.\n", - "\n", - "Currently, the extra galaxies API require that the centres of the light and mass profiles are fixed to the input centres\n", - "(but the other parameters of the light and mass profiles remain free). \n", - "\n", - "Therefore, in this example fits a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", - "\n", - " - The source galaxy's light is a Multi Gaussian Expansion [4 parameters].\n", - "\n", - " - Each extra galaxy's light is a Multi Gaussian expansion with fixed centre [2 extra galaxies x 2 parameters = 4 parameters].\n", - "\n", - " - Each extra galaxy's total mass distribution is a `IsothermalSph` profile with fixed \n", - " centre [2 extra galaxies x 1 parameters = 2 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=20.\n", - "\n", - "Extra galaxy mass profiles often to go unphysically high `einstein_radius` values, degrading the fit. The \n", - "`einstein_radius` parameter is set a `UniformPrior` with an upper limit of 0.1\" to prevent this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for extra_galaxy_centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Light\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=extra_galaxy_centre\n", - " )\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = extra_galaxy_centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source=source), extra_galaxies=extra_galaxies\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute confirms the model includes extra galaxies that we defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis__ \n", - "\n", - "The code below performs the normal steps to set up a model-fit.\n", - "\n", - "Given the extra model parameters due to the extra gaxies, we increase the number of live points to 200." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"extra_galaxies_model\",\n", - " unique_tag=dataset_name,\n", - " n_live=200,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=20000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", - "required by a model.\n", - "\n", - "Adding extra galaxies increases VRAM usage because each additional component adds calculations that JAX stores \n", - "in GPU memory. If you see VRAM warnings or errors, or if model fitting is slower than expected, you should \n", - "print the estimated VRAM usage and compare \n", - "\n", - "__Run Time__\n", - "\n", - "Adding extra galaxies to the model increases the likelihood evaluation times, because their light profiles need \n", - "their images evaluated and their mass profiles need their deflection angles computed. These calculations are pretty \n", - "fast, so only a small increase in time is expected.\n", - "\n", - "The bigger hit on run time is due to the extra free parameters, 2 free parameters for the `ell_comps` of each \n", - "multi Gaussian expansion of each extra galaxy and 1 `einstein_radius` for its mass. This increases the dimensionality \n", - "of non-linear parameter space. This means Nautilus takes longer to converge on the highest likelihood regions of \n", - "parameter space.\n", - "\n", - "The Source, Light and Mass (SLaM) pipelines support extra galaxies but in a way that reduces the number of free\n", - "parameters they add to the model. This is described in the `slam` examples. The `group` package, which models systems\n", - "with 10+ extra galaxies, introduces even more clever parameterizations which add 0 free parameters per extra galaxy,\n", - "so if your model has many extra galaxies you should check out the `group` package.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We can now begin the model-fit by passing the model and analysis object to the search, which performs a non-linear\n", - "search to find which models fit the data with the highest likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "By plotting the maximum log likelihood `FitImaging` object we can confirm the extra galaxies contribute to the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "These examples show how the results API can be extended to investigate extra galaxies in the results.\n", - "\n", - "__Approaches to Extra Galaxies__\n", - "\n", - "We illustrated two extremes of how to prevent the emission of extra galaxies impacting the model-fit:\n", - "\n", - "- **Noise Scaling**: We scaled the emission of the extra galaxies, such that their light did not impact the fit,\n", - " and ignored their mass entirely.\n", - "\n", - "- **Modeling**: We included the extra galaxies in the model, such that their light and mass profiles were fitted.\n", - "\n", - "There are approach that fall between these two, for example the light profiles could be omitted from the model\n", - "by applying an extra galaxies mask, but their mass profiles can still be included via the modeling API. You could also\n", - "include just the light profiles and not the mass profiles, or visa versa. You could also make the redshifts of the\n", - "extra galaxies free parameters in the model, or provide different light and mass profiles for each galaxy.\n", - "\n", - "Extending the modeling API should be straight forward given the above examples, and if anything is unclear then\n", - "checkout the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", - "\n", - "__Scaling Relations__\n", - "\n", - "The modeling API has full support for composing the extra galaxies such that their light and or mass follow scaling\n", - "relations. For example, you could assume that the mass of the extra galaxies is related to their luminosity via a\n", - "constant mass-to-light ratio.\n", - "\n", - "This is documented in the `autolens_workspace/*/imaging/features/scaling_relation` example.\n", - "\n", - "__Wrap Up__\n", - "\n", - "The extra galaxies API makes it straight forward for us to model galaxy-scale strong lenses with additional components for\n", - "the light and mass of nearby objects.\n", - "\n", - "The `autolens_workspace` includes a `group` package, for modeling group scale strong lenses which have multiple lens \n", - "galaxies. When you should use the extra galaxies API as shown here, and when you should use the group package? \n", - "\n", - "The distinction is as follows:\n", - "\n", - " - A galaxy scale lens is a system which can be modeled to a high level of accuracy using a single light and mass \n", - " distribution for the main lens galaxy. Including additional galaxies in the model via the extra galaxies API makes small \n", - " improvements on the lens model, but a good fit is possible without them. \n", - "\n", - " - A group scale lens is a system which cannot be modeled to a high level of accuracy using a single light and mass \n", - " distribution. Defining a 'main' lens galaxy is unclear and two or more main lens galaxies are required to fit an \n", - " accurate model. \n", - "\n", - "The `group` package also uses the extra galaxies API for model composition, but does so to compose and fit more complex lens \n", - "models." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Extra Galaxies\n", + "=================================\n", + "\n", + "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source\n", + "and whose mass may contribute to the ray-tracing and lens model.\n", + "\n", + "If their emission is significant, and close enough to the lens and source, we may simply mask it from the data\n", + "to ensure it does not impact the model-fit. In this script, we first illustrate how to do this, and outline two\n", + "different approaches to masking the emission of these extra galaxies which is appropriate for different lens models.\n", + "\n", + "Next, we consider a different approach which extends the modeling API to include these extra galaxies in the model-fit.\n", + "This includes both light profiles which fit and subtract their emission and mass profiles which account for their\n", + "contribution to the lensing of the source galaxy. The centres of each galaxy (e.g. their brightest pixels in the data)\n", + "are used as the centre of the light and mass profiles of these galaxies, in order to reduce model complexity.\n", + "\n", + "The second approach is more complex and computationally expensive, but if the emission of the extra galaxies blends\n", + "significantly with the lensed source emission, or if their mass is anticipated to contributed signficiantly, it is the\n", + "best approach to take.\n", + "\n", + "__Contents__\n", + "\n", + "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Extra Galaxies Over Sampling:** Over sampling is a numerical technique where the images of light profiles and galaxies are.\n", + "- **Extra Galaxies Noise Scaling:** To prevent extra galaxies from impacting the model-fit, we do not mask them entirely from the fit.\n", + "- **Extra Galaxies Dataset:** We are now going to model the dataset with extra galaxies included in the model, where these.\n", + "- **Extra Galaxies Centres:** To set up a lens model including each extra galaxy with light and / or mass profile, we input.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Extra Galaxies Model:** We now use the modeling API to create the model for the extra galaxies.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Approaches to Extra Galaxies:** We illustrated two extremes of how to prevent the emission of extra galaxies impacting the.\n", + "- **Scaling Relations:** The modeling API has full support for composing the extra galaxies such that their light and or.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Data Preparation__\n", + "\n", + "To perform modeling which accounts for extra galaxies, a mask of their emission or list of the centre of each extra\n", + "galaxy are used to set up the model-fit. For the example dataset used here, these tasks have already been performed and\n", + "the metadata (`mask_extra_galaxies.fits` and `extra_galaxies_centres.json` are already included in results folder.\n", + "\n", + "The tutorial `autolens_workspace/*/imaging/data_preparation/optional/extra_galaxies_centres.py`\n", + "describes how to create these centres and output them to a `.json` file.\n", + "\n", + "To mask the emission of extra galaxies and omit them from the fit, a `mask_extra_galaxies.fits` file is required.\n", + "The `data_preparation` tutorial `autolens_workspace/*/imaging/data_preparation/optional/mask_extra_galaxies.py`\n", + "describes how to create this mask.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `extra_galaxies` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Visualization of this dataset shows two galaxies either side of the lensed source. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define a bigger circular mask of 6.0\" than the 3.0\" masks used in other tutorials, to ensure the extra galaxy's \n", + "emission is included." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 6.0\n", + "\n", + "mask_main = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask_main)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets plot the masked imaging to make sure the extra galaxies are included." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Over Sampling__\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated \n", + "on a higher resolution grid than the image data to ensure the calculation is accurate. \n", + "\n", + "For a new user, the details of over-sampling are not important, therefore just be aware that below we make it so that \n", + "all calculations use an adaptive over sampling scheme which ensures high accuracy and precision.\n", + "\n", + "Crucially, this over sampling is applied at the centre of both extra galaxies, ensuring the light of both are over \n", + "sampled correctly.\n", + "\n", + "Once you are more experienced, you should read up on over-sampling in more detail via \n", + "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0), (1.0, 3.5), (-2.0, -3.5)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Noise Scaling__\n", + "\n", + "To prevent extra galaxies from impacting the model-fit, we do not mask them entirely from the fit, which\n", + "would be analogous to making the circular mask smaller or using a more refined mask. When pixels are masked and\n", + "removed entirely from the fit, their coordinates are not used when performing ray-tracing and the light of the\n", + "lens and source galaxies in these pixels not evaluated.\n", + "\n", + "Instead, the pixels are kept in the fit, but their data values are scaled to zero and their noise-map values\n", + "are increased to very large values. This means that during the model-fit, these pixels contribute negligibly to\n", + "the likelihood of the fit, and therefore do not impact the lens model.\n", + "\n", + "This approach is used because for certain types of modeling approaches, like a pixelized source reconstruction, \n", + "masking regions of the image in a way that removes their image pixels entirely from the fit can produce \n", + "discontinuities in the pixelixation. This can lead to unexpected systematics and unsatisfactory results\n", + "\n", + "In this case, applying the mask in a way where the image pixels are not removed from the fit, but their data and \n", + "noise-map values are scaled such that they contribute negligibly to the fit, is a better approach. \n", + "\n", + "We illustrate the API for doing this below, and show the subplot imaging where the extra galaxies mask has scaled\n", + "the data values to zeros, increasing the noise-map values to large values and in turn made the signal to noise\n", + "of its pixels effectively zero.\n", + "\n", + "We reload the dataset to ensure that the trimming performed for the FFT, based on the previous mask,\n", + "does not impact the noise-scaling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_extra_galaxies = al.Mask2D.from_fits(\n", + " file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"),\n", + " pixel_scales=0.1,\n", + " invert=True, # Note that we invert the mask here as `True` means a pixel is scaled.\n", + ")\n", + "\n", + "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native, pixel_scales=0.1, centre=(0.0, 0.0), radius=6.0\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now perform a model-fit using the standard API, where the extra galaxies are not included in the model.\n", + "\n", + "The mask we have applied ensures the extra galaxies do not impact the fit, and the model-fit returns a good fit to the\n", + "lensed source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + ")\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"extra_galaxies_noise_scaling\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " iterations_per_quick_update=20000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the `features/pixelization` example we perform a fit using this noise scaling scheme and a pixelization,\n", + "so check this out if you are interested in how to do this.\n", + "\n", + "__Extra Galaxies Dataset__\n", + "\n", + "We are now going to model the dataset with extra galaxies included in the model, where these galaxies include\n", + "both the light and mass profiles of the extra galaxies.\n", + "\n", + "We therefore reload the dataset and apply the 6.0\" circular mask to it, but do not use the extra galaxies mask\n", + "as the emission of the extra galaxies is included in the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_main = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=6.0\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask_main)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Centres__\n", + "\n", + "To set up a lens model including each extra galaxy with light and / or mass profile, we input manually the centres of\n", + "the extra galaxies.\n", + "\n", + "In principle, a lens model including the extra galaxies could be composed without these centres. For example, if \n", + "there were two extra galaxies in the data, we could simply add two additional light and mass profiles into the model. \n", + "The modeling API does support this, but we will not use it in this example.\n", + "\n", + "This is because models where the extra galaxies have free centres are often too complex to fit. It is likely the fit \n", + "will infer an inaccurate lens model and local maxima, because the parameter space is too complex.\n", + "\n", + "For example, a common problem is that one of the extra galaxy light profiles intended to model a nearby galaxy instead \n", + "fit one of the lensed source's multiple images. Alternatively, an extra galaxy's mass profile may recenter itself and \n", + "act as part of the main lens galaxy's mass distribution.\n", + "\n", + "Therefore, when modeling extra galaxies we input the centre of each, in order to fix their light and mass profile \n", + "centres or set up priors centre around these values.\n", + "\n", + "The `data_preparation` tutorial `autolens_workspace/*/imaging/data_preparation/examples/optional/extra_galaxies_centres.py` \n", + "describes how to create these centres. Using this script they have been output to the `.json` file we load below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_centres = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"extra_galaxies_centres.json\"))\n", + ")\n", + "\n", + "print(extra_galaxies_centres)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Over Sampling__\n", + "\n", + "Over sampling was discussed above, below we show how to apply it using the loaded centres of the extra galaxies.\n", + "\n", + "There is still a galaxy at the centre of the image so we include this in the `centre_list` with a centre \n", + "of (0.0\", 0.0\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)] + extra_galaxies_centres.in_list,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__ \n", + "\n", + "Perform the normal steps to set up the main model of the lens galaxy and source.\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + ")\n", + "mass = af.Model(al.mp.Isothermal)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Model__ \n", + "\n", + "We now use the modeling API to create the model for the extra galaxies.\n", + "\n", + "Currently, the extra galaxies API require that the centres of the light and mass profiles are fixed to the input centres\n", + "(but the other parameters of the light and mass profiles remain free). \n", + "\n", + "Therefore, in this example fits a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + "\n", + " - The source galaxy's light is a Multi Gaussian Expansion [4 parameters].\n", + "\n", + " - Each extra galaxy's light is a Multi Gaussian expansion with fixed centre [2 extra galaxies x 2 parameters = 4 parameters].\n", + "\n", + " - Each extra galaxy's total mass distribution is a `IsothermalSph` profile with fixed \n", + " centre [2 extra galaxies x 1 parameters = 2 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=20.\n", + "\n", + "Extra galaxy mass profiles often to go unphysically high `einstein_radius` values, degrading the fit. The \n", + "`einstein_radius` parameter is set a `UniformPrior` with an upper limit of 0.1\" to prevent this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for extra_galaxy_centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Light\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=extra_galaxy_centre\n", + " )\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = extra_galaxy_centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source=source), extra_galaxies=extra_galaxies\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute confirms the model includes extra galaxies that we defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis__ \n", + "\n", + "The code below performs the normal steps to set up a model-fit.\n", + "\n", + "Given the extra model parameters due to the extra gaxies, we increase the number of live points to 200." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"extra_galaxies_model\",\n", + " unique_tag=dataset_name,\n", + " n_live=200,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=20000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", + "required by a model.\n", + "\n", + "Adding extra galaxies increases VRAM usage because each additional component adds calculations that JAX stores \n", + "in GPU memory. If you see VRAM warnings or errors, or if model fitting is slower than expected, you should \n", + "print the estimated VRAM usage and compare \n", + "\n", + "__Run Time__\n", + "\n", + "Adding extra galaxies to the model increases the likelihood evaluation times, because their light profiles need \n", + "their images evaluated and their mass profiles need their deflection angles computed. These calculations are pretty \n", + "fast, so only a small increase in time is expected.\n", + "\n", + "The bigger hit on run time is due to the extra free parameters, 2 free parameters for the `ell_comps` of each \n", + "multi Gaussian expansion of each extra galaxy and 1 `einstein_radius` for its mass. This increases the dimensionality \n", + "of non-linear parameter space. This means Nautilus takes longer to converge on the highest likelihood regions of \n", + "parameter space.\n", + "\n", + "The Source, Light and Mass (SLaM) pipelines support extra galaxies but in a way that reduces the number of free\n", + "parameters they add to the model. This is described in the `slam` examples. The `group` package, which models systems\n", + "with 10+ extra galaxies, introduces even more clever parameterizations which add 0 free parameters per extra galaxy,\n", + "so if your model has many extra galaxies you should check out the `group` package.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We can now begin the model-fit by passing the model and analysis object to the search, which performs a non-linear\n", + "search to find which models fit the data with the highest likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "By plotting the maximum log likelihood `FitImaging` object we can confirm the extra galaxies contribute to the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "These examples show how the results API can be extended to investigate extra galaxies in the results.\n", + "\n", + "__Approaches to Extra Galaxies__\n", + "\n", + "We illustrated two extremes of how to prevent the emission of extra galaxies impacting the model-fit:\n", + "\n", + "- **Noise Scaling**: We scaled the emission of the extra galaxies, such that their light did not impact the fit,\n", + " and ignored their mass entirely.\n", + "\n", + "- **Modeling**: We included the extra galaxies in the model, such that their light and mass profiles were fitted.\n", + "\n", + "There are approach that fall between these two, for example the light profiles could be omitted from the model\n", + "by applying an extra galaxies mask, but their mass profiles can still be included via the modeling API. You could also\n", + "include just the light profiles and not the mass profiles, or visa versa. You could also make the redshifts of the\n", + "extra galaxies free parameters in the model, or provide different light and mass profiles for each galaxy.\n", + "\n", + "Extending the modeling API should be straight forward given the above examples, and if anything is unclear then\n", + "checkout the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", + "\n", + "__Scaling Relations__\n", + "\n", + "The modeling API has full support for composing the extra galaxies such that their light and or mass follow scaling\n", + "relations. For example, you could assume that the mass of the extra galaxies is related to their luminosity via a\n", + "constant mass-to-light ratio.\n", + "\n", + "This is documented in the `autolens_workspace/*/imaging/features/scaling_relation` example.\n", + "\n", + "__Wrap Up__\n", + "\n", + "The extra galaxies API makes it straight forward for us to model galaxy-scale strong lenses with additional components for\n", + "the light and mass of nearby objects.\n", + "\n", + "The `autolens_workspace` includes a `group` package, for modeling group scale strong lenses which have multiple lens \n", + "galaxies. When you should use the extra galaxies API as shown here, and when you should use the group package? \n", + "\n", + "The distinction is as follows:\n", + "\n", + " - A galaxy scale lens is a system which can be modeled to a high level of accuracy using a single light and mass \n", + " distribution for the main lens galaxy. Including additional galaxies in the model via the extra galaxies API makes small \n", + " improvements on the lens model, but a good fit is possible without them. \n", + "\n", + " - A group scale lens is a system which cannot be modeled to a high level of accuracy using a single light and mass \n", + " distribution. Defining a 'main' lens galaxy is unclear and two or more main lens galaxies are required to fit an \n", + " accurate model. \n", + "\n", + "The `group` package also uses the extra galaxies API for model composition, but does so to compose and fit more complex lens \n", + "models." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/extra_galaxies/simulator.ipynb b/notebooks/imaging/features/extra_galaxies/simulator.ipynb index 59d39a9fc..3e7c201e6 100644 --- a/notebooks/imaging/features/extra_galaxies/simulator.ipynb +++ b/notebooks/imaging/features/extra_galaxies/simulator.ipynb @@ -1,543 +1,580 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Extra Galaxies\n", - "=========================\n", - "\n", - "Certain lenses have small galaxies within their Einstein radius, or nearby the lensed source emission.\n", - "\n", - "The emission of these galaxies may overlap the lensed source emission, and their mass may contribute to the lensing\n", - "of the source.\n", - "\n", - "We therefore will need to mask the emission of these extra galaxies or include them in the model as light profiles which\n", - "fit and subtract the emission. We may also include these galaxies as mass profiles in the lens model, accounting for\n", - "their lensing effects via ray-tracing.\n", - "\n", - "This uses the modeling API, which is illustrated in\n", - "the script `autolens_workspace/*/features/extra_galaxies/modeling`.\n", - "\n", - "This script simulates an imaging dataset which includes extra galaxies near the lens and source\n", - "galaxies. This is used to illustrate the extra galaxies API in the script above.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Other Scripts:** This dataset is used in the following scripts.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Galaxies:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Extra Galaxies:** Includes two extra galaxies, which must be modeled or masked to ensure they do not impact the fit.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Mask Extra Galaxies:** Building and saving `mask_extra_galaxies.fits` so consumers can load it directly.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "- **Multiple Images:** Output the multiple image positions of the source galaxy which can help with lens modeling.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The source galaxy's light is an `Sersic`.\n", - " - There are two extra galaxies whose light is near the strong lens and their mass perturbs the lensed source's emission.\n", - "\n", - "__Other Scripts__\n", - "\n", - "This dataset is used in the following scripts:\n", - "\n", - " `autolens_workspace/*/imaging/data_preparation/examples/optional/scaled_dataset.ipynb`\n", - "\n", - "To illustrate how to subtract and remove the light of extra galaxies in real strong lensing data, so that it does\n", - "not impact the lens model.\n", - "\n", - " `autolens_workspace/*/imaging/data_preparation/examples/optional/extra_galaxies_centres.ipynb`\n", - "\n", - "To illustrate how mark extra galaxy centres on a dataset so they can be used in the lens model.\n", - "\n", - " `autolens_workspace/*/modeling/features/extra_galaxies.ipynb`\n", - "\n", - "To illustrate how compose and fit a lens model which includes the extra galaxies as light and mass profiles.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from pathlib import Path\n", - "\n", - "import numpy as np\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "\n", - "This simulated galaxy has additional galaxies and light profiles which are offset from the main galaxy centre \n", - "of (0.0\", 0.0\"). The adaptive over sampling grid has all centres input to account for this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(150, 150),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0), (1.0, 3.5), (-2.0, -3.5)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=1.0,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.1, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=1.0,\n", - " sersic_index=2.5,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "Includes two extra galaxies, which must be modeled or masked to ensure they do not impact the fit.\n", - "\n", - "Note that their redshift is the same as the main galaxy, which is not necessarily the case in real observations. \n", - "\n", - "If they are at a different redshift, the tools for masking or modeling the luminous emission of the extra galaxies \n", - "are equipped to handle this.\n", - "\n", - "For mass modeling, their redshifts being different to the main galaxy will lead to multi-plane ray-tracing being\n", - "performed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0_centre = (1.0, 3.5)\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " light=al.lp.ExponentialSph(\n", - " centre=extra_galaxy_0_centre, intensity=2.0, effective_radius=0.5\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=extra_galaxy_0_centre, einstein_radius=0.1),\n", - ")\n", - "\n", - "extra_galaxy_1_centre = (-2.0, -3.5)\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " light=al.lp.ExponentialSph(\n", - " centre=extra_galaxy_1_centre, intensity=2.0, effective_radius=0.8\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=extra_galaxy_1_centre, einstein_radius=0.2),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask Extra Galaxies__\n", - "\n", - "Build and output a `mask_extra_galaxies.fits` covering the two extra galaxy regions, so that downstream tutorials\n", - "which use this dataset (e.g. `imaging/features/extra_galaxies/modeling.py`,\n", - "`imaging/features/pixelization/modeling.py`) can load the mask directly without a separate data-preparation step.\n", - "\n", - "Each circle is sized to ~3x the galaxy's `effective_radius`, which comfortably covers the light extent for the\n", - "`ExponentialSph` profiles used above. The geometry is derived from the same centres + radii defined for the\n", - "extra galaxies in this script, so it stays in sync with any future tweak to those values.\n", - "\n", - "`Mask2D.circular` honours the `PYAUTO_SMALL_DATASETS=1` env var (caps to 15x15 at 0.6\"/px), so the mask\n", - "automatically shrinks alongside the small-dataset image and never raises an out-of-bounds error." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_mask = np.zeros(dataset.shape_native, dtype=bool)\n", - "\n", - "for centre, radius in [\n", - " (extra_galaxy_0_centre, 3.0 * 0.5),\n", - " (extra_galaxy_1_centre, 3.0 * 0.8),\n", - "]:\n", - " circle = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " centre=centre,\n", - " radius=radius,\n", - " invert=True, # True inside the circle (i.e. masked region)\n", - " )\n", - " extra_galaxies_mask = np.logical_or(extra_galaxies_mask, circle.native)\n", - "\n", - "mask_extra_galaxies = al.Mask2D(\n", - " mask=extra_galaxies_mask,\n", - " pixel_scales=dataset.pixel_scales,\n", - ")\n", - "\n", - "aplt.fits_array(\n", - " array=mask_extra_galaxies,\n", - " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multiple Images__\n", - "\n", - "Output the multiple image positions of the source galaxy which can help with lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " file_path=dataset_path / \"positions.json\",\n", - " obj=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Centres__\n", - "\n", - "Output the centres of the extra galaxies to a .json file, so that they can be used to set up the model\n", - "in the modeling scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(values=[extra_galaxy_0_centre, extra_galaxy_1_centre]),\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/extra_galaxies`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Extra Galaxies\n", + "=========================\n", + "\n", + "Certain lenses have small galaxies within their Einstein radius, or nearby the lensed source emission.\n", + "\n", + "The emission of these galaxies may overlap the lensed source emission, and their mass may contribute to the lensing\n", + "of the source.\n", + "\n", + "We therefore will need to mask the emission of these extra galaxies or include them in the model as light profiles which\n", + "fit and subtract the emission. We may also include these galaxies as mass profiles in the lens model, accounting for\n", + "their lensing effects via ray-tracing.\n", + "\n", + "This uses the modeling API, which is illustrated in\n", + "the script `autolens_workspace/*/features/extra_galaxies/modeling`.\n", + "\n", + "This script simulates an imaging dataset which includes extra galaxies near the lens and source\n", + "galaxies. This is used to illustrate the extra galaxies API in the script above.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Other Scripts:** This dataset is used in the following scripts.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Galaxies:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Extra Galaxies:** Includes two extra galaxies, which must be modeled or masked to ensure they do not impact the fit.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Mask Extra Galaxies:** Building and saving `mask_extra_galaxies.fits` so consumers can load it directly.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "- **Multiple Images:** Output the multiple image positions of the source galaxy which can help with lens modeling.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The source galaxy's light is an `Sersic`.\n", + " - There are two extra galaxies whose light is near the strong lens and their mass perturbs the lensed source's emission.\n", + "\n", + "__Other Scripts__\n", + "\n", + "This dataset is used in the following scripts:\n", + "\n", + " `autolens_workspace/*/imaging/data_preparation/examples/optional/scaled_dataset.ipynb`\n", + "\n", + "To illustrate how to subtract and remove the light of extra galaxies in real strong lensing data, so that it does\n", + "not impact the lens model.\n", + "\n", + " `autolens_workspace/*/imaging/data_preparation/examples/optional/extra_galaxies_centres.ipynb`\n", + "\n", + "To illustrate how mark extra galaxy centres on a dataset so they can be used in the lens model.\n", + "\n", + " `autolens_workspace/*/modeling/features/extra_galaxies.ipynb`\n", + "\n", + "To illustrate how compose and fit a lens model which includes the extra galaxies as light and mass profiles.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "\n", + "This simulated galaxy has additional galaxies and light profiles which are offset from the main galaxy centre \n", + "of (0.0\", 0.0\"). The adaptive over sampling grid has all centres input to account for this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(150, 150),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0), (1.0, 3.5), (-2.0, -3.5)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=1.0,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.1, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=1.0,\n", + " sersic_index=2.5,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "Includes two extra galaxies, which must be modeled or masked to ensure they do not impact the fit.\n", + "\n", + "Note that their redshift is the same as the main galaxy, which is not necessarily the case in real observations. \n", + "\n", + "If they are at a different redshift, the tools for masking or modeling the luminous emission of the extra galaxies \n", + "are equipped to handle this.\n", + "\n", + "For mass modeling, their redshifts being different to the main galaxy will lead to multi-plane ray-tracing being\n", + "performed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0_centre = (1.0, 3.5)\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " light=al.lp.ExponentialSph(\n", + " centre=extra_galaxy_0_centre, intensity=2.0, effective_radius=0.5\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=extra_galaxy_0_centre, einstein_radius=0.1),\n", + ")\n", + "\n", + "extra_galaxy_1_centre = (-2.0, -3.5)\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " light=al.lp.ExponentialSph(\n", + " centre=extra_galaxy_1_centre, intensity=2.0, effective_radius=0.8\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=extra_galaxy_1_centre, einstein_radius=0.2),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask Extra Galaxies__\n", + "\n", + "Build and output a `mask_extra_galaxies.fits` covering the two extra galaxy regions, so that downstream tutorials\n", + "which use this dataset (e.g. `imaging/features/extra_galaxies/modeling.py`,\n", + "`imaging/features/pixelization/modeling.py`) can load the mask directly without a separate data-preparation step.\n", + "\n", + "Each circle is sized to ~3x the galaxy's `effective_radius`, which comfortably covers the light extent for the\n", + "`ExponentialSph` profiles used above. The geometry is derived from the same centres + radii defined for the\n", + "extra galaxies in this script, so it stays in sync with any future tweak to those values.\n", + "\n", + "`Mask2D.circular` honours the `PYAUTO_SMALL_DATASETS=1` env var (caps to 15x15 at 0.6\"/px), so the mask\n", + "automatically shrinks alongside the small-dataset image and never raises an out-of-bounds error." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_mask = np.zeros(dataset.shape_native, dtype=bool)\n", + "\n", + "for centre, radius in [\n", + " (extra_galaxy_0_centre, 3.0 * 0.5),\n", + " (extra_galaxy_1_centre, 3.0 * 0.8),\n", + "]:\n", + " circle = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " centre=centre,\n", + " radius=radius,\n", + " invert=True, # True inside the circle (i.e. masked region)\n", + " )\n", + " extra_galaxies_mask = np.logical_or(extra_galaxies_mask, circle.native)\n", + "\n", + "mask_extra_galaxies = al.Mask2D(\n", + " mask=extra_galaxies_mask,\n", + " pixel_scales=dataset.pixel_scales,\n", + ")\n", + "\n", + "aplt.fits_array(\n", + " array=mask_extra_galaxies,\n", + " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multiple Images__\n", + "\n", + "Output the multiple image positions of the source galaxy which can help with lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " file_path=dataset_path / \"positions.json\",\n", + " obj=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Centres__\n", + "\n", + "Output the centres of the extra galaxies to a .json file, so that they can be used to set up the model\n", + "in the modeling scripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(values=[extra_galaxy_0_centre, extra_galaxy_1_centre]),\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/extra_galaxies`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/extra_galaxies/slam.ipynb b/notebooks/imaging/features/extra_galaxies/slam.ipynb index 5f522f361..5a6b18344 100644 --- a/notebooks/imaging/features/extra_galaxies/slam.ipynb +++ b/notebooks/imaging/features/extra_galaxies/slam.ipynb @@ -1,748 +1,785 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Extra Galaxies: SLaM\n", - "=====================\n", - "\n", - "This script uses the SLaM pipelines to fit a lens dataset that includes extra galaxies\n", - "surrounding the main lens, whose light and mass are both included in the model.\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. This script\n", - "only documents how the pipeline differs from that reference.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **Group SLaM:** This pipeline is designed for the galaxy-scale regime with a small number of nearby extra galaxies.\n", - "- **This Script:** Using SOURCE LP, SOURCE PIX, LIGHT LP and MASS TOTAL pipelines this script fits `Imaging` data.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except each extra galaxy is included in the model with a.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`, except extra-galaxy light is fixed from `source_lp` and their.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`, except extra galaxies are fully fixed as instances from.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except extra-galaxy mass is fixed from `source_pix[1]` and their.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except extra-galaxy light is fixed from `light[1]` and their.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Extra Galaxies Centres:** The centres of the extra galaxies are loaded from a `.json` file and used to set up the light and.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - "- **Extra Galaxies** (`features/extra_galaxies.ipynb`)\n", - "\n", - "__Group SLaM__\n", - "\n", - "This pipeline is designed for the galaxy-scale regime with a small number of nearby extra\n", - "galaxies. For systems with many companions or group-scale complexity, use the group SLaM\n", - "pipeline (`scripts/group/slam.py`), which models companion masses via a shared luminosity\n", - "scaling relation rather than individually.\n", - "\n", - "__This Script__\n", - "\n", - "Using SOURCE LP, SOURCE PIX, LIGHT LP and MASS TOTAL pipelines this script fits `Imaging`\n", - "data where in the final model:\n", - "\n", - " - The lens galaxy's light is an MGE bulge.\n", - " - The lens galaxy's total mass is a `PowerLaw` plus `ExternalShear`.\n", - " - The source galaxy's light is a `Pixelization`.\n", - " - Each extra galaxy has a spherical MGE light profile and an `IsothermalSph` mass.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except each extra galaxy is included in the model with a\n", - "spherical MGE light profile (`GaussianSph` basis) and a free `IsothermalSph` mass, both\n", - "centred on its known position from `extra_galaxies_centres`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search,\n", - " dataset,\n", - " mask_radius,\n", - " extra_galaxies_centres,\n", - " redshift_lens,\n", - " redshift_source,\n", - " n_batch=50,\n", - "):\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=30,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " extra_galaxies_list = []\n", - "\n", - " for centre in extra_galaxies_centres:\n", - " total_gaussians = 10\n", - "\n", - " log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - " gaussian_list = af.Collection(\n", - " af.Model(al.lp_linear.GaussianSph) for _ in range(total_gaussians)\n", - " )\n", - "\n", - " for i, gaussian in enumerate(gaussian_list):\n", - " gaussian.centre.centre_0 = centre[0]\n", - " gaussian.centre.centre_1 = centre[1]\n", - " gaussian.sigma = 10 ** log10_sigma_list[i]\n", - "\n", - " extra_galaxy_bulge = af.Model(\n", - " al.lp_basis.Basis, profile_list=list(gaussian_list)\n", - " )\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.1)\n", - "\n", - " extra_galaxies_list.append(\n", - " af.Model(\n", - " al.Galaxy, redshift=redshift_lens, bulge=extra_galaxy_bulge, mass=mass\n", - " )\n", - " )\n", - "\n", - " extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`, except extra-galaxy light is fixed from `source_lp` and\n", - "their `IsothermalSph` mass priors are carried forward as free parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search,\n", - " dataset,\n", - " source_lp_result,\n", - " mesh_shape,\n", - " n_batch=20,\n", - "):\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " extra_galaxies = source_lp_result.model.extra_galaxies\n", - " for galaxy, result_galaxy in zip(\n", - " extra_galaxies, source_lp_result.instance.extra_galaxies\n", - " ):\n", - " galaxy.bulge = result_galaxy.bulge\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_lp_result.model.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, except extra galaxies are fully fixed as instances\n", - "from `source_pix[1]`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search,\n", - " dataset,\n", - " source_lp_result,\n", - " source_pix_result_1,\n", - " mesh_shape,\n", - " n_batch=20,\n", - "):\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except extra-galaxy mass is fixed from `source_pix[1]`\n", - "and their light is a new free spherical MGE centred on the same position." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search,\n", - " dataset,\n", - " mask_radius,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " n_batch=20,\n", - "):\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " )\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " extra_galaxies = source_result_for_lens.model.extra_galaxies\n", - " for galaxy, result_galaxy in zip(\n", - " extra_galaxies, source_result_for_lens.instance.extra_galaxies\n", - " ):\n", - " galaxy.mass = result_galaxy.mass\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except extra-galaxy light is fixed from `light[1]` and\n", - "their `IsothermalSph` mass priors are carried forward as free parameters from `source_pix[1]`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search,\n", - " dataset,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " light_result,\n", - " n_batch=20,\n", - "):\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " extra_galaxies = source_result_for_lens.model.extra_galaxies\n", - " for galaxy, result_galaxy in zip(\n", - " extra_galaxies, light_result.instance.extra_galaxies\n", - " ):\n", - " galaxy.bulge = result_galaxy.bulge\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=light_result.instance.galaxies.lens.bulge,\n", - " disk=light_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Centres__\n", - "\n", - "The centres of the extra galaxies are loaded from a `.json` file and used to set up the\n", - "light and mass models for each companion galaxy inside the pipeline functions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_centres = al.Grid2DIrregular(\n", - " al.from_json(file_path=dataset_path / \"extra_galaxies_centres.json\")\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)] + list(extra_galaxies_centres),\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " extra_galaxies_centres=extra_galaxies_centres,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_shape=mesh_shape,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh_shape=mesh_shape,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Extra Galaxies: SLaM\n", + "=====================\n", + "\n", + "This script uses the SLaM pipelines to fit a lens dataset that includes extra galaxies\n", + "surrounding the main lens, whose light and mass are both included in the model.\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. This script\n", + "only documents how the pipeline differs from that reference.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **Group SLaM:** This pipeline is designed for the galaxy-scale regime with a small number of nearby extra galaxies.\n", + "- **This Script:** Using SOURCE LP, SOURCE PIX, LIGHT LP and MASS TOTAL pipelines this script fits `Imaging` data.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except each extra galaxy is included in the model with a.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`, except extra-galaxy light is fixed from `source_lp` and their.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`, except extra galaxies are fully fixed as instances from.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except extra-galaxy mass is fixed from `source_pix[1]` and their.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except extra-galaxy light is fixed from `light[1]` and their.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Extra Galaxies Centres:** The centres of the extra galaxies are loaded from a `.json` file and used to set up the light and.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + "- **Extra Galaxies** (`features/extra_galaxies.ipynb`)\n", + "\n", + "__Group SLaM__\n", + "\n", + "This pipeline is designed for the galaxy-scale regime with a small number of nearby extra\n", + "galaxies. For systems with many companions or group-scale complexity, use the group SLaM\n", + "pipeline (`scripts/group/slam.py`), which models companion masses via a shared luminosity\n", + "scaling relation rather than individually.\n", + "\n", + "__This Script__\n", + "\n", + "Using SOURCE LP, SOURCE PIX, LIGHT LP and MASS TOTAL pipelines this script fits `Imaging`\n", + "data where in the final model:\n", + "\n", + " - The lens galaxy's light is an MGE bulge.\n", + " - The lens galaxy's total mass is a `PowerLaw` plus `ExternalShear`.\n", + " - The source galaxy's light is a `Pixelization`.\n", + " - Each extra galaxy has a spherical MGE light profile and an `IsothermalSph` mass.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except each extra galaxy is included in the model with a\n", + "spherical MGE light profile (`GaussianSph` basis) and a free `IsothermalSph` mass, both\n", + "centred on its known position from `extra_galaxies_centres`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search,\n", + " dataset,\n", + " mask_radius,\n", + " extra_galaxies_centres,\n", + " redshift_lens,\n", + " redshift_source,\n", + " n_batch=50,\n", + "):\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=30,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " extra_galaxies_list = []\n", + "\n", + " for centre in extra_galaxies_centres:\n", + " total_gaussians = 10\n", + "\n", + " log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + " gaussian_list = af.Collection(\n", + " af.Model(al.lp_linear.GaussianSph) for _ in range(total_gaussians)\n", + " )\n", + "\n", + " for i, gaussian in enumerate(gaussian_list):\n", + " gaussian.centre.centre_0 = centre[0]\n", + " gaussian.centre.centre_1 = centre[1]\n", + " gaussian.sigma = 10 ** log10_sigma_list[i]\n", + "\n", + " extra_galaxy_bulge = af.Model(\n", + " al.lp_basis.Basis, profile_list=list(gaussian_list)\n", + " )\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.1)\n", + "\n", + " extra_galaxies_list.append(\n", + " af.Model(\n", + " al.Galaxy, redshift=redshift_lens, bulge=extra_galaxy_bulge, mass=mass\n", + " )\n", + " )\n", + "\n", + " extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`, except extra-galaxy light is fixed from `source_lp` and\n", + "their `IsothermalSph` mass priors are carried forward as free parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search,\n", + " dataset,\n", + " source_lp_result,\n", + " mesh_shape,\n", + " n_batch=20,\n", + "):\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " extra_galaxies = source_lp_result.model.extra_galaxies\n", + " for galaxy, result_galaxy in zip(\n", + " extra_galaxies, source_lp_result.instance.extra_galaxies\n", + " ):\n", + " galaxy.bulge = result_galaxy.bulge\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_lp_result.model.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, except extra galaxies are fully fixed as instances\n", + "from `source_pix[1]`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search,\n", + " dataset,\n", + " source_lp_result,\n", + " source_pix_result_1,\n", + " mesh_shape,\n", + " n_batch=20,\n", + "):\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except extra-galaxy mass is fixed from `source_pix[1]`\n", + "and their light is a new free spherical MGE centred on the same position." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search,\n", + " dataset,\n", + " mask_radius,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " n_batch=20,\n", + "):\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " )\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " extra_galaxies = source_result_for_lens.model.extra_galaxies\n", + " for galaxy, result_galaxy in zip(\n", + " extra_galaxies, source_result_for_lens.instance.extra_galaxies\n", + " ):\n", + " galaxy.mass = result_galaxy.mass\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except extra-galaxy light is fixed from `light[1]` and\n", + "their `IsothermalSph` mass priors are carried forward as free parameters from `source_pix[1]`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search,\n", + " dataset,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " light_result,\n", + " n_batch=20,\n", + "):\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " extra_galaxies = source_result_for_lens.model.extra_galaxies\n", + " for galaxy, result_galaxy in zip(\n", + " extra_galaxies, light_result.instance.extra_galaxies\n", + " ):\n", + " galaxy.bulge = result_galaxy.bulge\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=light_result.instance.galaxies.lens.bulge,\n", + " disk=light_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Centres__\n", + "\n", + "The centres of the extra galaxies are loaded from a `.json` file and used to set up the\n", + "light and mass models for each companion galaxy inside the pipeline functions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_centres = al.Grid2DIrregular(\n", + " al.from_json(file_path=dataset_path / \"extra_galaxies_centres.json\")\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)] + list(extra_galaxies_centres),\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " extra_galaxies_centres=extra_galaxies_centres,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_shape=mesh_shape,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh_shape=mesh_shape,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/linear_light_profiles/fit.ipynb b/notebooks/imaging/features/linear_light_profiles/fit.ipynb index 53a9caebe..d22b71616 100644 --- a/notebooks/imaging/features/linear_light_profiles/fit.ipynb +++ b/notebooks/imaging/features/linear_light_profiles/fit.ipynb @@ -1,399 +1,436 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Linear Light Profiles Fit\n", - "============================================\n", - "\n", - "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", - "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", - "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood)\n", - "given the light profile's other parameters.\n", - "\n", - "Based on the advantages below, we recommended you always use linear light profiles to fit models over standard\n", - "light profiles!\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Each light profile's `intensity` parameter is therefore not a free parameter in the model-fit.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Notes:** This script is identical to `modeling/start_here.py` except that the light profiles are switched to.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Intensities:** The fit contains the solved for intensity values.\n", - "- **Visualization:** Linear light profiles and objects containing them (e.g.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "Each light profile's `intensity` parameter is therefore not a free parameter in the model-fit, reducing the\n", - "dimensionality of non-linear parameter space by the number of light profiles (in this example by 2 dimensions).\n", - "\n", - "This also removes the degeneracies that occur between the `intensity` and other light profile parameters\n", - "(e.g. `effective_radius`, `sersic_index`), which are difficult degeneracies for the non-linear search to map out\n", - "accurately. This produces more reliable lens model results and the fit converges in fewer iterations, speeding up the\n", - "overall analysis.\n", - "\n", - "The inversion has a relatively small computational cost, thus we reduce the model complexity without much slow-down and\n", - "can therefore fit models more reliably and faster!\n", - "\n", - "__Disadvantages__\n", - "\n", - "Although the computation time of the inversion is small, it is not non-negligible. It is approximately 3-4x slower\n", - "than using a standard light profile.\n", - "\n", - "The gains in run times due to the simpler non-linear parameter space therefore are somewhat balanced by the slower\n", - "likelihood calculation.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and negative\n", - "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical.\n", - "\n", - "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is as fast\n", - "as positive-negative solvers. This ensures that all light profile intensities are positive and therefore physical.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a linear `Sersic` bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a linear `Sersic`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook.\n", - "\n", - "__Notes__\n", - "\n", - "This script is identical to `modeling/start_here.py` except that the light profiles are switched to linear light\n", - "profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We now illustrate how to perform a fit to the dataset using linear light profiles, using known light profile parameters.\n", - "\n", - "The API follows closely the standard use of a `FitImaging` object, but simply uses linear light profiles (via the\n", - "`lp_linear` module) instead of standard light profiles. \n", - "\n", - "Note that the linear light profiles below do not have `intensity` parameters input and we use the true input values\n", - "of all other parameters for illustrative purposes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the linear light profiles have solved for `intensity` values that give a good fit\n", - "to the image. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The fit contains the solved for intensity values.\n", - "\n", - "These are computed using a fit's `linear_light_profile_intensity_dict`, which maps each linear light profile \n", - "in the model parameterization above to its `intensity`.\n", - "\n", - "The code below shows how to use this dictionary, as an alternative to using the max_log_likelihood quantities above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_bulge = tracer.galaxies[0].bulge\n", - "source_bulge = tracer.galaxies[1].bulge\n", - "\n", - "print(fit.linear_light_profile_intensity_dict)\n", - "\n", - "print(\n", - " f\"\\n Intensity of lens galaxy's bulge = {fit.linear_light_profile_intensity_dict[lens_bulge]}\"\n", - ")\n", - "\n", - "print(\n", - " f\"\\n Intensity of source bulge (lp_linear.SersicCore) = {fit.linear_light_profile_intensity_dict[source_bulge]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the solved \n", - "for `intensity` values is also accessible from a fit.\n", - "\n", - "For example, the linear light profile `Sersic` of the `bulge` component above has a solved for `intensity` of ~0.75. \n", - "\n", - "The `tracer` created below instead has an ordinary light profile with an `intensity` of ~0.75.\n", - "\n", - "The benefit of using a tracer with standard light profiles is it can be visualized (linear light profiles cannot \n", - "by default because they do not have `intensity` values)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "print(tracer.galaxies[0].bulge.intensity)\n", - "print(tracer.galaxies[1].bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization__\n", - "\n", - "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they do not \n", - "have an `intensity` value.\n", - "\n", - "Therefore, the objects created above which replaces all linear light profiles with ordinary light profiles must be\n", - "used for visualization:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Image\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Linear Light Profiles Fit\n", + "============================================\n", + "\n", + "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", + "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", + "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood)\n", + "given the light profile's other parameters.\n", + "\n", + "Based on the advantages below, we recommended you always use linear light profiles to fit models over standard\n", + "light profiles!\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Each light profile's `intensity` parameter is therefore not a free parameter in the model-fit.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Notes:** This script is identical to `modeling/start_here.py` except that the light profiles are switched to.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Intensities:** The fit contains the solved for intensity values.\n", + "- **Visualization:** Linear light profiles and objects containing them (e.g.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "Each light profile's `intensity` parameter is therefore not a free parameter in the model-fit, reducing the\n", + "dimensionality of non-linear parameter space by the number of light profiles (in this example by 2 dimensions).\n", + "\n", + "This also removes the degeneracies that occur between the `intensity` and other light profile parameters\n", + "(e.g. `effective_radius`, `sersic_index`), which are difficult degeneracies for the non-linear search to map out\n", + "accurately. This produces more reliable lens model results and the fit converges in fewer iterations, speeding up the\n", + "overall analysis.\n", + "\n", + "The inversion has a relatively small computational cost, thus we reduce the model complexity without much slow-down and\n", + "can therefore fit models more reliably and faster!\n", + "\n", + "__Disadvantages__\n", + "\n", + "Although the computation time of the inversion is small, it is not non-negligible. It is approximately 3-4x slower\n", + "than using a standard light profile.\n", + "\n", + "The gains in run times due to the simpler non-linear parameter space therefore are somewhat balanced by the slower\n", + "likelihood calculation.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and negative\n", + "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical.\n", + "\n", + "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is as fast\n", + "as positive-negative solvers. This ensures that all light profile intensities are positive and therefore physical.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a linear `Sersic` bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a linear `Sersic`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook.\n", + "\n", + "__Notes__\n", + "\n", + "This script is identical to `modeling/start_here.py` except that the light profiles are switched to linear light\n", + "profiles." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We now illustrate how to perform a fit to the dataset using linear light profiles, using known light profile parameters.\n", + "\n", + "The API follows closely the standard use of a `FitImaging` object, but simply uses linear light profiles (via the\n", + "`lp_linear` module) instead of standard light profiles. \n", + "\n", + "Note that the linear light profiles below do not have `intensity` parameters input and we use the true input values\n", + "of all other parameters for illustrative purposes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the linear light profiles have solved for `intensity` values that give a good fit\n", + "to the image. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The fit contains the solved for intensity values.\n", + "\n", + "These are computed using a fit's `linear_light_profile_intensity_dict`, which maps each linear light profile \n", + "in the model parameterization above to its `intensity`.\n", + "\n", + "The code below shows how to use this dictionary, as an alternative to using the max_log_likelihood quantities above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_bulge = tracer.galaxies[0].bulge\n", + "source_bulge = tracer.galaxies[1].bulge\n", + "\n", + "print(fit.linear_light_profile_intensity_dict)\n", + "\n", + "print(\n", + " f\"\\n Intensity of lens galaxy's bulge = {fit.linear_light_profile_intensity_dict[lens_bulge]}\"\n", + ")\n", + "\n", + "print(\n", + " f\"\\n Intensity of source bulge (lp_linear.SersicCore) = {fit.linear_light_profile_intensity_dict[source_bulge]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the solved \n", + "for `intensity` values is also accessible from a fit.\n", + "\n", + "For example, the linear light profile `Sersic` of the `bulge` component above has a solved for `intensity` of ~0.75. \n", + "\n", + "The `tracer` created below instead has an ordinary light profile with an `intensity` of ~0.75.\n", + "\n", + "The benefit of using a tracer with standard light profiles is it can be visualized (linear light profiles cannot \n", + "by default because they do not have `intensity` values)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "print(tracer.galaxies[0].bulge.intensity)\n", + "print(tracer.galaxies[1].bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization__\n", + "\n", + "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they do not \n", + "have an `intensity` value.\n", + "\n", + "Therefore, the objects created above which replaces all linear light profiles with ordinary light profiles must be\n", + "used for visualization:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Image\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/linear_light_profiles/likelihood_function.ipynb b/notebooks/imaging/features/linear_light_profiles/likelihood_function.ipynb index 6745dd11d..4200c1409 100644 --- a/notebooks/imaging/features/linear_light_profiles/likelihood_function.ipynb +++ b/notebooks/imaging/features/linear_light_profiles/likelihood_function.ipynb @@ -1,1086 +1,1123 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Linear Light Profile__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data with\n", - "linear light profiles (e.g. a Sersic bulge and Exponential disk).\n", - "\n", - "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", - "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", - "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood)\n", - "given the light profile's other parameters.\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers, so that readers can understand the likelihood\n", - " function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", - "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", - "packages are called when the likelihood is evaluated.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** The likelihood function of a linear light profile builds on that used for standard light profiles.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Masked Image Grid:** To perform galaxy calculations we used a 2D image-plane grid of (y,x) coordinates, which evaluated.\n", - "- **Linear Light Profiles:** To use a linear light profile, whose `intensity` is computed via linear algebra, we simply use the.\n", - "- **LightProfileLinearObjFuncList:** For standard light profiles, we combined our linear light profiles into a single `Galaxies` object.\n", - "- **Mapping Matrix:** The `mapping_matrix` is a matrix where each column is an image of each linear light profiles.\n", - "- **Combining Matrices:** The linear algebra system solves for all light profile `intensity` values at once, so we need to.\n", - "- **Image Reconstruction:** Using the reconstructed `intensity` values we can map the reconstruction back to the image plane.\n", - "- **Likelihood Function:** We now quantify the goodness-of-fit of our galaxy model.\n", - "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", - "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", - "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function of a linear light profile builds on that used for standard light profiles,\n", - "therefore you must read the following notebooks before this script:\n", - "\n", - "- `light_profile/likelihood_function.ipynb`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Following the `light_profile/log_likelihood_function.py` script, we load and mask an `Imaging` dataset and\n", - "set oversampling to 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=masked_dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masked Image Grid__\n", - "\n", - "To perform galaxy calculations we used a 2D image-plane grid of (y,x) coordinates, which evaluated the\n", - "emission of galaxy light profiles created as `LightProfile` objects.\n", - "\n", - "The code below repeats that used in `light_profile/log_likelihood_function.py` to show how this was done." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=4.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "image = tracer.image_2d_from(grid=masked_dataset.grids.lp)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Light Profiles__\n", - "\n", - "To use a linear light profile, whose `intensity` is computed via linear algebra, we simply use the `lp_linear`\n", - "module instead of the `lp` module used throughout other example scripts. \n", - "\n", - "The `intensity` parameter of the light profile is no longer passed into the light profiles created via the\n", - "`lp_linear` module, as it is inferred via linear algebra.\n", - "\n", - "In this example, the lens galaxy has a linear `Sersic` bulge and the source galaxy has a linear `SersicCore` bulge.\n", - "Both are linear light profiles, so their `intensity` parameters are not free parameters of the model \u2014 they are\n", - "solved for via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.lp_linear.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Internally in the source code, linear light profiles have an `intensity` parameter, but its value is always set to \n", - "1.0. It will be clear why this is later in the script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Lens Bulge Internal Intensity:\")\n", - "print(lens_galaxy.bulge.intensity)\n", - "\n", - "print(\"Source Bulge Internal Intensity:\")\n", - "print(source_galaxy.bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Like standard light profiles, we can compute images of each linear light profile, but their overall\n", - "normalization is arbitrary given that the internal `intensity` value of 1.0 is used.\n", - "\n", - "Note the source bulge image computed below is its source-plane image (i.e. before lensing); its lensed image\n", - "that contributes to the data is computed inside `LightProfileLinearObjFuncList` further below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_2d_lens_bulge = lens_galaxy.bulge.image_2d_from(grid=masked_dataset.grid)\n", - "image_2d_source_bulge = source_galaxy.bulge.image_2d_from(grid=masked_dataset.grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now put them together in a `Tracer` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LightProfileLinearObjFuncList__\n", - "\n", - "For standard light profiles, we combined our linear light profiles into a single `Galaxies` object. The \n", - "galaxies object computed each individual light profile's image and added them together.\n", - "\n", - "This no longer occurs for linear light profiles, instead linear light profiles are passed into the \n", - "`LightProfileLinearObjFuncList` object, which acts as an interface between the linear light profiles and the\n", - "linear algebra used to compute their intensity via the inversion.\n", - "\n", - "The quantities used to compute the image, blurring image and blurred image of each light profiles (the\n", - "dataset grid, PSF, etc.) are passed to the `LightProfileLinearObjFuncList` object, because it internally uses these\n", - "to compute each linear light profile image to set up the linear algebra.\n", - "\n", - "For lensing, this means we have to use a different `LightProfileLinearObjFuncList` object for each plane, because\n", - "each plane has its own ray-traced grid of (y,x) coordinates. Below, we set up the first `LightProfileLinearObjFuncList`,\n", - "which uses the image-plane grid and lens galaxy bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lp_linear_func_lens = al.LightProfileLinearObjFuncList(\n", - " grid=masked_dataset.grids.lp,\n", - " blurring_grid=masked_dataset.grids.blurring,\n", - " psf=masked_dataset.psf,\n", - " light_profile_list=[tracer.galaxies[0].bulge],\n", - " regularization=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This has a property `params` which is the number of intensity values that are computed via the inversion,\n", - "which because we have 1 light profiles is equal to 1.\n", - "\n", - "The `params` defines the dimensions of many of the matrices used in the linear algebra we discuss below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Number of Parameters (Intensity Values) in Linear Algebra:\")\n", - "print(lp_linear_func_lens.params)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapping Matrix__\n", - "\n", - "The `mapping_matrix` is a matrix where each column is an image of each linear light profiles (assuming its \n", - "intensity is 1.0), not accounting for the PSF convolution.\n", - "\n", - "It has dimensions `(total_image_pixels, total_linear_light_profiles)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapping_matrix = lp_linear_func_lens.mapping_matrix" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing the first column of the mapping matrix shows the image of the lens bulge light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge_image = mapping_matrix[:, 0]\n", - "print(bulge_image)\n", - "print(image_2d_lens_bulge.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `mapping_matrix` shows each light profile image in 1D, which is a bit odd to look at but\n", - "is a good way to think about the linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now make the second `LightProfileLinearObjFuncList`, which uses the ray-traced source-plane grid and source\n", - "galaxy bulge light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grids_of_planes_list = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.lp\n", - ")\n", - "traced_blurring_grids_of_planes_list = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")\n", - "\n", - "lp_linear_func_source = al.LightProfileLinearObjFuncList(\n", - " grid=traced_grids_of_planes_list[-1],\n", - " blurring_grid=traced_blurring_grids_of_planes_list[1],\n", - " psf=masked_dataset.psf,\n", - " light_profile_list=[tracer.galaxies[1].bulge],\n", - " regularization=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing the first column of the mapping matrix shows the image of the source bulge light profile (the\n", - "mapping_matrix entry is in the image-plane, after ray-tracing through the lens; the unlensed source-plane\n", - "image is plotted below for comparison)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapping_matrix = lp_linear_func_source.mapping_matrix\n", - "\n", - "bulge_image = mapping_matrix[:, 0]\n", - "print(bulge_image)\n", - "print(image_2d_source_bulge.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Combining Matrices__\n", - "\n", - "The linear algebra system solves for all light profile `intensity` values at once, so we need to combine\n", - "each of their individual mapping matrices into a single matrix.\n", - "\n", - "This is done via `hstack`, which stacks the two matrices horizontally to create a single matrix with dimensions\n", - "`(total_image_pixels, total_linear_light_profiles)`, where the latter dimension is now 2 because we have combined\n", - "two linear light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapping_matrix = np.hstack(\n", - " [lp_linear_func_lens.mapping_matrix, lp_linear_func_source.mapping_matrix]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Blurred Mapping Matrix ($f$)__\n", - "\n", - "The `mapping_matrix` does not account for the blurring of the light profile images by the PSF and therefore \n", - "is not used directly to compute the likelihood.\n", - "\n", - "Instead, we create a `blurred_mapping_matrix` which does account for this blurring. This is computed by \n", - "convolving each light profile image with the PSF.\n", - "\n", - "The `blurred_mapping_matrix` is a matrix analogous to the mapping matrix, but where each column is the image of each\n", - "light profile after it has been blurred by the PSF.\n", - "\n", - "This operation does not change the dimensions of the mapping matrix, meaning the `blurred_mapping_matrix` also has\n", - "dimensions `(total_image_pixels, total_linear_light_profiles)`. \n", - "\n", - "The property is actually called `operated_mapping_matrix_override` for two reasons: \n", - "\n", - "1) The operated signifies that this matrix could have any operation applied to it, it just happens for imaging\n", - " data that this operation is a convolution with the PSF.\n", - "\n", - "2) The `override` signifies that in the source code is changes how the `operated_mapping_matrix` is computed internally. \n", - " This is important if you are looking at the source code, but not important for the description of the likelihood \n", - " function in this guide.\n", - " \n", - "We have two separate `LightProfileLinearObjFuncList` objects, one for the lens and one for the source, we combine\n", - "the `blurred_mapping_matrix` of each via `hstack` to create a single `blurred_mapping_matrix` that represents\n", - "the linear system that will be solved for both." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "blurred_mapping_matrix = np.hstack(\n", - " [\n", - " lp_linear_func_lens.operated_mapping_matrix_override,\n", - " lp_linear_func_source.operated_mapping_matrix_override,\n", - " ],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing the first column of the mapping matrix shows the blurred image of the bulge light profile, the\n", - "second the blurred image of the source light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(blurred_mapping_matrix[:, 0])\n", - "print(blurred_mapping_matrix[:, 1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `mapping_matrix` shows each light profile image in 1D, with a PSF convolution applied." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " blurred_mapping_matrix,\n", - " aspect=(blurred_mapping_matrix.shape[1] / blurred_mapping_matrix.shape[0]),\n", - ")\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) (hereafter WD03) introduce the linear inversion formalism \n", - "used to compute the intensity values of the linear light profiles. In WD03, the science case is centred around strong\n", - "gravitational lensing and the galaxy is reconstructed on a rectangular grid of pixels, as opposed to linear light \n", - "profiles.\n", - "\n", - "However, the mathematics of the WD03 linear inversion formalism is the same as that used here, therefore this guide \n", - "describes which quantities in the linear inversion formalism map to the equations given in WD03. The pixelized \n", - "reconstruction methods, available in the code but described in the `pixelization` likelihood function guide, \n", - "also follow the WD03 formalism.\n", - "\n", - "The `blurred_mapping_matrix` is denoted $f_{ij}$ where $i$ maps over all $I$ linear light profiles and $j$ maps \n", - "over all $J$ image pixels. \n", - "\n", - "For example: \n", - "\n", - " - $f_{0, 1} = 0.3$ indicates that image-pixel $2$ maps to linear light profile $1$ with an intensity in that image \n", - " pixel of $0.3$ after PSF convolution.\n", - "\n", - "The indexing of the `mapping_matrix` is reversed compared to the notation of WD03 (e.g. image pixels\n", - "are the first entry of `mapping_matrix` whereas for $f$ they are the second index)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"Mapping between image pixel 0 and linear light profile pixel 1 = {mapping_matrix[0, 1]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Data Vector (D)__\n", - "\n", - "To solve for the linear light profile intensities we now pose the problem as a linear inversion.\n", - "\n", - "This requires us to convert the `blurred_mapping_matrix` and our `data` and `noise map` into matrices of certain \n", - "dimensions. \n", - "\n", - "The `data_vector`, $D$, is the first matrix and it has dimensions `(total_linear_light_profiles,)`.\n", - "\n", - "In WD03 (https://arxiv.org/abs/astro-ph/0302587) the data vector is given by: \n", - "\n", - " $\\vec{D}_{i} = \\sum_{\\rm j=1}^{J}f_{ij}(d_{j} - b_{j})/\\sigma_{j}^2 \\, \\, .$\n", - "\n", - "Where:\n", - "\n", - " - $d_{\\rm j}$ are the image-pixel data flux values.\n", - " - $b_{\\rm j}$ are the image values of all standard light profiles (therefore $d_{\\rm j} - b_{\\rm j}$ is \n", - " the data minus any standard light profiles).\n", - " - $\\sigma{\\rm _j}^2$ are the statistical uncertainties of each image-pixel value.\n", - "\n", - "$i$ maps over all $I$ linear light profiles and $j$ maps over all $J$ image pixels. \n", - "\n", - "This equation highlights a first aspect of linear inversions, if we are combining standard light profiles (which\n", - "have an input `intensity` value) with linear light profiles, the inversion is performed on the data minus\n", - "the standard light profile images. In this example, we have no standard light profiles and therefore the data\n", - "vector uses the data directly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_vector = al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(\n", - " blurred_mapping_matrix=blurred_mapping_matrix,\n", - " image=np.array(masked_dataset.data),\n", - " noise_map=np.array(masked_dataset.noise_map),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$D$'s meaning is a bit abstract, it essentially weights each linear light profile's `intensity` based on how it\n", - "maps to the data, so that the linear algebra can compute the `intensity` values that best-fit the data.\n", - "\n", - "We can plot $D$ as a column vector:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dimensions of $D$ are the number of linear light profiles, which in this case is 2." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Data Vector:\")\n", - "print(data_vector)\n", - "print(data_vector.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Matrix (F)__\n", - "\n", - "The `curvature_matrix` $F$ is the second matrix and it has \n", - "dimensions `(total_linear_light_profiles, total_linear_light_profiles)`.\n", - "\n", - "In WD03 (https://arxiv.org/abs/astro-ph/0302587) the curvature matrix is a 2D matrix given by:\n", - "\n", - " ${F}_{ik} = \\sum_{\\rm j=1}^{J}f_{ij}f_{kj}/\\sigma_{j}^2 \\, \\, .$\n", - "\n", - "NOTE: this notation implicitly assumes a summation over $K$, where $k$ runs over all linear light profile indexes $K$.\n", - "\n", - "Note how summation over $J$ runs over $f$ twice, such that every entry of $F$ is the sum of the multiplication\n", - "between all values in every two columns of $f$.\n", - "\n", - "For example, $F_{0,1}$ is the sum of every blurred image pixels values in $f$ of linear light profile 0 multiplied by\n", - "every blurred image pixel value of linear light profile 1.\n", - "\n", - "$F$'s meaning is also a bit abstract, but it essentially quantifies how much each linear light profile's image\n", - "overlaps with every other linear light profile's image, weighted by the noise in the data. This is what combined with\n", - "the `data_vector` allows the inversion to compute the `intensity` values that best-fit the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix, noise_map=masked_dataset.noise_map\n", - ")\n", - "\n", - "plt.imshow(curvature_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction (Positive-Negative)__\n", - "\n", - "The following chi-squared is minimized when we perform the inversion and reconstruct the galaxy:\n", - "\n", - "$\\chi^2 = \\sum_{\\rm j=1}^{J} \\bigg[ \\frac{(\\sum_{\\rm i=1}^{I} s_{i} f_{ij}) + b_{j} - d_{j}}{\\sigma_{j}} \\bigg]$\n", - "\n", - "Where $s$ is the `intensity` values in all $I$ linear light profile images.\n", - "\n", - "The solution for $s$ is therefore given by (equation 5 WD03):\n", - "\n", - " $s = F^{-1} D$\n", - "\n", - "We can compute this using NumPy linear algebra and the `solve` function.\n", - "\n", - "However, this function allows for the solved `intensity` values to be negative. For linear light profiles which\n", - "are a good fit to the data, this is unlikely to happen and the `intensity` values will be positive. However, \n", - "for more complex models this may not be the case. Below, we describes how we can ensure the `intensity` values\n", - "are positive." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = np.linalg.solve(curvature_matrix, data_vector)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `reconstruction` is a 1D vector of length equal to the number of linear light profiles, which in this case is 2.\n", - "\n", - "Each value represents the intensity of the linear light profile.\n", - "\n", - "In this example, both values are positive, but remember that this is not guaranteed for all linear inversions\n", - "that are solve using this method." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Reconstruction (S) of Linear Light Profiles Intensity:\")\n", - "print(reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction (Positive Only)__\n", - "\n", - "The linear algebra can be solved for with the constraint that all solutions, and therefore all `intensity` values,\n", - "are positive. \n", - "\n", - "This could be achieved by using the `scipy` `nnls` non-negative least squares solver.\n", - "\n", - "The nnls poses the problem slightly different than the code above. It solves for the `intensity` values in an\n", - "iterative manner meaning that it is slower. It does not use `data_vector` $D$ and `curvature_matrix` $F$ but instead\n", - "works directly with the `blurred_mapping_matrix` $f$ and the data and noise-map.\n", - "\n", - "The `nnls` function is therefore computationally slow, especially for cases where there are many linear light profiles \n", - "or even more complex linear inversions like a pixelized reconstruction.\n", - "\n", - "The source code therefore uses a \"fast nnls\" algorithm, which is an adaptation of the algorithm found at\n", - "this URL: https://github.com/jvendrow/fnnls\n", - "\n", - "Unlike the scipy nnls function, the fnnls method uses the `data_vector` $D$ and `curvature_matrix` $F$ to solve for\n", - "the `intensity` values. This provides it with additional information about the linear algebra problem, which is\n", - "why it is faster.\n", - "\n", - "The function `reconstruction_positive_only_from` uses the `fnnls` algorithm to compute the `intensity` values\n", - "of the linear light profiles, ensuring they are positive." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", - " data_vector=data_vector,\n", - " curvature_reg_matrix=curvature_matrix, # ignore _reg_ tag in this guide\n", - ")\n", - "\n", - "print(reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image Reconstruction__\n", - "\n", - "Using the reconstructed `intensity` values we can map the reconstruction back to the image plane (via \n", - "the `blurred mapping_matrix`) and produce a reconstruction of the image data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_operated_data = (\n", - " al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix, reconstruction=reconstruction\n", - " )\n", - ")\n", - "\n", - "mapped_reconstructed_operated_data = al.Array2D(\n", - " values=mapped_reconstructed_operated_data, mask=mask\n", - ")\n", - "\n", - "aplt.plot_array(array=mapped_reconstructed_operated_data, title=\"\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our galaxy model.\n", - "\n", - "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", - "\n", - "The likelihood function for parametric galaxy modeling, even if linear light profiles are used, consists of two terms:\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "We now explain what each of these terms mean.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", - "\n", - " - `model_data` = `convolved_image_2d`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", - "\n", - "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", - "for, corresponding to a fit with a lower likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image = mapped_reconstructed_operated_data\n", - "\n", - "residual_map = masked_dataset.data - model_image\n", - "normalized_residual_map = residual_map / masked_dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(chi_squared)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", - "\n", - "aplt.plot_array(array=chi_squared_map, title=\"\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", - "\n", - "The final term in the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", - "of the log of every noise-map value squared. \n", - "\n", - "Given the `noise_map` is fixed, this term does not change during the galaxy modeling process and has no impact on the \n", - "model we infer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms computed above using\n", - "the likelihood function defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(figure_of_merit)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(\n", - " dataset=masked_dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_border_relocator=True),\n", - ")\n", - "fit_log_evidence = fit.log_evidence\n", - "print(fit_log_evidence)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit contains an `Inversion` object, which handles all the linear algebra we have covered in this script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.inversion)\n", - "print(fit.inversion.data_vector)\n", - "print(fit.inversion.curvature_matrix)\n", - "print(fit.inversion.reconstruction)\n", - "print(fit.inversion.mapped_reconstructed_operated_data)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Inversion` object can be computed from a tracer and a dataset, by passing them to the `TracerToInversion` object.\n", - "\n", - "This objects handles a lot of extra functionality that we have not covered in this script, such as:\n", - "\n", - "- Separating out the linear light profiles from the standard light profiles.\n", - "- Separating out objects which reconstruct the galaxy using a pixelized reconstruction, which are passed into\n", - " the `Inversion` object as well." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(\n", - " tracer=tracer,\n", - " dataset=masked_dataset,\n", - ")\n", - "\n", - "inversion = tracer_to_inversion.inversion" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "For linear light profiles, the reduced number of free parameters (e.g. the `intensity` values are solved for\n", - "via linear algebra and not a dimension of the non-linear parameter space) means that the sampler converges in fewer\n", - "iterations and is less likely to infer a local maximum.\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the linear light profile likelihood function, which uses \n", - "analytic light profiles to fit the galaxies light and solve for the `intensity` values via linear algebra.\n", - "\n", - "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", - "are described in additional notebooks found in the `guides` package:\n", - "\n", - " - `over_sampling`: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", - " ray-traced to the source-plane and used to evaluate the light profile more accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Linear Light Profile__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data with\n", + "linear light profiles (e.g. a Sersic bulge and Exponential disk).\n", + "\n", + "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", + "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", + "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood)\n", + "given the light profile's other parameters.\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers, so that readers can understand the likelihood\n", + " function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", + "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", + "packages are called when the likelihood is evaluated.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** The likelihood function of a linear light profile builds on that used for standard light profiles.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Masked Image Grid:** To perform galaxy calculations we used a 2D image-plane grid of (y,x) coordinates, which evaluated.\n", + "- **Linear Light Profiles:** To use a linear light profile, whose `intensity` is computed via linear algebra, we simply use the.\n", + "- **LightProfileLinearObjFuncList:** For standard light profiles, we combined our linear light profiles into a single `Galaxies` object.\n", + "- **Mapping Matrix:** The `mapping_matrix` is a matrix where each column is an image of each linear light profiles.\n", + "- **Combining Matrices:** The linear algebra system solves for all light profile `intensity` values at once, so we need to.\n", + "- **Image Reconstruction:** Using the reconstructed `intensity` values we can map the reconstruction back to the image plane.\n", + "- **Likelihood Function:** We now quantify the goodness-of-fit of our galaxy model.\n", + "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", + "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", + "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function of a linear light profile builds on that used for standard light profiles,\n", + "therefore you must read the following notebooks before this script:\n", + "\n", + "- `light_profile/likelihood_function.ipynb`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Following the `light_profile/log_likelihood_function.py` script, we load and mask an `Imaging` dataset and\n", + "set oversampling to 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=masked_dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masked Image Grid__\n", + "\n", + "To perform galaxy calculations we used a 2D image-plane grid of (y,x) coordinates, which evaluated the\n", + "emission of galaxy light profiles created as `LightProfile` objects.\n", + "\n", + "The code below repeats that used in `light_profile/log_likelihood_function.py` to show how this was done." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=4.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "image = tracer.image_2d_from(grid=masked_dataset.grids.lp)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Light Profiles__\n", + "\n", + "To use a linear light profile, whose `intensity` is computed via linear algebra, we simply use the `lp_linear`\n", + "module instead of the `lp` module used throughout other example scripts. \n", + "\n", + "The `intensity` parameter of the light profile is no longer passed into the light profiles created via the\n", + "`lp_linear` module, as it is inferred via linear algebra.\n", + "\n", + "In this example, the lens galaxy has a linear `Sersic` bulge and the source galaxy has a linear `SersicCore` bulge.\n", + "Both are linear light profiles, so their `intensity` parameters are not free parameters of the model \u2014 they are\n", + "solved for via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.lp_linear.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Internally in the source code, linear light profiles have an `intensity` parameter, but its value is always set to \n", + "1.0. It will be clear why this is later in the script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Lens Bulge Internal Intensity:\")\n", + "print(lens_galaxy.bulge.intensity)\n", + "\n", + "print(\"Source Bulge Internal Intensity:\")\n", + "print(source_galaxy.bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Like standard light profiles, we can compute images of each linear light profile, but their overall\n", + "normalization is arbitrary given that the internal `intensity` value of 1.0 is used.\n", + "\n", + "Note the source bulge image computed below is its source-plane image (i.e. before lensing); its lensed image\n", + "that contributes to the data is computed inside `LightProfileLinearObjFuncList` further below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_2d_lens_bulge = lens_galaxy.bulge.image_2d_from(grid=masked_dataset.grid)\n", + "image_2d_source_bulge = source_galaxy.bulge.image_2d_from(grid=masked_dataset.grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now put them together in a `Tracer` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LightProfileLinearObjFuncList__\n", + "\n", + "For standard light profiles, we combined our linear light profiles into a single `Galaxies` object. The \n", + "galaxies object computed each individual light profile's image and added them together.\n", + "\n", + "This no longer occurs for linear light profiles, instead linear light profiles are passed into the \n", + "`LightProfileLinearObjFuncList` object, which acts as an interface between the linear light profiles and the\n", + "linear algebra used to compute their intensity via the inversion.\n", + "\n", + "The quantities used to compute the image, blurring image and blurred image of each light profiles (the\n", + "dataset grid, PSF, etc.) are passed to the `LightProfileLinearObjFuncList` object, because it internally uses these\n", + "to compute each linear light profile image to set up the linear algebra.\n", + "\n", + "For lensing, this means we have to use a different `LightProfileLinearObjFuncList` object for each plane, because\n", + "each plane has its own ray-traced grid of (y,x) coordinates. Below, we set up the first `LightProfileLinearObjFuncList`,\n", + "which uses the image-plane grid and lens galaxy bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lp_linear_func_lens = al.LightProfileLinearObjFuncList(\n", + " grid=masked_dataset.grids.lp,\n", + " blurring_grid=masked_dataset.grids.blurring,\n", + " psf=masked_dataset.psf,\n", + " light_profile_list=[tracer.galaxies[0].bulge],\n", + " regularization=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This has a property `params` which is the number of intensity values that are computed via the inversion,\n", + "which because we have 1 light profiles is equal to 1.\n", + "\n", + "The `params` defines the dimensions of many of the matrices used in the linear algebra we discuss below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Number of Parameters (Intensity Values) in Linear Algebra:\")\n", + "print(lp_linear_func_lens.params)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapping Matrix__\n", + "\n", + "The `mapping_matrix` is a matrix where each column is an image of each linear light profiles (assuming its \n", + "intensity is 1.0), not accounting for the PSF convolution.\n", + "\n", + "It has dimensions `(total_image_pixels, total_linear_light_profiles)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapping_matrix = lp_linear_func_lens.mapping_matrix" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing the first column of the mapping matrix shows the image of the lens bulge light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge_image = mapping_matrix[:, 0]\n", + "print(bulge_image)\n", + "print(image_2d_lens_bulge.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `mapping_matrix` shows each light profile image in 1D, which is a bit odd to look at but\n", + "is a good way to think about the linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now make the second `LightProfileLinearObjFuncList`, which uses the ray-traced source-plane grid and source\n", + "galaxy bulge light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grids_of_planes_list = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.lp\n", + ")\n", + "traced_blurring_grids_of_planes_list = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")\n", + "\n", + "lp_linear_func_source = al.LightProfileLinearObjFuncList(\n", + " grid=traced_grids_of_planes_list[-1],\n", + " blurring_grid=traced_blurring_grids_of_planes_list[1],\n", + " psf=masked_dataset.psf,\n", + " light_profile_list=[tracer.galaxies[1].bulge],\n", + " regularization=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing the first column of the mapping matrix shows the image of the source bulge light profile (the\n", + "mapping_matrix entry is in the image-plane, after ray-tracing through the lens; the unlensed source-plane\n", + "image is plotted below for comparison)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapping_matrix = lp_linear_func_source.mapping_matrix\n", + "\n", + "bulge_image = mapping_matrix[:, 0]\n", + "print(bulge_image)\n", + "print(image_2d_source_bulge.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Combining Matrices__\n", + "\n", + "The linear algebra system solves for all light profile `intensity` values at once, so we need to combine\n", + "each of their individual mapping matrices into a single matrix.\n", + "\n", + "This is done via `hstack`, which stacks the two matrices horizontally to create a single matrix with dimensions\n", + "`(total_image_pixels, total_linear_light_profiles)`, where the latter dimension is now 2 because we have combined\n", + "two linear light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapping_matrix = np.hstack(\n", + " [lp_linear_func_lens.mapping_matrix, lp_linear_func_source.mapping_matrix]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Blurred Mapping Matrix ($f$)__\n", + "\n", + "The `mapping_matrix` does not account for the blurring of the light profile images by the PSF and therefore \n", + "is not used directly to compute the likelihood.\n", + "\n", + "Instead, we create a `blurred_mapping_matrix` which does account for this blurring. This is computed by \n", + "convolving each light profile image with the PSF.\n", + "\n", + "The `blurred_mapping_matrix` is a matrix analogous to the mapping matrix, but where each column is the image of each\n", + "light profile after it has been blurred by the PSF.\n", + "\n", + "This operation does not change the dimensions of the mapping matrix, meaning the `blurred_mapping_matrix` also has\n", + "dimensions `(total_image_pixels, total_linear_light_profiles)`. \n", + "\n", + "The property is actually called `operated_mapping_matrix_override` for two reasons: \n", + "\n", + "1) The operated signifies that this matrix could have any operation applied to it, it just happens for imaging\n", + " data that this operation is a convolution with the PSF.\n", + "\n", + "2) The `override` signifies that in the source code is changes how the `operated_mapping_matrix` is computed internally. \n", + " This is important if you are looking at the source code, but not important for the description of the likelihood \n", + " function in this guide.\n", + " \n", + "We have two separate `LightProfileLinearObjFuncList` objects, one for the lens and one for the source, we combine\n", + "the `blurred_mapping_matrix` of each via `hstack` to create a single `blurred_mapping_matrix` that represents\n", + "the linear system that will be solved for both." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "blurred_mapping_matrix = np.hstack(\n", + " [\n", + " lp_linear_func_lens.operated_mapping_matrix_override,\n", + " lp_linear_func_source.operated_mapping_matrix_override,\n", + " ],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing the first column of the mapping matrix shows the blurred image of the bulge light profile, the\n", + "second the blurred image of the source light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(blurred_mapping_matrix[:, 0])\n", + "print(blurred_mapping_matrix[:, 1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `mapping_matrix` shows each light profile image in 1D, with a PSF convolution applied." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " blurred_mapping_matrix,\n", + " aspect=(blurred_mapping_matrix.shape[1] / blurred_mapping_matrix.shape[0]),\n", + ")\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) (hereafter WD03) introduce the linear inversion formalism \n", + "used to compute the intensity values of the linear light profiles. In WD03, the science case is centred around strong\n", + "gravitational lensing and the galaxy is reconstructed on a rectangular grid of pixels, as opposed to linear light \n", + "profiles.\n", + "\n", + "However, the mathematics of the WD03 linear inversion formalism is the same as that used here, therefore this guide \n", + "describes which quantities in the linear inversion formalism map to the equations given in WD03. The pixelized \n", + "reconstruction methods, available in the code but described in the `pixelization` likelihood function guide, \n", + "also follow the WD03 formalism.\n", + "\n", + "The `blurred_mapping_matrix` is denoted $f_{ij}$ where $i$ maps over all $I$ linear light profiles and $j$ maps \n", + "over all $J$ image pixels. \n", + "\n", + "For example: \n", + "\n", + " - $f_{0, 1} = 0.3$ indicates that image-pixel $2$ maps to linear light profile $1$ with an intensity in that image \n", + " pixel of $0.3$ after PSF convolution.\n", + "\n", + "The indexing of the `mapping_matrix` is reversed compared to the notation of WD03 (e.g. image pixels\n", + "are the first entry of `mapping_matrix` whereas for $f$ they are the second index)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"Mapping between image pixel 0 and linear light profile pixel 1 = {mapping_matrix[0, 1]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Data Vector (D)__\n", + "\n", + "To solve for the linear light profile intensities we now pose the problem as a linear inversion.\n", + "\n", + "This requires us to convert the `blurred_mapping_matrix` and our `data` and `noise map` into matrices of certain \n", + "dimensions. \n", + "\n", + "The `data_vector`, $D$, is the first matrix and it has dimensions `(total_linear_light_profiles,)`.\n", + "\n", + "In WD03 (https://arxiv.org/abs/astro-ph/0302587) the data vector is given by: \n", + "\n", + " $\\vec{D}_{i} = \\sum_{\\rm j=1}^{J}f_{ij}(d_{j} - b_{j})/\\sigma_{j}^2 \\, \\, .$\n", + "\n", + "Where:\n", + "\n", + " - $d_{\\rm j}$ are the image-pixel data flux values.\n", + " - $b_{\\rm j}$ are the image values of all standard light profiles (therefore $d_{\\rm j} - b_{\\rm j}$ is \n", + " the data minus any standard light profiles).\n", + " - $\\sigma{\\rm _j}^2$ are the statistical uncertainties of each image-pixel value.\n", + "\n", + "$i$ maps over all $I$ linear light profiles and $j$ maps over all $J$ image pixels. \n", + "\n", + "This equation highlights a first aspect of linear inversions, if we are combining standard light profiles (which\n", + "have an input `intensity` value) with linear light profiles, the inversion is performed on the data minus\n", + "the standard light profile images. In this example, we have no standard light profiles and therefore the data\n", + "vector uses the data directly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_vector = al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(\n", + " blurred_mapping_matrix=blurred_mapping_matrix,\n", + " image=np.array(masked_dataset.data),\n", + " noise_map=np.array(masked_dataset.noise_map),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$D$'s meaning is a bit abstract, it essentially weights each linear light profile's `intensity` based on how it\n", + "maps to the data, so that the linear algebra can compute the `intensity` values that best-fit the data.\n", + "\n", + "We can plot $D$ as a column vector:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dimensions of $D$ are the number of linear light profiles, which in this case is 2." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Data Vector:\")\n", + "print(data_vector)\n", + "print(data_vector.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Matrix (F)__\n", + "\n", + "The `curvature_matrix` $F$ is the second matrix and it has \n", + "dimensions `(total_linear_light_profiles, total_linear_light_profiles)`.\n", + "\n", + "In WD03 (https://arxiv.org/abs/astro-ph/0302587) the curvature matrix is a 2D matrix given by:\n", + "\n", + " ${F}_{ik} = \\sum_{\\rm j=1}^{J}f_{ij}f_{kj}/\\sigma_{j}^2 \\, \\, .$\n", + "\n", + "NOTE: this notation implicitly assumes a summation over $K$, where $k$ runs over all linear light profile indexes $K$.\n", + "\n", + "Note how summation over $J$ runs over $f$ twice, such that every entry of $F$ is the sum of the multiplication\n", + "between all values in every two columns of $f$.\n", + "\n", + "For example, $F_{0,1}$ is the sum of every blurred image pixels values in $f$ of linear light profile 0 multiplied by\n", + "every blurred image pixel value of linear light profile 1.\n", + "\n", + "$F$'s meaning is also a bit abstract, but it essentially quantifies how much each linear light profile's image\n", + "overlaps with every other linear light profile's image, weighted by the noise in the data. This is what combined with\n", + "the `data_vector` allows the inversion to compute the `intensity` values that best-fit the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix, noise_map=masked_dataset.noise_map\n", + ")\n", + "\n", + "plt.imshow(curvature_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction (Positive-Negative)__\n", + "\n", + "The following chi-squared is minimized when we perform the inversion and reconstruct the galaxy:\n", + "\n", + "$\\chi^2 = \\sum_{\\rm j=1}^{J} \\bigg[ \\frac{(\\sum_{\\rm i=1}^{I} s_{i} f_{ij}) + b_{j} - d_{j}}{\\sigma_{j}} \\bigg]$\n", + "\n", + "Where $s$ is the `intensity` values in all $I$ linear light profile images.\n", + "\n", + "The solution for $s$ is therefore given by (equation 5 WD03):\n", + "\n", + " $s = F^{-1} D$\n", + "\n", + "We can compute this using NumPy linear algebra and the `solve` function.\n", + "\n", + "However, this function allows for the solved `intensity` values to be negative. For linear light profiles which\n", + "are a good fit to the data, this is unlikely to happen and the `intensity` values will be positive. However, \n", + "for more complex models this may not be the case. Below, we describes how we can ensure the `intensity` values\n", + "are positive." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = np.linalg.solve(curvature_matrix, data_vector)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `reconstruction` is a 1D vector of length equal to the number of linear light profiles, which in this case is 2.\n", + "\n", + "Each value represents the intensity of the linear light profile.\n", + "\n", + "In this example, both values are positive, but remember that this is not guaranteed for all linear inversions\n", + "that are solve using this method." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Reconstruction (S) of Linear Light Profiles Intensity:\")\n", + "print(reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction (Positive Only)__\n", + "\n", + "The linear algebra can be solved for with the constraint that all solutions, and therefore all `intensity` values,\n", + "are positive. \n", + "\n", + "This could be achieved by using the `scipy` `nnls` non-negative least squares solver.\n", + "\n", + "The nnls poses the problem slightly different than the code above. It solves for the `intensity` values in an\n", + "iterative manner meaning that it is slower. It does not use `data_vector` $D$ and `curvature_matrix` $F$ but instead\n", + "works directly with the `blurred_mapping_matrix` $f$ and the data and noise-map.\n", + "\n", + "The `nnls` function is therefore computationally slow, especially for cases where there are many linear light profiles \n", + "or even more complex linear inversions like a pixelized reconstruction.\n", + "\n", + "The source code therefore uses a \"fast nnls\" algorithm, which is an adaptation of the algorithm found at\n", + "this URL: https://github.com/jvendrow/fnnls\n", + "\n", + "Unlike the scipy nnls function, the fnnls method uses the `data_vector` $D$ and `curvature_matrix` $F$ to solve for\n", + "the `intensity` values. This provides it with additional information about the linear algebra problem, which is\n", + "why it is faster.\n", + "\n", + "The function `reconstruction_positive_only_from` uses the `fnnls` algorithm to compute the `intensity` values\n", + "of the linear light profiles, ensuring they are positive." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", + " data_vector=data_vector,\n", + " curvature_reg_matrix=curvature_matrix, # ignore _reg_ tag in this guide\n", + ")\n", + "\n", + "print(reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image Reconstruction__\n", + "\n", + "Using the reconstructed `intensity` values we can map the reconstruction back to the image plane (via \n", + "the `blurred mapping_matrix`) and produce a reconstruction of the image data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_operated_data = (\n", + " al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix, reconstruction=reconstruction\n", + " )\n", + ")\n", + "\n", + "mapped_reconstructed_operated_data = al.Array2D(\n", + " values=mapped_reconstructed_operated_data, mask=mask\n", + ")\n", + "\n", + "aplt.plot_array(array=mapped_reconstructed_operated_data, title=\"\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our galaxy model.\n", + "\n", + "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", + "\n", + "The likelihood function for parametric galaxy modeling, even if linear light profiles are used, consists of two terms:\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "We now explain what each of these terms mean.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", + "\n", + " - `model_data` = `convolved_image_2d`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", + "\n", + "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", + "for, corresponding to a fit with a lower likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image = mapped_reconstructed_operated_data\n", + "\n", + "residual_map = masked_dataset.data - model_image\n", + "normalized_residual_map = residual_map / masked_dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(chi_squared)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", + "\n", + "aplt.plot_array(array=chi_squared_map, title=\"\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", + "\n", + "The final term in the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", + "of the log of every noise-map value squared. \n", + "\n", + "Given the `noise_map` is fixed, this term does not change during the galaxy modeling process and has no impact on the \n", + "model we infer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms computed above using\n", + "the likelihood function defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(figure_of_merit)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(\n", + " dataset=masked_dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_border_relocator=True),\n", + ")\n", + "fit_log_evidence = fit.log_evidence\n", + "print(fit_log_evidence)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit contains an `Inversion` object, which handles all the linear algebra we have covered in this script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.inversion)\n", + "print(fit.inversion.data_vector)\n", + "print(fit.inversion.curvature_matrix)\n", + "print(fit.inversion.reconstruction)\n", + "print(fit.inversion.mapped_reconstructed_operated_data)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Inversion` object can be computed from a tracer and a dataset, by passing them to the `TracerToInversion` object.\n", + "\n", + "This objects handles a lot of extra functionality that we have not covered in this script, such as:\n", + "\n", + "- Separating out the linear light profiles from the standard light profiles.\n", + "- Separating out objects which reconstruct the galaxy using a pixelized reconstruction, which are passed into\n", + " the `Inversion` object as well." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(\n", + " tracer=tracer,\n", + " dataset=masked_dataset,\n", + ")\n", + "\n", + "inversion = tracer_to_inversion.inversion" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "For linear light profiles, the reduced number of free parameters (e.g. the `intensity` values are solved for\n", + "via linear algebra and not a dimension of the non-linear parameter space) means that the sampler converges in fewer\n", + "iterations and is less likely to infer a local maximum.\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the linear light profile likelihood function, which uses \n", + "analytic light profiles to fit the galaxies light and solve for the `intensity` values via linear algebra.\n", + "\n", + "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", + "are described in additional notebooks found in the `guides` package:\n", + "\n", + " - `over_sampling`: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", + " ray-traced to the source-plane and used to evaluate the light profile more accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/linear_light_profiles/modeling.ipynb b/notebooks/imaging/features/linear_light_profiles/modeling.ipynb index 29c1f69cc..f126b1afa 100644 --- a/notebooks/imaging/features/linear_light_profiles/modeling.ipynb +++ b/notebooks/imaging/features/linear_light_profiles/modeling.ipynb @@ -1,671 +1,708 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Linear Light Profiles\n", - "========================================\n", - "\n", - "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", - "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", - "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood)\n", - "given the light profile's other parameters.\n", - "\n", - "Based on the advantages below, we recommended you always use linear light profiles to fit models over standard\n", - "light profiles!\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Each light profile's `intensity` parameter is therefore not a free parameter in the model-fit.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Notes:** This script is identical to `modeling/start_here.py` except that the light profiles are switched to.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Intensities:** The intensities of linear light profiles are not a part of the model parameterization and therefore.\n", - "- **Visualization:** Linear light profiles and objects containing them (e.g.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "- **Max Likelihood Inversion:** As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit`, which contains.\n", - "\n", - "__Advantages__\n", - "\n", - "Each light profile's `intensity` parameter is therefore not a free parameter in the model-fit, reducing the\n", - "dimensionality of non-linear parameter space by the number of light profiles (in this example by 2 dimensions).\n", - "\n", - "This also removes the degeneracies that occur between the `intensity` and other light profile parameters\n", - "(e.g. `effective_radius`, `sersic_index`), which are difficult degeneracies for the non-linear search to map out\n", - "accurately. This produces more reliable lens model results and the fit converges in fewer iterations, speeding up the\n", - "overall analysis.\n", - "\n", - "The inversion has a relatively small computational cost, thus we reduce the model complexity without much slow-down and\n", - "can therefore fit models more reliably and faster!\n", - "\n", - "__Disadvantages__\n", - "\n", - "Although the computation time of the inversion is small, it is not non-negligible. It is approximately 3-4x slower\n", - "than using a standard light profile.\n", - "\n", - "The gains in run times due to the simpler non-linear parameter space therefore are somewhat balanced by the slower\n", - "likelihood calculation.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and negative\n", - "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical.\n", - "\n", - "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is as fast\n", - "as positive-negative solvers. This ensures that all light profile intensities are positive and therefore physical.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a linear `Sersic` bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a linear `Sersic`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook.\n", - "\n", - "__Notes__\n", - "\n", - "This script is identical to `modeling/start_here.py` except that the light profiles are switched to linear light\n", - "profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's light is a linear `Sersic` bulge [6 parameters].\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is a linear `Sersic` [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=19.\n", - "\n", - "Note how both the lens and source galaxies use linear light profiles, meaning that the `intensity` parameter of both\n", - "is no longer a free parameter in the fit.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the light profiles of the lens and source galaxies do not include an `intensity` parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", - "full description).\n", - "\n", - "In the `start_here.py` example 150 live points (`n_live=150`) were used to sample parameter space. For the linear\n", - "light profiles this is reduced to 100, as the simpler parameter space means we need fewer live points to map it out\n", - "accurately. This will lead to faster run times." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"linear_light_profiles\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to\n", - "print the estimated VRAM required by a model.\n", - "\n", - "For each linear light profile in the model extra is used VRAM. For 3-10 linear Sersic light profiles this is a tiny \n", - "amount of VRAM (e.g. < 10MB per batched likelihood). Even for large batch sizes (e.g. over 100) you probably \n", - "will not use enough VRAM to require monitoring.\n", - "\n", - "__Run Time__\n", - "\n", - "For standard light profiles, the log likelihood evaluation time is of order ~0.01 seconds for this dataset.\n", - "\n", - "For linear light profiles, the log likelihood evaluation increases to around ~0.05 seconds per likelihood evaluation.\n", - "This is still fast, but it does mean that the fit may take around five times longer to run.\n", - "\n", - "However, because two free parameters have been removed from the model (the `intensity` of the lens bulge and \n", - "source bulge), the total number of likelihood evaluations will reduce. Furthermore, the simpler parameter space\n", - "likely means that the fit will take less than 10000 likelihood evaluations per free parameter to converge. This is aided further\n", - "by the reduction in `n_live` to 100.\n", - "\n", - "Fits using standard light profiles and linear light profiles therefore take roughly the same time to run. However,\n", - "the simpler parameter space of linear light profiles means that the model-fit is more reliable, less susceptible to\n", - "converging to an incorrect solution and scales better if even more light profiles are included in the model.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that `intensity` parameters are not inferred by the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "The lens and source galaxies appear similar to those in the data, confirming that the `intensity` values inferred by\n", - "the inversion process are accurate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The intensities of linear light profiles are not a part of the model parameterization and therefore are not displayed\n", - "in the `model.results` file.\n", - "\n", - "To extract the `intensity` values of a specific component in the model, we use the `max_log_likelihood_tracer`,\n", - "which has already performed the inversion and therefore the galaxy light profiles have their solved for\n", - "`intensity`'s associated with them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "print(tracer.galaxies[0].bulge.intensity)\n", - "print(tracer.galaxies[1].bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Above, we access these values using the list index entry of each galaxy in the tracer. However, we may not be certain\n", - "of the order of the galaxies in the tracer, and therefore which galaxy index corresponds to the lens and source.\n", - "\n", - "We can therefore use the model composition API to access these values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(tracer.galaxies[0].bulge.intensity)\n", - "print(tracer.galaxies[-1].bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Tracer` contained in the `max_log_likelihood_fit` also has the solved for `intensity` values:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = result.max_log_likelihood_fit\n", - "\n", - "tracer = fit.tracer\n", - "\n", - "print(tracer.galaxies[0].bulge.intensity)\n", - "print(tracer.galaxies[-1].bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization__\n", - "\n", - "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they do not \n", - "have an `intensity` value.\n", - "\n", - "Therefore, the objects created above which replaces all linear light profiles with ordinary light profiles must be\n", - "used for visualization:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Image\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "__Result (Advanced)__\n", - "\n", - "The code belows shows all additional results that can be computed from a `Result` object following a fit with a\n", - "linear light profile.\n", - "\n", - "__Max Likelihood Inversion__\n", - "\n", - "As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit`, which contains the\n", - "`Inversion` object we need." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = result.max_log_likelihood_fit.inversion" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This `Inversion` can be used to plot the reconstructed image of specifically all linear light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Linear Objects (Internal Source Code)__\n", - "\n", - "An `Inversion` contains all of the linear objects used to reconstruct the data in its `linear_obj_list`.\n", - "\n", - "This list may include the following objects:\n", - "\n", - " - `LightProfileLinearObjFuncList`: Holds a list of linear light profiles and the functionality used to\n", - " reconstruct data in an inversion. It may contain a single light profile (e.g. `lp_linear.Sersic`) or\n", - " many light profiles combined in a `Basis` (e.g. `lp_basis.Basis`).\n", - "\n", - " - `Mapper`: The linear object used by a `Pixelization` to reconstruct data via an `Inversion`. The `Mapper`\n", - " is specific to the `Pixelization`'s `Mesh` (e.g. a `RectangularMapper` is used for a `RectangularAdaptDensity`\n", - " mesh).\n", - "\n", - "In this example, the model uses one linear `Sersic` for the lens galaxy's bulge and one linear `SersicCore`\n", - "for the source galaxy's bulge. The inversion therefore has two `LightProfileLinearObjFuncList` entries \u2014\n", - "one for each plane (lens at the image-plane grid, source at the ray-traced source-plane grid).\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.linear_obj_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To extract results from an inversion many quantities come in lists or require us to specify the linear object\n", - "we wish to use. Knowing what linear objects are in the `linear_obj_list`, and what indexes they correspond to,\n", - "is therefore important.\n", - "\n", - "The lens-plane entry comes first, followed by the source-plane entry." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"LightProfileLinearObjFuncList (Lens Sersic) = {inversion.linear_obj_list[0]}\"\n", - ")\n", - "print(\n", - " f\"LightProfileLinearObjFuncList (Source SersicCore) = {inversion.linear_obj_list[1]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Each `LightProfileLinearObjFuncList` contains a `light_profile_list`. For both entries in this example the\n", - "list has a single light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"Linear Light Profile list (Lens Sersic) = {inversion.linear_obj_list[0].light_profile_list}\"\n", - ")\n", - "print(\n", - " f\"Linear Light Profile list (Source SersicCore) = {inversion.linear_obj_list[1].light_profile_list}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Linear Light Profiles\n", + "========================================\n", + "\n", + "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", + "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", + "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood)\n", + "given the light profile's other parameters.\n", + "\n", + "Based on the advantages below, we recommended you always use linear light profiles to fit models over standard\n", + "light profiles!\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Each light profile's `intensity` parameter is therefore not a free parameter in the model-fit.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Notes:** This script is identical to `modeling/start_here.py` except that the light profiles are switched to.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Intensities:** The intensities of linear light profiles are not a part of the model parameterization and therefore.\n", + "- **Visualization:** Linear light profiles and objects containing them (e.g.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "- **Max Likelihood Inversion:** As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit`, which contains.\n", + "\n", + "__Advantages__\n", + "\n", + "Each light profile's `intensity` parameter is therefore not a free parameter in the model-fit, reducing the\n", + "dimensionality of non-linear parameter space by the number of light profiles (in this example by 2 dimensions).\n", + "\n", + "This also removes the degeneracies that occur between the `intensity` and other light profile parameters\n", + "(e.g. `effective_radius`, `sersic_index`), which are difficult degeneracies for the non-linear search to map out\n", + "accurately. This produces more reliable lens model results and the fit converges in fewer iterations, speeding up the\n", + "overall analysis.\n", + "\n", + "The inversion has a relatively small computational cost, thus we reduce the model complexity without much slow-down and\n", + "can therefore fit models more reliably and faster!\n", + "\n", + "__Disadvantages__\n", + "\n", + "Although the computation time of the inversion is small, it is not non-negligible. It is approximately 3-4x slower\n", + "than using a standard light profile.\n", + "\n", + "The gains in run times due to the simpler non-linear parameter space therefore are somewhat balanced by the slower\n", + "likelihood calculation.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and negative\n", + "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical.\n", + "\n", + "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is as fast\n", + "as positive-negative solvers. This ensures that all light profile intensities are positive and therefore physical.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a linear `Sersic` bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a linear `Sersic`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook.\n", + "\n", + "__Notes__\n", + "\n", + "This script is identical to `modeling/start_here.py` except that the light profiles are switched to linear light\n", + "profiles." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's light is a linear `Sersic` bulge [6 parameters].\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is a linear `Sersic` [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=19.\n", + "\n", + "Note how both the lens and source galaxies use linear light profiles, meaning that the `intensity` parameter of both\n", + "is no longer a free parameter in the fit.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the light profiles of the lens and source galaxies do not include an `intensity` parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", + "full description).\n", + "\n", + "In the `start_here.py` example 150 live points (`n_live=150`) were used to sample parameter space. For the linear\n", + "light profiles this is reduced to 100, as the simpler parameter space means we need fewer live points to map it out\n", + "accurately. This will lead to faster run times." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"linear_light_profiles\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to\n", + "print the estimated VRAM required by a model.\n", + "\n", + "For each linear light profile in the model extra is used VRAM. For 3-10 linear Sersic light profiles this is a tiny \n", + "amount of VRAM (e.g. < 10MB per batched likelihood). Even for large batch sizes (e.g. over 100) you probably \n", + "will not use enough VRAM to require monitoring.\n", + "\n", + "__Run Time__\n", + "\n", + "For standard light profiles, the log likelihood evaluation time is of order ~0.01 seconds for this dataset.\n", + "\n", + "For linear light profiles, the log likelihood evaluation increases to around ~0.05 seconds per likelihood evaluation.\n", + "This is still fast, but it does mean that the fit may take around five times longer to run.\n", + "\n", + "However, because two free parameters have been removed from the model (the `intensity` of the lens bulge and \n", + "source bulge), the total number of likelihood evaluations will reduce. Furthermore, the simpler parameter space\n", + "likely means that the fit will take less than 10000 likelihood evaluations per free parameter to converge. This is aided further\n", + "by the reduction in `n_live` to 100.\n", + "\n", + "Fits using standard light profiles and linear light profiles therefore take roughly the same time to run. However,\n", + "the simpler parameter space of linear light profiles means that the model-fit is more reliable, less susceptible to\n", + "converging to an incorrect solution and scales better if even more light profiles are included in the model.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that `intensity` parameters are not inferred by the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "The lens and source galaxies appear similar to those in the data, confirming that the `intensity` values inferred by\n", + "the inversion process are accurate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The intensities of linear light profiles are not a part of the model parameterization and therefore are not displayed\n", + "in the `model.results` file.\n", + "\n", + "To extract the `intensity` values of a specific component in the model, we use the `max_log_likelihood_tracer`,\n", + "which has already performed the inversion and therefore the galaxy light profiles have their solved for\n", + "`intensity`'s associated with them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "print(tracer.galaxies[0].bulge.intensity)\n", + "print(tracer.galaxies[1].bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Above, we access these values using the list index entry of each galaxy in the tracer. However, we may not be certain\n", + "of the order of the galaxies in the tracer, and therefore which galaxy index corresponds to the lens and source.\n", + "\n", + "We can therefore use the model composition API to access these values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(tracer.galaxies[0].bulge.intensity)\n", + "print(tracer.galaxies[-1].bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Tracer` contained in the `max_log_likelihood_fit` also has the solved for `intensity` values:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = result.max_log_likelihood_fit\n", + "\n", + "tracer = fit.tracer\n", + "\n", + "print(tracer.galaxies[0].bulge.intensity)\n", + "print(tracer.galaxies[-1].bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization__\n", + "\n", + "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they do not \n", + "have an `intensity` value.\n", + "\n", + "Therefore, the objects created above which replaces all linear light profiles with ordinary light profiles must be\n", + "used for visualization:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Image\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "__Result (Advanced)__\n", + "\n", + "The code belows shows all additional results that can be computed from a `Result` object following a fit with a\n", + "linear light profile.\n", + "\n", + "__Max Likelihood Inversion__\n", + "\n", + "As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit`, which contains the\n", + "`Inversion` object we need." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = result.max_log_likelihood_fit.inversion" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `Inversion` can be used to plot the reconstructed image of specifically all linear light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Linear Objects (Internal Source Code)__\n", + "\n", + "An `Inversion` contains all of the linear objects used to reconstruct the data in its `linear_obj_list`.\n", + "\n", + "This list may include the following objects:\n", + "\n", + " - `LightProfileLinearObjFuncList`: Holds a list of linear light profiles and the functionality used to\n", + " reconstruct data in an inversion. It may contain a single light profile (e.g. `lp_linear.Sersic`) or\n", + " many light profiles combined in a `Basis` (e.g. `lp_basis.Basis`).\n", + "\n", + " - `Mapper`: The linear object used by a `Pixelization` to reconstruct data via an `Inversion`. The `Mapper`\n", + " is specific to the `Pixelization`'s `Mesh` (e.g. a `RectangularMapper` is used for a `RectangularAdaptDensity`\n", + " mesh).\n", + "\n", + "In this example, the model uses one linear `Sersic` for the lens galaxy's bulge and one linear `SersicCore`\n", + "for the source galaxy's bulge. The inversion therefore has two `LightProfileLinearObjFuncList` entries \u2014\n", + "one for each plane (lens at the image-plane grid, source at the ray-traced source-plane grid).\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.linear_obj_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To extract results from an inversion many quantities come in lists or require us to specify the linear object\n", + "we wish to use. Knowing what linear objects are in the `linear_obj_list`, and what indexes they correspond to,\n", + "is therefore important.\n", + "\n", + "The lens-plane entry comes first, followed by the source-plane entry." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"LightProfileLinearObjFuncList (Lens Sersic) = {inversion.linear_obj_list[0]}\"\n", + ")\n", + "print(\n", + " f\"LightProfileLinearObjFuncList (Source SersicCore) = {inversion.linear_obj_list[1]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each `LightProfileLinearObjFuncList` contains a `light_profile_list`. For both entries in this example the\n", + "list has a single light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"Linear Light Profile list (Lens Sersic) = {inversion.linear_obj_list[0].light_profile_list}\"\n", + ")\n", + "print(\n", + " f\"Linear Light Profile list (Source SersicCore) = {inversion.linear_obj_list[1].light_profile_list}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/linear_light_profiles/slam.ipynb b/notebooks/imaging/features/linear_light_profiles/slam.ipynb index b62a72622..3e91944ef 100644 --- a/notebooks/imaging/features/linear_light_profiles/slam.ipynb +++ b/notebooks/imaging/features/linear_light_profiles/slam.ipynb @@ -1,663 +1,700 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Linear Light Profiles: SLaM\n", - "============================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines using linear light profiles\n", - "for the lens galaxy's light.\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that guide before\n", - "working through this example.\n", - "\n", - "The differences from `slam_start_here` are:\n", - "\n", - " - The SOURCE LP PIPELINE uses a `Sersic` linear light profile (`al.lp_linear.Sersic`) for the lens galaxy's\n", - " bulge instead of an MGE.\n", - " - The LIGHT LP PIPELINE uses a `Sersic` linear light profile for the lens galaxy's bulge instead of an MGE.\n", - "\n", - "Linear light profiles solve for the `intensity` analytically via linear algebra, removing it from the non-linear\n", - "parameter space. This reduces the dimensionality of the fit and eliminates intensity-shape degeneracies, resulting\n", - "in more reliable and faster inference.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "- **Linear Light Profiles** (`features/linear_light_profiles`)\n", - " How linear light profiles work and their advantages over standard light profiles.\n", - "\n", - "You can still run the script without fully understanding these guides, but reviewing them later will\n", - "make the structure and choices of the SLaM workflow clearer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light profile\n", - "(`al.lp_linear.Sersic`) instead of an MGE.\n", - "\n", - "The linear `Sersic` profile has fewer free parameters than an MGE, as `intensity` is solved analytically.\n", - "This makes the source LP search faster and more reliable." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = af.Model(al.lp_linear.Sersic)\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light profile\n", - "(`al.lp_linear.Sersic`) instead of an MGE.\n", - "\n", - "Using the linear `Sersic` in the LIGHT LP PIPELINE ensures the lens-light model is consistent with the\n", - "SOURCE LP PIPELINE and reduces the number of free parameters in this fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " )\n", - "\n", - " lens_bulge = af.Model(al.lp_linear.Sersic)\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " bulge = light_result.instance.galaxies.lens.bulge\n", - " disk = light_result.instance.galaxies.lens.disk\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=bulge,\n", - " disk=disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Linear Light Profiles: SLaM\n", + "============================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines using linear light profiles\n", + "for the lens galaxy's light.\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that guide before\n", + "working through this example.\n", + "\n", + "The differences from `slam_start_here` are:\n", + "\n", + " - The SOURCE LP PIPELINE uses a `Sersic` linear light profile (`al.lp_linear.Sersic`) for the lens galaxy's\n", + " bulge instead of an MGE.\n", + " - The LIGHT LP PIPELINE uses a `Sersic` linear light profile for the lens galaxy's bulge instead of an MGE.\n", + "\n", + "Linear light profiles solve for the `intensity` analytically via linear algebra, removing it from the non-linear\n", + "parameter space. This reduces the dimensionality of the fit and eliminates intensity-shape degeneracies, resulting\n", + "in more reliable and faster inference.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "- **Linear Light Profiles** (`features/linear_light_profiles`)\n", + " How linear light profiles work and their advantages over standard light profiles.\n", + "\n", + "You can still run the script without fully understanding these guides, but reviewing them later will\n", + "make the structure and choices of the SLaM workflow clearer." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light profile\n", + "(`al.lp_linear.Sersic`) instead of an MGE.\n", + "\n", + "The linear `Sersic` profile has fewer free parameters than an MGE, as `intensity` is solved analytically.\n", + "This makes the source LP search faster and more reliable." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = af.Model(al.lp_linear.Sersic)\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except the lens galaxy's bulge uses a linear `Sersic` light profile\n", + "(`al.lp_linear.Sersic`) instead of an MGE.\n", + "\n", + "Using the linear `Sersic` in the LIGHT LP PIPELINE ensures the lens-light model is consistent with the\n", + "SOURCE LP PIPELINE and reduces the number of free parameters in this fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " )\n", + "\n", + " lens_bulge = af.Model(al.lp_linear.Sersic)\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " bulge = light_result.instance.galaxies.lens.bulge\n", + " disk = light_result.instance.galaxies.lens.disk\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=bulge,\n", + " disk=disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/multi_gaussian_expansion/fit.ipynb b/notebooks/imaging/features/multi_gaussian_expansion/fit.ipynb index 1afea4852..d23b43ab0 100644 --- a/notebooks/imaging/features/multi_gaussian_expansion/fit.ipynb +++ b/notebooks/imaging/features/multi_gaussian_expansion/fit.ipynb @@ -1,536 +1,573 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Multi Gaussian Expansion\n", - "===========================================\n", - "\n", - "A multi Gaussian expansion (MGE) decomposes the lens light into ~15-100 Gaussians, where the `intensity` of every\n", - "Gaussian is solved for via a linear algebra using a process called an \"inversion\" (see the `linear_light_profiles.py`\n", - "feature for a full description of this).\n", - "\n", - "This script fits a lens light model which uses an MGE consisting of 60 Gaussians. It is fitted to simulated data\n", - "where the lens galaxy's light has asymmetric and irregular features, which fitted poorly by symmetric light\n", - "profiles like the `Sersic`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Symmetric light profiles (e.g.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **MGE Source Galaxy:** The MGE was designed to model the light of lens galaxies, because they are typically elliptical.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Basis:** We first build a `Basis`, which is built from multiple light profiles (in this case, Gaussians).\n", - "- **Gaussians:** The `Basis` is composed of many Gaussians, each with different sizes (the `sigma` value) and.\n", - "- **Linear Light Profiles:** We now show Composing a basis of multiple Gaussians and use them to fit the lens galaxy's light in.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Intensities:** The fit contains the solved for intensity values.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals, because they fail to capture\n", - "irregular and asymmetric morphological of galaxies (e.g. isophotal twists, an ellipticity which varies radially).\n", - "An MGE fully captures these features and can therefore much better represent the emission of complex lens galaxies.\n", - "\n", - "The MGE model can be composed in a way that has fewer non-linear parameters than an elliptical Sersic. In this example,\n", - "two separate groups of Gaussians are used to represent the `bulge` and `disk` of the lens, which in total correspond\n", - "to just N=6 non-linear parameters (a `bulge` and `disk` comprising two linear Sersics has N=10 parameters).\n", - "\n", - "The MGE model parameterization is also composed such that neither the `intensity` parameters or any of the\n", - "parameters controlling the size of the Gaussians (their `sigma` values) are non-linear parameters sampled by Nautilus.\n", - "This removes the most significant degeneracies in parameter space, making the model much more reliable and efficient\n", - "to fit.\n", - "\n", - "Therefore, not only does an MGE fit more complex galaxy morphologies, it does so using fewer non-linear parameters\n", - "in a much simpler non-linear parameter space which has far less significant parameter degeneracies!\n", - "\n", - "__Disadvantages__\n", - "\n", - "To fit an MGE model to the data, the light of the ~15-75 or more Gaussian in the MGE must be evaluated and compared\n", - "to the data. This is slower than evaluating the light of ~2-3 Sersic profiles, producing slower computational run\n", - "times (although the simpler non-linear parameter space will speed up the fit overall).\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and negative\n", - "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical. For an MGE, this produces a positive-negative \"ringing\", where the\n", - "Gaussians alternate between large positive and negative values. This is clearly undesirable and unphysical.\n", - "\n", - "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is as fast\n", - "as positive-negative solvers. This ensures that all light profile intensities are positive and therefore physical.\n", - "\n", - "__MGE Source Galaxy__\n", - "\n", - "The MGE was designed to model the light of lens galaxies, because they are typically elliptical galaxies whose\n", - "morphology is better represented as many Gaussians. Their complex features (e.g. isophotal twists,\n", - "an ellipticity which varies radially) are accurately captured by an MGE.\n", - "\n", - "The morphological features typically seen in source galaxies (e.g. disks, bars, clumps of star formation) are less\n", - "suited to an MGE. The source-plane of many lenses also often have multiple galaxies, whereas the MGE fitted\n", - "in this example assumes a single `centre`.\n", - "\n", - "However, despite these limitations, an MGE turns out to be an extremely powerful way to model the source galaxies\n", - "of strong lenses. This is because, even if it struggles to capture the source's morphology, the simplification of\n", - "non-linear parameter space and removal of degeneracies makes it much easier to obtain a reliable lens model.\n", - "This is driven by the removal of any non-linear parameters which change the size of the source's light profile,\n", - "which are otherwise the most degenerate with the lens's mass model.\n", - "\n", - "The second example in this script therefore uses an MGE source. We strongly recommend you read that example and adopt\n", - "MGE lens light models and source models, instead of the elliptical Sersic profiles, as soon as possible!\n", - "\n", - "To capture the irregular and asymmetric features of the source's morphology, or reconstruct multiple source galaxies,\n", - "we recommend using a pixelized source reconstruction (see `autolens_workspace/modeling/features/pixelization.py`).\n", - "Combining this with an MGE for the len's light can be a very powerful way to model strong lenses!\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's bulge is a super position of 60 `Gaussian` profiles.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a Multi Gaussian Expansion.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"lens_light_asymmetric\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/multi_gaussian_expansion/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Basis__\n", - "\n", - "We first build a `Basis`, which is built from multiple light profiles (in this case, Gaussians). \n", - "\n", - "Below, we make a `Basis` out of 30 elliptical Gaussian light profiles which: \n", - "\n", - " - All share the same centre and elliptical components.\n", - " - The `sigma` size of the Gaussians increases in log10 increments.\n", - "\n", - "Note that any light profile can be used to compose a Basis, but Gaussians are a good choice for lens galaxies\n", - "because they can capture the structure of elliptical galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "\n", - "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", - "\n", - "mask_radius = 3.0\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "# A list of linear light profile Gaussians will be input here, which will then be used to fit the data.\n", - "\n", - "bulge_gaussian_list = []\n", - "\n", - "# Iterate over every Gaussian and create it, with it centered at (0.0\", 0.0\") and assuming spherical symmetry.\n", - "\n", - "for i in range(total_gaussians):\n", - " gaussian = al.lp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=1.0,\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - "\n", - " bulge_gaussian_list.append(gaussian)\n", - "\n", - "# The Basis object groups many light profiles together into a single model component and is used to fit the data.\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Gaussians__\n", - "\n", - "The `Basis` is composed of many Gaussians, each with different sizes (the `sigma` value) and therefore capturing\n", - "emission on different scales.\n", - "\n", - "These Gaussians are visualized below using `subplot_basis_image`, which shows that the Gaussians expand in size as the\n", - "sigma value increases, in log10 increments." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "subplot_basis_image(basis=bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Light Profiles__\n", - "\n", - "We now show Composing a basis of multiple Gaussians and use them to fit the lens galaxy's light in data.\n", - "\n", - "This does not perform a model-fit via a non-linear search, and therefore requires us to manually specify and guess\n", - "suitable parameter values for the shapelets (e.g. the `centre`, `ell_comps`). However, Gaussians are\n", - "very flexible and will give us a decent looking lens fit even if we just guess sensible values\n", - "for each parameter. \n", - "\n", - "The one parameter that is tricky to guess is the `intensity` of each Gaussian. A wide range of positive `intensity` \n", - "values are required to decompose the lens galaxy's light accurately. We certainly cannot obtain a good solution by \n", - "guessing the `intensity` values by eye.\n", - "\n", - "We therefore use linear light profile Gaussians, which determine the optimal value for each Gaussian's `intensity` \n", - "via linear algebra. Linear light profiles are described in the `linear_light_profiles.py` example and you should\n", - "familiarize yourself with this example before using the multi-Gaussian expansion.\n", - "\n", - "We therefore again setup a `Basis` in an analogous fashion to the previous example, but this time we use linear\n", - "Gaussians (via the `lp_linear.linear` module)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 60\n", - "\n", - "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", - "\n", - "mask_radius = 3.0\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "# A list of linear light profile Gaussians will be input here, which will then be used to fit the data.\n", - "\n", - "bulge_gaussian_list = []\n", - "\n", - "# Iterate over every Gaussian and create it, with it centered at (0.0\", 0.0\") and assuming spherical symmetry.\n", - "\n", - "for i in range(total_gaussians):\n", - " gaussian = al.lp_linear.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - "\n", - " bulge_gaussian_list.append(gaussian)\n", - "\n", - "# The Basis object groups many light profiles together into a single model component and is used to fit the data.\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We now illustrate the API for performing an MGE using standard objects like the `Galaxy`, `Tracer`\n", - "and `FitImaging` \n", - "\n", - "Once we have a `Basis`, we can treat it like any other light profile in order to create a `Galaxy` and `Tracer` and \n", - "use it to fit data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, al.Galaxy(redshift=1.0)])\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the MGE does a reasonable job at capturing the appearance of the lens galaxy.\n", - "\n", - "The majority of residuals are due to the lensed source, which was not included in the model. There are faint\n", - "central residuals, which are due to the MGE not being a perfect fit to the lens galaxy's light. \n", - "\n", - "Given that there was no non-linear search to determine the optimal values of the Gaussians and the source galaxy\n", - "was omitted entirely, this is a pretty good fit!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can use `subplot_basis_image` to plot each individual Gaussian in the reconstructed basis.\n", - "\n", - "This plot shows each Gaussian has a unique positive `intensity` that was solved for via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer.galaxies[0].bulge, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The fit contains the solved for intensity values.\n", - "\n", - "These are computed using a fit's `linear_light_profile_intensity_dict`, which maps each linear light profile \n", - "in the model parameterization above to its `intensity`.\n", - "\n", - "The code below shows how to use this dictionary, as an alternative to using the max_log_likelihood quantities above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_bulge = fit.tracer.galaxies[0].bulge\n", - "\n", - "print(\n", - " f\"\\n Intensity of lens galaxy's first Gaussian in bulge = {fit.linear_light_profile_intensity_dict[lens_bulge.profile_list[0]]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the solved \n", - "for `intensity` values is also accessible from a fit.\n", - "\n", - "For example, the first linear light profile of the MGE `bulge` component above printed it solved for intensity value,\n", - "but it was still represented as a linear light profile. \n", - "\n", - "The `tracer` created below instead has a standard light profile with an `intensity` actually set.\n", - "\n", - "The benefit of using a tracer with standard light profiles is it can be visualized, as performed above (linear \n", - "light profiles cannot by default because they do not have `intensity` values)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "print(tracer.galaxies[0].bulge.profile_list[0].intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "A Multi Gaussian Expansion is a powerful tool for modeling the light of galaxies, and offers a compelling method to\n", - "fit complex light profiles with a small number of parameters.\n", - "\n", - "Now you are familiar with MGE modeling, it is recommended you adopt this as your default lens modeling approach. \n", - "However, it may not be suitable for lower resolution data, where the simpler Sersic profiles may be more appropriate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Multi Gaussian Expansion\n", + "===========================================\n", + "\n", + "A multi Gaussian expansion (MGE) decomposes the lens light into ~15-100 Gaussians, where the `intensity` of every\n", + "Gaussian is solved for via a linear algebra using a process called an \"inversion\" (see the `linear_light_profiles.py`\n", + "feature for a full description of this).\n", + "\n", + "This script fits a lens light model which uses an MGE consisting of 60 Gaussians. It is fitted to simulated data\n", + "where the lens galaxy's light has asymmetric and irregular features, which fitted poorly by symmetric light\n", + "profiles like the `Sersic`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Symmetric light profiles (e.g.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **MGE Source Galaxy:** The MGE was designed to model the light of lens galaxies, because they are typically elliptical.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Basis:** We first build a `Basis`, which is built from multiple light profiles (in this case, Gaussians).\n", + "- **Gaussians:** The `Basis` is composed of many Gaussians, each with different sizes (the `sigma` value) and.\n", + "- **Linear Light Profiles:** We now show Composing a basis of multiple Gaussians and use them to fit the lens galaxy's light in.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Intensities:** The fit contains the solved for intensity values.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals, because they fail to capture\n", + "irregular and asymmetric morphological of galaxies (e.g. isophotal twists, an ellipticity which varies radially).\n", + "An MGE fully captures these features and can therefore much better represent the emission of complex lens galaxies.\n", + "\n", + "The MGE model can be composed in a way that has fewer non-linear parameters than an elliptical Sersic. In this example,\n", + "two separate groups of Gaussians are used to represent the `bulge` and `disk` of the lens, which in total correspond\n", + "to just N=6 non-linear parameters (a `bulge` and `disk` comprising two linear Sersics has N=10 parameters).\n", + "\n", + "The MGE model parameterization is also composed such that neither the `intensity` parameters or any of the\n", + "parameters controlling the size of the Gaussians (their `sigma` values) are non-linear parameters sampled by Nautilus.\n", + "This removes the most significant degeneracies in parameter space, making the model much more reliable and efficient\n", + "to fit.\n", + "\n", + "Therefore, not only does an MGE fit more complex galaxy morphologies, it does so using fewer non-linear parameters\n", + "in a much simpler non-linear parameter space which has far less significant parameter degeneracies!\n", + "\n", + "__Disadvantages__\n", + "\n", + "To fit an MGE model to the data, the light of the ~15-75 or more Gaussian in the MGE must be evaluated and compared\n", + "to the data. This is slower than evaluating the light of ~2-3 Sersic profiles, producing slower computational run\n", + "times (although the simpler non-linear parameter space will speed up the fit overall).\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and negative\n", + "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical. For an MGE, this produces a positive-negative \"ringing\", where the\n", + "Gaussians alternate between large positive and negative values. This is clearly undesirable and unphysical.\n", + "\n", + "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is as fast\n", + "as positive-negative solvers. This ensures that all light profile intensities are positive and therefore physical.\n", + "\n", + "__MGE Source Galaxy__\n", + "\n", + "The MGE was designed to model the light of lens galaxies, because they are typically elliptical galaxies whose\n", + "morphology is better represented as many Gaussians. Their complex features (e.g. isophotal twists,\n", + "an ellipticity which varies radially) are accurately captured by an MGE.\n", + "\n", + "The morphological features typically seen in source galaxies (e.g. disks, bars, clumps of star formation) are less\n", + "suited to an MGE. The source-plane of many lenses also often have multiple galaxies, whereas the MGE fitted\n", + "in this example assumes a single `centre`.\n", + "\n", + "However, despite these limitations, an MGE turns out to be an extremely powerful way to model the source galaxies\n", + "of strong lenses. This is because, even if it struggles to capture the source's morphology, the simplification of\n", + "non-linear parameter space and removal of degeneracies makes it much easier to obtain a reliable lens model.\n", + "This is driven by the removal of any non-linear parameters which change the size of the source's light profile,\n", + "which are otherwise the most degenerate with the lens's mass model.\n", + "\n", + "The second example in this script therefore uses an MGE source. We strongly recommend you read that example and adopt\n", + "MGE lens light models and source models, instead of the elliptical Sersic profiles, as soon as possible!\n", + "\n", + "To capture the irregular and asymmetric features of the source's morphology, or reconstruct multiple source galaxies,\n", + "we recommend using a pixelized source reconstruction (see `autolens_workspace/modeling/features/pixelization.py`).\n", + "Combining this with an MGE for the len's light can be a very powerful way to model strong lenses!\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's bulge is a super position of 60 `Gaussian` profiles.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a Multi Gaussian Expansion.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"lens_light_asymmetric\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/multi_gaussian_expansion/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Basis__\n", + "\n", + "We first build a `Basis`, which is built from multiple light profiles (in this case, Gaussians). \n", + "\n", + "Below, we make a `Basis` out of 30 elliptical Gaussian light profiles which: \n", + "\n", + " - All share the same centre and elliptical components.\n", + " - The `sigma` size of the Gaussians increases in log10 increments.\n", + "\n", + "Note that any light profile can be used to compose a Basis, but Gaussians are a good choice for lens galaxies\n", + "because they can capture the structure of elliptical galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "\n", + "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", + "\n", + "mask_radius = 3.0\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "# A list of linear light profile Gaussians will be input here, which will then be used to fit the data.\n", + "\n", + "bulge_gaussian_list = []\n", + "\n", + "# Iterate over every Gaussian and create it, with it centered at (0.0\", 0.0\") and assuming spherical symmetry.\n", + "\n", + "for i in range(total_gaussians):\n", + " gaussian = al.lp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=1.0,\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + "\n", + " bulge_gaussian_list.append(gaussian)\n", + "\n", + "# The Basis object groups many light profiles together into a single model component and is used to fit the data.\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Gaussians__\n", + "\n", + "The `Basis` is composed of many Gaussians, each with different sizes (the `sigma` value) and therefore capturing\n", + "emission on different scales.\n", + "\n", + "These Gaussians are visualized below using `subplot_basis_image`, which shows that the Gaussians expand in size as the\n", + "sigma value increases, in log10 increments." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "subplot_basis_image(basis=bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Light Profiles__\n", + "\n", + "We now show Composing a basis of multiple Gaussians and use them to fit the lens galaxy's light in data.\n", + "\n", + "This does not perform a model-fit via a non-linear search, and therefore requires us to manually specify and guess\n", + "suitable parameter values for the shapelets (e.g. the `centre`, `ell_comps`). However, Gaussians are\n", + "very flexible and will give us a decent looking lens fit even if we just guess sensible values\n", + "for each parameter. \n", + "\n", + "The one parameter that is tricky to guess is the `intensity` of each Gaussian. A wide range of positive `intensity` \n", + "values are required to decompose the lens galaxy's light accurately. We certainly cannot obtain a good solution by \n", + "guessing the `intensity` values by eye.\n", + "\n", + "We therefore use linear light profile Gaussians, which determine the optimal value for each Gaussian's `intensity` \n", + "via linear algebra. Linear light profiles are described in the `linear_light_profiles.py` example and you should\n", + "familiarize yourself with this example before using the multi-Gaussian expansion.\n", + "\n", + "We therefore again setup a `Basis` in an analogous fashion to the previous example, but this time we use linear\n", + "Gaussians (via the `lp_linear.linear` module)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 60\n", + "\n", + "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", + "\n", + "mask_radius = 3.0\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "# A list of linear light profile Gaussians will be input here, which will then be used to fit the data.\n", + "\n", + "bulge_gaussian_list = []\n", + "\n", + "# Iterate over every Gaussian and create it, with it centered at (0.0\", 0.0\") and assuming spherical symmetry.\n", + "\n", + "for i in range(total_gaussians):\n", + " gaussian = al.lp_linear.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + "\n", + " bulge_gaussian_list.append(gaussian)\n", + "\n", + "# The Basis object groups many light profiles together into a single model component and is used to fit the data.\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We now illustrate the API for performing an MGE using standard objects like the `Galaxy`, `Tracer`\n", + "and `FitImaging` \n", + "\n", + "Once we have a `Basis`, we can treat it like any other light profile in order to create a `Galaxy` and `Tracer` and \n", + "use it to fit data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, al.Galaxy(redshift=1.0)])\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the MGE does a reasonable job at capturing the appearance of the lens galaxy.\n", + "\n", + "The majority of residuals are due to the lensed source, which was not included in the model. There are faint\n", + "central residuals, which are due to the MGE not being a perfect fit to the lens galaxy's light. \n", + "\n", + "Given that there was no non-linear search to determine the optimal values of the Gaussians and the source galaxy\n", + "was omitted entirely, this is a pretty good fit!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can use `subplot_basis_image` to plot each individual Gaussian in the reconstructed basis.\n", + "\n", + "This plot shows each Gaussian has a unique positive `intensity` that was solved for via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer.galaxies[0].bulge, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The fit contains the solved for intensity values.\n", + "\n", + "These are computed using a fit's `linear_light_profile_intensity_dict`, which maps each linear light profile \n", + "in the model parameterization above to its `intensity`.\n", + "\n", + "The code below shows how to use this dictionary, as an alternative to using the max_log_likelihood quantities above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_bulge = fit.tracer.galaxies[0].bulge\n", + "\n", + "print(\n", + " f\"\\n Intensity of lens galaxy's first Gaussian in bulge = {fit.linear_light_profile_intensity_dict[lens_bulge.profile_list[0]]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the solved \n", + "for `intensity` values is also accessible from a fit.\n", + "\n", + "For example, the first linear light profile of the MGE `bulge` component above printed it solved for intensity value,\n", + "but it was still represented as a linear light profile. \n", + "\n", + "The `tracer` created below instead has a standard light profile with an `intensity` actually set.\n", + "\n", + "The benefit of using a tracer with standard light profiles is it can be visualized, as performed above (linear \n", + "light profiles cannot by default because they do not have `intensity` values)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "print(tracer.galaxies[0].bulge.profile_list[0].intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "A Multi Gaussian Expansion is a powerful tool for modeling the light of galaxies, and offers a compelling method to\n", + "fit complex light profiles with a small number of parameters.\n", + "\n", + "Now you are familiar with MGE modeling, it is recommended you adopt this as your default lens modeling approach. \n", + "However, it may not be suitable for lower resolution data, where the simpler Sersic profiles may be more appropriate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/multi_gaussian_expansion/likelihood_function.ipynb b/notebooks/imaging/features/multi_gaussian_expansion/likelihood_function.ipynb index 021d20566..fca9a23ef 100644 --- a/notebooks/imaging/features/multi_gaussian_expansion/likelihood_function.ipynb +++ b/notebooks/imaging/features/multi_gaussian_expansion/likelihood_function.ipynb @@ -1,1276 +1,1313 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Multi Gaussian Expansion__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data with\n", - "a multi-Gaussian expansion (MGE), which is a superposition of multiple 2D Gaussian linear light profiles.\n", - "\n", - "You should be familiar with the `log_likelihood_function` of a linear light profile before reading this script,\n", - "which is described in the `log_likelihood_function/imaging/linear_light_profile/likelihood_function.ipynb` notebook.\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers, so that readers can understand the likelihood\n", - " function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", - "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", - "packages are called when the likelihood is evaluated.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** The likelihood function of a multi Gaussian expansion builds on that used for standard light.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Masked Image Grid:** To perform galaxy calculations we used a 2D image-plane grid of (y,x) coordinates, which evaluated.\n", - "- **Multiple Gaussians & Linear Light Profiles:** To use a linear light profile, whose `intensity` is computed via linear algebra, we simply use the.\n", - "- **Basis:** For a multi-Gaussian expansion (and other mdoels where the light profile is a superposition of.\n", - "- **Comparison To Linear Light Profiles Example:** The text below is nearly identical to the `linear_light_profile/likelihood_function.ipynb` example.\n", - "- **LightProfileLinearObjFuncList:** For standard light profiles, we combined our linear light profiles into a single `Galaxies` object.\n", - "- **Combining Matrices:** In the `linear_light_profile/log_likelihood_function.py` example, we used two.\n", - "- **Mapping Matrix:** The `mapping_matrix` is a matrix where each column is an image of each Gaussian linear light.\n", - "- **Image Reconstruction:** Using the reconstructed `intensity` values we can map the reconstruction back to the image plane.\n", - "- **Likelihood Function:** We now quantify the goodness-of-fit of our galaxy model.\n", - "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", - "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", - "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Galaxy Modeling:** To fit a galaxy model to data, the likelihood function illustrated in this tutorial is sampled.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function of a multi Gaussian expansion builds on that used for standard light profiles and\n", - "linear light profiles, therefore you must read the following notebooks before this script:\n", - "\n", - "- `light_profile/likelihood_function.ipynb`.\n", - "- `linear_light_profile/likelihood_function.ipynb`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from scipy.optimize import nnls\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Following the `linear_light_profile/log_likelihood_function.py` script, we load and mask an `Imaging` dataset and\n", - "set oversampling to 1.\n", - "\n", - "This example fits a simulated galaxy where galaxy has an asymmetric light distribution, which cannot be accurately \n", - "fitted with `Sersic` profile and therefore requires a multi-Gaussian expansion to fit accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"lens_light_asymmetric\"\n", - "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/multi_gaussian_expansion/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=masked_dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masked Image Grid__\n", - "\n", - "To perform galaxy calculations we used a 2D image-plane grid of (y,x) coordinates, which evaluated the\n", - "emission of galaxy light profiles created as `LightProfile` objects.\n", - "\n", - "The code below repeats that used in `light_profile/log_likelihood_function.py` to show how this was done." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=4.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "image = tracer.image_2d_from(grid=masked_dataset.grids.lp)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multiple Gaussians & Linear Light Profiles__\n", - "\n", - "To use a linear light profile, whose `intensity` is computed via linear algebra, we simply use the `lp_linear`\n", - "module instead of the `lp` module used throughout other example scripts. \n", - "\n", - "The `intensity` parameter of the light profile is no longer passed into the light profiles created via the\n", - "`lp_linear` module, as it is inferred via linear algebra.\n", - "\n", - "For a multi-Gaussian expansion, we use 30 linear light profile `Gaussian`'s, which is easily achieved by creating a\n", - "list of `Gaussian` objects via a for loop." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "\n", - "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", - "\n", - "mask_radius = 3.0\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "# A list of linear light profile Gaussians will be input here, which will then be used to fit the data.\n", - "\n", - "basis_gaussian_list = []\n", - "\n", - "# Iterate over every Gaussian and create it, with it centered at (0.0\", 0.0\") and assuming spherical symmetry.\n", - "\n", - "for i in range(total_gaussians):\n", - " gaussian = al.lp_linear.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - "\n", - " basis_gaussian_list.append(gaussian)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Basis__\n", - "\n", - "For a multi-Gaussian expansion (and other mdoels where the light profile is a superposition of multiple light profiles),\n", - "the list of linear light profiles is passed to the `Basis` class.\n", - "\n", - "The `Basis` basically stores all these light profiles into a single object such that they can collectively be used to\n", - "perform the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "basis = al.lp_basis.Basis(profile_list=basis_gaussian_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Basis` is composed of many Gaussians, each with different sizes (the `sigma` value) and therefore capturing\n", - "emission on different scales.\n", - "\n", - "These Gaussians are visualized below using `subplot_basis_image`, which shows that the Gaussians expand in size as the\n", - "sigma value increases, in log10 increments.\n", - "\n", - "This figure is a brilliant way to visualize the multi-Gaussian expansion, showing the 30 different Gaussian light\n", - "profiles that will be used perform the expansion on the data.\n", - "\n", - "Below, we will discuss how linear light profiles cannot be visualized (an exception is raised if you try). Therefore\n", - "below we make a separate `Basis` object of `Gaussians` using standard light profiles with input `intensity` values,\n", - "which we can visualize." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "basis_plot_gaussian_list = []\n", - "\n", - "for i in range(total_gaussians):\n", - " gaussian = al.lp.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(0.0, 0.0),\n", - " intensity=1.0,\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - "\n", - " basis_plot_gaussian_list.append(gaussian)\n", - "\n", - "basis_plot = al.lp_basis.Basis(profile_list=basis_plot_gaussian_list)\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", - "\n", - "subplot_basis_image(basis=basis_plot, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Internally in the source code, linear light profiles have an `intensity` parameter, but its value is always set to \n", - "1.0. \n", - "\n", - "This can be seen by printing the intensity of the first two Gaussians in the basis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Basis Internal Intensity Of First Gaussian:\")\n", - "print(basis.light_profile_list[0].intensity)\n", - "\n", - "print(\"Basis Internal Intensity Of Second Gaussian:\")\n", - "print(basis.light_profile_list[1].intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Like standard light profiles, we can compute images of each linear light profile in the basis, but their overall\n", - "normalization is arbitrary given that the internal `intensity` value of 1.0 is used." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_2d_basis_0 = basis.light_profile_list[0].image_2d_from(grid=masked_dataset.grid)\n", - "image_2d_basis_1 = basis.light_profile_list[1].image_2d_from(grid=masked_dataset.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we try and plot a linear light profile using a plotter, an exception is raised.\n", - "\n", - "This is to ensure that a user does not plot and interpret the intensity of a linear light profile, as it is not a\n", - "physical quantity. Plotting only works after a linear light profile has had its `intensity` computed via linear\n", - "algebra.\n", - "\n", - "Uncomment and run the code below to see the exception.\n", - "\n", - "Note that the `subplot_basis_image` used above did not raise an exception, because its intended purpose is to visualize\n", - "the basis light profiles and not the intensity of the light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"This will raise an exception\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now set up a `Tracer` using the MGE for the lens galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, bulge=basis, mass=mass, shear=shear)\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Comparison To Linear Light Profiles Example__\n", - "\n", - "The text below is nearly identical to the `linear_light_profile/likelihood_function.ipynb` example, because the \n", - "linear algebra and likelihood function of a multi-Gaussian expansion is essentially identical to that of a single \n", - "linear light profile.\n", - "\n", - "The key difference between the linear light profile and multi-Gaussian expansion calculation is essentially the\n", - "following:\n", - "\n", - "- The `mapping_matrix`, which for 2 linear light profiles had dimensions `(total_image_pixels, 2)`, now has dimensions\n", - " `(total_image_pixels, 30)`, corresponding to the 30 different Gaussian light profiles.\n", - " \n", - "- Each column of this `mapping_matrix` is the image of each Gaussian light profile, as opposed to the Sersic and\n", - " Exponential light profiles used in the previous example.\n", - "\n", - "- The use of the positive only solver for the reconstruction is more important for an MGE, because MGEs can otherwise\n", - " infer unphysical solutions where the Gaussians alternate between large positive and large negative values.\n", - "\n", - "Other than the above change, the calculation is performed in an identical manner to the linear light profile example,\n", - "with the `data_vector`, `curvature_matrix`, `reconstruction` and `log_likelihood` all computed in the same way\n", - "with the same dimensions. \n", - "\n", - "__LightProfileLinearObjFuncList__\n", - "\n", - "For standard light profiles, we combined our linear light profiles into a single `Galaxies` object. The \n", - "galaxies object computed each individual light profile's image and added them together.\n", - "\n", - "This no longer occurs for linear light profiles, instead linear light profiles are passed into the \n", - "`LightProfileLinearObjFuncList` object, which acts as an interface between the linear light profiles and the\n", - "linear algebra used to compute their intensity via the inversion.\n", - "\n", - "For an MGE, we input the whole `Basis` object into the `LightProfileLinearObjFuncList` object, which\n", - "contains all the Gaussian linmear light profiles.\n", - "\n", - "The quantities used to compute the image, blurring image and blurred image of each light profiles (the\n", - "dataset grid, PSF, etc.) are passed to the `LightProfileLinearObjFuncList` object, because it internally uses these\n", - "to compute each linear light profile image to set up the linear algebra.\n", - "\n", - "For lensing, this means we have to use a different `LightProfileLinearObjFuncList` object for each plane, because\n", - "each plane has its own ray-traced grid of (y,x) coordinates. Below, we set up the first `LightProfileLinearObjFuncList`,\n", - "which uses the image-plane grid and lens galaxy bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lp_linear_func_lens = al.LightProfileLinearObjFuncList(\n", - " grid=masked_dataset.grids.lp,\n", - " blurring_grid=masked_dataset.grids.blurring,\n", - " psf=masked_dataset.psf,\n", - " light_profile_list=basis.light_profile_list,\n", - " regularization=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This has a property `params` which is the number of intensity values that are computed via the inversion,\n", - "which because we have 30 Gaussian linear light profiles is equal to 30.\n", - "\n", - "The `params` defines the dimensions of many of the matrices used in the linear algebra we discuss below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Number of Parameters (Intensity Values) in Linear Algebra:\")\n", - "print(lp_linear_func_lens.params)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Combining Matrices__\n", - "\n", - "In the `linear_light_profile/log_likelihood_function.py` example, we used two `LightProfileLinearObjFuncList` to set\n", - "up the linear algebra for the different planes of the `Tracer`, which we do again below.\n", - "\n", - "In this example the source is a single Sersic linear light profile, but it could easily be an MGE itself." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grids_of_planes_list = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.lp\n", - ")\n", - "traced_blurring_grids_of_planes_list = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")\n", - "\n", - "lp_linear_func_source = al.LightProfileLinearObjFuncList(\n", - " grid=traced_grids_of_planes_list[-1],\n", - " blurring_grid=traced_blurring_grids_of_planes_list[1],\n", - " psf=masked_dataset.psf,\n", - " light_profile_list=[tracer.galaxies[1].bulge],\n", - " regularization=None,\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapping Matrix__\n", - "\n", - "The `mapping_matrix` is a matrix where each column is an image of each Gaussian linear light profiles (assuming its \n", - "intensity is 1.0), not accounting for the PSF convolution.\n", - "\n", - "We combine the `mapping_matrix` of the lens and source plane into a single matrix, which is used to compute the\n", - "`blurred_mapping_matrix` and the `data_vector` below.\n", - "\n", - "It has dimensions `(total_image_pixels, total_linear_light_profiles)` = `(total_image_pixels, 31)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapping_matrix = np.hstack(\n", - " [lp_linear_func_lens.mapping_matrix, lp_linear_func_source.mapping_matrix]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing the first column of the mapping matrix shows the image of the basis light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "basis_image = mapping_matrix[:, 0]\n", - "print(basis_image)\n", - "print(image_2d_basis_0.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `mapping_matrix` shows each light profile image in 1D, which is a bit odd to look at but\n", - "is a good way to think about the linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Blurred Mapping Matrix ($f$)__\n", - "\n", - "The `mapping_matrix` does not account for the blurring of the light profile images by the PSF and therefore \n", - "is not used directly to compute the likelihood.\n", - "\n", - "Instead, we create a `blurred_mapping_matrix` which does account for this blurring. This is computed by \n", - "convolving each light profile image with the PSF.\n", - "\n", - "The `blurred_mapping_matrix` is a matrix analogous to the mapping matrix, but where each column is the image of each\n", - "light profile after it has been blurred by the PSF.\n", - "\n", - "This operation does not change the dimensions of the mapping matrix, meaning the `blurred_mapping_matrix` also has\n", - "dimensions `(total_image_pixels, total_rectangular_pixels)`. \n", - "\n", - "The property is actually called `operated_mapping_matrix_override` for two reasons: \n", - "\n", - "1) The operated signifies that this matrix could have any operation applied to it, it just happens for imaging\n", - " data that this operation is a convolution with the PSF.\n", - "\n", - "2) The `override` signifies that in the source code is changes how the `operated_mapping_matrix` is computed internally. \n", - " This is important if you are looking at the source code, but not important for the description of the likelihood \n", - " function in this guide." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "blurred_mapping_matrix = np.hstack(\n", - " [\n", - " lp_linear_func_lens.operated_mapping_matrix_override,\n", - " lp_linear_func_source.operated_mapping_matrix_override,\n", - " ],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing the first column of the mapping matrix shows the blurred image of the basis light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "basis_image = blurred_mapping_matrix[:, 0]\n", - "print(basis_image)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `mapping_matrix` shows each light profile image in 1D, with a PSF convolution applied." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) (hereafter WD03) introduce the linear inversion formalism \n", - "used to compute the intensity values of the linear light profiles. In WD03, the science case is centred around strong\n", - "gravitational lensing and the galaxy is reconstructed on a rectangular grid of pixels, as opposed to linear light \n", - "profiles.\n", - "\n", - "However, the mathematics of the WD03 linear inversion formalism is the same as that used here, therefore this guide \n", - "describes which quantities in the linear inversion formalism map to the equations given in WD03. The pixelized \n", - "reconstruction methods, available in the code but described in the `pixelization` likelihood function guide, \n", - "also follow the WD03 formalism.\n", - "\n", - "The `blurred_mapping_matrix` is denoted $f_{ij}$ where $i$ maps over all $I$ linear light profiles and $j$ maps \n", - "over all $J$ image pixels. \n", - "\n", - "For example: \n", - "\n", - " - $f_{0, 1} = 0.3$ indicates that image-pixel $2$ maps to linear light profile $1$ with an intensity in that image \n", - " pixel of $0.3$ after PSF convolution.\n", - "\n", - "The indexing of the `mapping_matrix` is reversed compared to the notation of WD03 (e.g. image pixels\n", - "are the first entry of `mapping_matrix` whereas for $f$ they are the second index)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"Mapping between image pixel 0 and Gaussian linear light profile pixel 1 = {mapping_matrix[0, 1]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Data Vector (D)__\n", - "\n", - "To solve for the linear light profile intensities we now pose the problem as a linear inversion.\n", - "\n", - "This requires us to convert the `blurred_mapping_matrix` and our `data` and `noise map` into matrices of certain \n", - "dimensions. \n", - "\n", - "The `data_vector`, $D$, is the first matrix and it has dimensions `(total_linear_light_profiles,)`.\n", - "\n", - "In WD03 (https://arxiv.org/abs/astro-ph/0302587) the data vector is given by: \n", - "\n", - " $\\vec{D}_{i} = \\sum_{\\rm j=1}^{J}f_{ij}(d_{j} - b_{j})/\\sigma_{j}^2 \\, \\, .$\n", - "\n", - "Where:\n", - "\n", - " - $d_{\\rm j}$ are the image-pixel data flux values.\n", - " - $b_{\\rm j}$ are the image values of all standard light profiles (therefore $d_{\\rm j} - b_{\\rm j}$ is \n", - " the data minus any standard light profiles).\n", - " - $\\sigma{\\rm _j}^2$ are the statistical uncertainties of each image-pixel value.\n", - "\n", - "$i$ maps over all $I$ linear light profiles and $j$ maps over all $J$ image pixels. \n", - "\n", - "This equation highlights a first aspect of linear inversions, if we are combining standard light profiles (which\n", - "have an input `intensity` value) with linear light profiles, the inversion is performed on the data minus\n", - "the standard light profile images. In this example, we have no standard light profiles and therefore the data\n", - "vector uses the data directly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_vector = al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(\n", - " blurred_mapping_matrix=blurred_mapping_matrix,\n", - " image=np.array(masked_dataset.data),\n", - " noise_map=np.array(masked_dataset.noise_map),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$D$'s meaning is a bit abstract, it essentially weights each linear light profile's `intensity` based on how it\n", - "maps to the data, so that the linear algebra can compute the `intensity` values that best-fit the data.\n", - "\n", - "We can plot $D$ as a column vector:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dimensions of $D$ are the number of linear light profiles, which in this case is 2." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Data Vector:\")\n", - "print(data_vector)\n", - "print(data_vector.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Matrix (F)__\n", - "\n", - "The `curvature_matrix` $F$ is the second matrix and it has \n", - "dimensions `(total_linear_light_profiles, total_linear_light_profiles)`.\n", - "\n", - "In WD03 (https://arxiv.org/abs/astro-ph/0302587) the curvature matrix is a 2D matrix given by:\n", - "\n", - " ${F}_{ik} = \\sum_{\\rm j=1}^{J}f_{ij}f_{kj}/\\sigma_{j}^2 \\, \\, .$\n", - "\n", - "NOTE: this notation implicitly assumes a summation over $K$, where $k$ runs over all linear light profile indexes $K$.\n", - "\n", - "Note how summation over $J$ runs over $f$ twice, such that every entry of $F$ is the sum of the multiplication\n", - "between all values in every two columns of $f$.\n", - "\n", - "For example, $F_{0,1}$ is the sum of every blurred image pixels values in $f$ of linear light profile 0 multiplied by\n", - "every blurred image pixel value of linear light profile 1.\n", - "\n", - "$F$'s meaning is also a bit abstract, but it essentially quantifies how much each linear light profile's image\n", - "overlaps with every other linear light profile's image, weighted by the noise in the data. This is what combined with\n", - "the `data_vector` allows the inversion to compute the `intensity` values that best-fit the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix, noise_map=masked_dataset.noise_map\n", - ")\n", - "\n", - "plt.imshow(curvature_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction (Positive-Negative)__\n", - "\n", - "The following chi-squared is minimized when we perform the inversion and reconstruct the galaxy:\n", - "\n", - "$\\chi^2 = \\sum_{\\rm j=1}^{J} \\bigg[ \\frac{(\\sum_{\\rm i=1}^{I} s_{i} f_{ij}) + b_{j} - d_{j}}{\\sigma_{j}} \\bigg]$\n", - "\n", - "Where $s$ is the `intensity` values in all $I$ linear light profile images.\n", - "\n", - "The solution for $s$ is therefore given by (equation 5 WD03):\n", - "\n", - " $s = F^{-1} D$\n", - "\n", - "We can compute this using NumPy linear algebra and the `solve` function.\n", - "\n", - "However, this function allows for the solved `intensity` values to be negative, which are unphysical values for\n", - "describing the light profile of a galaxy. \n", - "\n", - "For a multi-Gaussian expansion, it is common for the inferred solution to contain negative `intensity` values. A common\n", - "solution is one where the Gaussians alternate between large positive and large negative values, creating an almost\n", - "\"ringing\" effect in the reconstruction. This is a very unphysical solution and one we want to avoid.\n", - "\n", - "We are able to illustrate this now, first by solving the linear algebra and then printing the `intensity` values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = np.linalg.solve(curvature_matrix, data_vector)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `reconstruction` is a 1D vector of length equal to the number of Gaussian linear light profiles, which in this case \n", - "is 30.\n", - "\n", - "Each value represents the solved for `intensity` of the Gaussian linear light profile.\n", - "\n", - "In this example, the values alternate between positive and negative, indicating a solution that is not physical\n", - "and one we must avoid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Reconstruction (S) of Linear Light Profiles Intensity:\")\n", - "print(reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction (Positive Only)__\n", - "\n", - "The linear algebra can be solved for with the constraint that all solutions, and therefore all `intensity` values,\n", - "are positive. \n", - "\n", - "This could be achieved by using the `scipy` `nnls` non-negative least squares solver.\n", - "\n", - "The nnls poses the problem slightly different than the code above. It solves for the `intensity` values in an\n", - "iterative manner meaning that it is slower. It does not use `data_vector` $D$ and `curvature_matrix` $F$ but instead\n", - "works directly with the `blurred_mapping_matrix` $f$ and the data and noise-map.\n", - "\n", - "The `nnls` function is therefore computationally slow, especially for cases where there are many linear light profiles \n", - "or even more complex linear inversions like a pixelized reconstruction.\n", - "\n", - "The source code therefore uses a \"fast nnls\" algorithm, which is an adaptation of the algorithm found at\n", - "this URL: https://github.com/jvendrow/fnnls\n", - "\n", - "Unlike the scipy nnls function, the fnnls method uses the `data_vector` $D$ and `curvature_matrix` $F$ to solve for\n", - "the `intensity` values. This provides it with additional information about the linear algebra problem, which is\n", - "why it is faster.\n", - "\n", - "The function `reconstruction_positive_only_from` uses the `fnnls` algorithm to compute the `intensity` values\n", - "of the linear light profiles, ensuring they are positive.\n", - "\n", - "However, the code below by itself actually produces a `LinAlgError` because the `curvature_matrix` is singular. Uncomment\n", - "the code below to see this error." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", - "# data_vector=data_vector,\n", - "# curvature_reg_matrix=curvature_matrix, # ignore _reg_ tag in this guide\n", - "# )\n", - "#\n", - "# print(reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To make the `curvature_matrix` non-singular, we simply add small numerical values to its diagonal elements. \n", - "\n", - "This is effectively add a small degree of \"zeroth order\" regularization to the inversion, which is sufficient to make\n", - "the matrix non-singular and ensure the inversion can be performed.\n", - "\n", - "There are a variety of ways to regularize the inversion, and these can be manually input into the `Basis` object.\n", - "However, for a multi-Gaussian expansion, testing has shown that adding a small degree of zeroth order regularization\n", - "in conjunction with a positive-only solution is sufficient to ensure the inversion is robust for all reasonable\n", - "science cases.\n", - "\n", - "In practise, the code only adds these small numerical values to the diagonal of the curvature matrix for elements\n", - "which have no other regularization applied to them. Therefore, in the function call below we input \n", - "`no_regularization_index_list=range(30)`, which tells the function to add small numerical values to all 30\n", - "diagonal values corresponding to the 30 Gaussian linear light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix,\n", - " noise_map=masked_dataset.noise_map,\n", - " add_to_curvature_diag=True,\n", - " no_regularization_index_list=list(range(30)),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `reconstruction` can now be computed successfully without a linear algebra error." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", - " data_vector=data_vector,\n", - " curvature_reg_matrix=curvature_matrix, # ignore _reg_ tag in this guide\n", - ")\n", - "\n", - "print(reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image Reconstruction__\n", - "\n", - "Using the reconstructed `intensity` values we can map the reconstruction back to the image plane (via \n", - "the `blurred mapping_matrix`) and produce a reconstruction of the image data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_operated_data = (\n", - " al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix, reconstruction=reconstruction\n", - " )\n", - ")\n", - "\n", - "mapped_reconstructed_operated_data = al.Array2D(\n", - " values=mapped_reconstructed_operated_data, mask=mask\n", - ")\n", - "\n", - "aplt.plot_array(array=mapped_reconstructed_operated_data, title=\"\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our galaxy model.\n", - "\n", - "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", - "\n", - "The likelihood function for parametric galaxy modeling, even if linear light profiles are used, consists of two terms:\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "We now explain what each of these terms mean.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", - "\n", - " - `model_data` = `convolved_image_2d`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", - "\n", - "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", - "for, corresponding to a fit with a lower likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image = mapped_reconstructed_operated_data\n", - "\n", - "residual_map = masked_dataset.data - model_image\n", - "normalized_residual_map = residual_map / masked_dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(chi_squared)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", - "\n", - "aplt.plot_array(array=chi_squared_map, title=\"\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", - "\n", - "The final term in the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", - "of the log of every noise-map value squared. \n", - "\n", - "Given the `noise_map` is fixed, this term does not change during the galaxy modeling process and has no impact on the \n", - "model we infer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms computed above using\n", - "the likelihood function defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(figure_of_merit)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " basis=basis,\n", - ")\n", - "\n", - "galaxies = al.Galaxies(galaxies=[galaxy])\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=masked_dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_border_relocator=True),\n", - ")\n", - "fit_log_evidence = fit.log_evidence\n", - "print(fit_log_evidence)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit contains an `Inversion` object, which handles all the linear algebra we have covered in this script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.inversion)\n", - "print(fit.inversion.data_vector)\n", - "print(fit.inversion.curvature_matrix)\n", - "print(fit.inversion.reconstruction)\n", - "print(fit.inversion.mapped_reconstructed_operated_data)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Inversion` object can be computed from a tracer and a dataset, by passing them to the `TracerToInversion` object.\n", - "\n", - "This objects handles a lot of extra functionality that we have not covered in this script, such as:\n", - "\n", - "- Separating out the linear light profiles from the standard light profiles.\n", - "- Separating out objects which reconstruct the galaxy using a pixelized reconstruction, which are passed into\n", - " the `Inversion` object as well." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(\n", - " tracer=tracer,\n", - " dataset=masked_dataset,\n", - ")\n", - "\n", - "inversion = tracer_to_inversion.inversion\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Modeling__\n", - "\n", - "To fit a galaxy model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "For an MGE, the reduced number of free parameters (e.g. the `intensity` values are solved for\n", - "via linear algebra and not a dimension of the non-linear parameter space) means that the sampler converges in fewer\n", - "iterations and is less likely to infer a local maximum. \n", - "\n", - "Furthermore, the size of the lens galaxy, controlled by the `sigma` values of the Gaussians, are also all fixed\n", - "and not non-linear free parameters. This further simplifies the non-linear parameter space.\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the multi Gaussian expansion likelihood function, which uses \n", - "many 2D Gaussians to fit the galaxy light and solve for the `intensity` values via linear algebra.\n", - "\n", - "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", - "are described in additional notebooks found in the `guides` package:\n", - "\n", - " - `over_sampling`: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", - " ray-traced to the source-plane and used to evaluate the light profile more accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Multi Gaussian Expansion__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data with\n", + "a multi-Gaussian expansion (MGE), which is a superposition of multiple 2D Gaussian linear light profiles.\n", + "\n", + "You should be familiar with the `log_likelihood_function` of a linear light profile before reading this script,\n", + "which is described in the `log_likelihood_function/imaging/linear_light_profile/likelihood_function.ipynb` notebook.\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers, so that readers can understand the likelihood\n", + " function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", + "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", + "packages are called when the likelihood is evaluated.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** The likelihood function of a multi Gaussian expansion builds on that used for standard light.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Masked Image Grid:** To perform galaxy calculations we used a 2D image-plane grid of (y,x) coordinates, which evaluated.\n", + "- **Multiple Gaussians & Linear Light Profiles:** To use a linear light profile, whose `intensity` is computed via linear algebra, we simply use the.\n", + "- **Basis:** For a multi-Gaussian expansion (and other mdoels where the light profile is a superposition of.\n", + "- **Comparison To Linear Light Profiles Example:** The text below is nearly identical to the `linear_light_profile/likelihood_function.ipynb` example.\n", + "- **LightProfileLinearObjFuncList:** For standard light profiles, we combined our linear light profiles into a single `Galaxies` object.\n", + "- **Combining Matrices:** In the `linear_light_profile/log_likelihood_function.py` example, we used two.\n", + "- **Mapping Matrix:** The `mapping_matrix` is a matrix where each column is an image of each Gaussian linear light.\n", + "- **Image Reconstruction:** Using the reconstructed `intensity` values we can map the reconstruction back to the image plane.\n", + "- **Likelihood Function:** We now quantify the goodness-of-fit of our galaxy model.\n", + "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", + "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", + "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Galaxy Modeling:** To fit a galaxy model to data, the likelihood function illustrated in this tutorial is sampled.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function of a multi Gaussian expansion builds on that used for standard light profiles and\n", + "linear light profiles, therefore you must read the following notebooks before this script:\n", + "\n", + "- `light_profile/likelihood_function.ipynb`.\n", + "- `linear_light_profile/likelihood_function.ipynb`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from scipy.optimize import nnls\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Following the `linear_light_profile/log_likelihood_function.py` script, we load and mask an `Imaging` dataset and\n", + "set oversampling to 1.\n", + "\n", + "This example fits a simulated galaxy where galaxy has an asymmetric light distribution, which cannot be accurately \n", + "fitted with `Sersic` profile and therefore requires a multi-Gaussian expansion to fit accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"lens_light_asymmetric\"\n", + "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/multi_gaussian_expansion/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=masked_dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masked Image Grid__\n", + "\n", + "To perform galaxy calculations we used a 2D image-plane grid of (y,x) coordinates, which evaluated the\n", + "emission of galaxy light profiles created as `LightProfile` objects.\n", + "\n", + "The code below repeats that used in `light_profile/log_likelihood_function.py` to show how this was done." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=4.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "image = tracer.image_2d_from(grid=masked_dataset.grids.lp)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multiple Gaussians & Linear Light Profiles__\n", + "\n", + "To use a linear light profile, whose `intensity` is computed via linear algebra, we simply use the `lp_linear`\n", + "module instead of the `lp` module used throughout other example scripts. \n", + "\n", + "The `intensity` parameter of the light profile is no longer passed into the light profiles created via the\n", + "`lp_linear` module, as it is inferred via linear algebra.\n", + "\n", + "For a multi-Gaussian expansion, we use 30 linear light profile `Gaussian`'s, which is easily achieved by creating a\n", + "list of `Gaussian` objects via a for loop." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "\n", + "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", + "\n", + "mask_radius = 3.0\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "# A list of linear light profile Gaussians will be input here, which will then be used to fit the data.\n", + "\n", + "basis_gaussian_list = []\n", + "\n", + "# Iterate over every Gaussian and create it, with it centered at (0.0\", 0.0\") and assuming spherical symmetry.\n", + "\n", + "for i in range(total_gaussians):\n", + " gaussian = al.lp_linear.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + "\n", + " basis_gaussian_list.append(gaussian)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Basis__\n", + "\n", + "For a multi-Gaussian expansion (and other mdoels where the light profile is a superposition of multiple light profiles),\n", + "the list of linear light profiles is passed to the `Basis` class.\n", + "\n", + "The `Basis` basically stores all these light profiles into a single object such that they can collectively be used to\n", + "perform the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "basis = al.lp_basis.Basis(profile_list=basis_gaussian_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Basis` is composed of many Gaussians, each with different sizes (the `sigma` value) and therefore capturing\n", + "emission on different scales.\n", + "\n", + "These Gaussians are visualized below using `subplot_basis_image`, which shows that the Gaussians expand in size as the\n", + "sigma value increases, in log10 increments.\n", + "\n", + "This figure is a brilliant way to visualize the multi-Gaussian expansion, showing the 30 different Gaussian light\n", + "profiles that will be used perform the expansion on the data.\n", + "\n", + "Below, we will discuss how linear light profiles cannot be visualized (an exception is raised if you try). Therefore\n", + "below we make a separate `Basis` object of `Gaussians` using standard light profiles with input `intensity` values,\n", + "which we can visualize." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "basis_plot_gaussian_list = []\n", + "\n", + "for i in range(total_gaussians):\n", + " gaussian = al.lp.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(0.0, 0.0),\n", + " intensity=1.0,\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + "\n", + " basis_plot_gaussian_list.append(gaussian)\n", + "\n", + "basis_plot = al.lp_basis.Basis(profile_list=basis_plot_gaussian_list)\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)\n", + "\n", + "subplot_basis_image(basis=basis_plot, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Internally in the source code, linear light profiles have an `intensity` parameter, but its value is always set to \n", + "1.0. \n", + "\n", + "This can be seen by printing the intensity of the first two Gaussians in the basis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Basis Internal Intensity Of First Gaussian:\")\n", + "print(basis.light_profile_list[0].intensity)\n", + "\n", + "print(\"Basis Internal Intensity Of Second Gaussian:\")\n", + "print(basis.light_profile_list[1].intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Like standard light profiles, we can compute images of each linear light profile in the basis, but their overall\n", + "normalization is arbitrary given that the internal `intensity` value of 1.0 is used." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_2d_basis_0 = basis.light_profile_list[0].image_2d_from(grid=masked_dataset.grid)\n", + "image_2d_basis_1 = basis.light_profile_list[1].image_2d_from(grid=masked_dataset.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we try and plot a linear light profile using a plotter, an exception is raised.\n", + "\n", + "This is to ensure that a user does not plot and interpret the intensity of a linear light profile, as it is not a\n", + "physical quantity. Plotting only works after a linear light profile has had its `intensity` computed via linear\n", + "algebra.\n", + "\n", + "Uncomment and run the code below to see the exception.\n", + "\n", + "Note that the `subplot_basis_image` used above did not raise an exception, because its intended purpose is to visualize\n", + "the basis light profiles and not the intensity of the light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"This will raise an exception\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now set up a `Tracer` using the MGE for the lens galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, bulge=basis, mass=mass, shear=shear)\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Comparison To Linear Light Profiles Example__\n", + "\n", + "The text below is nearly identical to the `linear_light_profile/likelihood_function.ipynb` example, because the \n", + "linear algebra and likelihood function of a multi-Gaussian expansion is essentially identical to that of a single \n", + "linear light profile.\n", + "\n", + "The key difference between the linear light profile and multi-Gaussian expansion calculation is essentially the\n", + "following:\n", + "\n", + "- The `mapping_matrix`, which for 2 linear light profiles had dimensions `(total_image_pixels, 2)`, now has dimensions\n", + " `(total_image_pixels, 30)`, corresponding to the 30 different Gaussian light profiles.\n", + " \n", + "- Each column of this `mapping_matrix` is the image of each Gaussian light profile, as opposed to the Sersic and\n", + " Exponential light profiles used in the previous example.\n", + "\n", + "- The use of the positive only solver for the reconstruction is more important for an MGE, because MGEs can otherwise\n", + " infer unphysical solutions where the Gaussians alternate between large positive and large negative values.\n", + "\n", + "Other than the above change, the calculation is performed in an identical manner to the linear light profile example,\n", + "with the `data_vector`, `curvature_matrix`, `reconstruction` and `log_likelihood` all computed in the same way\n", + "with the same dimensions. \n", + "\n", + "__LightProfileLinearObjFuncList__\n", + "\n", + "For standard light profiles, we combined our linear light profiles into a single `Galaxies` object. The \n", + "galaxies object computed each individual light profile's image and added them together.\n", + "\n", + "This no longer occurs for linear light profiles, instead linear light profiles are passed into the \n", + "`LightProfileLinearObjFuncList` object, which acts as an interface between the linear light profiles and the\n", + "linear algebra used to compute their intensity via the inversion.\n", + "\n", + "For an MGE, we input the whole `Basis` object into the `LightProfileLinearObjFuncList` object, which\n", + "contains all the Gaussian linmear light profiles.\n", + "\n", + "The quantities used to compute the image, blurring image and blurred image of each light profiles (the\n", + "dataset grid, PSF, etc.) are passed to the `LightProfileLinearObjFuncList` object, because it internally uses these\n", + "to compute each linear light profile image to set up the linear algebra.\n", + "\n", + "For lensing, this means we have to use a different `LightProfileLinearObjFuncList` object for each plane, because\n", + "each plane has its own ray-traced grid of (y,x) coordinates. Below, we set up the first `LightProfileLinearObjFuncList`,\n", + "which uses the image-plane grid and lens galaxy bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lp_linear_func_lens = al.LightProfileLinearObjFuncList(\n", + " grid=masked_dataset.grids.lp,\n", + " blurring_grid=masked_dataset.grids.blurring,\n", + " psf=masked_dataset.psf,\n", + " light_profile_list=basis.light_profile_list,\n", + " regularization=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This has a property `params` which is the number of intensity values that are computed via the inversion,\n", + "which because we have 30 Gaussian linear light profiles is equal to 30.\n", + "\n", + "The `params` defines the dimensions of many of the matrices used in the linear algebra we discuss below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Number of Parameters (Intensity Values) in Linear Algebra:\")\n", + "print(lp_linear_func_lens.params)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Combining Matrices__\n", + "\n", + "In the `linear_light_profile/log_likelihood_function.py` example, we used two `LightProfileLinearObjFuncList` to set\n", + "up the linear algebra for the different planes of the `Tracer`, which we do again below.\n", + "\n", + "In this example the source is a single Sersic linear light profile, but it could easily be an MGE itself." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grids_of_planes_list = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.lp\n", + ")\n", + "traced_blurring_grids_of_planes_list = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")\n", + "\n", + "lp_linear_func_source = al.LightProfileLinearObjFuncList(\n", + " grid=traced_grids_of_planes_list[-1],\n", + " blurring_grid=traced_blurring_grids_of_planes_list[1],\n", + " psf=masked_dataset.psf,\n", + " light_profile_list=[tracer.galaxies[1].bulge],\n", + " regularization=None,\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapping Matrix__\n", + "\n", + "The `mapping_matrix` is a matrix where each column is an image of each Gaussian linear light profiles (assuming its \n", + "intensity is 1.0), not accounting for the PSF convolution.\n", + "\n", + "We combine the `mapping_matrix` of the lens and source plane into a single matrix, which is used to compute the\n", + "`blurred_mapping_matrix` and the `data_vector` below.\n", + "\n", + "It has dimensions `(total_image_pixels, total_linear_light_profiles)` = `(total_image_pixels, 31)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapping_matrix = np.hstack(\n", + " [lp_linear_func_lens.mapping_matrix, lp_linear_func_source.mapping_matrix]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing the first column of the mapping matrix shows the image of the basis light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "basis_image = mapping_matrix[:, 0]\n", + "print(basis_image)\n", + "print(image_2d_basis_0.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `mapping_matrix` shows each light profile image in 1D, which is a bit odd to look at but\n", + "is a good way to think about the linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Blurred Mapping Matrix ($f$)__\n", + "\n", + "The `mapping_matrix` does not account for the blurring of the light profile images by the PSF and therefore \n", + "is not used directly to compute the likelihood.\n", + "\n", + "Instead, we create a `blurred_mapping_matrix` which does account for this blurring. This is computed by \n", + "convolving each light profile image with the PSF.\n", + "\n", + "The `blurred_mapping_matrix` is a matrix analogous to the mapping matrix, but where each column is the image of each\n", + "light profile after it has been blurred by the PSF.\n", + "\n", + "This operation does not change the dimensions of the mapping matrix, meaning the `blurred_mapping_matrix` also has\n", + "dimensions `(total_image_pixels, total_rectangular_pixels)`. \n", + "\n", + "The property is actually called `operated_mapping_matrix_override` for two reasons: \n", + "\n", + "1) The operated signifies that this matrix could have any operation applied to it, it just happens for imaging\n", + " data that this operation is a convolution with the PSF.\n", + "\n", + "2) The `override` signifies that in the source code is changes how the `operated_mapping_matrix` is computed internally. \n", + " This is important if you are looking at the source code, but not important for the description of the likelihood \n", + " function in this guide." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "blurred_mapping_matrix = np.hstack(\n", + " [\n", + " lp_linear_func_lens.operated_mapping_matrix_override,\n", + " lp_linear_func_source.operated_mapping_matrix_override,\n", + " ],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing the first column of the mapping matrix shows the blurred image of the basis light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "basis_image = blurred_mapping_matrix[:, 0]\n", + "print(basis_image)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `mapping_matrix` shows each light profile image in 1D, with a PSF convolution applied." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) (hereafter WD03) introduce the linear inversion formalism \n", + "used to compute the intensity values of the linear light profiles. In WD03, the science case is centred around strong\n", + "gravitational lensing and the galaxy is reconstructed on a rectangular grid of pixels, as opposed to linear light \n", + "profiles.\n", + "\n", + "However, the mathematics of the WD03 linear inversion formalism is the same as that used here, therefore this guide \n", + "describes which quantities in the linear inversion formalism map to the equations given in WD03. The pixelized \n", + "reconstruction methods, available in the code but described in the `pixelization` likelihood function guide, \n", + "also follow the WD03 formalism.\n", + "\n", + "The `blurred_mapping_matrix` is denoted $f_{ij}$ where $i$ maps over all $I$ linear light profiles and $j$ maps \n", + "over all $J$ image pixels. \n", + "\n", + "For example: \n", + "\n", + " - $f_{0, 1} = 0.3$ indicates that image-pixel $2$ maps to linear light profile $1$ with an intensity in that image \n", + " pixel of $0.3$ after PSF convolution.\n", + "\n", + "The indexing of the `mapping_matrix` is reversed compared to the notation of WD03 (e.g. image pixels\n", + "are the first entry of `mapping_matrix` whereas for $f$ they are the second index)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"Mapping between image pixel 0 and Gaussian linear light profile pixel 1 = {mapping_matrix[0, 1]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Data Vector (D)__\n", + "\n", + "To solve for the linear light profile intensities we now pose the problem as a linear inversion.\n", + "\n", + "This requires us to convert the `blurred_mapping_matrix` and our `data` and `noise map` into matrices of certain \n", + "dimensions. \n", + "\n", + "The `data_vector`, $D$, is the first matrix and it has dimensions `(total_linear_light_profiles,)`.\n", + "\n", + "In WD03 (https://arxiv.org/abs/astro-ph/0302587) the data vector is given by: \n", + "\n", + " $\\vec{D}_{i} = \\sum_{\\rm j=1}^{J}f_{ij}(d_{j} - b_{j})/\\sigma_{j}^2 \\, \\, .$\n", + "\n", + "Where:\n", + "\n", + " - $d_{\\rm j}$ are the image-pixel data flux values.\n", + " - $b_{\\rm j}$ are the image values of all standard light profiles (therefore $d_{\\rm j} - b_{\\rm j}$ is \n", + " the data minus any standard light profiles).\n", + " - $\\sigma{\\rm _j}^2$ are the statistical uncertainties of each image-pixel value.\n", + "\n", + "$i$ maps over all $I$ linear light profiles and $j$ maps over all $J$ image pixels. \n", + "\n", + "This equation highlights a first aspect of linear inversions, if we are combining standard light profiles (which\n", + "have an input `intensity` value) with linear light profiles, the inversion is performed on the data minus\n", + "the standard light profile images. In this example, we have no standard light profiles and therefore the data\n", + "vector uses the data directly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_vector = al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(\n", + " blurred_mapping_matrix=blurred_mapping_matrix,\n", + " image=np.array(masked_dataset.data),\n", + " noise_map=np.array(masked_dataset.noise_map),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$D$'s meaning is a bit abstract, it essentially weights each linear light profile's `intensity` based on how it\n", + "maps to the data, so that the linear algebra can compute the `intensity` values that best-fit the data.\n", + "\n", + "We can plot $D$ as a column vector:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dimensions of $D$ are the number of linear light profiles, which in this case is 2." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Data Vector:\")\n", + "print(data_vector)\n", + "print(data_vector.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Matrix (F)__\n", + "\n", + "The `curvature_matrix` $F$ is the second matrix and it has \n", + "dimensions `(total_linear_light_profiles, total_linear_light_profiles)`.\n", + "\n", + "In WD03 (https://arxiv.org/abs/astro-ph/0302587) the curvature matrix is a 2D matrix given by:\n", + "\n", + " ${F}_{ik} = \\sum_{\\rm j=1}^{J}f_{ij}f_{kj}/\\sigma_{j}^2 \\, \\, .$\n", + "\n", + "NOTE: this notation implicitly assumes a summation over $K$, where $k$ runs over all linear light profile indexes $K$.\n", + "\n", + "Note how summation over $J$ runs over $f$ twice, such that every entry of $F$ is the sum of the multiplication\n", + "between all values in every two columns of $f$.\n", + "\n", + "For example, $F_{0,1}$ is the sum of every blurred image pixels values in $f$ of linear light profile 0 multiplied by\n", + "every blurred image pixel value of linear light profile 1.\n", + "\n", + "$F$'s meaning is also a bit abstract, but it essentially quantifies how much each linear light profile's image\n", + "overlaps with every other linear light profile's image, weighted by the noise in the data. This is what combined with\n", + "the `data_vector` allows the inversion to compute the `intensity` values that best-fit the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix, noise_map=masked_dataset.noise_map\n", + ")\n", + "\n", + "plt.imshow(curvature_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction (Positive-Negative)__\n", + "\n", + "The following chi-squared is minimized when we perform the inversion and reconstruct the galaxy:\n", + "\n", + "$\\chi^2 = \\sum_{\\rm j=1}^{J} \\bigg[ \\frac{(\\sum_{\\rm i=1}^{I} s_{i} f_{ij}) + b_{j} - d_{j}}{\\sigma_{j}} \\bigg]$\n", + "\n", + "Where $s$ is the `intensity` values in all $I$ linear light profile images.\n", + "\n", + "The solution for $s$ is therefore given by (equation 5 WD03):\n", + "\n", + " $s = F^{-1} D$\n", + "\n", + "We can compute this using NumPy linear algebra and the `solve` function.\n", + "\n", + "However, this function allows for the solved `intensity` values to be negative, which are unphysical values for\n", + "describing the light profile of a galaxy. \n", + "\n", + "For a multi-Gaussian expansion, it is common for the inferred solution to contain negative `intensity` values. A common\n", + "solution is one where the Gaussians alternate between large positive and large negative values, creating an almost\n", + "\"ringing\" effect in the reconstruction. This is a very unphysical solution and one we want to avoid.\n", + "\n", + "We are able to illustrate this now, first by solving the linear algebra and then printing the `intensity` values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = np.linalg.solve(curvature_matrix, data_vector)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `reconstruction` is a 1D vector of length equal to the number of Gaussian linear light profiles, which in this case \n", + "is 30.\n", + "\n", + "Each value represents the solved for `intensity` of the Gaussian linear light profile.\n", + "\n", + "In this example, the values alternate between positive and negative, indicating a solution that is not physical\n", + "and one we must avoid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Reconstruction (S) of Linear Light Profiles Intensity:\")\n", + "print(reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction (Positive Only)__\n", + "\n", + "The linear algebra can be solved for with the constraint that all solutions, and therefore all `intensity` values,\n", + "are positive. \n", + "\n", + "This could be achieved by using the `scipy` `nnls` non-negative least squares solver.\n", + "\n", + "The nnls poses the problem slightly different than the code above. It solves for the `intensity` values in an\n", + "iterative manner meaning that it is slower. It does not use `data_vector` $D$ and `curvature_matrix` $F$ but instead\n", + "works directly with the `blurred_mapping_matrix` $f$ and the data and noise-map.\n", + "\n", + "The `nnls` function is therefore computationally slow, especially for cases where there are many linear light profiles \n", + "or even more complex linear inversions like a pixelized reconstruction.\n", + "\n", + "The source code therefore uses a \"fast nnls\" algorithm, which is an adaptation of the algorithm found at\n", + "this URL: https://github.com/jvendrow/fnnls\n", + "\n", + "Unlike the scipy nnls function, the fnnls method uses the `data_vector` $D$ and `curvature_matrix` $F$ to solve for\n", + "the `intensity` values. This provides it with additional information about the linear algebra problem, which is\n", + "why it is faster.\n", + "\n", + "The function `reconstruction_positive_only_from` uses the `fnnls` algorithm to compute the `intensity` values\n", + "of the linear light profiles, ensuring they are positive.\n", + "\n", + "However, the code below by itself actually produces a `LinAlgError` because the `curvature_matrix` is singular. Uncomment\n", + "the code below to see this error." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", + "# data_vector=data_vector,\n", + "# curvature_reg_matrix=curvature_matrix, # ignore _reg_ tag in this guide\n", + "# )\n", + "#\n", + "# print(reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To make the `curvature_matrix` non-singular, we simply add small numerical values to its diagonal elements. \n", + "\n", + "This is effectively add a small degree of \"zeroth order\" regularization to the inversion, which is sufficient to make\n", + "the matrix non-singular and ensure the inversion can be performed.\n", + "\n", + "There are a variety of ways to regularize the inversion, and these can be manually input into the `Basis` object.\n", + "However, for a multi-Gaussian expansion, testing has shown that adding a small degree of zeroth order regularization\n", + "in conjunction with a positive-only solution is sufficient to ensure the inversion is robust for all reasonable\n", + "science cases.\n", + "\n", + "In practise, the code only adds these small numerical values to the diagonal of the curvature matrix for elements\n", + "which have no other regularization applied to them. Therefore, in the function call below we input \n", + "`no_regularization_index_list=range(30)`, which tells the function to add small numerical values to all 30\n", + "diagonal values corresponding to the 30 Gaussian linear light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix,\n", + " noise_map=masked_dataset.noise_map,\n", + " add_to_curvature_diag=True,\n", + " no_regularization_index_list=list(range(30)),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `reconstruction` can now be computed successfully without a linear algebra error." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", + " data_vector=data_vector,\n", + " curvature_reg_matrix=curvature_matrix, # ignore _reg_ tag in this guide\n", + ")\n", + "\n", + "print(reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image Reconstruction__\n", + "\n", + "Using the reconstructed `intensity` values we can map the reconstruction back to the image plane (via \n", + "the `blurred mapping_matrix`) and produce a reconstruction of the image data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_operated_data = (\n", + " al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix, reconstruction=reconstruction\n", + " )\n", + ")\n", + "\n", + "mapped_reconstructed_operated_data = al.Array2D(\n", + " values=mapped_reconstructed_operated_data, mask=mask\n", + ")\n", + "\n", + "aplt.plot_array(array=mapped_reconstructed_operated_data, title=\"\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our galaxy model.\n", + "\n", + "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", + "\n", + "The likelihood function for parametric galaxy modeling, even if linear light profiles are used, consists of two terms:\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "We now explain what each of these terms mean.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", + "\n", + " - `model_data` = `convolved_image_2d`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", + "\n", + "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", + "for, corresponding to a fit with a lower likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image = mapped_reconstructed_operated_data\n", + "\n", + "residual_map = masked_dataset.data - model_image\n", + "normalized_residual_map = residual_map / masked_dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(chi_squared)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", + "\n", + "aplt.plot_array(array=chi_squared_map, title=\"\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", + "\n", + "The final term in the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", + "of the log of every noise-map value squared. \n", + "\n", + "Given the `noise_map` is fixed, this term does not change during the galaxy modeling process and has no impact on the \n", + "model we infer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms computed above using\n", + "the likelihood function defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(figure_of_merit)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " basis=basis,\n", + ")\n", + "\n", + "galaxies = al.Galaxies(galaxies=[galaxy])\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=masked_dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_border_relocator=True),\n", + ")\n", + "fit_log_evidence = fit.log_evidence\n", + "print(fit_log_evidence)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit contains an `Inversion` object, which handles all the linear algebra we have covered in this script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.inversion)\n", + "print(fit.inversion.data_vector)\n", + "print(fit.inversion.curvature_matrix)\n", + "print(fit.inversion.reconstruction)\n", + "print(fit.inversion.mapped_reconstructed_operated_data)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Inversion` object can be computed from a tracer and a dataset, by passing them to the `TracerToInversion` object.\n", + "\n", + "This objects handles a lot of extra functionality that we have not covered in this script, such as:\n", + "\n", + "- Separating out the linear light profiles from the standard light profiles.\n", + "- Separating out objects which reconstruct the galaxy using a pixelized reconstruction, which are passed into\n", + " the `Inversion` object as well." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(\n", + " tracer=tracer,\n", + " dataset=masked_dataset,\n", + ")\n", + "\n", + "inversion = tracer_to_inversion.inversion\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Modeling__\n", + "\n", + "To fit a galaxy model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "For an MGE, the reduced number of free parameters (e.g. the `intensity` values are solved for\n", + "via linear algebra and not a dimension of the non-linear parameter space) means that the sampler converges in fewer\n", + "iterations and is less likely to infer a local maximum. \n", + "\n", + "Furthermore, the size of the lens galaxy, controlled by the `sigma` values of the Gaussians, are also all fixed\n", + "and not non-linear free parameters. This further simplifies the non-linear parameter space.\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the multi Gaussian expansion likelihood function, which uses \n", + "many 2D Gaussians to fit the galaxy light and solve for the `intensity` values via linear algebra.\n", + "\n", + "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", + "are described in additional notebooks found in the `guides` package:\n", + "\n", + " - `over_sampling`: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", + " ray-traced to the source-plane and used to evaluate the light profile more accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/multi_gaussian_expansion/modeling.ipynb b/notebooks/imaging/features/multi_gaussian_expansion/modeling.ipynb index 9d886075c..7986dec8d 100644 --- a/notebooks/imaging/features/multi_gaussian_expansion/modeling.ipynb +++ b/notebooks/imaging/features/multi_gaussian_expansion/modeling.ipynb @@ -1,729 +1,766 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Multi Gaussian Expansion\n", - "===========================================\n", - "\n", - "A multi Gaussian expansion (MGE) decomposes the lens light into ~15-100 Gaussians, where the `intensity` of every\n", - "Gaussian is solved for via a linear algebra using a process called an \"inversion\" (see the `linear_light_profiles.py`\n", - "feature for a full description of this).\n", - "\n", - "This script performs lensing modeling using a lens light model which uses an MGE consisting of 60 Gaussians. It is\n", - "fitted to simulated data where the lens galaxy's light has asymmetric and irregular features, which fitted poorly by symmetric light\n", - "profiles like the `Sersic`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Symmetric light profiles (e.g.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **MGE Source Galaxy:** The MGE was designed to model the light of lens galaxies, because they are typically elliptical.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Source MGE:** As discussed at the beginning of this tutorial, an MGE is an effective way to model the light of a.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "- **Description:** There is one downside to `Basis` functions, we may compose a model with too much freedom.\n", - "\n", - "__Advantages__\n", - "\n", - "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals, because they fail to capture\n", - "irregular and asymmetric morphological of galaxies (e.g. isophotal twists, an ellipticity which varies radially).\n", - "An MGE fully captures these features and can therefore much better represent the emission of complex lens galaxies.\n", - "\n", - "The MGE model can be composed in a way that has fewer non-linear parameters than an elliptical Sersic. In this example,\n", - "two separate groups of Gaussians are used to represent the `bulge` and `disk` of the lens, which in total correspond\n", - "to just N=6 non-linear parameters (a `bulge` and `disk` comprising two linear Sersics has N=10 parameters).\n", - "\n", - "The MGE model parameterization is also composed such that neither the `intensity` parameters or any of the\n", - "parameters controlling the size of the Gaussians (their `sigma` values) are non-linear parameters sampled by Nautilus.\n", - "This removes the most significant degeneracies in parameter space, making the model much more reliable and efficient\n", - "to fit.\n", - "\n", - "Therefore, not only does an MGE fit more complex galaxy morphologies, it does so using fewer non-linear parameters\n", - "in a much simpler non-linear parameter space which has far less significant parameter degeneracies!\n", - "\n", - "__Disadvantages__\n", - "\n", - "To fit an MGE model to the data, the light of the ~15-75 or more Gaussian in the MGE must be evaluated and compared\n", - "to the data. This is slower than evaluating the light of ~2-3 Sersic profiles, producing slower computational run\n", - "times (although the simpler non-linear parameter space will speed up the fit overall).\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and negative\n", - "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical. For an MGE, this produces a positive-negative \"ringing\", where the\n", - "Gaussians alternate between large positive and negative values. This is clearly undesirable and unphysical.\n", - "\n", - "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is as fast\n", - "as positive-negative solvers. This ensures that all light profile intensities are positive and therefore physical.\n", - "\n", - "__MGE Source Galaxy__\n", - "\n", - "The MGE was designed to model the light of lens galaxies, because they are typically elliptical galaxies whose\n", - "morphology is better represented as many Gaussians. Their complex features (e.g. isophotal twists,\n", - "an ellipticity which varies radially) are accurately captured by an MGE.\n", - "\n", - "The morphological features typically seen in source galaxies (e.g. disks, bars, clumps of star formation) are less\n", - "suited to an MGE. The source-plane of many lenses also often have multiple galaxies, whereas the MGE fitted\n", - "in this example assumes a single `centre`.\n", - "\n", - "However, despite these limitations, an MGE turns out to be an extremely powerful way to model the source galaxies\n", - "of strong lenses. This is because, even if it struggles to capture the source's morphology, the simplification of\n", - "non-linear parameter space and removal of degeneracies makes it much easier to obtain a reliable lens model.\n", - "This is driven by the removal of any non-linear parameters which change the size of the source's light profile,\n", - "which are otherwise the most degenerate with the lens's mass model.\n", - "\n", - "The second example in this script therefore uses an MGE source. We strongly recommend you read that example and adopt\n", - "MGE lens light models and source models, instead of the elliptical Sersic profiles, as soon as possible!\n", - "\n", - "To capture the irregular and asymmetric features of the source's morphology, or reconstruct multiple source galaxies,\n", - "we recommend using a pixelized source reconstruction (see `autolens_workspace/modeling/features/pixelization.py`).\n", - "Combining this with an MGE for the len's light can be a very powerful way to model strong lenses!\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's bulge is a super position of 60 `Gaussian` profiles.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a linear `SersicCore`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"lens_light_asymmetric\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/multi_gaussian_expansion/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", - "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The galaxy's bulge is 60 linear `Gaussian` profiles [6 parameters]. \n", - " - The centres and elliptical components of the Gaussians are all linked together in two groups of 30.\n", - " - The `sigma` size of the Gaussians increases in log10 increments.\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is a linear `Sersic` [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=19.\n", - "\n", - "Unlike the examples above, our MGE now comprises 2 * 30 Gaussians (instead of 1 * 30). Each group of 30\n", - "have their own elliptical components meaning the lens's light is decomposed into two distinct elliptical components,\n", - "which could be viewed as a bulge and disk.\n", - "\n", - "Note that above we combine the MGE for the lens light with a linear light profile for the source, meaning these two\n", - "categories of models can be combined.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "gaussian_per_basis = 2\n", - "\n", - "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", - "mask_radius = 3.0\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "# By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.\n", - "\n", - "centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "centre_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "\n", - "bulge_gaussian_list = []\n", - "\n", - "for j in range(gaussian_per_basis):\n", - " # A list of Gaussian model components whose parameters are customized belows.\n", - "\n", - " gaussian_list = af.Collection(\n", - " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", - " )\n", - "\n", - " # Iterate over every Gaussian and customize its parameters.\n", - "\n", - " for i, gaussian in enumerate(gaussian_list):\n", - " gaussian.centre.centre_0 = centre_0 # All Gaussians have same y centre.\n", - " gaussian.centre.centre_1 = centre_1 # All Gaussians have same x centre.\n", - " gaussian.ell_comps = gaussian_list[\n", - " 0\n", - " ].ell_comps # All Gaussians have same elliptical components.\n", - " gaussian.sigma = (\n", - " 10 ** log10_sigma_list[i]\n", - " ) # All Gaussian sigmas are fixed to values above.\n", - "\n", - " bulge_gaussian_list += gaussian_list\n", - "\n", - "# The Basis object groups many light profiles together into a single model component.\n", - "\n", - "bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=bulge_gaussian_list,\n", - ")\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This shows every single Gaussian light profile in the model, which is a lot of parameters! However, the vast\n", - "majority of these parameters are fixed to the values we set above, so the model actually has far fewer free\n", - "parameters than it looks!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", - "full description).\n", - "\n", - "Owing to the simplicity of fitting an MGE we an use even fewer live points than other examples, reducing it to\n", - "75 live points, speeding up convergence of the non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"mge\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=2000000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", - "required by a model.\n", - "\n", - "For each linear Gaussian light profile, extra VRAM is used. For around 60 linear Gaussians this typically requires \n", - "a modest amount of VRAM (e.g. 10\u201350 MB per batched likelihood). Models that use hundreds of Gaussians, especially in \n", - "combination with a large batch size, may therefore exceed GBs of VRAM and require you to adjust the batch size to fit \n", - "within your GPU's VRAM.\n", - "\n", - "__Run Time__\n", - "\n", - "The likelihood evaluation time for a multi-Gaussian expansion is significantly slower than standard / linear \n", - "light profiles. This is because the image of every Gaussian must be computed and evaluated, and each must be blurred \n", - "with the PSF. In this example, the evaluation time is ~0.5s, compared to ~0.01 seconds for standard light profiles.\n", - "\n", - "Huge gains in the overall run-time however are made thanks to the models significantly reduced complexity and lower\n", - "number of free parameters. Furthermore, because there are not free parameters which scale the size of lens galaxy,\n", - "this produces significantly faster convergence by Nautilus that any other lens light model. We also use fewer live\n", - "points, further speeding up the model-fit.\n", - "\n", - "Overall, it is difficult to state which approach will be faster overall. However, the MGE's ability to fit the data\n", - "more accurately and the less complex parameter due to removing parameters that scale the lens galaxy make it the \n", - "superior approach.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this does \n", - "not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix this):\n", - "\n", - "This confirms there are many `Gaussian`' in the lens light model and it lists their inferred parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "In particular, checkout the results example `linear.py` which details how to extract all information about linear\n", - "light profiles from a fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source MGE__\n", - "\n", - "As discussed at the beginning of this tutorial, an MGE is an effective way to model the light of a source galaxy and \n", - "get an initial estimate of the lens mass model.\n", - "\n", - "This MGE source is used alongside the MGE lens light model, which offers a lot of flexibility in modeling the lens\n", - "and source galaxies. Note that the MGE source uses fewer Gaussians, as the MGE only needs to capture the main\n", - "structure of the source galaxy's light to obtain an accurate lens model.\n", - "\n", - "We compose the model below, recreating the Gaussian's line-by-line in Python, so you can easily copy and paste\n", - "and reuse the code below in your own scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "total_gaussians = 30\n", - "gaussian_per_basis = 2\n", - "\n", - "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", - "mask_radius = 3.0\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "# By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.\n", - "\n", - "centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "centre_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "\n", - "bulge_gaussian_list = []\n", - "\n", - "for j in range(gaussian_per_basis):\n", - " # A list of Gaussian model components whose parameters are customized belows.\n", - "\n", - " gaussian_list = af.Collection(\n", - " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", - " )\n", - "\n", - " # Iterate over every Gaussian and customize its parameters.\n", - "\n", - " for i, gaussian in enumerate(gaussian_list):\n", - " gaussian.centre.centre_0 = centre_0 # All Gaussians have same y centre.\n", - " gaussian.centre.centre_1 = centre_1 # All Gaussians have same x centre.\n", - " gaussian.ell_comps = gaussian_list[\n", - " 0\n", - " ].ell_comps # All Gaussians have same elliptical components.\n", - " gaussian.sigma = (\n", - " 10 ** log10_sigma_list[i]\n", - " ) # All Gaussian sigmas are fixed to values above.\n", - "\n", - " bulge_gaussian_list += gaussian_list\n", - "\n", - "# The Basis object groups many light profiles together into a single model component.\n", - "\n", - "bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=bulge_gaussian_list,\n", - ")\n", - "mass = af.Model(al.mp.Isothermal)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - "\n", - "# Source:\n", - "\n", - "total_gaussians = 20\n", - "gaussian_per_basis = 1\n", - "\n", - "# By defining the centre here, it creates two free parameters that are assigned to the source Gaussians.\n", - "\n", - "centre_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - "centre_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - "\n", - "log10_sigma_list = np.linspace(-2, np.log10(1.0), total_gaussians)\n", - "\n", - "bulge_gaussian_list = []\n", - "\n", - "for j in range(gaussian_per_basis):\n", - " gaussian_list = af.Collection(\n", - " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", - " )\n", - "\n", - " for i, gaussian in enumerate(gaussian_list):\n", - " gaussian.centre.centre_0 = centre_0\n", - " gaussian.centre.centre_1 = centre_1\n", - " gaussian.ell_comps = gaussian_list[0].ell_comps\n", - " gaussian.sigma = 10 ** log10_sigma_list[i]\n", - "\n", - " bulge_gaussian_list += gaussian_list\n", - "\n", - "source_bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=bulge_gaussian_list,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing the model info confirms the model has Gaussians for both the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now fit this model, which includes the MGE source and lens light models." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"mge_including_source\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The likelihood evaluation time for a multi-Gaussian expansion for both lens and source galaxies is not much slower\n", - "than when just the lens galaxy uses an MGE.\n", - "\n", - "However, the overall run-time will be even faster than before, as treating the source as an MGE further\n", - "reduces the complexity of parameter space ensuring Nautilus converges even faster.\n", - "\n", - "For initial model-fits where the lens model parameters are not known, a lens + source MGE is possibly the best\n", - "model one can use. \n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "A Multi Gaussian Expansion is a powerful tool for modeling the light of galaxies, and offers a compelling method to\n", - "fit complex light profiles with a small number of parameters.\n", - "\n", - "For **PyAutoLens**'s advanced search chaining feature, it is common use the MGE to initialize the lens light and source \n", - "models. The lens light model is then made more complex by using an MGE with more Gaussians and the source becomes a \n", - "pixelized reconstruction.\n", - "\n", - "Now you are familiar with MGE modeling, it is recommended you adopt this as your default lens modeling approach.\n", - "However, it may not be suitable for lower resolution data, where the simpler Sersic profiles may be more appropriate.\n", - "\n", - "__Basis Regularization (Advanced / Unused)__\n", - "\n", - "An MGE `Basis` can additionally carry a regularization term (e.g. `al.reg.Constant`) that penalises non-smooth\n", - "solutions for the Gaussian intensities. This is a research-only feature: it is not used by any production\n", - "scientific analysis, because the positive-only linear algebra solver used above already solves the\n", - "\"positive/negative ringing\" problem regularization was originally intended to address.\n", - "\n", - "The code and rationale for the regularization branch have been moved out of this user-facing script to keep it\n", - "focused on the supported MGE workflow. If you want to experiment with adding a regularization to a Basis, see:\n", - "\n", - " autolens_workspace_developer/basis_regularization/mge_lens.py\n", - "\n", - "That script is self-contained and runs the same regularized fit that used to live here.\n", - "\n", - "__Finish__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Multi Gaussian Expansion\n", + "===========================================\n", + "\n", + "A multi Gaussian expansion (MGE) decomposes the lens light into ~15-100 Gaussians, where the `intensity` of every\n", + "Gaussian is solved for via a linear algebra using a process called an \"inversion\" (see the `linear_light_profiles.py`\n", + "feature for a full description of this).\n", + "\n", + "This script performs lensing modeling using a lens light model which uses an MGE consisting of 60 Gaussians. It is\n", + "fitted to simulated data where the lens galaxy's light has asymmetric and irregular features, which fitted poorly by symmetric light\n", + "profiles like the `Sersic`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Symmetric light profiles (e.g.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **MGE Source Galaxy:** The MGE was designed to model the light of lens galaxies, because they are typically elliptical.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Source MGE:** As discussed at the beginning of this tutorial, an MGE is an effective way to model the light of a.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "- **Description:** There is one downside to `Basis` functions, we may compose a model with too much freedom.\n", + "\n", + "__Advantages__\n", + "\n", + "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals, because they fail to capture\n", + "irregular and asymmetric morphological of galaxies (e.g. isophotal twists, an ellipticity which varies radially).\n", + "An MGE fully captures these features and can therefore much better represent the emission of complex lens galaxies.\n", + "\n", + "The MGE model can be composed in a way that has fewer non-linear parameters than an elliptical Sersic. In this example,\n", + "two separate groups of Gaussians are used to represent the `bulge` and `disk` of the lens, which in total correspond\n", + "to just N=6 non-linear parameters (a `bulge` and `disk` comprising two linear Sersics has N=10 parameters).\n", + "\n", + "The MGE model parameterization is also composed such that neither the `intensity` parameters or any of the\n", + "parameters controlling the size of the Gaussians (their `sigma` values) are non-linear parameters sampled by Nautilus.\n", + "This removes the most significant degeneracies in parameter space, making the model much more reliable and efficient\n", + "to fit.\n", + "\n", + "Therefore, not only does an MGE fit more complex galaxy morphologies, it does so using fewer non-linear parameters\n", + "in a much simpler non-linear parameter space which has far less significant parameter degeneracies!\n", + "\n", + "__Disadvantages__\n", + "\n", + "To fit an MGE model to the data, the light of the ~15-75 or more Gaussian in the MGE must be evaluated and compared\n", + "to the data. This is slower than evaluating the light of ~2-3 Sersic profiles, producing slower computational run\n", + "times (although the simpler non-linear parameter space will speed up the fit overall).\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and negative\n", + "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical. For an MGE, this produces a positive-negative \"ringing\", where the\n", + "Gaussians alternate between large positive and negative values. This is clearly undesirable and unphysical.\n", + "\n", + "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is as fast\n", + "as positive-negative solvers. This ensures that all light profile intensities are positive and therefore physical.\n", + "\n", + "__MGE Source Galaxy__\n", + "\n", + "The MGE was designed to model the light of lens galaxies, because they are typically elliptical galaxies whose\n", + "morphology is better represented as many Gaussians. Their complex features (e.g. isophotal twists,\n", + "an ellipticity which varies radially) are accurately captured by an MGE.\n", + "\n", + "The morphological features typically seen in source galaxies (e.g. disks, bars, clumps of star formation) are less\n", + "suited to an MGE. The source-plane of many lenses also often have multiple galaxies, whereas the MGE fitted\n", + "in this example assumes a single `centre`.\n", + "\n", + "However, despite these limitations, an MGE turns out to be an extremely powerful way to model the source galaxies\n", + "of strong lenses. This is because, even if it struggles to capture the source's morphology, the simplification of\n", + "non-linear parameter space and removal of degeneracies makes it much easier to obtain a reliable lens model.\n", + "This is driven by the removal of any non-linear parameters which change the size of the source's light profile,\n", + "which are otherwise the most degenerate with the lens's mass model.\n", + "\n", + "The second example in this script therefore uses an MGE source. We strongly recommend you read that example and adopt\n", + "MGE lens light models and source models, instead of the elliptical Sersic profiles, as soon as possible!\n", + "\n", + "To capture the irregular and asymmetric features of the source's morphology, or reconstruct multiple source galaxies,\n", + "we recommend using a pixelized source reconstruction (see `autolens_workspace/modeling/features/pixelization.py`).\n", + "Combining this with an MGE for the len's light can be a very powerful way to model strong lenses!\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's bulge is a super position of 60 `Gaussian` profiles.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a linear `SersicCore`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"lens_light_asymmetric\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/multi_gaussian_expansion/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Apply adaptive over sampling to ensure the lens galaxy light calculation is accurate, you can read up on over-sampling \n", + "in more detail via the `autogalaxy_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The galaxy's bulge is 60 linear `Gaussian` profiles [6 parameters]. \n", + " - The centres and elliptical components of the Gaussians are all linked together in two groups of 30.\n", + " - The `sigma` size of the Gaussians increases in log10 increments.\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is a linear `Sersic` [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=19.\n", + "\n", + "Unlike the examples above, our MGE now comprises 2 * 30 Gaussians (instead of 1 * 30). Each group of 30\n", + "have their own elliptical components meaning the lens's light is decomposed into two distinct elliptical components,\n", + "which could be viewed as a bulge and disk.\n", + "\n", + "Note that above we combine the MGE for the lens light with a linear light profile for the source, meaning these two\n", + "categories of models can be combined.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "gaussian_per_basis = 2\n", + "\n", + "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", + "mask_radius = 3.0\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "# By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.\n", + "\n", + "centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "centre_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "\n", + "bulge_gaussian_list = []\n", + "\n", + "for j in range(gaussian_per_basis):\n", + " # A list of Gaussian model components whose parameters are customized belows.\n", + "\n", + " gaussian_list = af.Collection(\n", + " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", + " )\n", + "\n", + " # Iterate over every Gaussian and customize its parameters.\n", + "\n", + " for i, gaussian in enumerate(gaussian_list):\n", + " gaussian.centre.centre_0 = centre_0 # All Gaussians have same y centre.\n", + " gaussian.centre.centre_1 = centre_1 # All Gaussians have same x centre.\n", + " gaussian.ell_comps = gaussian_list[\n", + " 0\n", + " ].ell_comps # All Gaussians have same elliptical components.\n", + " gaussian.sigma = (\n", + " 10 ** log10_sigma_list[i]\n", + " ) # All Gaussian sigmas are fixed to values above.\n", + "\n", + " bulge_gaussian_list += gaussian_list\n", + "\n", + "# The Basis object groups many light profiles together into a single model component.\n", + "\n", + "bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=bulge_gaussian_list,\n", + ")\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This shows every single Gaussian light profile in the model, which is a lot of parameters! However, the vast\n", + "majority of these parameters are fixed to the values we set above, so the model actually has far fewer free\n", + "parameters than it looks!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", + "full description).\n", + "\n", + "Owing to the simplicity of fitting an MGE we an use even fewer live points than other examples, reducing it to\n", + "75 live points, speeding up convergence of the non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"mge\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=2000000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", + "required by a model.\n", + "\n", + "For each linear Gaussian light profile, extra VRAM is used. For around 60 linear Gaussians this typically requires \n", + "a modest amount of VRAM (e.g. 10\u201350 MB per batched likelihood). Models that use hundreds of Gaussians, especially in \n", + "combination with a large batch size, may therefore exceed GBs of VRAM and require you to adjust the batch size to fit \n", + "within your GPU's VRAM.\n", + "\n", + "__Run Time__\n", + "\n", + "The likelihood evaluation time for a multi-Gaussian expansion is significantly slower than standard / linear \n", + "light profiles. This is because the image of every Gaussian must be computed and evaluated, and each must be blurred \n", + "with the PSF. In this example, the evaluation time is ~0.5s, compared to ~0.01 seconds for standard light profiles.\n", + "\n", + "Huge gains in the overall run-time however are made thanks to the models significantly reduced complexity and lower\n", + "number of free parameters. Furthermore, because there are not free parameters which scale the size of lens galaxy,\n", + "this produces significantly faster convergence by Nautilus that any other lens light model. We also use fewer live\n", + "points, further speeding up the model-fit.\n", + "\n", + "Overall, it is difficult to state which approach will be faster overall. However, the MGE's ability to fit the data\n", + "more accurately and the less complex parameter due to removing parameters that scale the lens galaxy make it the \n", + "superior approach.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this does \n", + "not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix this):\n", + "\n", + "This confirms there are many `Gaussian`' in the lens light model and it lists their inferred parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "In particular, checkout the results example `linear.py` which details how to extract all information about linear\n", + "light profiles from a fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source MGE__\n", + "\n", + "As discussed at the beginning of this tutorial, an MGE is an effective way to model the light of a source galaxy and \n", + "get an initial estimate of the lens mass model.\n", + "\n", + "This MGE source is used alongside the MGE lens light model, which offers a lot of flexibility in modeling the lens\n", + "and source galaxies. Note that the MGE source uses fewer Gaussians, as the MGE only needs to capture the main\n", + "structure of the source galaxy's light to obtain an accurate lens model.\n", + "\n", + "We compose the model below, recreating the Gaussian's line-by-line in Python, so you can easily copy and paste\n", + "and reuse the code below in your own scripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "total_gaussians = 30\n", + "gaussian_per_basis = 2\n", + "\n", + "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", + "mask_radius = 3.0\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "# By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.\n", + "\n", + "centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "centre_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "\n", + "bulge_gaussian_list = []\n", + "\n", + "for j in range(gaussian_per_basis):\n", + " # A list of Gaussian model components whose parameters are customized belows.\n", + "\n", + " gaussian_list = af.Collection(\n", + " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", + " )\n", + "\n", + " # Iterate over every Gaussian and customize its parameters.\n", + "\n", + " for i, gaussian in enumerate(gaussian_list):\n", + " gaussian.centre.centre_0 = centre_0 # All Gaussians have same y centre.\n", + " gaussian.centre.centre_1 = centre_1 # All Gaussians have same x centre.\n", + " gaussian.ell_comps = gaussian_list[\n", + " 0\n", + " ].ell_comps # All Gaussians have same elliptical components.\n", + " gaussian.sigma = (\n", + " 10 ** log10_sigma_list[i]\n", + " ) # All Gaussian sigmas are fixed to values above.\n", + "\n", + " bulge_gaussian_list += gaussian_list\n", + "\n", + "# The Basis object groups many light profiles together into a single model component.\n", + "\n", + "bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=bulge_gaussian_list,\n", + ")\n", + "mass = af.Model(al.mp.Isothermal)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + "\n", + "# Source:\n", + "\n", + "total_gaussians = 20\n", + "gaussian_per_basis = 1\n", + "\n", + "# By defining the centre here, it creates two free parameters that are assigned to the source Gaussians.\n", + "\n", + "centre_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + "centre_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + "\n", + "log10_sigma_list = np.linspace(-2, np.log10(1.0), total_gaussians)\n", + "\n", + "bulge_gaussian_list = []\n", + "\n", + "for j in range(gaussian_per_basis):\n", + " gaussian_list = af.Collection(\n", + " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", + " )\n", + "\n", + " for i, gaussian in enumerate(gaussian_list):\n", + " gaussian.centre.centre_0 = centre_0\n", + " gaussian.centre.centre_1 = centre_1\n", + " gaussian.ell_comps = gaussian_list[0].ell_comps\n", + " gaussian.sigma = 10 ** log10_sigma_list[i]\n", + "\n", + " bulge_gaussian_list += gaussian_list\n", + "\n", + "source_bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=bulge_gaussian_list,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing the model info confirms the model has Gaussians for both the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now fit this model, which includes the MGE source and lens light models." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"mge_including_source\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The likelihood evaluation time for a multi-Gaussian expansion for both lens and source galaxies is not much slower\n", + "than when just the lens galaxy uses an MGE.\n", + "\n", + "However, the overall run-time will be even faster than before, as treating the source as an MGE further\n", + "reduces the complexity of parameter space ensuring Nautilus converges even faster.\n", + "\n", + "For initial model-fits where the lens model parameters are not known, a lens + source MGE is possibly the best\n", + "model one can use. \n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "A Multi Gaussian Expansion is a powerful tool for modeling the light of galaxies, and offers a compelling method to\n", + "fit complex light profiles with a small number of parameters.\n", + "\n", + "For **PyAutoLens**'s advanced search chaining feature, it is common use the MGE to initialize the lens light and source \n", + "models. The lens light model is then made more complex by using an MGE with more Gaussians and the source becomes a \n", + "pixelized reconstruction.\n", + "\n", + "Now you are familiar with MGE modeling, it is recommended you adopt this as your default lens modeling approach.\n", + "However, it may not be suitable for lower resolution data, where the simpler Sersic profiles may be more appropriate.\n", + "\n", + "__Basis Regularization (Advanced / Unused)__\n", + "\n", + "An MGE `Basis` can additionally carry a regularization term (e.g. `al.reg.Constant`) that penalises non-smooth\n", + "solutions for the Gaussian intensities. This is a research-only feature: it is not used by any production\n", + "scientific analysis, because the positive-only linear algebra solver used above already solves the\n", + "\"positive/negative ringing\" problem regularization was originally intended to address.\n", + "\n", + "The code and rationale for the regularization branch have been moved out of this user-facing script to keep it\n", + "focused on the supported MGE workflow. If you want to experiment with adding a regularization to a Basis, see:\n", + "\n", + " autolens_workspace_developer/basis_regularization/mge_lens.py\n", + "\n", + "That script is self-contained and runs the same regularized fit that used to live here.\n", + "\n", + "__Finish__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/multi_gaussian_expansion/simulator.ipynb b/notebooks/imaging/features/multi_gaussian_expansion/simulator.ipynb index d21518ed9..2e76fbe51 100644 --- a/notebooks/imaging/features/multi_gaussian_expansion/simulator.ipynb +++ b/notebooks/imaging/features/multi_gaussian_expansion/simulator.ipynb @@ -1,482 +1,519 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Lens Light Asymmetric\n", - "================================\n", - "\n", - "The morphological of massive elliptical galaxies which act as strong lens are often asymmetric and irregular, with\n", - "features such as isophotal twists or radially varying elliptical components.\n", - "\n", - "This script uses a basis of 14 elliptical Gaussians to simulate the light of a massive elliptical galaxy which has\n", - "these irregular and asymmetric features. The parameters of the Gaussian basis are derived from a Multi-Gaussian\n", - "fit to a real strong lens.\n", - "\n", - "This dataset is used in the `modeling/features/multi_gaussian_expansion.py` script to illustrate how to fit these\n", - "features using a Multi-Gaussian Expansion (MGE).\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's light is a superposition of 14 `Gaussian` profiles.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an `Sersic`.\n", - "\n", - "The lens galaxy's light is derived from a Multi-Gaussian Expansion (MGE) fit to a massive elliptical galaxy.\n", - "\n", - "The simulated galaxy has irregular and asymmetric features in the galaxy, including a twist in the isophotes of its\n", - "emission.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"lens_light_asymmetric\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "\n", - "The lens galaxy uses 14 elliptical Gaussians, which represent a complex galaxy morphology with irregular and\n", - "asymmetric features such as an isophotal twist (which symmetric profiles like a Sersic cannot capture)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "centre_y_list = [\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - " -0.00361289,\n", - "]\n", - "\n", - "centre_x_list = [\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - " -0.00636064,\n", - "]\n", - "\n", - "ell_comps_0_list = [\n", - " 0.05843285,\n", - " 0.0,\n", - " 0.05368621,\n", - " 0.05090395,\n", - " 0.0,\n", - " 0.25367341,\n", - " 0.01677313,\n", - " 0.03626733,\n", - " 0.15887384,\n", - " 0.02790297,\n", - " 0.12368768,\n", - " 0.38624915,\n", - " -0.10490247,\n", - " 0.0385585,\n", - "]\n", - "\n", - "ell_comps_1_list = [\n", - " 0.05932136,\n", - " 0.0,\n", - " 0.04267542,\n", - " -0.06920487,\n", - " -0.0,\n", - " -0.15141799,\n", - " 0.01464508,\n", - " 0.03084128,\n", - " -0.17983965,\n", - " 0.02215257,\n", - " -0.16271084,\n", - " -0.15945967,\n", - " -0.3969543,\n", - " -0.03808391,\n", - "]\n", - "\n", - "intensity_list = [\n", - " 0.52107394,\n", - " 4.2933716,\n", - " 2.40608609,\n", - " 4.98902608,\n", - " 2.72773562,\n", - " 1.10429021,\n", - " 1.08190372,\n", - " 0.30007753,\n", - " 0.6462658,\n", - " 0.15766566,\n", - " 0.24687923,\n", - " 0.04815128,\n", - " 0.02559108,\n", - " 0.06763223,\n", - "]\n", - "\n", - "sigma_list = [\n", - " 0.01607907,\n", - " 0.04039063,\n", - " 0.06734373,\n", - " 0.08471335,\n", - " 0.16048498,\n", - " 0.13531624,\n", - " 0.25649938,\n", - " 0.46096968,\n", - " 0.34492195,\n", - " 0.92418119,\n", - " 0.71803244,\n", - " 1.23547346,\n", - " 1.2574071,\n", - " 2.69979461,\n", - "]\n", - "\n", - "gaussians = []\n", - "\n", - "for gaussian_index in range(len(centre_x_list)):\n", - " gaussian = al.lp.Gaussian(\n", - " centre=(centre_y_list[gaussian_index], centre_x_list[gaussian_index]),\n", - " ell_comps=(\n", - " ell_comps_0_list[gaussian_index],\n", - " ell_comps_1_list[gaussian_index],\n", - " ),\n", - " intensity=intensity_list[gaussian_index],\n", - " sigma=sigma_list[gaussian_index],\n", - " )\n", - "\n", - " gaussians.append(gaussian)\n", - "\n", - "basis = al.lp_basis.Basis(profile_list=gaussians)\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, bulge=basis, mass=mass)\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/lens_light_asymmetric`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Lens Light Asymmetric\n", + "================================\n", + "\n", + "The morphological of massive elliptical galaxies which act as strong lens are often asymmetric and irregular, with\n", + "features such as isophotal twists or radially varying elliptical components.\n", + "\n", + "This script uses a basis of 14 elliptical Gaussians to simulate the light of a massive elliptical galaxy which has\n", + "these irregular and asymmetric features. The parameters of the Gaussian basis are derived from a Multi-Gaussian\n", + "fit to a real strong lens.\n", + "\n", + "This dataset is used in the `modeling/features/multi_gaussian_expansion.py` script to illustrate how to fit these\n", + "features using a Multi-Gaussian Expansion (MGE).\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's light is a superposition of 14 `Gaussian` profiles.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an `Sersic`.\n", + "\n", + "The lens galaxy's light is derived from a Multi-Gaussian Expansion (MGE) fit to a massive elliptical galaxy.\n", + "\n", + "The simulated galaxy has irregular and asymmetric features in the galaxy, including a twist in the isophotes of its\n", + "emission.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"lens_light_asymmetric\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "\n", + "The lens galaxy uses 14 elliptical Gaussians, which represent a complex galaxy morphology with irregular and\n", + "asymmetric features such as an isophotal twist (which symmetric profiles like a Sersic cannot capture)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "centre_y_list = [\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + " -0.00361289,\n", + "]\n", + "\n", + "centre_x_list = [\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + " -0.00636064,\n", + "]\n", + "\n", + "ell_comps_0_list = [\n", + " 0.05843285,\n", + " 0.0,\n", + " 0.05368621,\n", + " 0.05090395,\n", + " 0.0,\n", + " 0.25367341,\n", + " 0.01677313,\n", + " 0.03626733,\n", + " 0.15887384,\n", + " 0.02790297,\n", + " 0.12368768,\n", + " 0.38624915,\n", + " -0.10490247,\n", + " 0.0385585,\n", + "]\n", + "\n", + "ell_comps_1_list = [\n", + " 0.05932136,\n", + " 0.0,\n", + " 0.04267542,\n", + " -0.06920487,\n", + " -0.0,\n", + " -0.15141799,\n", + " 0.01464508,\n", + " 0.03084128,\n", + " -0.17983965,\n", + " 0.02215257,\n", + " -0.16271084,\n", + " -0.15945967,\n", + " -0.3969543,\n", + " -0.03808391,\n", + "]\n", + "\n", + "intensity_list = [\n", + " 0.52107394,\n", + " 4.2933716,\n", + " 2.40608609,\n", + " 4.98902608,\n", + " 2.72773562,\n", + " 1.10429021,\n", + " 1.08190372,\n", + " 0.30007753,\n", + " 0.6462658,\n", + " 0.15766566,\n", + " 0.24687923,\n", + " 0.04815128,\n", + " 0.02559108,\n", + " 0.06763223,\n", + "]\n", + "\n", + "sigma_list = [\n", + " 0.01607907,\n", + " 0.04039063,\n", + " 0.06734373,\n", + " 0.08471335,\n", + " 0.16048498,\n", + " 0.13531624,\n", + " 0.25649938,\n", + " 0.46096968,\n", + " 0.34492195,\n", + " 0.92418119,\n", + " 0.71803244,\n", + " 1.23547346,\n", + " 1.2574071,\n", + " 2.69979461,\n", + "]\n", + "\n", + "gaussians = []\n", + "\n", + "for gaussian_index in range(len(centre_x_list)):\n", + " gaussian = al.lp.Gaussian(\n", + " centre=(centre_y_list[gaussian_index], centre_x_list[gaussian_index]),\n", + " ell_comps=(\n", + " ell_comps_0_list[gaussian_index],\n", + " ell_comps_1_list[gaussian_index],\n", + " ),\n", + " intensity=intensity_list[gaussian_index],\n", + " sigma=sigma_list[gaussian_index],\n", + " )\n", + "\n", + " gaussians.append(gaussian)\n", + "\n", + "basis = al.lp_basis.Basis(profile_list=gaussians)\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, bulge=basis, mass=mass)\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/lens_light_asymmetric`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/multi_gaussian_expansion/slam.ipynb b/notebooks/imaging/features/multi_gaussian_expansion/slam.ipynb index 6a7d66468..7ddc12121 100644 --- a/notebooks/imaging/features/multi_gaussian_expansion/slam.ipynb +++ b/notebooks/imaging/features/multi_gaussian_expansion/slam.ipynb @@ -1,470 +1,507 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Multi Gaussian Expansion: SLaM\n", - "==============================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for fitting a\n", - "lens model where the source is modeled using a Multi Gaussian Expansion (MGE).\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", - "guide before working through this example.\n", - "\n", - "This example only provides documentation specific to the use of an MGE source, describing how the pipeline\n", - "differs from the standard SLaM pipelines described in the SLaM start here guide.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py` with `gaussian_per_basis=1` for both the lens and source MGE.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "You can still run the script without fully understanding the guide, but reviewing it later will\n", - "make the structure and choices of the SLaM workflow clearer.\n", - "\n", - "__Model__\n", - "\n", - "Using a SOURCE LP PIPELINE, LIGHT PIPELINE and a MASS TOTAL PIPELINE this SLaM script fits a strong\n", - "lens system, where in the final model:\n", - "\n", - " - The lens galaxy's light is a bulge with an MGE.\n", - " - The lens galaxy's total mass distribution is an `PowerLaw`.\n", - " - The source galaxy's light is an MGE.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `slam_start_here` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py` with `gaussian_per_basis=1` for both the lens and source MGE models.\n", - "\n", - "Because the source is parametric (an MGE), the SOURCE PIX PIPELINE is skipped. The LIGHT LP and MASS TOTAL\n", - "pipelines use `source_lp_result` directly as both the lens and source initialization result." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " disk=None,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except:\n", - "\n", - " - `source_result_for_lens` and `source_result_for_source` are both set to `source_lp_result`, because\n", - " there is no SOURCE PIX PIPELINE for an MGE source.\n", - " - The analysis does not use adapt images (not required for a parametric MGE source)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_lp_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_lp_result, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_lp_result.instance.galaxies.lens.mass,\n", - " shear=source_lp_result.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except:\n", - "\n", - " - `source_result_for_lens` and `source_result_for_source` are both set to `source_lp_result`, because\n", - " there is no SOURCE PIX PIPELINE for an MGE source.\n", - " - The analysis does not use adapt images or positions likelihood (not required for a parametric MGE source)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " bulge = light_result.instance.galaxies.lens.bulge\n", - " disk = light_result.instance.galaxies.lens.disk\n", - "\n", - " source = al.util.chaining.source_from(result=source_lp_result)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=bulge,\n", - " disk=disk,\n", - " mass=mass,\n", - " shear=source_lp_result.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_lp_result=source_lp_result,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Multi Gaussian Expansion: SLaM\n", + "==============================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for fitting a\n", + "lens model where the source is modeled using a Multi Gaussian Expansion (MGE).\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", + "guide before working through this example.\n", + "\n", + "This example only provides documentation specific to the use of an MGE source, describing how the pipeline\n", + "differs from the standard SLaM pipelines described in the SLaM start here guide.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py` with `gaussian_per_basis=1` for both the lens and source MGE.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`, except.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "You can still run the script without fully understanding the guide, but reviewing it later will\n", + "make the structure and choices of the SLaM workflow clearer.\n", + "\n", + "__Model__\n", + "\n", + "Using a SOURCE LP PIPELINE, LIGHT PIPELINE and a MASS TOTAL PIPELINE this SLaM script fits a strong\n", + "lens system, where in the final model:\n", + "\n", + " - The lens galaxy's light is a bulge with an MGE.\n", + " - The lens galaxy's total mass distribution is an `PowerLaw`.\n", + " - The source galaxy's light is an MGE.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `slam_start_here` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py` with `gaussian_per_basis=1` for both the lens and source MGE models.\n", + "\n", + "Because the source is parametric (an MGE), the SOURCE PIX PIPELINE is skipped. The LIGHT LP and MASS TOTAL\n", + "pipelines use `source_lp_result` directly as both the lens and source initialization result." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " disk=None,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except:\n", + "\n", + " - `source_result_for_lens` and `source_result_for_source` are both set to `source_lp_result`, because\n", + " there is no SOURCE PIX PIPELINE for an MGE source.\n", + " - The analysis does not use adapt images (not required for a parametric MGE source)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_lp_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_lp_result, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_lp_result.instance.galaxies.lens.mass,\n", + " shear=source_lp_result.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except:\n", + "\n", + " - `source_result_for_lens` and `source_result_for_source` are both set to `source_lp_result`, because\n", + " there is no SOURCE PIX PIPELINE for an MGE source.\n", + " - The analysis does not use adapt images or positions likelihood (not required for a parametric MGE source)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " bulge = light_result.instance.galaxies.lens.bulge\n", + " disk = light_result.instance.galaxies.lens.disk\n", + "\n", + " source = al.util.chaining.source_from(result=source_lp_result)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=bulge,\n", + " disk=disk,\n", + " mass=mass,\n", + " shear=source_lp_result.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_lp_result=source_lp_result,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/multi_gaussian_expansion/source_science.ipynb b/notebooks/imaging/features/multi_gaussian_expansion/source_science.ipynb index b02d010fe..6e6436a8a 100644 --- a/notebooks/imaging/features/multi_gaussian_expansion/source_science.ipynb +++ b/notebooks/imaging/features/multi_gaussian_expansion/source_science.ipynb @@ -1,538 +1,575 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Source Science\n", - "==============\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", - "size of the source.\n", - "\n", - "This example shows how to perform these calculations using Multi Gaussian Expansion (MGE) sources on imaging data,\n", - "which is conceptually the simplest case for source science calculations and a good introduction to the topic.\n", - "\n", - "__Contents__\n", - "\n", - "- **Simulated Dataset:** We load and plot the `simple__no_lens_light` example dataset, which is simulated imaging of a.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Lens:** The simulated dataset was created using a lens galaxy with an Isothermal mass profile and External.\n", - "- **Multi Gaussian Expansion Source:** The default workspace source model is a Multi Gaussian Expansion (MGE) profile, which is a.\n", - "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", - "- **Source Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", - "- **Tracer:** Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to.\n", - "- **Parametric Source Models:** If your lens modeling uses a parametric source model (e.g.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulated Dataset__\n", - "\n", - "We load and plot the `simple__no_lens_light` example dataset, which is simulated imaging of a strong lens\n", - "that we will use to demonstrate source science caluculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We apply a 3.0 arcsecond circular mask and apply it to the `Imaging` object.\n", - "\n", - "Source science calculations are typically performed on masked datasets to ensure only the lensed source is used\n", - "in the calculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens__\n", - "\n", - "The simulated dataset was created using a lens galaxy with an Isothermal mass profile and External Shear,\n", - "which we now define." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multi Gaussian Expansion Source__\n", - "\n", - "The default workspace source model is a Multi Gaussian Expansion (MGE) profile, which is a superposition of 20\n", - "Gaussians whose `intensity` values are solved for via linear alegbra. \n", - "\n", - "These features are described in the `autolens_workspace/*/imaging/features/multi_gaussian_expansion`\n", - "and `autolens_workspace/*/imaging/features/linear_light_profiles` examples, but you do not need a full\n", - "understanding of these to follow this example on source science calculations.\n", - "\n", - "We now set up a source galaxy using an MGE, use it to compute the source flux and magnification and compare\n", - "them to the true values computed above. We will then consider whether making different assumptions about the source\n", - "model (e.g. MGE versus Sersic) changes the inferred source science calculations.\n", - "\n", - "We first set up a source galaxy using an MGE made up of 20 Gaussians whose `sigma` values span 0.01\" to the mask \n", - "radius of 3.0\"." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 20\n", - "\n", - "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", - "\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "# A list of linear light profile Gaussians will be input here, which will then be used to fit the data.\n", - "\n", - "bulge_gaussian_list = []\n", - "\n", - "# Iterate over every Gaussian and create it, with it centered at (0.0\", 0.0\") and assuming spherical symmetry.\n", - "\n", - "for i in range(total_gaussians):\n", - " gaussian = al.lp_linear.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - "\n", - " bulge_gaussian_list.append(gaussian)\n", - "\n", - "# The Basis object groups many light profiles together into a single model component and is used to fit the data.\n", - "\n", - "bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, bulge=bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now create the tracer using the lens galaxy and MGE source galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Each of the Gaussians above does not have a manually input intensity, instead their intensities are solved for via\n", - "linear algebra when we use the tracer to fit strong lens imaging data.\n", - "\n", - "We therefore create a fit of the tracer to the simulated dataset, which solves for the combination of Gaussian intensities\n", - "that best reconstruct the lensed source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From the fit, we can extract a tracer where the Gaussian intensities have been set to their best-fit values.\n", - "\n", - "We therefore now have the tracer which we need to perform source science calculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "However, lets first make a few plots confirming that the MGE source gives a good fit, and showing the individual \n", - "Gaussians that make up the MGE (which uses a lower sized grid to make the Gaussians more visible)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "grid_basis_plot = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "subplot_basis_image(basis=tracer.galaxies[1].bulge, grid=grid_basis_plot)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Flux__\n", - "\n", - "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", - "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", - "\n", - "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", - "we compute the image of the light profile, and then sum the image. \n", - "\n", - "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", - "we will assume everything is in electrons per second (`e- s^-1`), which is typical for Hubble Space Telescope imaging data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Source Galaxy's Intensity {source_galaxy.bulge.intensity} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The total flux, in units of `e- s^-1` , is computed by summing the image of the light profile over all pixels.\n", - "\n", - "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", - "and high enough resolution that it captures all the light from the light profile.\n", - "\n", - "Note that we are using the source galaxy's true light profile, which corresponds to its emission in the source-plane.\n", - "For real datasets, we have to infer this via lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "source_galaxy = tracer.galaxies[1]\n", - "\n", - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels\n", - "\n", - "print(f\"Total Source Flux: {total_flux} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Below, we will compare how this true source flux compares to the inferred source fluxes we compute using different\n", - "source modeling techniques (e.g. parametric and pixelized source models). Converting the flux to magnitudes or\n", - "other quantities used for tasks like SED fitting is described in the `autolens_workspace/*/guides/units/flux` example.\n", - "\n", - "__Source Magnification__\n", - "\n", - "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", - "total surface brightness in the source-plane.\n", - "\n", - "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", - "We therefore explicitly mention how area folds into the calculation below.\n", - "\n", - "To ensure the magnification is stable and that we resolve all source emission in both the image-plane and source-plane \n", - "we use a very high resolution grid, higher than we used to compute the total flux above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We repeat our calculation of the source's total flux in the source-plane using this higher resolution grid, note\n", - "that we do not take the area into account, the reason for this is explained below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_source_plane_flux = np.sum(image) # in units e- s^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now need the total flux of the lensed source in the image-plane, that is how much flux we measure after\n", - "gravitational lensing.\n", - "\n", - "To calculation this, we first ray-trace the grid above from the image-plane to the source-plane using the tracer\n", - "and then pass it to the source galaxy's light profile to compute the lensed image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "source_plane_grid = traced_grid_list[1]\n", - "\n", - "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(\n", - " lensed_source_image\n", - ") # in units e- s^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now take the ratio of the total image-plane flux to source-plane flux to estimate the magnification.\n", - "\n", - "Because both fluxes were computed on grids with the same total area and area per pixel, we do not need to\n", - "explicitly account for area in this calculation. This is because the area terms cancel out when taking the ratio.\n", - "Were the grid areas different, we would need to include area terms in the calculation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Magnification: {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to compute\n", - "source science calculations for real datasets.\n", - "\n", - "The code below shows how using a tracer, composed of any combination of lens and source galaxies, we can\n", - "compute the source flux and magnification. It reproduces the calculations above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "image_plane_grid = traced_grid_list[0]\n", - "source_plane_grid = traced_grid_list[1]\n", - "\n", - "lensed_source_image = tracer.planes[1].image_2d_from(grid=source_plane_grid)\n", - "source_plane_image = tracer.planes[1].image_2d_from(grid=image_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(lensed_source_image)\n", - "total_source_plane_flux = np.sum(source_plane_image)\n", - "\n", - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Plane Total Flux via Tracer: {total_source_plane_flux} e- s^-1\")\n", - "print(f\"Source Magnification via Tracer: {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parametric Source Models__\n", - "\n", - "If your lens modeling uses a parametric source model (e.g. Sersic, Multi Gaussian Expansion), the only object\n", - "you need to perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", - "\n", - "Alternatively, as done above, you can manually set up a tracer using the lens and source galaxies inferred\n", - "by lens modeling.\n", - "\n", - "Therefore, you may now wish to go to your results, extract the `max_log_likelihood_tracer`, and use it to compute\n", - "the source flux and magnification as shown above.\n", - "\n", - "__Wrap Up__\n", - "\n", - "In this example, the MGE Source gave very similar results for the source flux and magnification as did the Sersic\n", - "example illustrated in the `autolens_workspace/*/imaging/features/source_science.py` example.\n", - "\n", - "In real lenses, this is not always the case, and if you are really interested in precise source science calculations\n", - "which estimate systematic uncertainties, you should explore how different source models impact these calculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Source Science\n", + "==============\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", + "size of the source.\n", + "\n", + "This example shows how to perform these calculations using Multi Gaussian Expansion (MGE) sources on imaging data,\n", + "which is conceptually the simplest case for source science calculations and a good introduction to the topic.\n", + "\n", + "__Contents__\n", + "\n", + "- **Simulated Dataset:** We load and plot the `simple__no_lens_light` example dataset, which is simulated imaging of a.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Lens:** The simulated dataset was created using a lens galaxy with an Isothermal mass profile and External.\n", + "- **Multi Gaussian Expansion Source:** The default workspace source model is a Multi Gaussian Expansion (MGE) profile, which is a.\n", + "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", + "- **Source Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", + "- **Tracer:** Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to.\n", + "- **Parametric Source Models:** If your lens modeling uses a parametric source model (e.g.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulated Dataset__\n", + "\n", + "We load and plot the `simple__no_lens_light` example dataset, which is simulated imaging of a strong lens\n", + "that we will use to demonstrate source science caluculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We apply a 3.0 arcsecond circular mask and apply it to the `Imaging` object.\n", + "\n", + "Source science calculations are typically performed on masked datasets to ensure only the lensed source is used\n", + "in the calculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens__\n", + "\n", + "The simulated dataset was created using a lens galaxy with an Isothermal mass profile and External Shear,\n", + "which we now define." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multi Gaussian Expansion Source__\n", + "\n", + "The default workspace source model is a Multi Gaussian Expansion (MGE) profile, which is a superposition of 20\n", + "Gaussians whose `intensity` values are solved for via linear alegbra. \n", + "\n", + "These features are described in the `autolens_workspace/*/imaging/features/multi_gaussian_expansion`\n", + "and `autolens_workspace/*/imaging/features/linear_light_profiles` examples, but you do not need a full\n", + "understanding of these to follow this example on source science calculations.\n", + "\n", + "We now set up a source galaxy using an MGE, use it to compute the source flux and magnification and compare\n", + "them to the true values computed above. We will then consider whether making different assumptions about the source\n", + "model (e.g. MGE versus Sersic) changes the inferred source science calculations.\n", + "\n", + "We first set up a source galaxy using an MGE made up of 20 Gaussians whose `sigma` values span 0.01\" to the mask \n", + "radius of 3.0\"." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 20\n", + "\n", + "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0\".\n", + "\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "# A list of linear light profile Gaussians will be input here, which will then be used to fit the data.\n", + "\n", + "bulge_gaussian_list = []\n", + "\n", + "# Iterate over every Gaussian and create it, with it centered at (0.0\", 0.0\") and assuming spherical symmetry.\n", + "\n", + "for i in range(total_gaussians):\n", + " gaussian = al.lp_linear.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + "\n", + " bulge_gaussian_list.append(gaussian)\n", + "\n", + "# The Basis object groups many light profiles together into a single model component and is used to fit the data.\n", + "\n", + "bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, bulge=bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now create the tracer using the lens galaxy and MGE source galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each of the Gaussians above does not have a manually input intensity, instead their intensities are solved for via\n", + "linear algebra when we use the tracer to fit strong lens imaging data.\n", + "\n", + "We therefore create a fit of the tracer to the simulated dataset, which solves for the combination of Gaussian intensities\n", + "that best reconstruct the lensed source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the fit, we can extract a tracer where the Gaussian intensities have been set to their best-fit values.\n", + "\n", + "We therefore now have the tracer which we need to perform source science calculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "However, lets first make a few plots confirming that the MGE source gives a good fit, and showing the individual \n", + "Gaussians that make up the MGE (which uses a lower sized grid to make the Gaussians more visible)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "grid_basis_plot = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "subplot_basis_image(basis=tracer.galaxies[1].bulge, grid=grid_basis_plot)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Flux__\n", + "\n", + "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", + "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", + "\n", + "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", + "we compute the image of the light profile, and then sum the image. \n", + "\n", + "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", + "we will assume everything is in electrons per second (`e- s^-1`), which is typical for Hubble Space Telescope imaging data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Source Galaxy's Intensity {source_galaxy.bulge.intensity} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The total flux, in units of `e- s^-1` , is computed by summing the image of the light profile over all pixels.\n", + "\n", + "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", + "and high enough resolution that it captures all the light from the light profile.\n", + "\n", + "Note that we are using the source galaxy's true light profile, which corresponds to its emission in the source-plane.\n", + "For real datasets, we have to infer this via lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "source_galaxy = tracer.galaxies[1]\n", + "\n", + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels\n", + "\n", + "print(f\"Total Source Flux: {total_flux} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below, we will compare how this true source flux compares to the inferred source fluxes we compute using different\n", + "source modeling techniques (e.g. parametric and pixelized source models). Converting the flux to magnitudes or\n", + "other quantities used for tasks like SED fitting is described in the `autolens_workspace/*/guides/units/flux` example.\n", + "\n", + "__Source Magnification__\n", + "\n", + "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", + "total surface brightness in the source-plane.\n", + "\n", + "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", + "We therefore explicitly mention how area folds into the calculation below.\n", + "\n", + "To ensure the magnification is stable and that we resolve all source emission in both the image-plane and source-plane \n", + "we use a very high resolution grid, higher than we used to compute the total flux above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We repeat our calculation of the source's total flux in the source-plane using this higher resolution grid, note\n", + "that we do not take the area into account, the reason for this is explained below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_source_plane_flux = np.sum(image) # in units e- s^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now need the total flux of the lensed source in the image-plane, that is how much flux we measure after\n", + "gravitational lensing.\n", + "\n", + "To calculation this, we first ray-trace the grid above from the image-plane to the source-plane using the tracer\n", + "and then pass it to the source galaxy's light profile to compute the lensed image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "source_plane_grid = traced_grid_list[1]\n", + "\n", + "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(\n", + " lensed_source_image\n", + ") # in units e- s^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now take the ratio of the total image-plane flux to source-plane flux to estimate the magnification.\n", + "\n", + "Because both fluxes were computed on grids with the same total area and area per pixel, we do not need to\n", + "explicitly account for area in this calculation. This is because the area terms cancel out when taking the ratio.\n", + "Were the grid areas different, we would need to include area terms in the calculation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Magnification: {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to compute\n", + "source science calculations for real datasets.\n", + "\n", + "The code below shows how using a tracer, composed of any combination of lens and source galaxies, we can\n", + "compute the source flux and magnification. It reproduces the calculations above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "image_plane_grid = traced_grid_list[0]\n", + "source_plane_grid = traced_grid_list[1]\n", + "\n", + "lensed_source_image = tracer.planes[1].image_2d_from(grid=source_plane_grid)\n", + "source_plane_image = tracer.planes[1].image_2d_from(grid=image_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(lensed_source_image)\n", + "total_source_plane_flux = np.sum(source_plane_image)\n", + "\n", + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Plane Total Flux via Tracer: {total_source_plane_flux} e- s^-1\")\n", + "print(f\"Source Magnification via Tracer: {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parametric Source Models__\n", + "\n", + "If your lens modeling uses a parametric source model (e.g. Sersic, Multi Gaussian Expansion), the only object\n", + "you need to perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", + "\n", + "Alternatively, as done above, you can manually set up a tracer using the lens and source galaxies inferred\n", + "by lens modeling.\n", + "\n", + "Therefore, you may now wish to go to your results, extract the `max_log_likelihood_tracer`, and use it to compute\n", + "the source flux and magnification as shown above.\n", + "\n", + "__Wrap Up__\n", + "\n", + "In this example, the MGE Source gave very similar results for the source flux and magnification as did the Sersic\n", + "example illustrated in the `autolens_workspace/*/imaging/features/source_science.py` example.\n", + "\n", + "In real lenses, this is not always the case, and if you are really interested in precise source science calculations\n", + "which estimate systematic uncertainties, you should explore how different source models impact these calculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/no_lens_light/modeling.ipynb b/notebooks/imaging/features/no_lens_light/modeling.ipynb index 626ee45c5..12b4b9963 100644 --- a/notebooks/imaging/features/no_lens_light/modeling.ipynb +++ b/notebooks/imaging/features/no_lens_light/modeling.ipynb @@ -1,469 +1,506 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: No Lens Light\n", - "================================\n", - "\n", - "CCD imaging data of a strong lens may not have lens galaxy light emission present, for example if the lens galaxy light\n", - "has already been subtracted from the image.\n", - "\n", - "This example illustrates how to fit a lens model to data where the lens galaxy's light is not present.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** The main advantage of fitting data without lens light is the reduction in the number of free.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "The main advantage of fitting data without lens light is the reduction in the number of free parameters in the\n", - "model-fit. In the `start_here.py` example, the lens galaxy's light profile was composed of a `bulge` and `disk` which\n", - "were both Sersic profiles, which contributed an extra N=12 free parameters. Without the lens light we therefore\n", - "reduce the dimensionality of non-linear parameter space by 12, which makes the model-fit much faster and more efficient.\n", - "\n", - "__Disadvantages__\n", - "\n", - "If the lens light was simply not present in the observed data (e.g. it was too faint at the wavelength of the\n", - "observation) then there is no option but to fit the data without lens light, and therefore there is also no disadvantage.\n", - "\n", - "If you are fitting data where the lens light was subtracted before lens modeling (e.g. using a Sersic\n", - "subtraction or GUI / b-splines based fitting), there are disadvantages.\n", - "\n", - "The lens light subtraction process may not be clean. The lens and source light are blended with one another in the\n", - "data (e.g. due to PSF convlution). The process may over subtract source light in certain regions of the data and\n", - "fail to subtract lens light in other regions. This will distort the lensed source emission input into the lens\n", - "modeling and therefore potentially bias the inferred mass model.\n", - "\n", - "Performing lens and source modeling simultanoeusly means that the model-fit can find the optimal deblending of the two\n", - "components whilst fully propagating the errors on the inferred model parameters forward.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a Multi Gaussian Expansion.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple__no_lens_light` via .fits files" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Whereas we normally apply adaptive over sampling for the lens light, when its not present we do not need to.\n", - "\n", - "__Fit__\n", - "\n", - "This is to illustrate the API for performing a fit without lens light using standard autolens objects like \n", - "the `Galaxy`, `Tracer` and `FitImaging`.\n", - "\n", - "We simply do not input a `bulge` with a light profile into the `lens`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitImaging(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the lens light is not included in the model fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is a Multi Gaussian Expansion [6 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14.\n", - "\n", - "The lens galaxy does not include a light profile `bulge` or `disk` component, and thus its emission is not fitted for.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the lens galaxy's light is omitted from the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"no_lens_light\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", - "required by a model.\n", - "\n", - "Removing lens light from the model reduces VRAM use modestly, but likely wont have a noticeable impact on overall use.\n", - "\n", - "__Run Time__\n", - "\n", - "The likelihood evaluation time for fits to data without lens light are only small bit faster than fits to data with\n", - "lens light. This is because the most computationally expensive steps (e.g. computing the deflection angles, blurring\n", - "the image with the PSF) are performed for both model-fits.\n", - "\n", - "However, the overall run-time will be faster than before, as the removal of the lens light reduces the dimensionality\n", - "of non-linear parameter space by 7 or more parameters. This means that the non-linear search will more efficiently\n", - "converge on the highest likelihood regions of parameter space.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms there is no lens galaxy light in the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "__Wrap Up__\n", - "\n", - "This script shows how to fit a lens model to data where the lens galaxy's light is not present.\n", - "\n", - "It was a straightforward extension to the modeling API illustrated in `start_here.ipynb`, where one simply removed\n", - "the light profiles from the lens galaxy's model.\n", - "\n", - "Models where the source has no light, or other components of the model are omitted can also be easily composed using\n", - "the same API manipulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: No Lens Light\n", + "================================\n", + "\n", + "CCD imaging data of a strong lens may not have lens galaxy light emission present, for example if the lens galaxy light\n", + "has already been subtracted from the image.\n", + "\n", + "This example illustrates how to fit a lens model to data where the lens galaxy's light is not present.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** The main advantage of fitting data without lens light is the reduction in the number of free.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "The main advantage of fitting data without lens light is the reduction in the number of free parameters in the\n", + "model-fit. In the `start_here.py` example, the lens galaxy's light profile was composed of a `bulge` and `disk` which\n", + "were both Sersic profiles, which contributed an extra N=12 free parameters. Without the lens light we therefore\n", + "reduce the dimensionality of non-linear parameter space by 12, which makes the model-fit much faster and more efficient.\n", + "\n", + "__Disadvantages__\n", + "\n", + "If the lens light was simply not present in the observed data (e.g. it was too faint at the wavelength of the\n", + "observation) then there is no option but to fit the data without lens light, and therefore there is also no disadvantage.\n", + "\n", + "If you are fitting data where the lens light was subtracted before lens modeling (e.g. using a Sersic\n", + "subtraction or GUI / b-splines based fitting), there are disadvantages.\n", + "\n", + "The lens light subtraction process may not be clean. The lens and source light are blended with one another in the\n", + "data (e.g. due to PSF convlution). The process may over subtract source light in certain regions of the data and\n", + "fail to subtract lens light in other regions. This will distort the lensed source emission input into the lens\n", + "modeling and therefore potentially bias the inferred mass model.\n", + "\n", + "Performing lens and source modeling simultanoeusly means that the model-fit can find the optimal deblending of the two\n", + "components whilst fully propagating the errors on the inferred model parameters forward.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a Multi Gaussian Expansion.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple__no_lens_light` via .fits files" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Whereas we normally apply adaptive over sampling for the lens light, when its not present we do not need to.\n", + "\n", + "__Fit__\n", + "\n", + "This is to illustrate the API for performing a fit without lens light using standard autolens objects like \n", + "the `Galaxy`, `Tracer` and `FitImaging`.\n", + "\n", + "We simply do not input a `bulge` with a light profile into the `lens`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitImaging(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the lens light is not included in the model fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is a Multi Gaussian Expansion [6 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14.\n", + "\n", + "The lens galaxy does not include a light profile `bulge` or `disk` component, and thus its emission is not fitted for.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the lens galaxy's light is omitted from the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"no_lens_light\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM\n", + "required by a model.\n", + "\n", + "Removing lens light from the model reduces VRAM use modestly, but likely wont have a noticeable impact on overall use.\n", + "\n", + "__Run Time__\n", + "\n", + "The likelihood evaluation time for fits to data without lens light are only small bit faster than fits to data with\n", + "lens light. This is because the most computationally expensive steps (e.g. computing the deflection angles, blurring\n", + "the image with the PSF) are performed for both model-fits.\n", + "\n", + "However, the overall run-time will be faster than before, as the removal of the lens light reduces the dimensionality\n", + "of non-linear parameter space by 7 or more parameters. This means that the non-linear search will more efficiently\n", + "converge on the highest likelihood regions of parameter space.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms there is no lens galaxy light in the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "__Wrap Up__\n", + "\n", + "This script shows how to fit a lens model to data where the lens galaxy's light is not present.\n", + "\n", + "It was a straightforward extension to the modeling API illustrated in `start_here.ipynb`, where one simply removed\n", + "the light profiles from the lens galaxy's model.\n", + "\n", + "Models where the source has no light, or other components of the model are omitted can also be easily composed using\n", + "the same API manipulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/no_lens_light/simulator.ipynb b/notebooks/imaging/features/no_lens_light/simulator.ipynb index 2bfabcd33..d51816961 100644 --- a/notebooks/imaging/features/no_lens_light/simulator.ipynb +++ b/notebooks/imaging/features/no_lens_light/simulator.ipynb @@ -1,479 +1,516 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: No Lens Light\n", - "========================\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' which is identical to the `simple` simulated in the `start_here.py`\n", - "script, but where the lens galaxy's light is omitted.\n", - "\n", - "It is used in `autolens_workspace/notebooks/modeling/features/no_lens_light.ipynb` to illustrate how to fit a\n", - "lens model to data where the lens galaxy's light is not present (e.g. because it is too faint to be detected).\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Mask Extra Galaxies:** Save a `mask_extra_galaxies.fits` covering the extra galaxy for noise-scaling tutorials.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an `Sersic`.\n", - " - A faint extra galaxy is included offset from the lens, whose emission must be removed via noise scaling\n", - " (a `mask_extra_galaxies.fits` covering it is written below).\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The centre of a faint extra galaxy, placed inside the 3.0\" modeling mask but clear of the lensed source arcs\n", - "(Einstein radius ~1.6\"). It is reused for over-sampling, the galaxy itself and the `mask_extra_galaxies.fits`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_centre = (2.2, 1.6)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0), extra_galaxy_centre],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "\n", - "the `lens_galaxy` below does not include a `bulge` or `disk` component and therefore has no lens light." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A single faint extra galaxy offset from the lens, representing a nearby contaminating object. It has a light\n", - "profile only (no mass), so the lensed source arcs are unchanged; its emission is removed in the fit examples via\n", - "the `__Extra Galaxies Noise Scaling__` step using the `mask_extra_galaxies.fits` written below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " light=al.lp.ExponentialSph(\n", - " centre=extra_galaxy_centre, intensity=1.0, effective_radius=0.3\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, extra_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask Extra Galaxies__\n", - "\n", - "Build and output a `mask_extra_galaxies.fits` covering the extra galaxy, so the fit example (`imaging/fit.py`)\n", - "and the pixelization tutorials that load this dataset (`imaging/features/pixelization/modeling.py`,\n", - "`imaging/features/pixelization/fit.py`) can demonstrate the noise-scaling API on a real contaminant.\n", - "\n", - "The circle is sized to ~3x the galaxy's `effective_radius`, derived from the same `extra_galaxy_centre` defined\n", - "above so it stays in sync. The mask shape tracks `dataset.shape_native`, so `PYAUTO_SMALL_DATASETS=1` is honoured\n", - "automatically." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_extra_galaxies = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " centre=extra_galaxy_centre,\n", - " radius=3.0 * 0.3,\n", - " invert=True, # `True` inside the circle, i.e. the region whose noise is scaled.\n", - ")\n", - "\n", - "aplt.fits_array(\n", - " array=mask_extra_galaxies,\n", - " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future.\n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multiple Images__\n", - "\n", - "Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x) \n", - "coordinates of multiple images of a source galaxy to positions which are far apart from one another \n", - "in the source plane are penalized in the lens model's overall likelihood.\n", - "\n", - "This speeds up lens modeling, helps the non-linear search avoid local maxima and is vital for inferred \n", - "accurate solutions when using pixelized source reconstructions.\n", - "\n", - "For real data, the multiple image positions are determined by eye from the data, for example\n", - "using a Graphical User Interface (GUI) to mark them with mouse clicks. For simulated data, we can save\n", - "ourselves time by using the `PointSolver` to determine the multiple image positions automatically and\n", - "output to a .json file.\n", - "\n", - "If you have not looked in the `point_source` package, the point solver is the core tool used to find\n", - "multiple image positions for point source lens modeling (e.g. lensed quasars)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " file_path=dataset_path / \"positions.json\",\n", - " obj=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: No Lens Light\n", + "========================\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' which is identical to the `simple` simulated in the `start_here.py`\n", + "script, but where the lens galaxy's light is omitted.\n", + "\n", + "It is used in `autolens_workspace/notebooks/modeling/features/no_lens_light.ipynb` to illustrate how to fit a\n", + "lens model to data where the lens galaxy's light is not present (e.g. because it is too faint to be detected).\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Mask Extra Galaxies:** Save a `mask_extra_galaxies.fits` covering the extra galaxy for noise-scaling tutorials.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an `Sersic`.\n", + " - A faint extra galaxy is included offset from the lens, whose emission must be removed via noise scaling\n", + " (a `mask_extra_galaxies.fits` covering it is written below).\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The centre of a faint extra galaxy, placed inside the 3.0\" modeling mask but clear of the lensed source arcs\n", + "(Einstein radius ~1.6\"). It is reused for over-sampling, the galaxy itself and the `mask_extra_galaxies.fits`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_centre = (2.2, 1.6)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0), extra_galaxy_centre],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "\n", + "the `lens_galaxy` below does not include a `bulge` or `disk` component and therefore has no lens light." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A single faint extra galaxy offset from the lens, representing a nearby contaminating object. It has a light\n", + "profile only (no mass), so the lensed source arcs are unchanged; its emission is removed in the fit examples via\n", + "the `__Extra Galaxies Noise Scaling__` step using the `mask_extra_galaxies.fits` written below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " light=al.lp.ExponentialSph(\n", + " centre=extra_galaxy_centre, intensity=1.0, effective_radius=0.3\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, extra_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask Extra Galaxies__\n", + "\n", + "Build and output a `mask_extra_galaxies.fits` covering the extra galaxy, so the fit example (`imaging/fit.py`)\n", + "and the pixelization tutorials that load this dataset (`imaging/features/pixelization/modeling.py`,\n", + "`imaging/features/pixelization/fit.py`) can demonstrate the noise-scaling API on a real contaminant.\n", + "\n", + "The circle is sized to ~3x the galaxy's `effective_radius`, derived from the same `extra_galaxy_centre` defined\n", + "above so it stays in sync. The mask shape tracks `dataset.shape_native`, so `PYAUTO_SMALL_DATASETS=1` is honoured\n", + "automatically." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_extra_galaxies = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " centre=extra_galaxy_centre,\n", + " radius=3.0 * 0.3,\n", + " invert=True, # `True` inside the circle, i.e. the region whose noise is scaled.\n", + ")\n", + "\n", + "aplt.fits_array(\n", + " array=mask_extra_galaxies,\n", + " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future.\n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multiple Images__\n", + "\n", + "Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x) \n", + "coordinates of multiple images of a source galaxy to positions which are far apart from one another \n", + "in the source plane are penalized in the lens model's overall likelihood.\n", + "\n", + "This speeds up lens modeling, helps the non-linear search avoid local maxima and is vital for inferred \n", + "accurate solutions when using pixelized source reconstructions.\n", + "\n", + "For real data, the multiple image positions are determined by eye from the data, for example\n", + "using a Graphical User Interface (GUI) to mark them with mouse clicks. For simulated data, we can save\n", + "ourselves time by using the `PointSolver` to determine the multiple image positions automatically and\n", + "output to a .json file.\n", + "\n", + "If you have not looked in the `point_source` package, the point solver is the core tool used to find\n", + "multiple image positions for point source lens modeling (e.g. lensed quasars)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " file_path=dataset_path / \"positions.json\",\n", + " obj=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/no_lens_light/slam.ipynb b/notebooks/imaging/features/no_lens_light/slam.ipynb index 3525ac24a..81c39251b 100644 --- a/notebooks/imaging/features/no_lens_light/slam.ipynb +++ b/notebooks/imaging/features/no_lens_light/slam.ipynb @@ -1,560 +1,597 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "No Lens Light: SLaM\n", - "====================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for fitting a\n", - "lens model where there is no lens light observed in the imaging data.\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", - "guide before working through this example.\n", - "\n", - "This example only provides documentation specific to the use of SLaM for data without lens light, describing\n", - "how the pipeline differs from the standard SLaM pipelines described in the SLaM start here guide.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py` with `lens_bulge=None` to omit lens light from the model.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py` except no lens light is included in the model.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "You can still run the script without fully understanding the guide, but reviewing it later will\n", - "make the structure and choices of the SLaM workflow clearer.\n", - "\n", - "__Model__\n", - "\n", - "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE and a MASS TOTAL PIPELINE this SLaM script fits `Imaging`\n", - "of a strong lens system, where in the final model:\n", - "\n", - " - The lens galaxy's light is omitted from the data and model.\n", - " - The lens galaxy's total mass distribution is an `PowerLaw`.\n", - " - The source galaxy is reconstructed using a `RectangularAdaptImage` mesh and `Adapt` regularization scheme.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `slam_start_here` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py` with `lens_bulge=None` to omit lens light from the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`. Because no lens light components were included in `source_lp`, they\n", - "are automatically omitted from the lens galaxy model here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`. Note that the LIGHT LP PIPELINE from `slam_start_here` is not included\n", - "as there is no lens light to model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py` except no lens light is included in the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "No Lens Light: SLaM\n", + "====================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for fitting a\n", + "lens model where there is no lens light observed in the imaging data.\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", + "guide before working through this example.\n", + "\n", + "This example only provides documentation specific to the use of SLaM for data without lens light, describing\n", + "how the pipeline differs from the standard SLaM pipelines described in the SLaM start here guide.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py` with `lens_bulge=None` to omit lens light from the model.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py` except no lens light is included in the model.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "You can still run the script without fully understanding the guide, but reviewing it later will\n", + "make the structure and choices of the SLaM workflow clearer.\n", + "\n", + "__Model__\n", + "\n", + "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE and a MASS TOTAL PIPELINE this SLaM script fits `Imaging`\n", + "of a strong lens system, where in the final model:\n", + "\n", + " - The lens galaxy's light is omitted from the data and model.\n", + " - The lens galaxy's total mass distribution is an `PowerLaw`.\n", + " - The source galaxy is reconstructed using a `RectangularAdaptImage` mesh and `Adapt` regularization scheme.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `slam_start_here` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py` with `lens_bulge=None` to omit lens light from the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`. Because no lens light components were included in `source_lp`, they\n", + "are automatically omitted from the lens galaxy model here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`. Note that the LIGHT LP PIPELINE from `slam_start_here` is not included\n", + "as there is no lens light to model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py` except no lens light is included in the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/pixelization/adaptive.ipynb b/notebooks/imaging/features/pixelization/adaptive.ipynb index 33e82a7c9..066b96742 100644 --- a/notebooks/imaging/features/pixelization/adaptive.ipynb +++ b/notebooks/imaging/features/pixelization/adaptive.ipynb @@ -1,692 +1,729 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: Adaptive\n", - "======================\n", - "\n", - "Non-linear search chaining is an advanced model-fitting approach which breaks the model-fitting procedure down into\n", - "multiple non-linear searches, using the results of the initial searches to initialization parameter sampling in\n", - "subsequent searches. This contrasts the `modeling` example which fits a single lens model-fit using one non-linear search.\n", - "\n", - "An overview of search chaining is provided in the `autolens_workspace/*/guides/modeling/chaining` script, make\n", - "sure to read that before reading this script!\n", - "\n", - "This script introduces adaptive pixdelizations features, which use search chainig to pass the results of previous\n", - "model-fits performed by earlier searches to searches performed later in the chain, in order to adapt the pixelizaiton's\n", - "mesh and regularization to the source's unlensed properties. It also uses the results of previous searches to\n", - "calculate and pass the multiple image-plane positions of the lensed source to later searches, which resamples\n", - "bad mass models removing demagnified source reconstructions.\n", - "\n", - "This script illustrates using the `RectangularAdaptImage` mesh and `Adapt` regularization\n", - "scheme to adapt the source reconstruction to the source galaxy's morphology (as opposed to the methods used in other\n", - "examplesw hich adapt to the mass model magnification and apply a constant regularization scheme).\n", - "\n", - "This script illustrates the API used for adaptive pixelizations, but does not go into the details of how they\n", - "work. This is described in chapter 4 of the **HowToLens** lectures.\n", - "\n", - "__Why Chain?__\n", - "\n", - "There are a number of benefits of chaining a linear source model and a pixelized source, as opposed to fitting the\n", - "pixelization in one search:\n", - "\n", - " - Parametric sources are computationally faster to fit. Therefore, even though the MGE has more\n", - " parameters for the search to fit than a pixelized source, the model-fit is faster overall.\n", - "\n", - " - pixelizations often go to unphysical solutions where the mass model goes to high / low normalization_list and the source\n", - " is reconstructed as a demagnified version of the image. (see Chapter 4, tutorial 6 for a complete description of\n", - " this effect). This does not occur for a linear source, therefore the mass model can be initialized using a\n", - " parametric source, which sets up the search which fits a pixelization so as to not sample these unphysical solutions.\n", - "\n", - " - The positions and positions threshold can be updated to further ensure these unphysical solutions do not bias the\n", - " model-fit. The updated positions use the maximum log likelihood mass model of the first search to determine the\n", - " image-plane position of the lensed source. In the second search, we then require that a mass model must trace these\n", - " positions within a threshold arc-secoond value of one another in the source-plane, removing these unphysical solutions.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** This script chains three searches to fit `Imaging` data of a 'galaxy-scale' strong lens with a model where the lens galaxy's total mass distribution is an `Isothermal` and the source galaxy's light uses an MGE followed by a pixelization.\n", - "- **Dataset + Masking + Positions:** Load, plot and mask the `Imaging` data.\n", - "- **Paths:** The path the results of all chained searches are output.\n", - "- **Model (Search 1):** Search 1 fits a lens model where the lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` and the source galaxy's light is an MGE.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling.\n", - "- **Analysis + Position Likelihood:** We add a penalty term to the likelihood function, which penalizes models where the brightest multiple images of the lensed source galaxy do not trace close to one another in the source plane.\n", - "- **Brief Description:** In this example we update the positions between searches, where the positions correspond to the (y,x) locations of the lensed source's multiple images.\n", - "- **Adaptive Pixelization:** Search 3 uses two adaptive pixelization classes: `RectangularAdaptImage` mesh and `Adapt` regularization, which adapt the source reconstruction to the source galaxy's morphology.\n", - "- **Adapt Images:** When we create the analysis, we pass it an `adapt_images`, which contains a dictionary mapping each galaxy name to the corresponding lens subtracted image of the source galaxy from the result of a previous search.\n", - "- **SLaM Pipelines:** The API above allows you to write modeling code using adaptive features yourself, but it is recommended you use the Source, Light and Mass (SLaM) pipeline.\n", - "\n", - "__Model__\n", - "\n", - "This script chains three searches to fit `Imaging` data of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a multi-Gaussian expansion (MGE) in search 1 and a pixelization in searches 2 and 3.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking + Positions__ \n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "if not (dataset_path / \"positions.json\").exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Paths__\n", - "\n", - "The path the results of all chained searches are output:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "path_prefix = Path(\"imaging\") / \"pixelization\" / \"adaptive\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 1)__\n", - "\n", - "Search 1 fits a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=11.\n", - "\n", - "The benefit of using an MGE in search 1 is that it is computationally fast to fit, allowing the\n", - "non-linear search to quickly converge to a reasonable lens model. This lens model is then used \n", - "to set up the adaptive pixelization and multiple image positions in search 2." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_1.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", - "provides a reasonably accurate lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_1 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[1]__parametric\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_like_max=500,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "To use adapt features, we require a model image of the lensed source galaxy, which is what the code will adapt the\n", - "analysis too.\n", - "\n", - "When we begin a fit, we do not have such an image, and thus cannot use the adaptive features. This is why search chaining\n", - "is required, it allows us to perform an initial model-fit which gives us the source image, which we can then use to\n", - "perform a subsequent model-fit which adapts the analysis to the source's properties.\n", - "\n", - "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In the first\n", - "search our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", - " \n", - " - The source galaxy's light uses no image-mesh (only used for Delaunay meshes) [0 parameters].\n", - " \n", - " - The source-galaxy's light uses a 20 x 20 `RectangularAdaptDensity` mesh [0 parameters].\n", - "\n", - " - This pixelization is regularized using a `Constant` scheme [1 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=8.\n", - "\n", - "**Chaining API:** The term `model` below passes the source model as model-components that are to be fitted for by the \n", - "non-linear search. We pass the `lens` as a `model`, so that we can use the mass model inferred by search 1. The source\n", - "does not use any priors from the result of search 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = result_1.model.galaxies.lens\n", - "\n", - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", - " regularization=al.reg.Constant,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "model_2 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model_2.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis + Position Likelihood__\n", - "\n", - "We add a penalty term ot the likelihood function, which penalizes models where the brightest multiple images of\n", - "the lensed source galaxy do not trace close to one another in the source plane. This removes \"demagnified source\n", - "solutions\" from the source pixelization, which one is likely to infer without this penalty.\n", - "\n", - "A comprehensive description of why we do this is given at the following readthedocs page. I strongly recommend you \n", - "read this page in full if you are not familiar with the positions likelihood penalty and demagnified source reconstructions:\n", - "\n", - " https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", - "\n", - "__Brief Description__\n", - "\n", - "In this example we update the positions between searches, where the positions correspond to the (y,x) locations of the \n", - "lensed source's multiple images. When a model-fit uses positions, it requires them to trace within a threshold value of \n", - "one another for every mass model sampled by the non-linear search. If they do not, a penalty term is added to the\n", - "likelihood penalizing that solution \n", - "\n", - "Below, we use the results of the first search to compute the lensed source positions that are input into search 2. The\n", - "code below uses the maximum log likelihood model mass model and source galaxy centre, to determine where the source\n", - "positions are located in the image-plane. \n", - "\n", - "We also use this result to set the `threshold`, whereby the threshold value is based on how close these positions \n", - "trace to one another in the source-plane (using the best-fit mass model again). This threshold is multiplied by \n", - "a `factor` to ensure it is not too small (and thus does not remove plausible mass models). If, after this \n", - "multiplication, the threshold is below the `minimum_threshold`, it is rounded up to this minimum value." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_2 = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[\n", - " result_1.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis + Model-Fit (Search 1)__\n", - "\n", - "We now create the non-linear search, analysis and perform the model-fit using this model.\n", - "\n", - "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", - "provides a reasonably accurate lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[2]__adaptive_pixelization_setup\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_like_max=500,\n", - ")\n", - "\n", - "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Adaptive Pixelization__\n", - "\n", - "Search 3 uses two adaptive pixelization classes that have not been used elsewhere in the workspace:\n", - "\n", - " - `RectangularAdaptImage` mesh: adapts the rectangular source-pixel upsampling to the source's unlensed morphology. This \n", - " means that more rectangular pixels will be used where the source is located, even if its far away from the caustic\n", - " and therefore in lower magnification regions.\n", - "\n", - " - `Adapt` regularization: adapts the regularization coefficient to the source's\n", - " unlensed morphology. This means that the source's brightest regions are regularized less than its faintest regions, \n", - " ensuring that the bright central regions of the source are not over-smoothed.\n", - " \n", - "This adaptive mesh and regularization produces a significantly better lens analysis and reconstruction of the source \n", - "galaxy than other schemes used throughout the workspace. Now you are familiar with them, you should\n", - "never use anything else!\n", - "\n", - "It is recommend that the parameters governing these features are always fitted using a fixed lens light and\n", - "mass model. This ensures the adaptation is performed quickly, and removes degeneracies in the lens model that\n", - "are difficult to sample. Extensive testing has shown that this does not reduce the accuracy of the lens model.\n", - "\n", - "For this reason, search 2 fixes the lens galaxy's light and mass model to the best-fit model of search 1. A third\n", - "search will then fit for the lens galaxy's light and mass model using these adaptive features.\n", - "\n", - "The details of how the above features work is not provided here, but is given at the end of chapter 4 of the HowToLens\n", - "lecture series.\n", - "\n", - "__Model (Search 3)__\n", - "\n", - "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In \n", - "the second search our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` with fixed parameters from \n", - " search 1 [0 parameters].\n", - " \n", - " - The source galaxy's light uses no image-mesh (only used for Delaunay meshes) [0 parameters].\n", - " \n", - " - The source-galaxy's light uses a 20 x 20 `RectangularAdaptImage` mesh [0 parameters].\n", - "\n", - " - This pixelization is regularized using a `Adapt` scheme [2 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=4.\n", - "\n", - "**Chaining API:** The term `instance` below passes the lens model as fixed model-components that are not to be\n", - "fitted for by the non-linear search. We pass the `lens` as an `instance`, so that its parameters are fixed to \n", - "the best-fit values of search 2. The source" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = result_2.instance.galaxies.lens\n", - "\n", - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.RectangularAdaptImage(shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - ")\n", - "\n", - "source = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " pixelization=pixelization,\n", - ")\n", - "\n", - "model_3 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis (Search 2)__\n", - "\n", - "We now create the analysis for the second search.\n", - "\n", - "__Adapt Images__\n", - "\n", - "When we create the analysis, we pass it an `adapt_images`, which contains a dictionary mapping each galaxy name \n", - "(e.g. galaxies.source) to the corresponding lens subtracted image of the source galaxy from the result of search 1. \n", - "\n", - "This is telling the `Analysis` class to use the lens subtracted images of this fit to guide the `Adapt` \n", - "regularization for the source galaxy. Specifically, it uses the lens subtracted signal to noise map of the lensed \n", - "source in order to adapt the location of the source-pixels to the source's brightest regions and lower the \n", - "regularization coefficient in these regions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result_2)\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "analysis_3 = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " result_2.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Model-Fit (Search 3)__\n", - "\n", - "We now create the non-linear search and perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_3 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[3]__adaptive_pixelization\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - " n_like_max=500,\n", - ")\n", - "\n", - "result_3 = search_3.fit(model=model_3, analysis=analysis_3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result (Search 3)__\n", - "\n", - "If you inspect and compare the results of searches 2 and 3, you'll note how the model-fits of search 3 have a much\n", - "higher likelihood than search 2 and how the source reconstruction has congregated it pixels to the bright central\n", - "regions of the source. This indicates that a much better result has been achieved.\n", - "\n", - "__Model + Search + Analysis + Model-Fit (Search 4)__\n", - "\n", - "We now perform a final search which uses the `Adapt` regularization with their parameter fixed to the \n", - "results of search 2.\n", - "\n", - "The lens mass model is free to vary.\n", - "\n", - "The analysis class still uses the adapt images from search 2, because this is what the adaptive features adapted\n", - "to in search 3.\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=7." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", - ")\n", - "\n", - "source = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " pixelization=result_3.instance.galaxies.source.pixelization,\n", - ")\n", - "\n", - "model_4 = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "search_4 = af.Nautilus(\n", - " path_prefix=path_prefix,\n", - " name=\"search[4]__adapt\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_4 = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - ")\n", - "\n", - "result_4 = search_4.fit(model=model_4, analysis=analysis_4)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipelines__\n", - "\n", - "The API above allows you to write modeling code using adaptive features yourself.\n", - "\n", - "However, it is recommend you use the Source, Light and Mass (SLaM) pipeline, whcih are carefully crafted to automate \n", - "lens modeling of large samples whilst ensuring models of the highest complexity can be reliably fitted.\n", - "\n", - "The SLaM pipelines are built around the use of these adaptive pixelization features, with the Source pipeline first \n", - "so that these features are set up robustly before more complex lens light and mass models are fitted.\n", - "\n", - "The example `guides/modeling/slam_start_here` provides a full run through of how to use the SLaM pipelines with \n", - "adaptive pixelizations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: Adaptive\n", + "======================\n", + "\n", + "Non-linear search chaining is an advanced model-fitting approach which breaks the model-fitting procedure down into\n", + "multiple non-linear searches, using the results of the initial searches to initialization parameter sampling in\n", + "subsequent searches. This contrasts the `modeling` example which fits a single lens model-fit using one non-linear search.\n", + "\n", + "An overview of search chaining is provided in the `autolens_workspace/*/guides/modeling/chaining` script, make\n", + "sure to read that before reading this script!\n", + "\n", + "This script introduces adaptive pixdelizations features, which use search chainig to pass the results of previous\n", + "model-fits performed by earlier searches to searches performed later in the chain, in order to adapt the pixelizaiton's\n", + "mesh and regularization to the source's unlensed properties. It also uses the results of previous searches to\n", + "calculate and pass the multiple image-plane positions of the lensed source to later searches, which resamples\n", + "bad mass models removing demagnified source reconstructions.\n", + "\n", + "This script illustrates using the `RectangularAdaptImage` mesh and `Adapt` regularization\n", + "scheme to adapt the source reconstruction to the source galaxy's morphology (as opposed to the methods used in other\n", + "examplesw hich adapt to the mass model magnification and apply a constant regularization scheme).\n", + "\n", + "This script illustrates the API used for adaptive pixelizations, but does not go into the details of how they\n", + "work. This is described in chapter 4 of the **HowToLens** lectures.\n", + "\n", + "__Why Chain?__\n", + "\n", + "There are a number of benefits of chaining a linear source model and a pixelized source, as opposed to fitting the\n", + "pixelization in one search:\n", + "\n", + " - Parametric sources are computationally faster to fit. Therefore, even though the MGE has more\n", + " parameters for the search to fit than a pixelized source, the model-fit is faster overall.\n", + "\n", + " - pixelizations often go to unphysical solutions where the mass model goes to high / low normalization_list and the source\n", + " is reconstructed as a demagnified version of the image. (see Chapter 4, tutorial 6 for a complete description of\n", + " this effect). This does not occur for a linear source, therefore the mass model can be initialized using a\n", + " parametric source, which sets up the search which fits a pixelization so as to not sample these unphysical solutions.\n", + "\n", + " - The positions and positions threshold can be updated to further ensure these unphysical solutions do not bias the\n", + " model-fit. The updated positions use the maximum log likelihood mass model of the first search to determine the\n", + " image-plane position of the lensed source. In the second search, we then require that a mass model must trace these\n", + " positions within a threshold arc-secoond value of one another in the source-plane, removing these unphysical solutions.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** This script chains three searches to fit `Imaging` data of a 'galaxy-scale' strong lens with a model where the lens galaxy's total mass distribution is an `Isothermal` and the source galaxy's light uses an MGE followed by a pixelization.\n", + "- **Dataset + Masking + Positions:** Load, plot and mask the `Imaging` data.\n", + "- **Paths:** The path the results of all chained searches are output.\n", + "- **Model (Search 1):** Search 1 fits a lens model where the lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` and the source galaxy's light is an MGE.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling.\n", + "- **Analysis + Position Likelihood:** We add a penalty term to the likelihood function, which penalizes models where the brightest multiple images of the lensed source galaxy do not trace close to one another in the source plane.\n", + "- **Brief Description:** In this example we update the positions between searches, where the positions correspond to the (y,x) locations of the lensed source's multiple images.\n", + "- **Adaptive Pixelization:** Search 3 uses two adaptive pixelization classes: `RectangularAdaptImage` mesh and `Adapt` regularization, which adapt the source reconstruction to the source galaxy's morphology.\n", + "- **Adapt Images:** When we create the analysis, we pass it an `adapt_images`, which contains a dictionary mapping each galaxy name to the corresponding lens subtracted image of the source galaxy from the result of a previous search.\n", + "- **SLaM Pipelines:** The API above allows you to write modeling code using adaptive features yourself, but it is recommended you use the Source, Light and Mass (SLaM) pipeline.\n", + "\n", + "__Model__\n", + "\n", + "This script chains three searches to fit `Imaging` data of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a multi-Gaussian expansion (MGE) in search 1 and a pixelization in searches 2 and 3.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/chaining.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking + Positions__ \n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "if not (dataset_path / \"positions.json\").exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Paths__\n", + "\n", + "The path the results of all chained searches are output:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "path_prefix = Path(\"imaging\") / \"pixelization\" / \"adaptive\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 1)__\n", + "\n", + "Search 1 fits a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=11.\n", + "\n", + "The benefit of using an MGE in search 1 is that it is computationally fast to fit, allowing the\n", + "non-linear search to quickly converge to a reasonable lens model. This lens model is then used \n", + "to set up the adaptive pixelization and multiple image positions in search 2." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_1.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", + "provides a reasonably accurate lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_1 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[1]__parametric\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_like_max=500,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "To use adapt features, we require a model image of the lensed source galaxy, which is what the code will adapt the\n", + "analysis too.\n", + "\n", + "When we begin a fit, we do not have such an image, and thus cannot use the adaptive features. This is why search chaining\n", + "is required, it allows us to perform an initial model-fit which gives us the source image, which we can then use to\n", + "perform a subsequent model-fit which adapts the analysis to the source's properties.\n", + "\n", + "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In the first\n", + "search our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", + " \n", + " - The source galaxy's light uses no image-mesh (only used for Delaunay meshes) [0 parameters].\n", + " \n", + " - The source-galaxy's light uses a 20 x 20 `RectangularAdaptDensity` mesh [0 parameters].\n", + "\n", + " - This pixelization is regularized using a `Constant` scheme [1 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=8.\n", + "\n", + "**Chaining API:** The term `model` below passes the source model as model-components that are to be fitted for by the \n", + "non-linear search. We pass the `lens` as a `model`, so that we can use the mass model inferred by search 1. The source\n", + "does not use any priors from the result of search 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = result_1.model.galaxies.lens\n", + "\n", + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", + " regularization=al.reg.Constant,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "model_2 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model, including how parameters and priors were passed from `result_1`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model_2.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis + Position Likelihood__\n", + "\n", + "We add a penalty term ot the likelihood function, which penalizes models where the brightest multiple images of\n", + "the lensed source galaxy do not trace close to one another in the source plane. This removes \"demagnified source\n", + "solutions\" from the source pixelization, which one is likely to infer without this penalty.\n", + "\n", + "A comprehensive description of why we do this is given at the following readthedocs page. I strongly recommend you \n", + "read this page in full if you are not familiar with the positions likelihood penalty and demagnified source reconstructions:\n", + "\n", + " https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", + "\n", + "__Brief Description__\n", + "\n", + "In this example we update the positions between searches, where the positions correspond to the (y,x) locations of the \n", + "lensed source's multiple images. When a model-fit uses positions, it requires them to trace within a threshold value of \n", + "one another for every mass model sampled by the non-linear search. If they do not, a penalty term is added to the\n", + "likelihood penalizing that solution \n", + "\n", + "Below, we use the results of the first search to compute the lensed source positions that are input into search 2. The\n", + "code below uses the maximum log likelihood model mass model and source galaxy centre, to determine where the source\n", + "positions are located in the image-plane. \n", + "\n", + "We also use this result to set the `threshold`, whereby the threshold value is based on how close these positions \n", + "trace to one another in the source-plane (using the best-fit mass model again). This threshold is multiplied by \n", + "a `factor` to ensure it is not too small (and thus does not remove plausible mass models). If, after this \n", + "multiplication, the threshold is below the `minimum_threshold`, it is rounded up to this minimum value." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_2 = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[\n", + " result_1.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis + Model-Fit (Search 1)__\n", + "\n", + "We now create the non-linear search, analysis and perform the model-fit using this model.\n", + "\n", + "You may wish to inspect the results of the search 1 model-fit to ensure a fast non-linear search has been provided that \n", + "provides a reasonably accurate lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[2]__adaptive_pixelization_setup\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_like_max=500,\n", + ")\n", + "\n", + "analysis_2 = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Adaptive Pixelization__\n", + "\n", + "Search 3 uses two adaptive pixelization classes that have not been used elsewhere in the workspace:\n", + "\n", + " - `RectangularAdaptImage` mesh: adapts the rectangular source-pixel upsampling to the source's unlensed morphology. This \n", + " means that more rectangular pixels will be used where the source is located, even if its far away from the caustic\n", + " and therefore in lower magnification regions.\n", + "\n", + " - `Adapt` regularization: adapts the regularization coefficient to the source's\n", + " unlensed morphology. This means that the source's brightest regions are regularized less than its faintest regions, \n", + " ensuring that the bright central regions of the source are not over-smoothed.\n", + " \n", + "This adaptive mesh and regularization produces a significantly better lens analysis and reconstruction of the source \n", + "galaxy than other schemes used throughout the workspace. Now you are familiar with them, you should\n", + "never use anything else!\n", + "\n", + "It is recommend that the parameters governing these features are always fitted using a fixed lens light and\n", + "mass model. This ensures the adaptation is performed quickly, and removes degeneracies in the lens model that\n", + "are difficult to sample. Extensive testing has shown that this does not reduce the accuracy of the lens model.\n", + "\n", + "For this reason, search 2 fixes the lens galaxy's light and mass model to the best-fit model of search 1. A third\n", + "search will then fit for the lens galaxy's light and mass model using these adaptive features.\n", + "\n", + "The details of how the above features work is not provided here, but is given at the end of chapter 4 of the HowToLens\n", + "lecture series.\n", + "\n", + "__Model (Search 3)__\n", + "\n", + "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In \n", + "the second search our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` with fixed parameters from \n", + " search 1 [0 parameters].\n", + " \n", + " - The source galaxy's light uses no image-mesh (only used for Delaunay meshes) [0 parameters].\n", + " \n", + " - The source-galaxy's light uses a 20 x 20 `RectangularAdaptImage` mesh [0 parameters].\n", + "\n", + " - This pixelization is regularized using a `Adapt` scheme [2 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=4.\n", + "\n", + "**Chaining API:** The term `instance` below passes the lens model as fixed model-components that are not to be\n", + "fitted for by the non-linear search. We pass the `lens` as an `instance`, so that its parameters are fixed to \n", + "the best-fit values of search 2. The source" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = result_2.instance.galaxies.lens\n", + "\n", + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.RectangularAdaptImage(shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + ")\n", + "\n", + "source = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " pixelization=pixelization,\n", + ")\n", + "\n", + "model_3 = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis (Search 2)__\n", + "\n", + "We now create the analysis for the second search.\n", + "\n", + "__Adapt Images__\n", + "\n", + "When we create the analysis, we pass it an `adapt_images`, which contains a dictionary mapping each galaxy name \n", + "(e.g. galaxies.source) to the corresponding lens subtracted image of the source galaxy from the result of search 1. \n", + "\n", + "This is telling the `Analysis` class to use the lens subtracted images of this fit to guide the `Adapt` \n", + "regularization for the source galaxy. Specifically, it uses the lens subtracted signal to noise map of the lensed \n", + "source in order to adapt the location of the source-pixels to the source's brightest regions and lower the \n", + "regularization coefficient in these regions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result_2)\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "analysis_3 = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " result_2.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Model-Fit (Search 3)__\n", + "\n", + "We now create the non-linear search and perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_3 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[3]__adaptive_pixelization\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + " n_like_max=500,\n", + ")\n", + "\n", + "result_3 = search_3.fit(model=model_3, analysis=analysis_3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result (Search 3)__\n", + "\n", + "If you inspect and compare the results of searches 2 and 3, you'll note how the model-fits of search 3 have a much\n", + "higher likelihood than search 2 and how the source reconstruction has congregated it pixels to the bright central\n", + "regions of the source. This indicates that a much better result has been achieved.\n", + "\n", + "__Model + Search + Analysis + Model-Fit (Search 4)__\n", + "\n", + "We now perform a final search which uses the `Adapt` regularization with their parameter fixed to the \n", + "results of search 2.\n", + "\n", + "The lens mass model is free to vary.\n", + "\n", + "The analysis class still uses the adapt images from search 2, because this is what the adaptive features adapted\n", + "to in search 3.\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=7." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", + ")\n", + "\n", + "source = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " pixelization=result_3.instance.galaxies.source.pixelization,\n", + ")\n", + "\n", + "model_4 = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "search_4 = af.Nautilus(\n", + " path_prefix=path_prefix,\n", + " name=\"search[4]__adapt\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_4 = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + ")\n", + "\n", + "result_4 = search_4.fit(model=model_4, analysis=analysis_4)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipelines__\n", + "\n", + "The API above allows you to write modeling code using adaptive features yourself.\n", + "\n", + "However, it is recommend you use the Source, Light and Mass (SLaM) pipeline, whcih are carefully crafted to automate \n", + "lens modeling of large samples whilst ensuring models of the highest complexity can be reliably fitted.\n", + "\n", + "The SLaM pipelines are built around the use of these adaptive pixelization features, with the Source pipeline first \n", + "so that these features are set up robustly before more complex lens light and mass models are fitted.\n", + "\n", + "The example `guides/modeling/slam_start_here` provides a full run through of how to use the SLaM pipelines with \n", + "adaptive pixelizations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/pixelization/cpu_fast_modeling.ipynb b/notebooks/imaging/features/pixelization/cpu_fast_modeling.ipynb index 8f4c07cc2..4bfa3faeb 100644 --- a/notebooks/imaging/features/pixelization/cpu_fast_modeling.ipynb +++ b/notebooks/imaging/features/pixelization/cpu_fast_modeling.ipynb @@ -1,792 +1,829 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: CPU Fast Modeling\n", - "===============================\n", - "\n", - "This example demonstrates how to achieve **fast pixelization performance on a CPU without JAX**, by combining:\n", - "\n", - "- `numba` for optimized numerical routines, and\n", - "- Python `multiprocessing` to exploit multiple CPU cores.\n", - "\n", - "On machines with many CPU cores (e.g. HPC clusters with >10 CPUs), this method can **outperform JAX GPU acceleration**\n", - "for pixelized source modeling. The advantage arises because pixelizations rely heavily on **sparse linear algebra**,\n", - "which is not currently optimized in JAX.\n", - "\n", - "> Note: This performance advantage applies **only to pixelized sources**.\n", - "> For parametric sources or multi-Gaussian models, JAX (especially with a GPU) is significantly faster, and even JAX\n", - "> on a CPU outperforms the `numba` approach shown here.\n", - "\n", - "__Contents__\n", - "\n", - "- **Sparse Operators:** Pixelized source modeling requires dense linear algebra operations.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **SLaM Pipeline:** The example `guides/modeling//slam_start_here.ipynb` introduces the SLaM (Source, Light and Mass).\n", - "- **SLaM Pipeline Functions:** Overview of slam pipeline functions for this example.\n", - "- **CPU Fast SLaM Pipelines:** The SLaM pipeline is mostly identical to other examples, but via the `SettingsSearch` it uses a.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "try:\n", - " import numba\n", - "except ModuleNotFoundError:\n", - " input(\n", - " \"##################\\n\"\n", - " \"##### NUMBA ######\\n\"\n", - " \"##################\\n\\n\"\n", - " \"\"\"\n", - " Numba is not currently installed.\n", - "\n", - " Numba is a library which makes PyAutoLens run a lot faster. Certain functionality is disabled without numba\n", - " and will raise an exception if it is used.\n", - "\n", - " If you have not tried installing numba, I recommend you try and do so now by running the following \n", - " commands in your command line / bash terminal now:\n", - "\n", - " pip install --upgrade pip\n", - " pip install numba\n", - "\n", - " If your numba installation raises an error and fails, you should go ahead and use PyAutoLens without numba to \n", - " decide if it is the right software for you. If it is, you should then commit time to bug-fixing the numba\n", - " installation. Feel free to raise an issue on GitHub for support with installing numba.\n", - "\n", - " A warning will crop up throughout your *PyAutoLens** use until you install numba, to remind you to do so.\n", - "\n", - " [Press Enter to continue]\n", - " \"\"\"\n", - " )\n", - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import matplotlib.pyplot as plt\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking + Positions__ \n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "if not (dataset_path / \"positions.json\").exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n", - "\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "Pixelized source modeling requires dense linear algebra operations. These calculations can be greatly accelerated\n", - "using an alternative mathematical approach called the **sparse operator formalism**.\n", - "\n", - "You do not need to understand the full details of the method, but the key point is:\n", - "\n", - "- It exploits the **sparsity** of the matrices used in pixelized source reconstruction, reducing memory usage.\n", - "- This leads to a **significant speed-up on CPUs**.\n", - "- The current implementation does **not support JAX**, and therefore does not benefit from GPU acceleration.\n", - "\n", - "To enable this feature, we call `apply_sparse_operator()` on the `Imaging` dataset. This computes and stores operator\n", - "matrices, which are then reused in all subsequent pixelized source fits.\n", - "\n", - "- Computing the operator matrices takes anywhere from a few seconds to a few minutes, depending on the dataset size.\n", - "- After it is computed once, every model-fit using pixelization becomes substantially faster." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_sparse_operator_cpu()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "In the example `imaging/features/pixelization/fit.py`, we demonstrated fitting imaging data using a\n", - "pixelized source with a rectangular mesh.\n", - "\n", - "Below, we perform a similar fit using the **same pixelization**, but this time accelerated on the **CPU**\n", - "using `numba` and sparse operations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", - "regularization = al.reg.Constant(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", - "\n", - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - ")\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We now perform a full model-fit using the sparse operator formalism on the CPU.\n", - "\n", - "There are two key differences from the earlier JAX-based pixelization examples:\n", - "\n", - "- **JAX is disabled** \n", - " The `AnalysisImaging` class is created with `use_jax=False`, preventing JAX compilation and ensuring\n", - " that all computations run on the CPU.\n", - "\n", - "- **CPU parallelization** \n", - " The non-linear search is given a `number_of_cores` parameter, which parallelizes likelihood evaluations\n", - " using Python's `multiprocessing`. \n", - " In practice, this provides a speed-up of roughly half the number of CPU cores used \n", - " (e.g., 4 cores \u2192 ~2\u00d7 speed-up, 8 cores \u2192 ~4\u00d7 speed-up)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", - ")\n", - "\n", - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", - " regularization=al.reg.Constant,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"features\"),\n", - " name=\"cpu_fast_modeling\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " number_of_cores=2, # CPU specific code\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " use_jax=False, # CPU specific code\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The example `guides/modeling//slam_start_here.ipynb` introduces the SLaM (Source, Light and Mass) pipelines for\n", - "automated lens modeling of large samples of strong lenses.\n", - "\n", - "We finish this example by showing how to run the SLaM pipelines using CPU acceleration with sparse operators, similar \n", - "to the model-fit above.\n", - "\n", - "Note that the first pipeline, SOURCE LP, uses JAX acceleration as in previous examples and therefore does not \n", - "pass `use_jax=False` or a `number_of_cores` parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# %%\n", - "'''\n", - "__SLaM Pipeline Functions__\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search,\n", - " analysis,\n", - " lens_bulge,\n", - " source_bulge,\n", - " redshift_lens,\n", - " redshift_source,\n", - " mass_centre=(0.0, 0.0),\n", - " n_batch=50,\n", - "):\n", - " \"\"\"\n", - " SOURCE LP PIPELINE: fits an initial lens model using a parametric source to establish a robust\n", - " lens light, mass and source model before pixelized source fitting.\n", - " \"\"\"\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = mass_centre\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def source_pix_1(\n", - " settings_search,\n", - " analysis,\n", - " source_lp_result,\n", - " mesh_shape,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SOURCE PIX PIPELINE 1: initializes a pixelized source model with mass priors from SOURCE LP PIPELINE,\n", - " run on CPU using sparse operators and multiprocessing.\n", - " \"\"\"\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_lp_result.model.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def source_pix_2(\n", - " settings_search,\n", - " analysis,\n", - " source_lp_result,\n", - " source_pix_result_1,\n", - " mesh_shape,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SOURCE PIX PIPELINE 2: fits an improved pixelized source using adapt images from SOURCE PIX PIPELINE 1,\n", - " with fixed lens mass, run on CPU using sparse operators and multiprocessing.\n", - " \"\"\"\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def light_lp(\n", - " settings_search,\n", - " analysis,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " lens_bulge,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " LIGHT LP PIPELINE: fits the lens galaxy light with mass and source fixed from SOURCE PIX PIPELINE,\n", - " run on CPU using sparse operators and multiprocessing.\n", - " \"\"\"\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def mass_total(\n", - " settings_search,\n", - " analysis,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " light_result,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " MASS TOTAL PIPELINE: fits a PowerLaw total mass model with priors from SOURCE PIX PIPELINE and\n", - " lens light fixed from LIGHT LP PIPELINE, run on CPU using sparse operators and multiprocessing.\n", - " \"\"\"\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=light_result.instance.galaxies.lens.bulge,\n", - " disk=light_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "redshift_lens = 0.5\n", - "redshift_source = 1.0\n", - "\n", - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam_cpu_fast\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")\n", - "\n", - "# Lens Light\n", - "\n", - "lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "# Source Light\n", - "\n", - "source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=False)\n", - "\n", - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " lens_bulge=lens_bulge,\n", - " source_bulge=source_bulge,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__CPU Fast SLaM Pipelines__\n", - "\n", - "The SLaM pipeline is mostly identical to other examples, but via the `SettingsSearch` it\n", - "uses a `number_of_cores` parameter to parallelize the likelihood evaluations on the CPU\n", - "and disables JAX compilation for each `AnalysisImaging` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam_cpu_fast\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - " number_of_cores=2,\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " use_jax=False, # CPU specific code\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_lp_result=source_lp_result,\n", - " mesh_shape=mesh_shape,\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=False, # CPU specific code\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh_shape=mesh_shape,\n", - ")\n", - "\n", - "lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=False, # CPU specific code\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " lens_bulge=lens_bulge,\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_2.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " use_jax=False, # CPU specific code\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This example has demonstrated how to perform fast pixelized source modeling on a CPU without JAX, by combining\n", - "`numba` and Python `multiprocessing`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: CPU Fast Modeling\n", + "===============================\n", + "\n", + "This example demonstrates how to achieve **fast pixelization performance on a CPU without JAX**, by combining:\n", + "\n", + "- `numba` for optimized numerical routines, and\n", + "- Python `multiprocessing` to exploit multiple CPU cores.\n", + "\n", + "On machines with many CPU cores (e.g. HPC clusters with >10 CPUs), this method can **outperform JAX GPU acceleration**\n", + "for pixelized source modeling. The advantage arises because pixelizations rely heavily on **sparse linear algebra**,\n", + "which is not currently optimized in JAX.\n", + "\n", + "> Note: This performance advantage applies **only to pixelized sources**.\n", + "> For parametric sources or multi-Gaussian models, JAX (especially with a GPU) is significantly faster, and even JAX\n", + "> on a CPU outperforms the `numba` approach shown here.\n", + "\n", + "__Contents__\n", + "\n", + "- **Sparse Operators:** Pixelized source modeling requires dense linear algebra operations.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **SLaM Pipeline:** The example `guides/modeling//slam_start_here.ipynb` introduces the SLaM (Source, Light and Mass).\n", + "- **SLaM Pipeline Functions:** Overview of slam pipeline functions for this example.\n", + "- **CPU Fast SLaM Pipelines:** The SLaM pipeline is mostly identical to other examples, but via the `SettingsSearch` it uses a.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "try:\n", + " import numba\n", + "except ModuleNotFoundError:\n", + " input(\n", + " \"##################\\n\"\n", + " \"##### NUMBA ######\\n\"\n", + " \"##################\\n\\n\"\n", + " \"\"\"\n", + " Numba is not currently installed.\n", + "\n", + " Numba is a library which makes PyAutoLens run a lot faster. Certain functionality is disabled without numba\n", + " and will raise an exception if it is used.\n", + "\n", + " If you have not tried installing numba, I recommend you try and do so now by running the following \n", + " commands in your command line / bash terminal now:\n", + "\n", + " pip install --upgrade pip\n", + " pip install numba\n", + "\n", + " If your numba installation raises an error and fails, you should go ahead and use PyAutoLens without numba to \n", + " decide if it is the right software for you. If it is, you should then commit time to bug-fixing the numba\n", + " installation. Feel free to raise an issue on GitHub for support with installing numba.\n", + "\n", + " A warning will crop up throughout your *PyAutoLens** use until you install numba, to remind you to do so.\n", + "\n", + " [Press Enter to continue]\n", + " \"\"\"\n", + " )\n", + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking + Positions__ \n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "if not (dataset_path / \"positions.json\").exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n", + "\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "Pixelized source modeling requires dense linear algebra operations. These calculations can be greatly accelerated\n", + "using an alternative mathematical approach called the **sparse operator formalism**.\n", + "\n", + "You do not need to understand the full details of the method, but the key point is:\n", + "\n", + "- It exploits the **sparsity** of the matrices used in pixelized source reconstruction, reducing memory usage.\n", + "- This leads to a **significant speed-up on CPUs**.\n", + "- The current implementation does **not support JAX**, and therefore does not benefit from GPU acceleration.\n", + "\n", + "To enable this feature, we call `apply_sparse_operator()` on the `Imaging` dataset. This computes and stores operator\n", + "matrices, which are then reused in all subsequent pixelized source fits.\n", + "\n", + "- Computing the operator matrices takes anywhere from a few seconds to a few minutes, depending on the dataset size.\n", + "- After it is computed once, every model-fit using pixelization becomes substantially faster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_sparse_operator_cpu()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "In the example `imaging/features/pixelization/fit.py`, we demonstrated fitting imaging data using a\n", + "pixelized source with a rectangular mesh.\n", + "\n", + "Below, we perform a similar fit using the **same pixelization**, but this time accelerated on the **CPU**\n", + "using `numba` and sparse operations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", + "regularization = al.reg.Constant(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", + "\n", + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + ")\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We now perform a full model-fit using the sparse operator formalism on the CPU.\n", + "\n", + "There are two key differences from the earlier JAX-based pixelization examples:\n", + "\n", + "- **JAX is disabled** \n", + " The `AnalysisImaging` class is created with `use_jax=False`, preventing JAX compilation and ensuring\n", + " that all computations run on the CPU.\n", + "\n", + "- **CPU parallelization** \n", + " The non-linear search is given a `number_of_cores` parameter, which parallelizes likelihood evaluations\n", + " using Python's `multiprocessing`. \n", + " In practice, this provides a speed-up of roughly half the number of CPU cores used \n", + " (e.g., 4 cores \u2192 ~2\u00d7 speed-up, 8 cores \u2192 ~4\u00d7 speed-up)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", + ")\n", + "\n", + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", + " regularization=al.reg.Constant,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"features\"),\n", + " name=\"cpu_fast_modeling\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " number_of_cores=2, # CPU specific code\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " use_jax=False, # CPU specific code\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The example `guides/modeling//slam_start_here.ipynb` introduces the SLaM (Source, Light and Mass) pipelines for\n", + "automated lens modeling of large samples of strong lenses.\n", + "\n", + "We finish this example by showing how to run the SLaM pipelines using CPU acceleration with sparse operators, similar \n", + "to the model-fit above.\n", + "\n", + "Note that the first pipeline, SOURCE LP, uses JAX acceleration as in previous examples and therefore does not \n", + "pass `use_jax=False` or a `number_of_cores` parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# %%\n", + "'''\n", + "__SLaM Pipeline Functions__\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search,\n", + " analysis,\n", + " lens_bulge,\n", + " source_bulge,\n", + " redshift_lens,\n", + " redshift_source,\n", + " mass_centre=(0.0, 0.0),\n", + " n_batch=50,\n", + "):\n", + " \"\"\"\n", + " SOURCE LP PIPELINE: fits an initial lens model using a parametric source to establish a robust\n", + " lens light, mass and source model before pixelized source fitting.\n", + " \"\"\"\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = mass_centre\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def source_pix_1(\n", + " settings_search,\n", + " analysis,\n", + " source_lp_result,\n", + " mesh_shape,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SOURCE PIX PIPELINE 1: initializes a pixelized source model with mass priors from SOURCE LP PIPELINE,\n", + " run on CPU using sparse operators and multiprocessing.\n", + " \"\"\"\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_lp_result.model.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def source_pix_2(\n", + " settings_search,\n", + " analysis,\n", + " source_lp_result,\n", + " source_pix_result_1,\n", + " mesh_shape,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SOURCE PIX PIPELINE 2: fits an improved pixelized source using adapt images from SOURCE PIX PIPELINE 1,\n", + " with fixed lens mass, run on CPU using sparse operators and multiprocessing.\n", + " \"\"\"\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def light_lp(\n", + " settings_search,\n", + " analysis,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " lens_bulge,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " LIGHT LP PIPELINE: fits the lens galaxy light with mass and source fixed from SOURCE PIX PIPELINE,\n", + " run on CPU using sparse operators and multiprocessing.\n", + " \"\"\"\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def mass_total(\n", + " settings_search,\n", + " analysis,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " light_result,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " MASS TOTAL PIPELINE: fits a PowerLaw total mass model with priors from SOURCE PIX PIPELINE and\n", + " lens light fixed from LIGHT LP PIPELINE, run on CPU using sparse operators and multiprocessing.\n", + " \"\"\"\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=light_result.instance.galaxies.lens.bulge,\n", + " disk=light_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "redshift_lens = 0.5\n", + "redshift_source = 1.0\n", + "\n", + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam_cpu_fast\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")\n", + "\n", + "# Lens Light\n", + "\n", + "lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "# Source Light\n", + "\n", + "source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=False)\n", + "\n", + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " lens_bulge=lens_bulge,\n", + " source_bulge=source_bulge,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__CPU Fast SLaM Pipelines__\n", + "\n", + "The SLaM pipeline is mostly identical to other examples, but via the `SettingsSearch` it\n", + "uses a `number_of_cores` parameter to parallelize the likelihood evaluations on the CPU\n", + "and disables JAX compilation for each `AnalysisImaging` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam_cpu_fast\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + " number_of_cores=2,\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " use_jax=False, # CPU specific code\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_lp_result=source_lp_result,\n", + " mesh_shape=mesh_shape,\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=False, # CPU specific code\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh_shape=mesh_shape,\n", + ")\n", + "\n", + "lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=False, # CPU specific code\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " lens_bulge=lens_bulge,\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_2.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " use_jax=False, # CPU specific code\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This example has demonstrated how to perform fast pixelized source modeling on a CPU without JAX, by combining\n", + "`numba` and Python `multiprocessing`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/pixelization/delaunay.ipynb b/notebooks/imaging/features/pixelization/delaunay.ipynb index f89c84eac..838e5118d 100644 --- a/notebooks/imaging/features/pixelization/delaunay.ipynb +++ b/notebooks/imaging/features/pixelization/delaunay.ipynb @@ -1,2000 +1,2037 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: Delaunay\n", - "======================\n", - "\n", - "The majority of pixelized source reconstructions in the workspace use a rectangular mesh to reconstruct\n", - "the source's surface brightness.\n", - "\n", - "This example illustrates an alternative pixelization that uses a Delaunay triangulation mesh to reconstruct the\n", - "source.\n", - "\n", - "The approach is distinct from the rectangular mesh and has a number of traits which are unique to it:\n", - "\n", - "- `Adaptive Mesh`: In the source plane, the Delaunay mesh uses irregularly shaped triangles to reconstruct the\n", - " source, as opposed to uniform rectangular pixels. This allows the mesh to better adapt to irregular and\n", - " asymmetric source morphologies and change the distribution of source pixels to better match the source's\n", - " surface brightness.\n", - "\n", - "- `Image Mesh`: The vertexes of the Delaunay triangles are computed by overlaying a coarse uniform grid in the\n", - " image-plane and ray-tracing these coordinates to the source-plane. This is unlike the rectangular mesh, which\n", - " simply overlays a uniform grid in the source-plane. This again helps the Delaunay mesh to better adapt to the\n", - " source's surface brightness.\n", - "\n", - "- `Interpolation`: The Delaunay mesh uses a different interpolation scheme to the rectangular mesh, which is\n", - " barycentric interpolation within each triangle. This is different to the rectangular mesh, which uses bilinear\n", - " interpolation within each rectangular pixel.\n", - "\n", - "- `Regularization`: The Delaunay mesh provides different approaches to regularization, with the default being\n", - " one which uses the barycentric coordinates of the triangles to compute how source pixels are regularized with\n", - " their neighbors.\n", - "\n", - "Currently it is not expected that the Delaunay is better or worse than the rectangular mesh, it is simply a different\n", - "approach to pixelization that may work better for certain datasets.\n", - "\n", - "__JAX + GPU__\n", - "\n", - "Generating a Delaunay mesh supports JAX and GPU acceleration, however certain operations (e.g. generating the Delaunay\n", - "triangulation itself) do not run on the GPU because they cannot be easily converted to JAX.\n", - "\n", - "Instead, JAX sends them to a CPU, runs them there, and then sends the results back to the GPU. This process is\n", - "very efficient, because these operations run very fast on a CPU and the data being sent back and forth is small.\n", - "Current benchmarking suggests the Delaunay runs less than twice as long as the same fit using a rectangular mesh,\n", - "but scientfically offers better results in many cases.\n", - "\n", - "If you do want to run only on CPU, you can use fast CPU method described in\n", - "example `imaging/features/pixelization/cpu_fast_modeling` with the Delaunay mesh.\n", - "\n", - "__Source Science (Magnification, Flux and More)__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", - "size of the source.\n", - "\n", - "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", - "including examples using a Delaunay source reconstruction. Once you have completed lens modeling using a Delaunay mesh,\n", - "you can jump to that example to study the source galaxy.\n", - "\n", - "__Contents__\n", - "\n", - "- **Image Mesh:** For a Delaunay mesh, the vertices of the triangles are defined by (y, x) coordinates in the.\n", - "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Source Science:** Source science focuses on studying the highly magnified properties of the background lensed source.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **JAX CPU:** On CPU, the Delaunay mesh computation via JAX can lead to slow down or indefinite freezing.\n", - "- **VRAM:** The `pixelization/modeling` example explains how VRAM use is an important consideration for.\n", - "- **Adaptive Delaunay:** The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to.\n", - "- **SLaM Pipelines:** The API above allows you to use adaptive features yourself, and you should go ahead an explore them.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "- **Likelihood Function:** The example `imaging/features/pixelization/likelihood_function.py` provides a step-by-step.\n", - "- **Source Galaxy Pixelization and Regularization:** The source galaxy is reconstructed using a pixel-grid, in this example a Delaunay mesh, which.\n", - "- **Lens Light:** Overview of lens light for this example.\n", - "- **Source Pixel Centre Calculation:** In order to reconstruct the source galaxy using a Delaunay mesh, we need to determine the centres.\n", - "- **Ray Tracing:** Overview of ray tracing for this example.\n", - "- **Border Relocation:** Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace.\n", - "- **Delaunay Mesh:** The relocated mesh grid is used to create the `Pixelization`'s Delaunay mesh using the.\n", - "- **Interpolation:** Overview of interpolation for this example.\n", - "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", - "- **Sub Gridding:** The calculation above uses a `Grid2D` object, with a `sub-size=1`, meaning it does not perform.\n", - "- **Sourrce Plane Interpolation:** For the `Delaunay` mesh used in this example, every image-sub pixel maps to a single source Voronoi.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "import matplotlib.pyplot as plt\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking + Positions__ \n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image Mesh__\n", - "\n", - "For a Delaunay mesh, the vertices of the triangles are defined by (y, x) coordinates in the image-plane. These\n", - "coordinates are then ray-traced into the source-plane for each mass model sampled during the non-linear search.\n", - "This `image_plane_mesh_grid` must be computed before lens modeling.\n", - "\n", - "We compute this `image_plane_mesh_grid` using an `Overlay` image-mesh, which places a regular grid of\n", - "(y, x) points across the image-plane. This has a mild adaptive effect: regions of high lens magnification receive\n", - "more source pixels once they are ray-traced. Later in this example, we switch to a `Hilbert` image-mesh, which adapts\n", - "the pixel distribution more strongly to the source\u2019s surface brightness.\n", - "\n", - "The `Delaunay` mesh has an input number of `pixels`, which is the number of source pixels used to reconstruct the \n", - "source. The number of `pixels` must be equal to the number of coordinates in the `image_plane_mesh_grid`. \n", - "\n", - "Like for the `mesh_shape` rectangular mesh, `pixels` must be fixed for lens modeling because JAX uses the \n", - "number of `pixels` to determine static array shapes. \n", - "\n", - "To pass the `image_plane_mesh_grid` to the modeling, we use the `AdaptImages` object below, which pairs\n", - "the `image_plane_mesh_grid` to the source galaxy. For double source plane lenses, this means we can\n", - "attach an `image_plane_mesh_grid` to each source galaxy and use adaptive meshes for each source plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Edge Zeroing__\n", - "\n", - "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", - "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", - "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", - "\n", - "For a rectangular mesh, the source code computes edge pixels internally using the known pixels at the edge of the mesh,\n", - "requiring no input from the user. \n", - "\n", - "For the `Delaunay` mesh, we use the `append_with_circle_edge_points` function to manually setup the Delaunay image \n", - "mesh to include a ring of edge pixels and then input the total number into the mesh to perform zeroing. \n", - "\n", - "These points are added to the edge of the image-plane mesh, ray-traced to the source-plane during lens modeling, \n", - "included in the Delaunay triangulation but zeroed during the inversion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "edge_pixels_total = 30\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "In the example `imaging/features/pixelization/fit.py`, we illustrate how to use a pixelized source\n", - "with a rectangular mesh to fit imaging data.\n", - "\n", - "Below, we use a Delaunay mesh to perform a fit using the Delaunay source reconstruction.\n", - "\n", - "The API is nearly identical to the rectangular mesh example, noting that the inputs to the `Delaunay` \n", - "mesh are different to the rectangular mesh and use image mesh quantities computed above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh = al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - ")\n", - "\n", - "regularization = al.reg.ConstantSplit(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", - "\n", - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "adapt_images = al.AdaptImages(\n", - " galaxy_image_plane_mesh_grid_dict={source: image_plane_mesh_grid}\n", - ")\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " adapt_images=adapt_images,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the Delaunay source does a good job at capturing the appearance of the source galaxy\n", - "using adaptive triangular pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Science__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source pixelization, we can compute key quantities such as the magnification, total flux, and\n", - "intrinsic size of the source.\n", - "\n", - "For rectangular meshes, the example `autolens_workspace/*/imaging/features/source_science` gives a complete overview of \n", - "how to this. \n", - "\n", - "For the Delaunay, specific functionality which manipulates the Delaunay triangles is required to perform these \n", - "calculations. We show below how to interpolate the Delaunay source reconstruction to a regular grid.\n", - "\n", - "When an inversion is performed on a Delaunay mesh, the reconstructed values are defined only at the irregular\n", - "mesh vertices (the Delaunay nodes). In order to visualise the source reconstruction as a regular 2D image\n", - "(e.g. for plotting or saving to a FITS file), we must interpolate these values onto a uniform grid.\n", - "\n", - "This is done using SciPy's Delaunay-based linear interpolation:\n", - "\n", - "- A `scipy.spatial.Delaunay` triangulation is constructed from the mesh node coordinates.\n", - "\n", - "- Within each triangle, interpolation is performed using *barycentric coordinates*, meaning the value at any point\n", - " inside the triangle is computed as a weighted linear combination of the values at the triangle\u2019s three vertices.\n", - "\n", - "- Points outside the convex hull of the triangulation are assigned a fallback value (here `fill_value=0.0`).\n", - "\n", - "__Important Coordinate Convention (Delaunay(__\n", - "\n", - "SciPy expects all Delaunay coordinates to be provided in **(x, y)** order.\n", - "\n", - "However, grids are stored internally in **(y, x)** order.\n", - "\n", - "Therefore:\n", - "\n", - "- The mesh coordinates must be converted to (x, y) before building the triangulation.\n", - "- The interpolation grid must also be flipped to (x, y) before evaluating the interpolator.\n", - "\n", - "This coordinate flip is essential for correct interpolation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.spatial import Delaunay\n", - "from scipy.interpolate import LinearNDInterpolator\n", - "\n", - "inversion = fit.inversion\n", - "\n", - "mapper = inversion.cls_list_from(cls=al.Mapper)[0]\n", - "\n", - "reconstruction = inversion.reconstruction\n", - "\n", - "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", - "\n", - "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", - "\n", - "mesh_grid_xy = np.stack([source_plane_mesh_grid[:, 0], source_plane_mesh_grid[:, 1]]).T\n", - "\n", - "# Uses find simplex so recomputes delaunay internally\n", - "delaunay = Delaunay(mesh_grid_xy)\n", - "\n", - "interp = LinearNDInterpolator(delaunay, reconstruction, fill_value=0.0)\n", - "\n", - "interpolation_grid_xy = np.asarray(interpolation_grid)[:, ::-1] # (y,x)->(x,y)\n", - "interpolated_reconstruction = interp(interpolation_grid_xy)\n", - "\n", - "print(f\"Brightest Interpolated Source Pixel: {np.max(interpolated_reconstruction)}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We now perform lens modeling using the Delaunay pixelization with the Overlay image-mesh.\n", - "\n", - "The code below is a simple adaptive modeling example using the Delaunay mesh, which mirrors the\n", - "API used in other pixelization modeling examples.\n", - "\n", - "The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to\n", - "adapt the rectangular mesh and its regularization to the source's surface brightness. In particular, an image\n", - "of the lensed source is passed to the modeling via the `AdaptImages` object, in order to adapt\n", - "the mesh and regularization during the model-fit.\n", - "\n", - "You have already seen this used once above, but we set up the adapt images again to remind you\n", - "of the API." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "adapt_images = al.AdaptImages(\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In the first\n", - "search our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", - " \n", - " - The source galaxy's light uses an `Overlay` image-mesh with fixed resolution 30 x 30 pixels [0 parameters].\n", - " \n", - " - The source-galaxy's light uses a `Delaunay` mesh [0 parameters].\n", - "\n", - " - This pixelization is regularized using a `Constant` scheme [1 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=8.\n", - "\n", - "__JAX CPU__\n", - "\n", - "On CPU, the Delaunay mesh computation via JAX can lead to slow down or indefinite freezing. \n", - "\n", - "You may find better performance if you set `use_jax_vmap=False` in the `Nautilus` search below, which\n", - "disables JAX's vectorization of certain computations.\n", - "\n", - "This does not impact GPU performance, which should always use `use_jax_vmap=True`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", - ")\n", - "\n", - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.ConstantSplit,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "search_1 = af.Nautilus(\n", - " path_prefix=Path(\"features\"),\n", - " name=\"delaunay\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=10,\n", - " # use_jax_vmap=False # Set to False if CPU performance is slow or hangs\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `pixelization/modeling` example explains how VRAM use is an important consideration for pixelization models\n", - "and how it depends on image resolution, number of source pixels and batch size.\n", - "\n", - "This is true for the Delaunay mesh, therefore we print out the estimated VRAM required for this model-fit.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_1.print_vram_use(model=model_1, batch_size=search_1.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit (Search 1)__\n", - "\n", - "Perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Adaptive Delaunay__\n", - "\n", - "The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to\n", - "adapt the rectangular mesh and its regularization to the source's surface brightness.\n", - "\n", - "The image-mesh has a special adaptive variant called the `Hilbert` image-mesh, which adapts the distribution \n", - "of source-pixels to the source's unlensed morphology. This means that the source's brightest regions are \n", - "reconstructed using significantly more source pixels than seen for the `Overlay` image mesh. \n", - "Conversely, the source's faintest regions are reconstructed using significantly fewer source pixels.\n", - "\n", - "Unlike the adaptive rectangular mesh, the Hilbert image-plane mesh is computed before modeling, passed\n", - "to the `AdaptImages` object, and remains fixed during the model-fit.\n", - "\n", - "It is recommend that the parameters governing these features are always fitted from using a fixed lens light and\n", - "mass model. This ensures the adaptation is performed quickly, and removes degeneracies in the lens model that\n", - "are difficult to sample. Given the Hilbert mesh is fixed, this modeling only fits for the regularization coefficients\n", - "of the adaptive regularization scheme.\n", - "\n", - "For this reason, search 2 fixes the lens galaxy's light and mass model to the best-fit model of search 1. A third\n", - "search will then fit for the lens galaxy's light and mass model using these adaptive features.\n", - "\n", - "The details of how the above features work is not provided here, but is given at the end of chapter 4 of the HowToLens\n", - "lecture series." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result_1)\n", - "\n", - "image_mesh = al.image_mesh.Hilbert(pixels=1000, weight_power=3.5, weight_floor=0.01)\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask, adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", - ")\n", - "\n", - "# Repeat edge zeroing set up describe above.\n", - "\n", - "edge_pixels_total = 30\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=mask.mask_centre,\n", - " radius=mask_radius + mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In \n", - "the second search our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` with fixed parameters from \n", - " search 1 [0 parameters].\n", - " \n", - " - The source galaxy's light uses a `Hilbert` image-mesh with fixed resolution 1000 pixels [2 parameters].\n", - " \n", - " - The source-galaxy's light uses a `Delaunay` mesh [0 parameters].\n", - "\n", - " - This pixelization is regularized using a `AdaptSplit` scheme [2 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=4." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.AdaptSplit,\n", - ")\n", - "\n", - "source = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " pixelization=pixelization,\n", - ")\n", - "\n", - "model_2 = af.Collection(\n", - " galaxies=af.Collection(lens=result_1.instance.galaxies.lens, source=source)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis (Search 2)__\n", - "\n", - "We now create the analysis for the second search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_2 = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Model-Fit (Search 2)__\n", - "\n", - "We now create the non-linear search and perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=Path(\"features\"),\n", - " name=\"delaunay_adapt\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - " # use_jax_vmap=False # Set to False if CPU performance is slow or hangs\n", - ")\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We could perform a third fit where we free all lens model parameters and fit them using the adaptive \n", - "image mesh and regularization.\n", - "\n", - "However, it is better to use all of these features with the Delaunay via the\n", - "SLaM pipelines, which we jump to immediately below.\n", - "\n", - "__SLaM Pipelines__\n", - "\n", - "The API above allows you to use adaptive features yourself, and you should go ahead an explore them on datasets you\n", - "are familiar with.\n", - "\n", - "However, you may also wish to use the Source, Light and Mass (SLaM) pipelines, which are pipelines that\n", - "have been carefully crafted to automate lens modeling of large samples whilst ensuring models of the highest\n", - "complexity can be reliably fitted.\n", - "\n", - "These pipelines are built around the use of adaptive features -- for example the Source pipeline comes first so that\n", - "these features are set up robustly before more complex lens light and mass models are fitted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", - "\n", - "The `source_pix_1` search uses an `Overlay` image-mesh to place the initial Delaunay mesh pixels, with\n", - "additional edge points added around the mask boundary to ensure full coverage." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_lp_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask,\n", - " )\n", - "\n", - " edge_pixels_total = 30\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=dataset.mask.mask_centre,\n", - " radius=mask_radius + dataset.mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0],\n", - " zeroed_pixels=edge_pixels_total,\n", - " ),\n", - " regularization=al.reg.AdaptSplit,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", - "\n", - "The `source_pix_2` search uses a `Hilbert` image-mesh to place the final Delaunay mesh pixels, which adapts\n", - "the mesh to the source morphology using the high-quality adapt images from search 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Hilbert(pixels=1000, weight_power=3.5, weight_floor=0.01)\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask, adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", - " )\n", - "\n", - " edge_pixels_total = 30\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=dataset.mask.mask_centre,\n", - " radius=mask_radius + dataset.mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0],\n", - " zeroed_pixels=edge_pixels_total,\n", - " ),\n", - " regularization=al.reg.AdaptSplit,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict=(\n", - " source_result_for_source.analysis.adapt_images.galaxy_name_image_plane_mesh_grid_dict\n", - " ),\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict=(\n", - " source_result_for_source.analysis.adapt_images.galaxy_name_image_plane_mesh_grid_dict\n", - " ),\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " bulge = light_result.instance.galaxies.lens.bulge\n", - " disk = light_result.instance.galaxies.lens.disk\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=bulge,\n", - " disk=disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam_delaunay\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - " # use_jax_vmap=False # Set to False if CPU performance is slow or hangs\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_lp_result=source_lp_result,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "The example `imaging/features/pixelization/likelihood_function.py` provides a step-by-step description of how\n", - "a likelihood evaluation is performed for imaging data using a pixelized source reconstruction with a rectangular\n", - "mesh.\n", - "\n", - "We now give the same step-by-step description for a pixelized source reconstruction using a Delaunay mesh and\n", - "adaptive features.\n", - "\n", - "We only describe code which is specific to Delaunay meshes and adaptive features -- for all other aspects of the likelihood\n", - "evaluation, refer to rectangular mesh example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=masked_dataset)\n", - "\n", - "masked_dataset = masked_dataset.apply_over_sampling(\n", - " over_sample_size_lp=1,\n", - " over_sample_size_pixelization=1,\n", - ")\n", - "\n", - "aplt.plot_grid(grid=masked_dataset.grids.pixelization, title=\"\")\n", - "\n", - "bulge = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Pixelization and Regularization__\n", - "\n", - "The source galaxy is reconstructed using a pixel-grid, in this example a Delaunay mesh, which accounts for \n", - "irregularities and asymmetries in the source's surface brightness. \n", - "\n", - "A constant regularization scheme is applied which applies a smoothness prior on the reconstruction. \n", - "\n", - "One of the biggest differences between a Delaunay mesh and rectangular mesh is how the centres of the mesh pixels\n", - "in the source-plane are computed. \n", - "\n", - "For the rectangular mesh, the pixel centres are computed by overlaying a uniform grid over the source-plane.\n", - "\n", - "For a Delaunay mesh, the uniform grid is instead laid over the image-plane to create a course grid of (y,x) coordinates.\n", - "These are then ray-traced to the source-plane and are used as the vertexes of the Delaunay triangles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=masked_dataset.grids.pixelization.shape[0],\n", - " zeroed_pixels=edge_pixels_total,\n", - " ),\n", - " regularization=al.reg.ConstantSplit(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Light__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", - "\n", - "\n", - "blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n", - "\n", - "\n", - "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", - " image=image, blurring_image=blurring_image_2d\n", - ")\n", - "\n", - "aplt.plot_array(array=convolved_image_2d, title=\"\")\n", - "\n", - "\n", - "lens_subtracted_image = masked_dataset.data - convolved_image_2d\n", - "\n", - "aplt.plot_array(array=lens_subtracted_image, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Pixel Centre Calculation__\n", - "\n", - "In order to reconstruct the source galaxy using a Delaunay mesh, we need to determine the centres of the Delaunay\n", - "source pixels.\n", - "\n", - "The image-mesh `Overlay` object computes the source-pixel centres in the image-plane (which are ray-traced to the\n", - "source-plane below). The source pixelization therefore adapts to the lens model magnification, because more\n", - "source pixels will congregate in higher magnification regions.\n", - "\n", - "This calculation is performed by overlaying a uniform regular grid with an `pixelization_shape_2d` over the image\n", - "mask and retaining all pixels that fall within the mask. This uses a `Grid2DSparse` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_mesh = al.image_mesh.Overlay(shape=(30, 30)) # Specific to Delaunay\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=masked_dataset.mask,\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(\n", - " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid},\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting this grid shows a sparse grid of (y,x) coordinates within the mask, which will form our source pixel centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=masked_dataset.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The source code gets quite complex when handling grids for a pixelization, but it is all handled in\n", - "the `TracerToInversion` objects.\n", - "\n", - "The plots at the bottom of this cell show the traced grids used by the source pixelization, showing\n", - "how the Delaunay mesh and traced image pixels are constructed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(\n", - " tracer=tracer, dataset=masked_dataset, adapt_images=adapt_images\n", - ")\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", - "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.pixelization\n", - ")[-1]\n", - "\n", - "# This functions a bit weird - it returns a list of lists of ndarrays. Best not to worry about it for now!\n", - "traced_mesh_grid = tracer_to_inversion.traced_mesh_grid_pg_list[-1][-1]\n", - "\n", - "\n", - "aplt.plot_grid(grid=traced_grid_pixelization, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We have also ray-traced the coarse grid of image-pixel coordinates used to form the source pixelization's\n", - "Delaunay mesh, which we can also plot." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=traced_mesh_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Border Relocation__\n", - "\n", - "Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace to far outskirts of\n", - "the source-plane. \n", - "\n", - "Border relocation is performed on both the traced image-pixel grid and traced mesh pixels, therefore ensuring that\n", - "the vertexes of the Delaunay triangles are not at the extreme outskirts of the source-plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", - "\n", - "border_relocator = BorderRelocator(mask=masked_dataset.mask, sub_size=1)\n", - "\n", - "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", - "\n", - "relocated_mesh_grid = border_relocator.relocated_mesh_grid_from(\n", - " grid=traced_grid_pixelization, mesh_grid=traced_mesh_grid\n", - ")\n", - "\n", - "\n", - "aplt.plot_grid(grid=relocated_grid, title=\"\")\n", - "\n", - "aplt.plot_grid(grid=relocated_mesh_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Delaunay Mesh__\n", - "\n", - "The relocated mesh grid is used to create the `Pixelization`'s Delaunay mesh using the `scipy.spatial` library." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "interpolator = al.InterpolatorDelaunay(\n", - " mesh=pixelization.mesh,\n", - " mesh_grid=relocated_mesh_grid,\n", - " data_grid=relocated_grid,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting the Delaunay mesh shows that the source-plane and been discretized into a grid of irregular Delaunay pixels.\n", - "\n", - "(To plot the Delaunay mesh, we have to convert it to a `Mapper` object, which is described in the next likelihood step).\n", - "\n", - "Below, we plot the Delaunay mesh without the traced image-grid pixels (for clarity) and with them as black dots in order\n", - "to show how each set of image-pixels fall within a Delaunay pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = al.Mapper(\n", - " interpolator=interpolator,\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - ")\n", - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolation__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pix_indexes_for_sub_slim_index = mapper.pix_indexes_for_sub_slim_index\n", - "\n", - "print(pix_indexes_for_sub_slim_index[0:9])\n", - "\n", - "\n", - "aplt.plot_array(\n", - " array=lens_subtracted_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", - ")\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "pix_indexes = [[200]]\n", - "\n", - "indexes = mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)\n", - "\n", - "\n", - "aplt.plot_array(\n", - " array=lens_subtracted_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", - ")\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", - " pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,\n", - " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index, # unused for Delaunay\n", - " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index, # unused for Delaunay\n", - " pixels=mapper.pixels,\n", - " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", - " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", - " sub_fraction=mapper.over_sampler.sub_fraction,\n", - ")\n", - "\n", - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "indexes_source_pix_200 = np.nonzero(mapping_matrix[:, 200])\n", - "\n", - "print(indexes_source_pix_200[0])\n", - "\n", - "array_2d = al.Array2D(values=mapping_matrix[:, 200], mask=masked_dataset.mask)\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")\n", - "\n", - "blurred_mapping_matrix = masked_dataset.psf.convolved_mapping_matrix_from(\n", - " mapping_matrix=mapping_matrix, mask=masked_dataset.mask\n", - ")\n", - "\n", - "plt.imshow(\n", - " blurred_mapping_matrix,\n", - " aspect=(blurred_mapping_matrix.shape[1] / blurred_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "indexes_source_pix_200 = np.nonzero(blurred_mapping_matrix[:, 200])\n", - "\n", - "print(indexes_source_pix_200[0])\n", - "\n", - "array_2d = al.Array2D(values=blurred_mapping_matrix[:, 200], mask=masked_dataset.mask)\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")\n", - "\n", - "print(f\"Mapping between image pixel 0 and source pixel 2 = {mapping_matrix[0, 2]}\")\n", - "\n", - "data_vector = al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(\n", - " blurred_mapping_matrix=blurred_mapping_matrix,\n", - " image=np.array(lens_subtracted_image),\n", - " noise_map=np.array(masked_dataset.noise_map),\n", - ")\n", - "\n", - "plt.imshow(\n", - " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix, noise_map=masked_dataset.noise_map\n", - ")\n", - "\n", - "plt.imshow(curvature_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "source_pixel_0 = 0\n", - "source_pixel_1 = 1\n", - "\n", - "print(curvature_matrix[source_pixel_0, source_pixel_1])\n", - "\n", - "array_2d = al.Array2D(\n", - " values=blurred_mapping_matrix[:, source_pixel_0], mask=masked_dataset.mask\n", - ")\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")\n", - "\n", - "array_2d = al.Array2D(\n", - " values=blurred_mapping_matrix[:, source_pixel_1], mask=masked_dataset.mask\n", - ")\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")\n", - "\n", - "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", - " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", - " neighbors=mapper.neighbors,\n", - " neighbors_sizes=mapper.neighbors.sizes,\n", - ")\n", - "\n", - "plt.imshow(regularization_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)\n", - "\n", - "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)\n", - "\n", - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "mapped_reconstructed_operated_data = (\n", - " al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix, reconstruction=reconstruction\n", - " )\n", - ")\n", - "\n", - "mapped_reconstructed_operated_data = al.Array2D(\n", - " values=mapped_reconstructed_operated_data, mask=mask\n", - ")\n", - "\n", - "aplt.plot_array(array=mapped_reconstructed_operated_data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image = convolved_image_2d + mapped_reconstructed_operated_data\n", - "\n", - "residual_map = masked_dataset.data - model_image\n", - "normalized_residual_map = residual_map / masked_dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(chi_squared)\n", - "\n", - "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", - "\n", - "aplt.plot_array(array=chi_squared_map, title=\"\")\n", - "\n", - "regularization_term = np.matmul(\n", - " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", - ")\n", - "\n", - "print(regularization_term)\n", - "\n", - "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", - "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]\n", - "\n", - "print(log_curvature_reg_matrix_term)\n", - "print(log_regularization_matrix_term)\n", - "\n", - "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))\n", - "\n", - "log_evidence = float(\n", - " -0.5\n", - " * (\n", - " chi_squared\n", - " + regularization_term\n", - " + log_curvature_reg_matrix_term\n", - " - log_regularization_matrix_term\n", - " + noise_normalization\n", - " )\n", - ")\n", - "\n", - "print(log_evidence)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(\n", - " dataset=masked_dataset,\n", - " tracer=tracer,\n", - " adapt_images=adapt_images,\n", - " settings=al.Settings(use_border_relocator=True),\n", - ")\n", - "fit_log_evidence = fit.log_evidence\n", - "print(fit_log_evidence)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `Nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "__Sub Gridding__\n", - "\n", - "The calculation above uses a `Grid2D` object, with a `sub-size=1`, meaning it does not perform oversampling to\n", - "evaluate the light profile flux at every image pixel.\n", - "\n", - "**PyAutoLens** has alternative methods of computing the lens galaxy images above, which uses a grid whose sub-size\n", - "adaptively increases depending on a required fractional accuracy of the light profile.\n", - "\n", - " https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/structures/grids/two_d/grid_iterate.py\n", - "\n", - "__Sourrce Plane Interpolation__\n", - "\n", - "For the `Delaunay` mesh used in this example, every image-sub pixel maps to a single source Voronoi\n", - "pixel. Therefore, the plural use of `pix_indexes` is not required. However, for other pixelizations each sub-pixel\n", - "can map to multiple source pixels with an interpolation weight (e.g. `Delaunay` triangulation or a `Voronoi` mesh\n", - "which uses natural neighbor interpolation).\n", - "\n", - "`MapperVoronoiNoInterp.pix_index_for_sub_slim_index`:\n", - "https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/inversion/mappers/voronoi.py\n", - "\n", - "`pixelization_index_for_voronoi_sub_slim_index_from`:\n", - " https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/util/mapper_util.py\n", - "\n", - "The number of pixels that each sub-pixel maps too is also stored and extracted. This is used for speeding up\n", - "the calculation of the `mapping_matrix` described next.\n", - "\n", - "As discussed above, because for the `VoronoiNoInterp` pixelization where every sub-pixel maps to one source pixel,\n", - "every entry of this array will be equal to 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# pix_sizes_for_sub_slim_index = mapper.pix_sizes_for_sub_slim_index" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When each sub-pixel maps to multiple source pixels, the mappings are described via an interpolation weight. For \n", - "example, for a `Delaunay` triangulation, every sub-pixel maps to 3 Delaunay triangles based on which triangle\n", - "it lands in.\n", - "\n", - "For the `VoronoiNoInterp` pixelization where every sub-pixel maps to a single source pixel without inteprolation,\n", - "every entry of this weight array is 1.0." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# pix_weights_for_sub_slim_index = mapper.pix_weights_for_sub_slim_index" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the **PyAutoLens** likelihood function, which uses a pixelization, \n", - "regularization scheme and inversion to reconstruct the source galaxy.\n", - "\n", - "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", - "are described in additional notebooks found in this package. In brief, these describe:\n", - "\n", - " - **Sub-gridding**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", - " ray-traced to the source-plane and paired fractionally with each source pixel.\n", - " \n", - " - **Source-plane Interpolation**: Using a Delaunay triangulation or Delaunay mesh with natural neighbor interpolation\n", - " to pair each image (sub-)pixel to multiple source-plane pixels with interpolation weights.\n", - " \n", - " - **Source Morphology Pixelization Adaption**: Adapting the pixelization such that is congregates source pixels around\n", - " the source's brightest regions, as opposed to the magnification-based pixelization used here.\n", - " \n", - " - **Luminosity Weighted Regularization**: Using an adaptive regularization coefficient which adapts the level of \n", - " regularization applied to the source based on its luminosity." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: Delaunay\n", + "======================\n", + "\n", + "The majority of pixelized source reconstructions in the workspace use a rectangular mesh to reconstruct\n", + "the source's surface brightness.\n", + "\n", + "This example illustrates an alternative pixelization that uses a Delaunay triangulation mesh to reconstruct the\n", + "source.\n", + "\n", + "The approach is distinct from the rectangular mesh and has a number of traits which are unique to it:\n", + "\n", + "- `Adaptive Mesh`: In the source plane, the Delaunay mesh uses irregularly shaped triangles to reconstruct the\n", + " source, as opposed to uniform rectangular pixels. This allows the mesh to better adapt to irregular and\n", + " asymmetric source morphologies and change the distribution of source pixels to better match the source's\n", + " surface brightness.\n", + "\n", + "- `Image Mesh`: The vertexes of the Delaunay triangles are computed by overlaying a coarse uniform grid in the\n", + " image-plane and ray-tracing these coordinates to the source-plane. This is unlike the rectangular mesh, which\n", + " simply overlays a uniform grid in the source-plane. This again helps the Delaunay mesh to better adapt to the\n", + " source's surface brightness.\n", + "\n", + "- `Interpolation`: The Delaunay mesh uses a different interpolation scheme to the rectangular mesh, which is\n", + " barycentric interpolation within each triangle. This is different to the rectangular mesh, which uses bilinear\n", + " interpolation within each rectangular pixel.\n", + "\n", + "- `Regularization`: The Delaunay mesh provides different approaches to regularization, with the default being\n", + " one which uses the barycentric coordinates of the triangles to compute how source pixels are regularized with\n", + " their neighbors.\n", + "\n", + "Currently it is not expected that the Delaunay is better or worse than the rectangular mesh, it is simply a different\n", + "approach to pixelization that may work better for certain datasets.\n", + "\n", + "__JAX + GPU__\n", + "\n", + "Generating a Delaunay mesh supports JAX and GPU acceleration, however certain operations (e.g. generating the Delaunay\n", + "triangulation itself) do not run on the GPU because they cannot be easily converted to JAX.\n", + "\n", + "Instead, JAX sends them to a CPU, runs them there, and then sends the results back to the GPU. This process is\n", + "very efficient, because these operations run very fast on a CPU and the data being sent back and forth is small.\n", + "Current benchmarking suggests the Delaunay runs less than twice as long as the same fit using a rectangular mesh,\n", + "but scientfically offers better results in many cases.\n", + "\n", + "If you do want to run only on CPU, you can use fast CPU method described in\n", + "example `imaging/features/pixelization/cpu_fast_modeling` with the Delaunay mesh.\n", + "\n", + "__Source Science (Magnification, Flux and More)__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", + "size of the source.\n", + "\n", + "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", + "including examples using a Delaunay source reconstruction. Once you have completed lens modeling using a Delaunay mesh,\n", + "you can jump to that example to study the source galaxy.\n", + "\n", + "__Contents__\n", + "\n", + "- **Image Mesh:** For a Delaunay mesh, the vertices of the triangles are defined by (y, x) coordinates in the.\n", + "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Source Science:** Source science focuses on studying the highly magnified properties of the background lensed source.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **JAX CPU:** On CPU, the Delaunay mesh computation via JAX can lead to slow down or indefinite freezing.\n", + "- **VRAM:** The `pixelization/modeling` example explains how VRAM use is an important consideration for.\n", + "- **Adaptive Delaunay:** The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to.\n", + "- **SLaM Pipelines:** The API above allows you to use adaptive features yourself, and you should go ahead an explore them.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "- **Likelihood Function:** The example `imaging/features/pixelization/likelihood_function.py` provides a step-by-step.\n", + "- **Source Galaxy Pixelization and Regularization:** The source galaxy is reconstructed using a pixel-grid, in this example a Delaunay mesh, which.\n", + "- **Lens Light:** Overview of lens light for this example.\n", + "- **Source Pixel Centre Calculation:** In order to reconstruct the source galaxy using a Delaunay mesh, we need to determine the centres.\n", + "- **Ray Tracing:** Overview of ray tracing for this example.\n", + "- **Border Relocation:** Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace.\n", + "- **Delaunay Mesh:** The relocated mesh grid is used to create the `Pixelization`'s Delaunay mesh using the.\n", + "- **Interpolation:** Overview of interpolation for this example.\n", + "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", + "- **Sub Gridding:** The calculation above uses a `Grid2D` object, with a `sub-size=1`, meaning it does not perform.\n", + "- **Sourrce Plane Interpolation:** For the `Delaunay` mesh used in this example, every image-sub pixel maps to a single source Voronoi.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking + Positions__ \n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image Mesh__\n", + "\n", + "For a Delaunay mesh, the vertices of the triangles are defined by (y, x) coordinates in the image-plane. These\n", + "coordinates are then ray-traced into the source-plane for each mass model sampled during the non-linear search.\n", + "This `image_plane_mesh_grid` must be computed before lens modeling.\n", + "\n", + "We compute this `image_plane_mesh_grid` using an `Overlay` image-mesh, which places a regular grid of\n", + "(y, x) points across the image-plane. This has a mild adaptive effect: regions of high lens magnification receive\n", + "more source pixels once they are ray-traced. Later in this example, we switch to a `Hilbert` image-mesh, which adapts\n", + "the pixel distribution more strongly to the source\u2019s surface brightness.\n", + "\n", + "The `Delaunay` mesh has an input number of `pixels`, which is the number of source pixels used to reconstruct the \n", + "source. The number of `pixels` must be equal to the number of coordinates in the `image_plane_mesh_grid`. \n", + "\n", + "Like for the `mesh_shape` rectangular mesh, `pixels` must be fixed for lens modeling because JAX uses the \n", + "number of `pixels` to determine static array shapes. \n", + "\n", + "To pass the `image_plane_mesh_grid` to the modeling, we use the `AdaptImages` object below, which pairs\n", + "the `image_plane_mesh_grid` to the source galaxy. For double source plane lenses, this means we can\n", + "attach an `image_plane_mesh_grid` to each source galaxy and use adaptive meshes for each source plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Edge Zeroing__\n", + "\n", + "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", + "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", + "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", + "\n", + "For a rectangular mesh, the source code computes edge pixels internally using the known pixels at the edge of the mesh,\n", + "requiring no input from the user. \n", + "\n", + "For the `Delaunay` mesh, we use the `append_with_circle_edge_points` function to manually setup the Delaunay image \n", + "mesh to include a ring of edge pixels and then input the total number into the mesh to perform zeroing. \n", + "\n", + "These points are added to the edge of the image-plane mesh, ray-traced to the source-plane during lens modeling, \n", + "included in the Delaunay triangulation but zeroed during the inversion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "edge_pixels_total = 30\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "In the example `imaging/features/pixelization/fit.py`, we illustrate how to use a pixelized source\n", + "with a rectangular mesh to fit imaging data.\n", + "\n", + "Below, we use a Delaunay mesh to perform a fit using the Delaunay source reconstruction.\n", + "\n", + "The API is nearly identical to the rectangular mesh example, noting that the inputs to the `Delaunay` \n", + "mesh are different to the rectangular mesh and use image mesh quantities computed above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh = al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + ")\n", + "\n", + "regularization = al.reg.ConstantSplit(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", + "\n", + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "adapt_images = al.AdaptImages(\n", + " galaxy_image_plane_mesh_grid_dict={source: image_plane_mesh_grid}\n", + ")\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " adapt_images=adapt_images,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the Delaunay source does a good job at capturing the appearance of the source galaxy\n", + "using adaptive triangular pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Science__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source pixelization, we can compute key quantities such as the magnification, total flux, and\n", + "intrinsic size of the source.\n", + "\n", + "For rectangular meshes, the example `autolens_workspace/*/imaging/features/source_science` gives a complete overview of \n", + "how to this. \n", + "\n", + "For the Delaunay, specific functionality which manipulates the Delaunay triangles is required to perform these \n", + "calculations. We show below how to interpolate the Delaunay source reconstruction to a regular grid.\n", + "\n", + "When an inversion is performed on a Delaunay mesh, the reconstructed values are defined only at the irregular\n", + "mesh vertices (the Delaunay nodes). In order to visualise the source reconstruction as a regular 2D image\n", + "(e.g. for plotting or saving to a FITS file), we must interpolate these values onto a uniform grid.\n", + "\n", + "This is done using SciPy's Delaunay-based linear interpolation:\n", + "\n", + "- A `scipy.spatial.Delaunay` triangulation is constructed from the mesh node coordinates.\n", + "\n", + "- Within each triangle, interpolation is performed using *barycentric coordinates*, meaning the value at any point\n", + " inside the triangle is computed as a weighted linear combination of the values at the triangle\u2019s three vertices.\n", + "\n", + "- Points outside the convex hull of the triangulation are assigned a fallback value (here `fill_value=0.0`).\n", + "\n", + "__Important Coordinate Convention (Delaunay(__\n", + "\n", + "SciPy expects all Delaunay coordinates to be provided in **(x, y)** order.\n", + "\n", + "However, grids are stored internally in **(y, x)** order.\n", + "\n", + "Therefore:\n", + "\n", + "- The mesh coordinates must be converted to (x, y) before building the triangulation.\n", + "- The interpolation grid must also be flipped to (x, y) before evaluating the interpolator.\n", + "\n", + "This coordinate flip is essential for correct interpolation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.spatial import Delaunay\n", + "from scipy.interpolate import LinearNDInterpolator\n", + "\n", + "inversion = fit.inversion\n", + "\n", + "mapper = inversion.cls_list_from(cls=al.Mapper)[0]\n", + "\n", + "reconstruction = inversion.reconstruction\n", + "\n", + "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", + "\n", + "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", + "\n", + "mesh_grid_xy = np.stack([source_plane_mesh_grid[:, 0], source_plane_mesh_grid[:, 1]]).T\n", + "\n", + "# Uses find simplex so recomputes delaunay internally\n", + "delaunay = Delaunay(mesh_grid_xy)\n", + "\n", + "interp = LinearNDInterpolator(delaunay, reconstruction, fill_value=0.0)\n", + "\n", + "interpolation_grid_xy = np.asarray(interpolation_grid)[:, ::-1] # (y,x)->(x,y)\n", + "interpolated_reconstruction = interp(interpolation_grid_xy)\n", + "\n", + "print(f\"Brightest Interpolated Source Pixel: {np.max(interpolated_reconstruction)}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We now perform lens modeling using the Delaunay pixelization with the Overlay image-mesh.\n", + "\n", + "The code below is a simple adaptive modeling example using the Delaunay mesh, which mirrors the\n", + "API used in other pixelization modeling examples.\n", + "\n", + "The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to\n", + "adapt the rectangular mesh and its regularization to the source's surface brightness. In particular, an image\n", + "of the lensed source is passed to the modeling via the `AdaptImages` object, in order to adapt\n", + "the mesh and regularization during the model-fit.\n", + "\n", + "You have already seen this used once above, but we set up the adapt images again to remind you\n", + "of the API." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "adapt_images = al.AdaptImages(\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In the first\n", + "search our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", + " \n", + " - The source galaxy's light uses an `Overlay` image-mesh with fixed resolution 30 x 30 pixels [0 parameters].\n", + " \n", + " - The source-galaxy's light uses a `Delaunay` mesh [0 parameters].\n", + "\n", + " - This pixelization is regularized using a `Constant` scheme [1 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=8.\n", + "\n", + "__JAX CPU__\n", + "\n", + "On CPU, the Delaunay mesh computation via JAX can lead to slow down or indefinite freezing. \n", + "\n", + "You may find better performance if you set `use_jax_vmap=False` in the `Nautilus` search below, which\n", + "disables JAX's vectorization of certain computations.\n", + "\n", + "This does not impact GPU performance, which should always use `use_jax_vmap=True`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", + ")\n", + "\n", + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.ConstantSplit,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "search_1 = af.Nautilus(\n", + " path_prefix=Path(\"features\"),\n", + " name=\"delaunay\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=10,\n", + " # use_jax_vmap=False # Set to False if CPU performance is slow or hangs\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `pixelization/modeling` example explains how VRAM use is an important consideration for pixelization models\n", + "and how it depends on image resolution, number of source pixels and batch size.\n", + "\n", + "This is true for the Delaunay mesh, therefore we print out the estimated VRAM required for this model-fit.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_1.print_vram_use(model=model_1, batch_size=search_1.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit (Search 1)__\n", + "\n", + "Perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Adaptive Delaunay__\n", + "\n", + "The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to\n", + "adapt the rectangular mesh and its regularization to the source's surface brightness.\n", + "\n", + "The image-mesh has a special adaptive variant called the `Hilbert` image-mesh, which adapts the distribution \n", + "of source-pixels to the source's unlensed morphology. This means that the source's brightest regions are \n", + "reconstructed using significantly more source pixels than seen for the `Overlay` image mesh. \n", + "Conversely, the source's faintest regions are reconstructed using significantly fewer source pixels.\n", + "\n", + "Unlike the adaptive rectangular mesh, the Hilbert image-plane mesh is computed before modeling, passed\n", + "to the `AdaptImages` object, and remains fixed during the model-fit.\n", + "\n", + "It is recommend that the parameters governing these features are always fitted from using a fixed lens light and\n", + "mass model. This ensures the adaptation is performed quickly, and removes degeneracies in the lens model that\n", + "are difficult to sample. Given the Hilbert mesh is fixed, this modeling only fits for the regularization coefficients\n", + "of the adaptive regularization scheme.\n", + "\n", + "For this reason, search 2 fixes the lens galaxy's light and mass model to the best-fit model of search 1. A third\n", + "search will then fit for the lens galaxy's light and mass model using these adaptive features.\n", + "\n", + "The details of how the above features work is not provided here, but is given at the end of chapter 4 of the HowToLens\n", + "lecture series." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result_1)\n", + "\n", + "image_mesh = al.image_mesh.Hilbert(pixels=1000, weight_power=3.5, weight_floor=0.01)\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask, adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", + ")\n", + "\n", + "# Repeat edge zeroing set up describe above.\n", + "\n", + "edge_pixels_total = 30\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=mask.mask_centre,\n", + " radius=mask_radius + mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In \n", + "the second search our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` with fixed parameters from \n", + " search 1 [0 parameters].\n", + " \n", + " - The source galaxy's light uses a `Hilbert` image-mesh with fixed resolution 1000 pixels [2 parameters].\n", + " \n", + " - The source-galaxy's light uses a `Delaunay` mesh [0 parameters].\n", + "\n", + " - This pixelization is regularized using a `AdaptSplit` scheme [2 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=4." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.AdaptSplit,\n", + ")\n", + "\n", + "source = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " pixelization=pixelization,\n", + ")\n", + "\n", + "model_2 = af.Collection(\n", + " galaxies=af.Collection(lens=result_1.instance.galaxies.lens, source=source)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis (Search 2)__\n", + "\n", + "We now create the analysis for the second search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_2 = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Model-Fit (Search 2)__\n", + "\n", + "We now create the non-linear search and perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=Path(\"features\"),\n", + " name=\"delaunay_adapt\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + " # use_jax_vmap=False # Set to False if CPU performance is slow or hangs\n", + ")\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We could perform a third fit where we free all lens model parameters and fit them using the adaptive \n", + "image mesh and regularization.\n", + "\n", + "However, it is better to use all of these features with the Delaunay via the\n", + "SLaM pipelines, which we jump to immediately below.\n", + "\n", + "__SLaM Pipelines__\n", + "\n", + "The API above allows you to use adaptive features yourself, and you should go ahead an explore them on datasets you\n", + "are familiar with.\n", + "\n", + "However, you may also wish to use the Source, Light and Mass (SLaM) pipelines, which are pipelines that\n", + "have been carefully crafted to automate lens modeling of large samples whilst ensuring models of the highest\n", + "complexity can be reliably fitted.\n", + "\n", + "These pipelines are built around the use of adaptive features -- for example the Source pipeline comes first so that\n", + "these features are set up robustly before more complex lens light and mass models are fitted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", + "\n", + "The `source_pix_1` search uses an `Overlay` image-mesh to place the initial Delaunay mesh pixels, with\n", + "additional edge points added around the mask boundary to ensure full coverage." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_lp_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask,\n", + " )\n", + "\n", + " edge_pixels_total = 30\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=dataset.mask.mask_centre,\n", + " radius=mask_radius + dataset.mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0],\n", + " zeroed_pixels=edge_pixels_total,\n", + " ),\n", + " regularization=al.reg.AdaptSplit,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", + "\n", + "The `source_pix_2` search uses a `Hilbert` image-mesh to place the final Delaunay mesh pixels, which adapts\n", + "the mesh to the source morphology using the high-quality adapt images from search 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Hilbert(pixels=1000, weight_power=3.5, weight_floor=0.01)\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask, adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", + " )\n", + "\n", + " edge_pixels_total = 30\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=dataset.mask.mask_centre,\n", + " radius=mask_radius + dataset.mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0],\n", + " zeroed_pixels=edge_pixels_total,\n", + " ),\n", + " regularization=al.reg.AdaptSplit,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict=(\n", + " source_result_for_source.analysis.adapt_images.galaxy_name_image_plane_mesh_grid_dict\n", + " ),\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict=(\n", + " source_result_for_source.analysis.adapt_images.galaxy_name_image_plane_mesh_grid_dict\n", + " ),\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " bulge = light_result.instance.galaxies.lens.bulge\n", + " disk = light_result.instance.galaxies.lens.disk\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=bulge,\n", + " disk=disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam_delaunay\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + " # use_jax_vmap=False # Set to False if CPU performance is slow or hangs\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_lp_result=source_lp_result,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "The example `imaging/features/pixelization/likelihood_function.py` provides a step-by-step description of how\n", + "a likelihood evaluation is performed for imaging data using a pixelized source reconstruction with a rectangular\n", + "mesh.\n", + "\n", + "We now give the same step-by-step description for a pixelized source reconstruction using a Delaunay mesh and\n", + "adaptive features.\n", + "\n", + "We only describe code which is specific to Delaunay meshes and adaptive features -- for all other aspects of the likelihood\n", + "evaluation, refer to rectangular mesh example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=masked_dataset)\n", + "\n", + "masked_dataset = masked_dataset.apply_over_sampling(\n", + " over_sample_size_lp=1,\n", + " over_sample_size_pixelization=1,\n", + ")\n", + "\n", + "aplt.plot_grid(grid=masked_dataset.grids.pixelization, title=\"\")\n", + "\n", + "bulge = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Pixelization and Regularization__\n", + "\n", + "The source galaxy is reconstructed using a pixel-grid, in this example a Delaunay mesh, which accounts for \n", + "irregularities and asymmetries in the source's surface brightness. \n", + "\n", + "A constant regularization scheme is applied which applies a smoothness prior on the reconstruction. \n", + "\n", + "One of the biggest differences between a Delaunay mesh and rectangular mesh is how the centres of the mesh pixels\n", + "in the source-plane are computed. \n", + "\n", + "For the rectangular mesh, the pixel centres are computed by overlaying a uniform grid over the source-plane.\n", + "\n", + "For a Delaunay mesh, the uniform grid is instead laid over the image-plane to create a course grid of (y,x) coordinates.\n", + "These are then ray-traced to the source-plane and are used as the vertexes of the Delaunay triangles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=masked_dataset.grids.pixelization.shape[0],\n", + " zeroed_pixels=edge_pixels_total,\n", + " ),\n", + " regularization=al.reg.ConstantSplit(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Light__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n", + "\n", + "\n", + "blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n", + "\n", + "\n", + "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", + " image=image, blurring_image=blurring_image_2d\n", + ")\n", + "\n", + "aplt.plot_array(array=convolved_image_2d, title=\"\")\n", + "\n", + "\n", + "lens_subtracted_image = masked_dataset.data - convolved_image_2d\n", + "\n", + "aplt.plot_array(array=lens_subtracted_image, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Pixel Centre Calculation__\n", + "\n", + "In order to reconstruct the source galaxy using a Delaunay mesh, we need to determine the centres of the Delaunay\n", + "source pixels.\n", + "\n", + "The image-mesh `Overlay` object computes the source-pixel centres in the image-plane (which are ray-traced to the\n", + "source-plane below). The source pixelization therefore adapts to the lens model magnification, because more\n", + "source pixels will congregate in higher magnification regions.\n", + "\n", + "This calculation is performed by overlaying a uniform regular grid with an `pixelization_shape_2d` over the image\n", + "mask and retaining all pixels that fall within the mask. This uses a `Grid2DSparse` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_mesh = al.image_mesh.Overlay(shape=(30, 30)) # Specific to Delaunay\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=masked_dataset.mask,\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(\n", + " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid},\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting this grid shows a sparse grid of (y,x) coordinates within the mask, which will form our source pixel centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=masked_dataset.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The source code gets quite complex when handling grids for a pixelization, but it is all handled in\n", + "the `TracerToInversion` objects.\n", + "\n", + "The plots at the bottom of this cell show the traced grids used by the source pixelization, showing\n", + "how the Delaunay mesh and traced image pixels are constructed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(\n", + " tracer=tracer, dataset=masked_dataset, adapt_images=adapt_images\n", + ")\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", + "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.pixelization\n", + ")[-1]\n", + "\n", + "# This functions a bit weird - it returns a list of lists of ndarrays. Best not to worry about it for now!\n", + "traced_mesh_grid = tracer_to_inversion.traced_mesh_grid_pg_list[-1][-1]\n", + "\n", + "\n", + "aplt.plot_grid(grid=traced_grid_pixelization, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We have also ray-traced the coarse grid of image-pixel coordinates used to form the source pixelization's\n", + "Delaunay mesh, which we can also plot." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=traced_mesh_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Border Relocation__\n", + "\n", + "Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace to far outskirts of\n", + "the source-plane. \n", + "\n", + "Border relocation is performed on both the traced image-pixel grid and traced mesh pixels, therefore ensuring that\n", + "the vertexes of the Delaunay triangles are not at the extreme outskirts of the source-plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", + "\n", + "border_relocator = BorderRelocator(mask=masked_dataset.mask, sub_size=1)\n", + "\n", + "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", + "\n", + "relocated_mesh_grid = border_relocator.relocated_mesh_grid_from(\n", + " grid=traced_grid_pixelization, mesh_grid=traced_mesh_grid\n", + ")\n", + "\n", + "\n", + "aplt.plot_grid(grid=relocated_grid, title=\"\")\n", + "\n", + "aplt.plot_grid(grid=relocated_mesh_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Delaunay Mesh__\n", + "\n", + "The relocated mesh grid is used to create the `Pixelization`'s Delaunay mesh using the `scipy.spatial` library." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "interpolator = al.InterpolatorDelaunay(\n", + " mesh=pixelization.mesh,\n", + " mesh_grid=relocated_mesh_grid,\n", + " data_grid=relocated_grid,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting the Delaunay mesh shows that the source-plane and been discretized into a grid of irregular Delaunay pixels.\n", + "\n", + "(To plot the Delaunay mesh, we have to convert it to a `Mapper` object, which is described in the next likelihood step).\n", + "\n", + "Below, we plot the Delaunay mesh without the traced image-grid pixels (for clarity) and with them as black dots in order\n", + "to show how each set of image-pixels fall within a Delaunay pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = al.Mapper(\n", + " interpolator=interpolator,\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + ")\n", + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolation__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pix_indexes_for_sub_slim_index = mapper.pix_indexes_for_sub_slim_index\n", + "\n", + "print(pix_indexes_for_sub_slim_index[0:9])\n", + "\n", + "\n", + "aplt.plot_array(\n", + " array=lens_subtracted_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", + ")\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "pix_indexes = [[200]]\n", + "\n", + "indexes = mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)\n", + "\n", + "\n", + "aplt.plot_array(\n", + " array=lens_subtracted_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", + ")\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", + " pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,\n", + " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index, # unused for Delaunay\n", + " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index, # unused for Delaunay\n", + " pixels=mapper.pixels,\n", + " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", + " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", + " sub_fraction=mapper.over_sampler.sub_fraction,\n", + ")\n", + "\n", + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "indexes_source_pix_200 = np.nonzero(mapping_matrix[:, 200])\n", + "\n", + "print(indexes_source_pix_200[0])\n", + "\n", + "array_2d = al.Array2D(values=mapping_matrix[:, 200], mask=masked_dataset.mask)\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")\n", + "\n", + "blurred_mapping_matrix = masked_dataset.psf.convolved_mapping_matrix_from(\n", + " mapping_matrix=mapping_matrix, mask=masked_dataset.mask\n", + ")\n", + "\n", + "plt.imshow(\n", + " blurred_mapping_matrix,\n", + " aspect=(blurred_mapping_matrix.shape[1] / blurred_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "indexes_source_pix_200 = np.nonzero(blurred_mapping_matrix[:, 200])\n", + "\n", + "print(indexes_source_pix_200[0])\n", + "\n", + "array_2d = al.Array2D(values=blurred_mapping_matrix[:, 200], mask=masked_dataset.mask)\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")\n", + "\n", + "print(f\"Mapping between image pixel 0 and source pixel 2 = {mapping_matrix[0, 2]}\")\n", + "\n", + "data_vector = al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(\n", + " blurred_mapping_matrix=blurred_mapping_matrix,\n", + " image=np.array(lens_subtracted_image),\n", + " noise_map=np.array(masked_dataset.noise_map),\n", + ")\n", + "\n", + "plt.imshow(\n", + " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix, noise_map=masked_dataset.noise_map\n", + ")\n", + "\n", + "plt.imshow(curvature_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "source_pixel_0 = 0\n", + "source_pixel_1 = 1\n", + "\n", + "print(curvature_matrix[source_pixel_0, source_pixel_1])\n", + "\n", + "array_2d = al.Array2D(\n", + " values=blurred_mapping_matrix[:, source_pixel_0], mask=masked_dataset.mask\n", + ")\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")\n", + "\n", + "array_2d = al.Array2D(\n", + " values=blurred_mapping_matrix[:, source_pixel_1], mask=masked_dataset.mask\n", + ")\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")\n", + "\n", + "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", + " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", + " neighbors=mapper.neighbors,\n", + " neighbors_sizes=mapper.neighbors.sizes,\n", + ")\n", + "\n", + "plt.imshow(regularization_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)\n", + "\n", + "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)\n", + "\n", + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "mapped_reconstructed_operated_data = (\n", + " al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix, reconstruction=reconstruction\n", + " )\n", + ")\n", + "\n", + "mapped_reconstructed_operated_data = al.Array2D(\n", + " values=mapped_reconstructed_operated_data, mask=mask\n", + ")\n", + "\n", + "aplt.plot_array(array=mapped_reconstructed_operated_data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image = convolved_image_2d + mapped_reconstructed_operated_data\n", + "\n", + "residual_map = masked_dataset.data - model_image\n", + "normalized_residual_map = residual_map / masked_dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(chi_squared)\n", + "\n", + "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", + "\n", + "aplt.plot_array(array=chi_squared_map, title=\"\")\n", + "\n", + "regularization_term = np.matmul(\n", + " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", + ")\n", + "\n", + "print(regularization_term)\n", + "\n", + "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", + "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]\n", + "\n", + "print(log_curvature_reg_matrix_term)\n", + "print(log_regularization_matrix_term)\n", + "\n", + "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))\n", + "\n", + "log_evidence = float(\n", + " -0.5\n", + " * (\n", + " chi_squared\n", + " + regularization_term\n", + " + log_curvature_reg_matrix_term\n", + " - log_regularization_matrix_term\n", + " + noise_normalization\n", + " )\n", + ")\n", + "\n", + "print(log_evidence)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(\n", + " dataset=masked_dataset,\n", + " tracer=tracer,\n", + " adapt_images=adapt_images,\n", + " settings=al.Settings(use_border_relocator=True),\n", + ")\n", + "fit_log_evidence = fit.log_evidence\n", + "print(fit_log_evidence)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `Nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "__Sub Gridding__\n", + "\n", + "The calculation above uses a `Grid2D` object, with a `sub-size=1`, meaning it does not perform oversampling to\n", + "evaluate the light profile flux at every image pixel.\n", + "\n", + "**PyAutoLens** has alternative methods of computing the lens galaxy images above, which uses a grid whose sub-size\n", + "adaptively increases depending on a required fractional accuracy of the light profile.\n", + "\n", + " https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/structures/grids/two_d/grid_iterate.py\n", + "\n", + "__Sourrce Plane Interpolation__\n", + "\n", + "For the `Delaunay` mesh used in this example, every image-sub pixel maps to a single source Voronoi\n", + "pixel. Therefore, the plural use of `pix_indexes` is not required. However, for other pixelizations each sub-pixel\n", + "can map to multiple source pixels with an interpolation weight (e.g. `Delaunay` triangulation or a `Voronoi` mesh\n", + "which uses natural neighbor interpolation).\n", + "\n", + "`MapperVoronoiNoInterp.pix_index_for_sub_slim_index`:\n", + "https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/inversion/mappers/voronoi.py\n", + "\n", + "`pixelization_index_for_voronoi_sub_slim_index_from`:\n", + " https://github.com/PyAutoLabs/PyAutoArray/blob/main/autoarray/util/mapper_util.py\n", + "\n", + "The number of pixels that each sub-pixel maps too is also stored and extracted. This is used for speeding up\n", + "the calculation of the `mapping_matrix` described next.\n", + "\n", + "As discussed above, because for the `VoronoiNoInterp` pixelization where every sub-pixel maps to one source pixel,\n", + "every entry of this array will be equal to 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# pix_sizes_for_sub_slim_index = mapper.pix_sizes_for_sub_slim_index" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When each sub-pixel maps to multiple source pixels, the mappings are described via an interpolation weight. For \n", + "example, for a `Delaunay` triangulation, every sub-pixel maps to 3 Delaunay triangles based on which triangle\n", + "it lands in.\n", + "\n", + "For the `VoronoiNoInterp` pixelization where every sub-pixel maps to a single source pixel without inteprolation,\n", + "every entry of this weight array is 1.0." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# pix_weights_for_sub_slim_index = mapper.pix_weights_for_sub_slim_index" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the **PyAutoLens** likelihood function, which uses a pixelization, \n", + "regularization scheme and inversion to reconstruct the source galaxy.\n", + "\n", + "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", + "are described in additional notebooks found in this package. In brief, these describe:\n", + "\n", + " - **Sub-gridding**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", + " ray-traced to the source-plane and paired fractionally with each source pixel.\n", + " \n", + " - **Source-plane Interpolation**: Using a Delaunay triangulation or Delaunay mesh with natural neighbor interpolation\n", + " to pair each image (sub-)pixel to multiple source-plane pixels with interpolation weights.\n", + " \n", + " - **Source Morphology Pixelization Adaption**: Adapting the pixelization such that is congregates source pixels around\n", + " the source's brightest regions, as opposed to the magnification-based pixelization used here.\n", + " \n", + " - **Luminosity Weighted Regularization**: Using an adaptive regularization coefficient which adapts the level of \n", + " regularization applied to the source based on its luminosity." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/pixelization/fit.ipynb b/notebooks/imaging/features/pixelization/fit.ipynb index 3020873e2..286de2f39 100644 --- a/notebooks/imaging/features/pixelization/fit.ipynb +++ b/notebooks/imaging/features/pixelization/fit.ipynb @@ -1,867 +1,904 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Pixelization Fit\n", - "==========================\n", - "\n", - "A pixelization reconstructs the source's light using a pixel-grid, which is regularized using a prior that forces\n", - "the solution to have a degree of smoothness.\n", - "\n", - "This is important for reconstructing complex and irregular source morphologies that cannot be well represented by\n", - "light profiles like a Sersic or shapelets. It is vital to ensuring the infferred lens mass model is accurate and unbiased.\n", - "\n", - "This script fits a source galaxy in a way which uses a pixelization to reconstruct the source's light. This uses a\n", - "rectangular mesh and constant regularization scheme are, which are the simplest forms of each, both providing\n", - "computationally fast and accurate solutions.\n", - "\n", - "For simplicity, the lens galaxy's light is omitted from the model and is not present in the simulated data. It is\n", - "straightforward to include the lens galaxy's light in the model.\n", - "\n", - "Pixelizations are covered in detail in chapter 4 of the **HowToLens** lectures.\n", - "\n", - "__CPU Users__\n", - "\n", - "Matrices must be set up for a pixelized source reconstruction which speed up the linear algebra. On GPU, this takes\n", - "seconds, or at most a minute for datasets with tens of millions, or more visibities. On CPU, this can be a lot slower,\n", - "taking over hours. If you are on CPU, the `feature/pixelization/many_visibilities_preparation` explains how this\n", - "initial setup can be performed before lens modeling and saved to hard disk for fast loading before the model fit.\n", - "\n", - "__Contents__\n", - "\n", - "- **CPU Users:** Matrices must be set up for a pixelized source reconstruction which speed up the linear algebra.\n", - "- **Advantages & Disadvantages:** Many strongly lensed source galaxies are complex, and have asymmetric and irregular morphologies.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", - "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", - "- **Pixelization:** We create a `Pixelization` object to perform the pixelized source reconstruction, which is made up.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Mask Extra Galaxies:** There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "- **Linear Objects:** An `Inversion` contains all of the linear objects used to reconstruct the data in its.\n", - "- **Grids:** The role of a mapper is to map between the image-plane and source-plane.\n", - "- **Reconstruction:** The source reconstruction is also available as a 1D numpy array of values representative of the.\n", - "- **Mapped Reconstructed Images:** The source reconstruction(s) are mapped to the image-plane in order to fit the lens model.\n", - "- **Simulated Imaging:** We load the source galaxy image from the pixelized inversion of a previous fit, which was performed.\n", - "\n", - "__Advantages__\n", - "\n", - "Many strongly lensed source galaxies are complex, and have asymmetric and irregular morphologies. These morphologies\n", - "cannot be well approximated by a light profiles like a Sersic, or many Sersics, and thus a pixelization\n", - "is required to reconstruct the source's irregular light.\n", - "\n", - "Even basis functions like shapelets or a multi-Gaussian expansion cannot reconstruct a source-plane accurately\n", - "if there are multiple source galaxies, or if the source galaxy has a very complex morphology.\n", - "\n", - "To infer detailed components of a lens mass model (e.g. its density slope, whether there's a dark matter subhalo, etc.)\n", - "then pixelized source models are required, to ensure the mass model is fitting all of the lensed source light.\n", - "\n", - "There are also many science cases where one wants to study the highly magnified light of the source galaxy in detail,\n", - "to learnt about distant and faint galaxies. A pixelization reconstructs the source's unlensed emission and thus\n", - "enables this.\n", - "\n", - "__Disadvantages__\n", - "\n", - "Pixelizations are computationally slow and run times are typically longer than a parametric source model. It is not\n", - "uncommon for lens models using a pixelization to take hours to fit high resolution imaging data (e.g. Hubble Space\n", - "Telescope imaging), albeit on modern GPUs run times are often closer to < 20 minutes.\n", - "\n", - "Lens modeling with pixelizations is also more complex than parametric source models, with there being more things\n", - "that can go wrong. For example, there are solutions where a demagnified version of the lensed source galaxy is\n", - "reconstructed, using a mass model which effectively has no mass or too much mass. These are described in detail below.\n", - "\n", - "It will take you longer to learn how to successfully fit lens models with a pixelization than other methods illustrated\n", - "in the workspace!\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algabra solver which allows for positive and negative\n", - "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical. For a pixelizaiton, this often produces negative source pixels which over-fit\n", - "the data, producing unphysical solutions.\n", - "\n", - "All pixelized source reconstructions use a positive-only solver, meaning that every source-pixel is only allowed\n", - "to reconstruct positive flux values. This ensures that the source reconstruction is physical and that we don't\n", - "reconstruct negative flux values that don't exist in the real source galaxy (a common systematic solution in lens\n", - "analysis).\n", - "\n", - "Enforcing positive reconstructions efficiently requires non-trivial linear algebra, so a bespoke JAX fast non-negative\n", - "solver was developed; many methods in the literature omit this and therefore allow unphysical negative solutions that\n", - "can degrade lens modeling results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autoarray.inversion.plot.inversion_plots import subplot_of_mapper, subplot_mappings" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple__no_lens_light` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "A pixelization uses a separate grid for ray tracing, with its own over sampling scheme, which below we set to a \n", - "uniform grid of values of 4. \n", - "\n", - "The pixelization only reconstructs the source galaxy, therefore the adaptive over sampling used for the lens galaxy's \n", - "light in other examples is not applied to the pixelization. \n", - "\n", - "This example does not model lens light, for examples which combine lens light and a pixelization both over sampling \n", - "schemes should be used, with the lens light adaptive and the pixelization uniform.\n", - "\n", - "Note that the over sampling is input into the `over_sample_size_pixelization` because we are using a `Pixelization`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_over_sampling(\n", - " over_sample_size_pixelization=4,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", - "set below to 28 x 28. \n", - "\n", - "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", - "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", - "mesh, the same number of pixels must be used in the y and x directions.\n", - "\n", - "__Edge Zeroing__\n", - "\n", - "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", - "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", - "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", - "\n", - "For a rectangular mesh, the source code computes edge pixels internally using the known\n", - "pixels at the edge of the mesh. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixelization__\n", - "\n", - "We create a `Pixelization` object to perform the pixelized source reconstruction, which is made up of three\n", - "components:\n", - "\n", - "- `mesh:` Different types of mesh can be used to perform the source reconstruction, where the mesh changes the\n", - "details of how the source is reconstructed (e.g. interpolation weights). In this example, we use a rectangular mesh,\n", - "where the centres computed by overlayiong a rectangular mesh over the source plane.\n", - "\n", - "- `regularization:` A pixelization uses many pixels to reconstructed the source, which will often lead to over fitting\n", - "of the noise in the data and an unrealistically complex and structured source. Regularization smooths the source\n", - "reconstruction solution by penalizing solutions where neighboring pixels have\n", - "large flux differences." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", - "regularization = al.reg.Constant(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This is to illustrate the API for performing a fit via a pixelization using standard autolens objects like \n", - "the `Galaxy`, `Tracer` and `FitImaging` \n", - "\n", - "We simply create a `Pixelization` and pass it to the source galaxy, which then gets input into the tracer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the pixelized source does a good job at capturing the appearance of the source galaxy\n", - "and fitting the data to roughly the noise level." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelizations have bespoke visualizations which show more details about the source-reconstruction, image-mesh\n", - "and other quantities.\n", - "\n", - "The `subplot_of_mapper` function produces a comprehensive diagnostic subplot for the inversion. The\n", - "`subplot_mappings` overlays colored circles in the image and source planes that map to one another, thereby\n", - "allowing one to assess how the mass model ray-traces image-pixels and therefore to assess how the source\n", - "reconstruction maps to the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "\n", - "subplot_of_mapper(inversion=inversion, mapper_index=0)\n", - "subplot_mappings(inversion=inversion, pixelization_index=0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Science (Magnification, Flux and More)__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", - "size of the source.\n", - "\n", - "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", - "including examples using a pixelized source reconstruction. Now you know how to fit a pixelization, go check it out!\n", - "\n", - "__Mask Extra Galaxies__\n", - "\n", - "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source.\n", - "\n", - "If their emission is significant, and close enough to the lens and source, we may simply remove it from the data\n", - "to ensure it does not impact the model-fit. A standard masking approach would be to remove the image pixels containing\n", - "the emission of these galaxies altogether. This is analogous to what the circular masks used throughout the examples\n", - "does.\n", - "\n", - "For fits using a pixelization, masking regions of the image in a way that removes their image pixels entirely from\n", - "the fit. This can produce discontinuities in the pixelixation used to reconstruct the source and produce unexpected\n", - "systematics and unsatisfactory results. In this case, applying the mask in a way where the image pixels are not\n", - "removed from the fit, but their data and noise-map values are scaled such that they contribute negligibly to the fit,\n", - "is a better approach.\n", - "\n", - "We illustrate the API for doing this below, using the `extra_galaxies` dataset which has extra galaxies whose emission\n", - "needs to be removed via scaling in this way. We apply the scaling and show the subplot imaging where the extra\n", - "galaxies mask has scaled the data values to zeros, increasing the noise-map values to large values and in turn made\n", - "the signal to noise of its pixels effectively zero." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_extra_galaxies = al.Mask2D.from_fits(\n", - " file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"),\n", - " pixel_scales=0.1,\n", - " invert=True, # Note that we invert the mask here as `True` means a pixel is scaled.\n", - ")\n", - "\n", - "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native, pixel_scales=0.1, centre=(0.0, 0.0), radius=6.0\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We do not explictly fit this data, for the sake of brevity, however if your data has these nearby galaxies you should\n", - "apply the mask as above before fitting the data.\n", - "\n", - "__Wrap Up__\n", - "\n", - "Pixelizations are the most complex but also most powerful way to model a source galaxy.\n", - "\n", - "Whether you need to use them or not depends on the science you are doing. If you are only interested in measuring a\n", - "simple quantity like the Einstein radius of a lens, you can get away with using light profiles like a Sersic, MGE or \n", - "shapelets to model the source. Low resolution data also means that using a pixelization is not necessary, as the\n", - "complex structure of the source galaxy is not resolved anyway.\n", - "\n", - "However, fitting complex mass models (e.g. a power-law, stellar / dark model or dark matter substructure) requires \n", - "this level of complexity in the source model. Furthermore, if you are interested in studying the properties of the\n", - "source itself, you won't find a better way to do this than using a pixelization.\n", - "\n", - "__Linear Objects__\n", - "\n", - "An `Inversion` contains all of the linear objects used to reconstruct the data in its `linear_obj_list`. \n", - "\n", - "This list may include the following objects:\n", - "\n", - " - `LightProfileLinearObjFuncList`: This object contains lists of linear light profiles and the functionality used\n", - " by them to reconstruct data in an inversion. For example it may only contain a list with a single light profile\n", - " (e.g. `lp_linear.Sersic`) or many light profiles combined in a `Basis` (e.g. `lp_basis.Basis`).\n", - "\n", - "- `Mapper`: The linear objected used by a `Pixelization` to reconstruct data via an `Inversion`, where the `Mapper` \n", - "is specific to the `Pixelization`'s `Mesh` (e.g. a `RectnagularMapper` is used for a `Voronoi` mesh).\n", - "\n", - "In this example, the only linear object used to fit the data was a `Pixelization`, thus the `linear_obj_list`\n", - "contains just one entry corresponding to a `Mapper`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.linear_obj_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To extract results from an inversion many quantities will come in lists or require that we specific the linear object\n", - "we with to use. \n", - "\n", - "Thus, knowing what linear objects are contained in the `linear_obj_list` and what indexes they correspond to\n", - "is important." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Mapper = {inversion.linear_obj_list[0]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grids__\n", - "\n", - "The role of a mapper is to map between the image-plane and source-plane. \n", - "\n", - "This includes mapping grids corresponding to the data grid (e.g. the centers of each image-pixel in the image and\n", - "source plane) and the pixelization grid (e.g. the centre of the Delaunay triangulation in the image-plane and \n", - "source-plane).\n", - "\n", - "All grids are available in a mapper via its `mapper` property." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = inversion.linear_obj_list[0]\n", - "\n", - "# Centre of each masked image pixel in the image-plane.\n", - "print(mapper.image_plane_data_grid)\n", - "\n", - "# Centre of each source pixel in the source-plane.\n", - "print(mapper.source_plane_data_grid)\n", - "\n", - "# Centre of each pixelization pixel in the image-plane (the `Overlay` image_mesh computes these in the image-plane\n", - "# and maps to the source-plane).\n", - "print(mapper.image_plane_mesh_grid)\n", - "\n", - "# Centre of each pixelization pixel in the source-plane.\n", - "print(mapper.source_plane_mesh_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction__\n", - "\n", - "The source reconstruction is also available as a 1D numpy array of values representative of the source pixelization\n", - "itself (in this example, the reconstructed source values at the vertexes of each Voronoi triangle)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The (y,x) grid of coordinates associated with these values is given by the `Inversion`'s `Mapper` (which are \n", - "described in chapter 4 of **HowToLens**." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = inversion.linear_obj_list[0]\n", - "print(mapper.source_plane_mesh_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The mapper also contains the (y,x) grid of coordinates that correspond to the ray-traced image sub-pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(mapper.source_plane_data_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapped Reconstructed Images__\n", - "\n", - "The source reconstruction(s) are mapped to the image-plane in order to fit the lens model.\n", - "\n", - "These mapped reconstructed images are also accessible via the `Inversion`. \n", - "\n", - "Note that any light profiles in the lens model (e.g. the `bulge` and `disk` of a lens galaxy) are not \n", - "included in this image -- it only contains the source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.mapped_reconstructed_operated_data.native)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Algebra Matrices (Advanced)__\n", - "\n", - "To perform an `Inversion` a number of matrices are constructed which use linear algebra to perform the reconstruction.\n", - "\n", - "These are accessible in the inversion object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.curvature_matrix)\n", - "print(inversion.regularization_matrix)\n", - "print(inversion.curvature_reg_matrix)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Evidence Terms (Advanced)__\n", - "\n", - "In **HowToLens** and the papers below, we cover how an `Inversion` uses a Bayesian evidence to quantify the goodness\n", - "of fit:\n", - "\n", - "https://arxiv.org/abs/1708.07377\n", - "https://arxiv.org/abs/astro-ph/0601493\n", - "\n", - "This evidence balances solutions which fit the data accurately, without using an overly complex regularization source.\n", - "\n", - "The individual terms of the evidence and accessed via the following properties:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.regularization_term)\n", - "print(inversion.log_det_regularization_matrix_term)\n", - "print(inversion.log_det_curvature_reg_matrix_term)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulated Imaging__\n", - "\n", - "We load the source galaxy image from the pixelized inversion of a previous fit, which was performed on an irregular \n", - "RectangularAdaptDensity. \n", - "\n", - "Since irregular meshes cannot be directly used to simulate lensed images, we interpolate the source onto a uniform \n", - "grid with shape `interpolated_pixelized_shape`. This grid should have a high resolution (e.g., 1000 \u00d7 1000) to preserve \n", - "all resolved structure from the original mesh. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.interpolate import griddata\n", - "\n", - "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", - "\n", - "reconstruction = inversion.reconstruction\n", - "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", - "\n", - "interpolated_reconstruction = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", - ")\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_ndarray,\n", - " pixel_scales=interpolation_grid.pixel_scales,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To create the lensed image, we ray-trace image pixels to the source plane and interpolate them onto the source \n", - "galaxy image. \n", - "\n", - "This requires an image-plane grid of (y, x) coordinates. In this example, we use a grid with the same \n", - "resolution as the `Imaging` dataset, but without applying a mask. \n", - "\n", - "To ensure accurate ray-tracing, we apply an 8\u00d78 oversampling scheme. This means that for each pixel in the \n", - "image-plane grid, an 8\u00d78 sub-pixel grid is ray-traced. This approach fully resolves how light is distributed \n", - "across each simulated image pixel, given the source pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=mask.shape_native,\n", - " pixel_scales=mask.pixel_scales,\n", - " over_sample_size=8,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We create a tracer to generate the lensed grid onto which we overlay the interpolated source galaxy image, \n", - "producing the lensed source galaxy image. \n", - "\n", - "The source-plane requires a source galaxy with a defined `redshift` for the tracer to function. Since the source\u2019s \n", - "emission is entirely determined by the source galaxy image, this galaxy has no light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[\n", - " lens,\n", - " al.Galaxy(redshift=source.redshift),\n", - " ]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Using the tracer, we generate the lensed source galaxy image on the image-plane grid. This process incorporates \n", - "the `interpolated_reconstruction`, preserving the irregular and asymmetric morphological features captured by the source reconstruction. \n", - "\n", - "Next, we configure the grid, PSF, and simulator settings to match the signal-to-noise ratio (S/N) and noise properties \n", - "of the observed data used for sensitivity mapping. \n", - "\n", - "The `SimulatorImaging` takes the generated strong lens image and convolves it with the PSF before adding noise. To \n", - "prevent edge effects, the image is padded before convolution and then trimmed to restore its original `shape_native`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=dataset.psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=False,\n", - " noise_seed=1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - "# dataset = simulator.via_interpolated_reconstruction_from(\n", - "# tracer=tracer, grid=grid, interpolated_reconstruction=interpolated_reconstruction\n", - "# )\n", - "#\n", - "#\n", - "#\n", - "# dataset=dataset\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Future Ideas / Contributions__\n", - "\n", - "Here are a list of things I would like to add to this tutorial but haven't found the time. If you are interested\n", - "in having a go at adding them contact me on SLACK! :)\n", - "\n", - "- More magnification calculations.\n", - "- Source gradient calculations.\n", - "- A calculation which shows differential lensing effects (e.g. magnification across the source plane)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Pixelization Fit\n", + "==========================\n", + "\n", + "A pixelization reconstructs the source's light using a pixel-grid, which is regularized using a prior that forces\n", + "the solution to have a degree of smoothness.\n", + "\n", + "This is important for reconstructing complex and irregular source morphologies that cannot be well represented by\n", + "light profiles like a Sersic or shapelets. It is vital to ensuring the infferred lens mass model is accurate and unbiased.\n", + "\n", + "This script fits a source galaxy in a way which uses a pixelization to reconstruct the source's light. This uses a\n", + "rectangular mesh and constant regularization scheme are, which are the simplest forms of each, both providing\n", + "computationally fast and accurate solutions.\n", + "\n", + "For simplicity, the lens galaxy's light is omitted from the model and is not present in the simulated data. It is\n", + "straightforward to include the lens galaxy's light in the model.\n", + "\n", + "Pixelizations are covered in detail in chapter 4 of the **HowToLens** lectures.\n", + "\n", + "__CPU Users__\n", + "\n", + "Matrices must be set up for a pixelized source reconstruction which speed up the linear algebra. On GPU, this takes\n", + "seconds, or at most a minute for datasets with tens of millions, or more visibities. On CPU, this can be a lot slower,\n", + "taking over hours. If you are on CPU, the `feature/pixelization/many_visibilities_preparation` explains how this\n", + "initial setup can be performed before lens modeling and saved to hard disk for fast loading before the model fit.\n", + "\n", + "__Contents__\n", + "\n", + "- **CPU Users:** Matrices must be set up for a pixelized source reconstruction which speed up the linear algebra.\n", + "- **Advantages & Disadvantages:** Many strongly lensed source galaxies are complex, and have asymmetric and irregular morphologies.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", + "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", + "- **Pixelization:** We create a `Pixelization` object to perform the pixelized source reconstruction, which is made up.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Mask Extra Galaxies:** There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "- **Linear Objects:** An `Inversion` contains all of the linear objects used to reconstruct the data in its.\n", + "- **Grids:** The role of a mapper is to map between the image-plane and source-plane.\n", + "- **Reconstruction:** The source reconstruction is also available as a 1D numpy array of values representative of the.\n", + "- **Mapped Reconstructed Images:** The source reconstruction(s) are mapped to the image-plane in order to fit the lens model.\n", + "- **Simulated Imaging:** We load the source galaxy image from the pixelized inversion of a previous fit, which was performed.\n", + "\n", + "__Advantages__\n", + "\n", + "Many strongly lensed source galaxies are complex, and have asymmetric and irregular morphologies. These morphologies\n", + "cannot be well approximated by a light profiles like a Sersic, or many Sersics, and thus a pixelization\n", + "is required to reconstruct the source's irregular light.\n", + "\n", + "Even basis functions like shapelets or a multi-Gaussian expansion cannot reconstruct a source-plane accurately\n", + "if there are multiple source galaxies, or if the source galaxy has a very complex morphology.\n", + "\n", + "To infer detailed components of a lens mass model (e.g. its density slope, whether there's a dark matter subhalo, etc.)\n", + "then pixelized source models are required, to ensure the mass model is fitting all of the lensed source light.\n", + "\n", + "There are also many science cases where one wants to study the highly magnified light of the source galaxy in detail,\n", + "to learnt about distant and faint galaxies. A pixelization reconstructs the source's unlensed emission and thus\n", + "enables this.\n", + "\n", + "__Disadvantages__\n", + "\n", + "Pixelizations are computationally slow and run times are typically longer than a parametric source model. It is not\n", + "uncommon for lens models using a pixelization to take hours to fit high resolution imaging data (e.g. Hubble Space\n", + "Telescope imaging), albeit on modern GPUs run times are often closer to < 20 minutes.\n", + "\n", + "Lens modeling with pixelizations is also more complex than parametric source models, with there being more things\n", + "that can go wrong. For example, there are solutions where a demagnified version of the lensed source galaxy is\n", + "reconstructed, using a mass model which effectively has no mass or too much mass. These are described in detail below.\n", + "\n", + "It will take you longer to learn how to successfully fit lens models with a pixelization than other methods illustrated\n", + "in the workspace!\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algabra solver which allows for positive and negative\n", + "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical. For a pixelizaiton, this often produces negative source pixels which over-fit\n", + "the data, producing unphysical solutions.\n", + "\n", + "All pixelized source reconstructions use a positive-only solver, meaning that every source-pixel is only allowed\n", + "to reconstruct positive flux values. This ensures that the source reconstruction is physical and that we don't\n", + "reconstruct negative flux values that don't exist in the real source galaxy (a common systematic solution in lens\n", + "analysis).\n", + "\n", + "Enforcing positive reconstructions efficiently requires non-trivial linear algebra, so a bespoke JAX fast non-negative\n", + "solver was developed; many methods in the literature omit this and therefore allow unphysical negative solutions that\n", + "can degrade lens modeling results." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autoarray.inversion.plot.inversion_plots import subplot_of_mapper, subplot_mappings" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple__no_lens_light` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "A pixelization uses a separate grid for ray tracing, with its own over sampling scheme, which below we set to a \n", + "uniform grid of values of 4. \n", + "\n", + "The pixelization only reconstructs the source galaxy, therefore the adaptive over sampling used for the lens galaxy's \n", + "light in other examples is not applied to the pixelization. \n", + "\n", + "This example does not model lens light, for examples which combine lens light and a pixelization both over sampling \n", + "schemes should be used, with the lens light adaptive and the pixelization uniform.\n", + "\n", + "Note that the over sampling is input into the `over_sample_size_pixelization` because we are using a `Pixelization`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_over_sampling(\n", + " over_sample_size_pixelization=4,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", + "set below to 28 x 28. \n", + "\n", + "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", + "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", + "mesh, the same number of pixels must be used in the y and x directions.\n", + "\n", + "__Edge Zeroing__\n", + "\n", + "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", + "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", + "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", + "\n", + "For a rectangular mesh, the source code computes edge pixels internally using the known\n", + "pixels at the edge of the mesh. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixelization__\n", + "\n", + "We create a `Pixelization` object to perform the pixelized source reconstruction, which is made up of three\n", + "components:\n", + "\n", + "- `mesh:` Different types of mesh can be used to perform the source reconstruction, where the mesh changes the\n", + "details of how the source is reconstructed (e.g. interpolation weights). In this example, we use a rectangular mesh,\n", + "where the centres computed by overlayiong a rectangular mesh over the source plane.\n", + "\n", + "- `regularization:` A pixelization uses many pixels to reconstructed the source, which will often lead to over fitting\n", + "of the noise in the data and an unrealistically complex and structured source. Regularization smooths the source\n", + "reconstruction solution by penalizing solutions where neighboring pixels have\n", + "large flux differences." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", + "regularization = al.reg.Constant(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This is to illustrate the API for performing a fit via a pixelization using standard autolens objects like \n", + "the `Galaxy`, `Tracer` and `FitImaging` \n", + "\n", + "We simply create a `Pixelization` and pass it to the source galaxy, which then gets input into the tracer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the pixelized source does a good job at capturing the appearance of the source galaxy\n", + "and fitting the data to roughly the noise level." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelizations have bespoke visualizations which show more details about the source-reconstruction, image-mesh\n", + "and other quantities.\n", + "\n", + "The `subplot_of_mapper` function produces a comprehensive diagnostic subplot for the inversion. The\n", + "`subplot_mappings` overlays colored circles in the image and source planes that map to one another, thereby\n", + "allowing one to assess how the mass model ray-traces image-pixels and therefore to assess how the source\n", + "reconstruction maps to the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "\n", + "subplot_of_mapper(inversion=inversion, mapper_index=0)\n", + "subplot_mappings(inversion=inversion, pixelization_index=0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Science (Magnification, Flux and More)__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", + "size of the source.\n", + "\n", + "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", + "including examples using a pixelized source reconstruction. Now you know how to fit a pixelization, go check it out!\n", + "\n", + "__Mask Extra Galaxies__\n", + "\n", + "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source.\n", + "\n", + "If their emission is significant, and close enough to the lens and source, we may simply remove it from the data\n", + "to ensure it does not impact the model-fit. A standard masking approach would be to remove the image pixels containing\n", + "the emission of these galaxies altogether. This is analogous to what the circular masks used throughout the examples\n", + "does.\n", + "\n", + "For fits using a pixelization, masking regions of the image in a way that removes their image pixels entirely from\n", + "the fit. This can produce discontinuities in the pixelixation used to reconstruct the source and produce unexpected\n", + "systematics and unsatisfactory results. In this case, applying the mask in a way where the image pixels are not\n", + "removed from the fit, but their data and noise-map values are scaled such that they contribute negligibly to the fit,\n", + "is a better approach.\n", + "\n", + "We illustrate the API for doing this below, using the `extra_galaxies` dataset which has extra galaxies whose emission\n", + "needs to be removed via scaling in this way. We apply the scaling and show the subplot imaging where the extra\n", + "galaxies mask has scaled the data values to zeros, increasing the noise-map values to large values and in turn made\n", + "the signal to noise of its pixels effectively zero." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_extra_galaxies = al.Mask2D.from_fits(\n", + " file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"),\n", + " pixel_scales=0.1,\n", + " invert=True, # Note that we invert the mask here as `True` means a pixel is scaled.\n", + ")\n", + "\n", + "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native, pixel_scales=0.1, centre=(0.0, 0.0), radius=6.0\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We do not explictly fit this data, for the sake of brevity, however if your data has these nearby galaxies you should\n", + "apply the mask as above before fitting the data.\n", + "\n", + "__Wrap Up__\n", + "\n", + "Pixelizations are the most complex but also most powerful way to model a source galaxy.\n", + "\n", + "Whether you need to use them or not depends on the science you are doing. If you are only interested in measuring a\n", + "simple quantity like the Einstein radius of a lens, you can get away with using light profiles like a Sersic, MGE or \n", + "shapelets to model the source. Low resolution data also means that using a pixelization is not necessary, as the\n", + "complex structure of the source galaxy is not resolved anyway.\n", + "\n", + "However, fitting complex mass models (e.g. a power-law, stellar / dark model or dark matter substructure) requires \n", + "this level of complexity in the source model. Furthermore, if you are interested in studying the properties of the\n", + "source itself, you won't find a better way to do this than using a pixelization.\n", + "\n", + "__Linear Objects__\n", + "\n", + "An `Inversion` contains all of the linear objects used to reconstruct the data in its `linear_obj_list`. \n", + "\n", + "This list may include the following objects:\n", + "\n", + " - `LightProfileLinearObjFuncList`: This object contains lists of linear light profiles and the functionality used\n", + " by them to reconstruct data in an inversion. For example it may only contain a list with a single light profile\n", + " (e.g. `lp_linear.Sersic`) or many light profiles combined in a `Basis` (e.g. `lp_basis.Basis`).\n", + "\n", + "- `Mapper`: The linear objected used by a `Pixelization` to reconstruct data via an `Inversion`, where the `Mapper` \n", + "is specific to the `Pixelization`'s `Mesh` (e.g. a `RectnagularMapper` is used for a `Voronoi` mesh).\n", + "\n", + "In this example, the only linear object used to fit the data was a `Pixelization`, thus the `linear_obj_list`\n", + "contains just one entry corresponding to a `Mapper`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.linear_obj_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To extract results from an inversion many quantities will come in lists or require that we specific the linear object\n", + "we with to use. \n", + "\n", + "Thus, knowing what linear objects are contained in the `linear_obj_list` and what indexes they correspond to\n", + "is important." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Mapper = {inversion.linear_obj_list[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grids__\n", + "\n", + "The role of a mapper is to map between the image-plane and source-plane. \n", + "\n", + "This includes mapping grids corresponding to the data grid (e.g. the centers of each image-pixel in the image and\n", + "source plane) and the pixelization grid (e.g. the centre of the Delaunay triangulation in the image-plane and \n", + "source-plane).\n", + "\n", + "All grids are available in a mapper via its `mapper` property." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = inversion.linear_obj_list[0]\n", + "\n", + "# Centre of each masked image pixel in the image-plane.\n", + "print(mapper.image_plane_data_grid)\n", + "\n", + "# Centre of each source pixel in the source-plane.\n", + "print(mapper.source_plane_data_grid)\n", + "\n", + "# Centre of each pixelization pixel in the image-plane (the `Overlay` image_mesh computes these in the image-plane\n", + "# and maps to the source-plane).\n", + "print(mapper.image_plane_mesh_grid)\n", + "\n", + "# Centre of each pixelization pixel in the source-plane.\n", + "print(mapper.source_plane_mesh_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction__\n", + "\n", + "The source reconstruction is also available as a 1D numpy array of values representative of the source pixelization\n", + "itself (in this example, the reconstructed source values at the vertexes of each Voronoi triangle)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The (y,x) grid of coordinates associated with these values is given by the `Inversion`'s `Mapper` (which are \n", + "described in chapter 4 of **HowToLens**." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = inversion.linear_obj_list[0]\n", + "print(mapper.source_plane_mesh_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The mapper also contains the (y,x) grid of coordinates that correspond to the ray-traced image sub-pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(mapper.source_plane_data_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapped Reconstructed Images__\n", + "\n", + "The source reconstruction(s) are mapped to the image-plane in order to fit the lens model.\n", + "\n", + "These mapped reconstructed images are also accessible via the `Inversion`. \n", + "\n", + "Note that any light profiles in the lens model (e.g. the `bulge` and `disk` of a lens galaxy) are not \n", + "included in this image -- it only contains the source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.mapped_reconstructed_operated_data.native)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Algebra Matrices (Advanced)__\n", + "\n", + "To perform an `Inversion` a number of matrices are constructed which use linear algebra to perform the reconstruction.\n", + "\n", + "These are accessible in the inversion object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.curvature_matrix)\n", + "print(inversion.regularization_matrix)\n", + "print(inversion.curvature_reg_matrix)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Evidence Terms (Advanced)__\n", + "\n", + "In **HowToLens** and the papers below, we cover how an `Inversion` uses a Bayesian evidence to quantify the goodness\n", + "of fit:\n", + "\n", + "https://arxiv.org/abs/1708.07377\n", + "https://arxiv.org/abs/astro-ph/0601493\n", + "\n", + "This evidence balances solutions which fit the data accurately, without using an overly complex regularization source.\n", + "\n", + "The individual terms of the evidence and accessed via the following properties:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.regularization_term)\n", + "print(inversion.log_det_regularization_matrix_term)\n", + "print(inversion.log_det_curvature_reg_matrix_term)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulated Imaging__\n", + "\n", + "We load the source galaxy image from the pixelized inversion of a previous fit, which was performed on an irregular \n", + "RectangularAdaptDensity. \n", + "\n", + "Since irregular meshes cannot be directly used to simulate lensed images, we interpolate the source onto a uniform \n", + "grid with shape `interpolated_pixelized_shape`. This grid should have a high resolution (e.g., 1000 \u00d7 1000) to preserve \n", + "all resolved structure from the original mesh. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.interpolate import griddata\n", + "\n", + "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", + "\n", + "reconstruction = inversion.reconstruction\n", + "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", + "\n", + "interpolated_reconstruction = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", + ")\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_ndarray,\n", + " pixel_scales=interpolation_grid.pixel_scales,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To create the lensed image, we ray-trace image pixels to the source plane and interpolate them onto the source \n", + "galaxy image. \n", + "\n", + "This requires an image-plane grid of (y, x) coordinates. In this example, we use a grid with the same \n", + "resolution as the `Imaging` dataset, but without applying a mask. \n", + "\n", + "To ensure accurate ray-tracing, we apply an 8\u00d78 oversampling scheme. This means that for each pixel in the \n", + "image-plane grid, an 8\u00d78 sub-pixel grid is ray-traced. This approach fully resolves how light is distributed \n", + "across each simulated image pixel, given the source pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=mask.shape_native,\n", + " pixel_scales=mask.pixel_scales,\n", + " over_sample_size=8,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We create a tracer to generate the lensed grid onto which we overlay the interpolated source galaxy image, \n", + "producing the lensed source galaxy image. \n", + "\n", + "The source-plane requires a source galaxy with a defined `redshift` for the tracer to function. Since the source\u2019s \n", + "emission is entirely determined by the source galaxy image, this galaxy has no light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[\n", + " lens,\n", + " al.Galaxy(redshift=source.redshift),\n", + " ]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Using the tracer, we generate the lensed source galaxy image on the image-plane grid. This process incorporates \n", + "the `interpolated_reconstruction`, preserving the irregular and asymmetric morphological features captured by the source reconstruction. \n", + "\n", + "Next, we configure the grid, PSF, and simulator settings to match the signal-to-noise ratio (S/N) and noise properties \n", + "of the observed data used for sensitivity mapping. \n", + "\n", + "The `SimulatorImaging` takes the generated strong lens image and convolves it with the PSF before adding noise. To \n", + "prevent edge effects, the image is padded before convolution and then trimmed to restore its original `shape_native`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=dataset.psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=False,\n", + " noise_seed=1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + "# dataset = simulator.via_interpolated_reconstruction_from(\n", + "# tracer=tracer, grid=grid, interpolated_reconstruction=interpolated_reconstruction\n", + "# )\n", + "#\n", + "#\n", + "#\n", + "# dataset=dataset\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Future Ideas / Contributions__\n", + "\n", + "Here are a list of things I would like to add to this tutorial but haven't found the time. If you are interested\n", + "in having a go at adding them contact me on SLACK! :)\n", + "\n", + "- More magnification calculations.\n", + "- Source gradient calculations.\n", + "- A calculation which shows differential lensing effects (e.g. magnification across the source plane)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/pixelization/likelihood_function.ipynb b/notebooks/imaging/features/pixelization/likelihood_function.ipynb index 195073c28..afe8db9ec 100644 --- a/notebooks/imaging/features/pixelization/likelihood_function.ipynb +++ b/notebooks/imaging/features/pixelization/likelihood_function.ipynb @@ -1,1600 +1,1637 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Pixelization__\n", - "\n", - "This script provides a step-by-step guide of the **PyAutoLens** `log_likelihood_function` which is used to fit\n", - "`Imaging` data with a pixelization (`RectangularUniform` mesh and `Constant` regularization scheme`).\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can understand the\n", - " likelihood function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - " - To make pixelized reconstructions less of a \"black-box\" to users.\n", - "\n", - "__Contents__\n", - "\n", - "- **Simplifications:** This example uses a `RectangularUniform` mesh, where all rectangular source pixels have the same.\n", - "- **Prerequisites:** The likelihood function of a pixelization builds on that used for standard light profiles and.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Masked Image Grid:** To perform lensing calculations we first must define the 2D image-plane (y,x) coordinates used in.\n", - "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", - "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", - "- **Lens Galaxy:** We set up a lens galaxy with the lens light and mass, which we will use to demonstrate a pixelized.\n", - "- **Source Galaxy Pixelization and Regularization:** The source galaxy is reconstructed using a pixel-grid, in this example a RectangularUniform mesh.\n", - "- **Lens Light:** Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the.\n", - "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the.\n", - "- **Border Relocation:** Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace.\n", - "- **Source Pixel Centre Calculation:** In order to reconstruct the source galaxy using a RectangularUniform mesh, we need to determine the.\n", - "- **Interpolation:** We now combine grids computed above to create an `Interpolator`, which describes how image grid.\n", - "- **Mapper:** We now use the interpolator to create a `Mapper`, which describes the mapping between every image.\n", - "- **Alternative Meshes:** We can briefly consider how this step differs for other mesh types.\n", - "- **Mapping Matrix:** The `mapping_matrix` represents the image-pixel to source-pixel mappings above in a 2D matrix.\n", - "- **Image Reconstruction:** Using the reconstructed source pixel fluxes we can map the source reconstruction back to the image.\n", - "- **Likelihood Function:** We now quantify the goodness-of-fit of our lens model and source reconstruction.\n", - "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", - "- **Regularization Term:** The second term, $s^{T} H s$, corresponds to the $\\lambda $G_{\\rm L}$ regularization term we added.\n", - "- **Complexity Terms:** Up to this point, it is unclear why we chose a value of `regularization_coefficient=1.0`.\n", - "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", - "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the lens model, by combining the five terms.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Simplifications__\n", - "\n", - "This example uses a `RectangularUniform` mesh, where all rectangular source pixels have the same size. Most\n", - "pixelization examples use a `RectangularAdaptDensity` mesh, which adapts the size of source pixels to the\n", - "density of points in the source-plane (e.g. the caustic).\n", - "\n", - "The `RectangularUniform` mesh is used here because it is simpler to explain the likelihood function\n", - "and illustrate the key steps in the calculation. The same principles apply to other mesh types, which this\n", - "example will explain where relevant.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function of a pixelization builds on that used for standard light profiles and\n", - "linear light profiles, therefore you must read the following notebooks before this script:\n", - "\n", - "- `imaging/likelihood_function.ipynb`.\n", - "- `imaging/linear_light_profile/likelihood_function.ipynb`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "In order to perform a likelihood evaluation, we first load a dataset.\n", - "\n", - "This example fits a simulated strong lens which is simulated using a 0.1 arcsecond-per-pixel resolution (this is lower\n", - "resolution than the best quality Hubble Space Telescope imaging and close to that of the Euclid space satellite)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "I will use in-built visualization tools for plotting. \n", - "\n", - "For example, using the `aplt.subplot_imaging_dataset` I can plot the imaging dataset we performed a likelihood evaluation on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "The likelihood is only evaluated using image pixels contained within a 2D mask, which we choose before performing\n", - "lens modeling.\n", - "\n", - "Below, we define a 2D circular mask with a 3.0\" radius." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When we plot the masked imaging, only the circular masked region is shown." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=masked_dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling evaluates a light profile using multiple samples of its intensity per image-pixel.\n", - "\n", - "For simplicity, we disable over sampling in this guide by setting `sub_size=1`. \n", - "\n", - "a full description of over sampling and how to use it is given in `autolens_workspace/*/guides/over_sampling.py`.\n", - "\n", - "This example will explain how sub-gridding changes the linear algebra at the relevent section." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_dataset = masked_dataset.apply_over_sampling(\n", - " over_sample_size_lp=1,\n", - " over_sample_size_pixelization=1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masked Image Grid__\n", - "\n", - "To perform lensing calculations we first must define the 2D image-plane (y,x) coordinates used in the calculation.\n", - "\n", - "For light profiles these are given by `masked_dataset.lp`, which is a uniform grid of (y,x) Cartesian coordinates\n", - "which have had the 3.0\" circular mask applied.\n", - "\n", - "A pixelization uses a separate grid of (y,x) coordinates, called `masked_dataset.grids.pixelization`, which is\n", - "identical to the light profile grid but may of had a different over-sampling scale applied (but in this example\n", - "does not).\n", - "\n", - "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", - "used to construct a pixelization there is a straight forward mapping between the image data and pixelization pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=masked_dataset.grids.pixelization, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", - "set below to 28 x 28. \n", - "\n", - "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", - "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", - "mesh, the same number of pixels must be used in the y and x directions.\n", - "\n", - "__Edge Zeroing__\n", - "\n", - "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", - "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", - "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", - "\n", - "For a rectangular mesh, the source code computes edge pixels internally using the known\n", - "pixels at the edge of the mesh. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "We set up a lens galaxy with the lens light and mass, which we will use to demonstrate a pixelized source\n", - "reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Pixelization and Regularization__\n", - "\n", - "The source galaxy is reconstructed using a pixel-grid, in this example a RectangularUniform mesh, which accounts for \n", - "irregularities and asymmetries in the source's surface brightness. \n", - "\n", - "A constant regularization scheme is applied which applies a smoothness prior on the reconstruction. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.RectangularUniform(shape=mesh_shape),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Light__\n", - "\n", - "Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the `Sersic` \n", - "bulge). \n", - "\n", - "This computes the `image` of each `LightProfile` and adds them together. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To convolve the lens's 2D image with the imaging data's PSF, we need its `blurring_image`. This represents all flux \n", - "values not within the mask, which are close enough to it that their flux blurs into the mask after PSF convolution.\n", - "\n", - "To compute this, a `blurring_mask` and `blurring_grid` are used, corresponding to these pixels near the edge of the \n", - "actual mask whose light blurs into the image:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Light Convolution + Subtraction__\n", - "\n", - "Convolve the 2D lens light images above with the PSF using a `Convolver`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", - " image=image, blurring_image=blurring_image_2d\n", - ")\n", - "\n", - "aplt.plot_array(array=convolved_image_2d, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now subtract this image from the observed image to produce a `lens_subtracted_image`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_subtracted_image = masked_dataset.data - convolved_image_2d\n", - "\n", - "aplt.plot_array(array=lens_subtracted_image, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the image-plane to its (y,x) \n", - "source-plane coordinate $\\beta$ using the summed deflection angles $\\alpha$ of the mass profiles:\n", - "\n", - " $\\beta = \\theta - \\alpha(\\theta)$\n", - "\n", - "The likelihood function of a pixelized source reconstruction ray-traces one grid from the image-plane to the source-plane:\n", - "\n", - " 1) A 2D grid of (y,x) coordinates aligned with the imaging data's image-pixels.\n", - " \n", - "The function below computes the 2D deflection angles of the tracer's lens galaxies and subtracts them from the \n", - "image-plane 2D (y,x) coordinates $\\theta$ of each grid, thus ray-tracing their coordinates to the source plane to \n", - "compute their $\\beta$ values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The source code gets quite complex when handling grids for a pixelization, but it is all handled in\n", - "the `TracerToInversion` objects.\n", - "\n", - "The plots at the bottom of this cell show the traced grids used by the source pixelization, showing\n", - "how the traced image pixels are constructed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(tracer=tracer, dataset=masked_dataset)\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", - "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.pixelization\n", - ")[-1]\n", - "\n", - "\n", - "aplt.plot_grid(grid=traced_grid_pixelization, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Border Relocation__\n", - "\n", - "Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace to far outskirts of\n", - "the source-plane. \n", - "\n", - "We relocate these pixels (for both grids above) to the edge of the source-plane border (defined via the border of the \n", - "image-plane mask). This is detailed in **HowToLens chapter 4 tutorial 5** and figure 2 of https://arxiv.org/abs/1708.07377." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", - "\n", - "border_relocator = BorderRelocator(mask=masked_dataset.mask, sub_size=1)\n", - "\n", - "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", - "\n", - "\n", - "aplt.plot_grid(grid=relocated_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Pixel Centre Calculation__\n", - "\n", - "In order to reconstruct the source galaxy using a RectangularUniform mesh, we need to determine the centres of \n", - "the rectangular mesh's source pixels.\n", - "\n", - "We do this by overlying a rectangular grid on the relocated traced image-plane grid computed above.\n", - "\n", - "This distributes the rectangular mesh so it fully overlaps the region of the source-plane containing the traced \n", - "image-pixels without having edge pixels that extend beyond this region." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoarray.inversion.mesh.mesh.rectangular_adapt_density import overlay_grid_from\n", - "\n", - "mesh_grid = overlay_grid_from(\n", - " shape_native=mesh_shape, grid=al.Grid2DIrregular(relocated_grid)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolation__\n", - "\n", - "We now combine grids computed above to create an `Interpolator`, which describes how image grid pixel maps to\n", - "every rectangular mesh pixel. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "interpolator = pixelization.mesh.interpolator_from(\n", - " source_plane_data_grid=relocated_grid,\n", - " source_plane_mesh_grid=mesh_grid,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For the rectangular mesh, the interpolation scheme is called bilinear interpolation, which means that every image \n", - "pixel maps to the rectangular pixel it lands in and the three neighboring rectangular pixels. \n", - "\n", - "The weight of each mapping is determined by the bilinear interpolation scheme, which is a function of how close the \n", - "image pixel is to the centre of the rectangular pixel it lands in and the three neighboring rectangular pixels.\n", - "\n", - "We can print the mappings and weights, for example of the first image pixel, to confirm this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(interpolator.mappings[0])\n", - "print(interpolator.weights[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapper__\n", - "\n", - "We now use the interpolator to create a `Mapper`, which describes the mapping between every image pixel and every \n", - "rectangular pixel, based on the interpolation scheme above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = al.Mapper(interpolator=interpolator)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting the RectangularUniform mesh shows that the source-plane and been discretized into a grid of rectangular pixels.\n", - "\n", - "Below, we plot the RectangularUniform mesh without the traced image-grid pixels (for clarity) and with them as black \n", - "dots in order to show how each set of image-pixels fall within a RectangularUniform pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Mapper` contains:\n", - "\n", - " 1) `source_plane_data_grid`: the traced grid of (y,x) image-pixel coordinate centres (`relocated_grid`).\n", - " 2) `source_plane_mesh_grid`: The RectangularUniform mesh of traced (y,x) source-pixel coordinates (`mesh_grid`).\n", - "\n", - "We have therefore discretized the source-plane into a rectangular mesh, and can pair every traced image-pixel \n", - "coordinate with the corresponding source pixel it lands in.\n", - "\n", - "These quantities are both in the source-plane, and do not by themselves describe the mapping between the image and \n", - "source planes. The mapping is described by the `pix_indexes_for_sub_slim_index`, which maps every image-pixel index to \n", - "every pixelization pixel index.\n", - "\n", - "`pix_indexes` refers to the pixelization pixel indexes (e.g. rectangular pixel 0, 1, 2 etc.) and `sub_slim_index` \n", - "refers to the index of an image pixel (e.g. image-pixel 0, 1, 2 etc.). \n", - "\n", - "For example, printing the first ten entries of `pix_indexes_for_sub_slim_index` shows the first ten source-pixel\n", - "indexes these image sub-pixels map too." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pix_indexes_for_sub_slim_index = mapper.pix_indexes_for_sub_slim_index\n", - "\n", - "print(pix_indexes_for_sub_slim_index[0:9])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This array can be used to visualize how an input list of image-pixel indexes map to the source-plane.\n", - "\n", - "It also shows that image-pixel indexing begins from the top-left and goes rightwards and downwards, accounting for \n", - "all image-pixels which are not masked." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(\n", - " array=lens_subtracted_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", - ")\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The reverse mappings of source-pixels to image-pixels can also be used.\n", - "\n", - "If we choose the right source-pixel index, we can see that multiple imaging occur whereby image-pixels in different\n", - "regions of the image-plane are grouped into the same source-pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pix_indexes = [[200]]\n", - "\n", - "indexes = mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)\n", - "\n", - "\n", - "aplt.plot_array(\n", - " array=lens_subtracted_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", - ")\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolation__\n", - "\n", - "The right hand plot shows more laying over source pixel 200 than its retangular black lines. Pixels further \n", - "out than the pixel appear to be mapped to this source pixel. \n", - "\n", - "This is because the mesh uses an interpolation mapping scheme whereby each image pixels is paired with four source \n", - "pixels. For a rectangular mesh, this scheme is called bilinear interpolation, and it means that every pixel maps\n", - "not only to the rectangular source pixel it lands in, but also the three neighbouring source pixels. Interpolation is \n", - "key to ensuring that the pixelization can reconstruct smooth source morphologies.\n", - "\n", - "We can confirm that every image pixel maps to four source pixels by printing \n", - "the `pix_sizes_for_sub_slim_index`, which gives the number of mapped source pixels for every image pixel.\n", - "\n", - "We can also confirm that the interpolation introduces weights to each mapping by printing the \n", - "`pix_weights_for_sub_slim_index`, which gives the weight of each mapping for every image pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(mapper.pix_sizes_for_sub_slim_index[0:9])\n", - "print(mapper.pix_weights_for_sub_slim_index[0:9])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Lets quickly think about what happens when we use over sampling in the pixelization (e.g. `sub_size>1`). For\n", - "the `sub_size=1` case above, each image pixel maps to 4 source pixels (due to bilinear interpolation)\n", - "with a weight determined from the bilinear interpolation scheme.\n", - "\n", - "However, the default over sampling for a pixelization is `sub_size=4`, meaning each image pixel is divided\n", - "into a 4x4 grid of sub-pixels (16 sub-pixels in total). Each of these sub-pixels maps to 4 source pixels\n", - "(due to bilinear interpolation), where the weight of each mapping is determined by the bilinear interpolation\n", - "scheme divided by 16 (because there are 16 sub-pixels).\n", - "\n", - "This example therefore used a `sub_size=1` to keep the explanation of the likelihood, visualization of the\n", - "arrays above and understanding of the mapping scheme as simple as possible. You can manually increase the\n", - "`sub_size` above and re-run the notebook to see how this changes the mapping scheme.\n", - "\n", - "__Alternative Meshes__\n", - "\n", - "We can briefly consider how this step differs for other mesh types. Above, we simply overlaid a uniform rectangular\n", - "grid to define the source pixel centres and then mapped image pixels to these source pixels.\n", - "\n", - "The `RectangularAdaptDensity` mesh pretty much works exactly the same, its just that a calculation (which we don't\n", - "describe here) works out how to make a grid of rectangular pixels that adapt to the source-plane density and thus\n", - "vary in size. \n", - "\n", - "There is also a `RectangularAdaptImage` mesh which uses the image of the lensed source to adapt\n", - "the rectangular pixel sizes. This often puts even smaller pixels in the brightest regions of the source,\n", - "even if it lies offset or away from the caustic.\n", - "\n", - "There is also a `Delaunay` mesh which uses a Delaunay triangulation to define an irregular grid of source pixels.\n", - "This is described fully in the `delaunay` example including a likelihood function guide.\n", - "\n", - "__Mapping Matrix__\n", - "\n", - "The `mapping_matrix` represents the image-pixel to source-pixel mappings above in a 2D matrix. \n", - "\n", - "It has dimensions `(total_image_pixels, total_source_pixels)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", - " pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,\n", - " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,\n", - " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,\n", - " pixels=mapper.pixels,\n", - " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", - " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", - " sub_fraction=mapper.over_sampler.sub_fraction,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `mapping_matrix` shows of all image-source pixel mappings.\n", - "\n", - "No row of pixels has more than one non-zero entry. It is not possible for two image pixels to map to the same source \n", - "pixel (meaning that there are no correlated pixels in the mapping matrix)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Each column of the `mapping_matrix` can therefore be used to show all image-pixels it maps too. \n", - "\n", - "For example, above, we plotted all image-pixels of source-pixel 200 (as well as 202 and 204). We can extract all\n", - "image-pixel indexes of source pixels 200 using the `mapping_matrix` and use them to plot the image of this\n", - "source-pixel (which corresponds to only values of zeros or ones)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "indexes_source_pix_200 = np.nonzero(mapping_matrix[:, 200])\n", - "\n", - "print(indexes_source_pix_200[0])\n", - "\n", - "array_2d = al.Array2D(values=mapping_matrix[:, 200], mask=masked_dataset.mask)\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Blurred Mapping Matrix ($f$)__\n", - "\n", - "Each source-pixel can therefore be thought of as an image (where all entries of this image are zeros and ones). \n", - "\n", - "To incorporate the imaging data's PSF, we simply blur each one of these source-pixel images with the imaging data's \n", - "Point Spread Function (PSF) via 2D convolution.\n", - "\n", - "This operation does not change the dimensions of the mapping matrix, meaning the `blurred_mapping_matrix` also has\n", - "dimensions `(total_image_pixels, total_source_pixels)`. It turns the values of zeros and ones into \n", - "non-integer values which have been blurred by the PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "blurred_mapping_matrix = masked_dataset.psf.convolved_mapping_matrix_from(\n", - " mapping_matrix=mapping_matrix, mask=masked_dataset.mask\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `blurred_mapping_matrix` shows all image-source pixel mappings including PSF blurring.\n", - "\n", - "Note how, unlike for the `mapping_matrix`, every row of image-pixels now has multiple non-zero entries. It is now \n", - "possible for two image pixels to map to the same source pixel, because they become correlated by PSF convolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " blurred_mapping_matrix,\n", - " aspect=(blurred_mapping_matrix.shape[1] / blurred_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Each column of the `blurred_mapping_matrix` shows all image-pixels it maps to after PSF blurring. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "indexes_source_pix_200 = np.nonzero(blurred_mapping_matrix[:, 200])\n", - "\n", - "print(indexes_source_pix_200[0])\n", - "\n", - "array_2d = al.Array2D(values=blurred_mapping_matrix[:, 200], mask=masked_dataset.mask)\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) the `blurred_mapping_matrix` is denoted $f_{ij}$\n", - "where $i$ maps over all $I$ source pixels and $j$ maps over all $J$ image pixels. \n", - "\n", - "For example: \n", - "\n", - " - $f_{0, 2} = 0.3$ indicates that image-pixel $2$ maps to source-pixel $0$ with a weight of $0.3$ after PSF convolution.\n", - " - $f_{4, 8} = 0$ indicates that image-pixel $8$ does not map to source-pixel $4$, even after PSF convolution.\n", - "\n", - "The indexing of the `mapping_matrix` is reversed compared to the notation of WD03 (e.g. image pixels\n", - "are the first entry of `mapping_matrix` whereas for $f$ they are the second index)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Mapping between image pixel 0 and source pixel 2 = {mapping_matrix[0, 2]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Data Vector (D)__\n", - "\n", - "To solve for the source pixel fluxes we now pose the problem as a linear inversion.\n", - "\n", - "This requires us to convert the `blurred_mapping_matrix` and our `data` and `noise map` into matrices of certain dimensions. \n", - "\n", - "The `data_vector`, $D$, is the first matrix and it has dimensions `(total_source_pixels,)`.\n", - "\n", - "In WD03 (https://arxiv.org/abs/astro-ph/0302587) and N15 (https://arxiv.org/abs/1412.7436) the data vector \n", - "is give by: \n", - "\n", - " $\\vec{D}_{i} = \\sum_{\\rm j=1}^{J}f_{ij}(d_{j} - b_{j})/\\sigma_{j}^2 \\, \\, .$\n", - "\n", - "Where:\n", - "\n", - " - $d_{\\rm j}$ are the image-pixel data flux values.\n", - " - $b_{\\rm j}$ are the brightness values of the lens light model (therefore $d_{\\rm j} - b_{\\rm j}$ is the lens light\n", - " subtracted image).\n", - " - $\\sigma{\\rm _j}^2$ are the statistical uncertainties of each image-pixel value.\n", - "\n", - "$i$ maps over all $I$ source pixels and $j$ maps over all $J$ image pixels. \n", - "\n", - "NOTE: WD03 assume the data is already lens subtracted thus $b_{j}$ is omitted (e.g. all values are zero)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_vector = al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(\n", - " blurred_mapping_matrix=blurred_mapping_matrix,\n", - " image=np.array(lens_subtracted_image),\n", - " noise_map=np.array(masked_dataset.noise_map),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$D$ describes which deconvolved source pixels trace to which image-plane pixels. This ensures the source reconstruction\n", - "fully accounts for the PSF when fitting the data.\n", - "\n", - "We can plot $D$ as a column vector:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Matrix (F)__\n", - "\n", - "The `curvature_matrix` $F$ is the second matrix and it has dimensions `(total_source_pixels, total_source_pixels)`.\n", - "\n", - "In WD03 / N15 (https://arxiv.org/abs/astro-ph/0302587) the curvature matrix is a 2D matrix given by:\n", - "\n", - " ${F}_{ik} = \\sum_{\\rm j=1}^{J}f_{ij}f_{kj}/\\sigma_{j}^2 \\, \\, .$\n", - "\n", - "NOTE: this notation implicitly assumes a summation over $K$, where $k$ runs over all source-pixel indexes $K$.\n", - "\n", - "Note how summation over $J$ runs over $f$ twice, such that every entry of $F$ is the sum of the multiplication\n", - "between all values in every two columns of $f$.\n", - "\n", - "For example, $F_{0,1}$ is the sum of every blurred image pixels values in $f$ of source pixel 0 multiplied by\n", - "every blurred image pixel value of source pixel 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix, noise_map=masked_dataset.noise_map\n", - ")\n", - "\n", - "plt.imshow(curvature_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For $F_{ik}$ to be non-zero, this requires that the images of source pixels $i$ and $k$ share at least one\n", - "image-pixel, which we saw above is only possible due to PSF blurring.\n", - "\n", - "For example, we can see a non-zero entry for $F_{100,101}$ and plotting their images\n", - "show overlap." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_pixel_0 = 0\n", - "source_pixel_1 = 1\n", - "\n", - "print(curvature_matrix[source_pixel_0, source_pixel_1])\n", - "\n", - "array_2d = al.Array2D(\n", - " values=blurred_mapping_matrix[:, source_pixel_0], mask=masked_dataset.mask\n", - ")\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")\n", - "\n", - "array_2d = al.Array2D(\n", - " values=blurred_mapping_matrix[:, source_pixel_1], mask=masked_dataset.mask\n", - ")\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The following chi-squared is minimized when we perform the inversion and reconstruct the source:\n", - "\n", - "$\\chi^2 = \\sum_{\\rm j=1}^{J} \\bigg[ \\frac{(\\sum_{\\rm i=1}^{I} s_{i} f_{ij}) + b_{j} - d_{j}}{\\sigma_{j}} \\bigg]$\n", - "\n", - "Where $s$ is the reconstructed source pixel fluxes in all $I$ source pixels.\n", - "\n", - "The solution for $s$ is therefore given by (equation 5 WD03):\n", - "\n", - " $s = F^{-1} D$\n", - " \n", - "We can compute this using NumPy linear algebra:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# Because we are no using regularizartion (see below) it is common for the curvature matrix to be singular and lead\n", - "# to a LinAlgException. The loop below mitigates this -- you can ignore it as it is not important for understanding\n", - "# the PyAutoLens likelihood function.\n", - "\n", - "for i in range(curvature_matrix.shape[0]):\n", - " curvature_matrix[i, i] += 1e-8\n", - "\n", - "reconstruction = np.linalg.solve(curvature_matrix, data_vector)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot this source reconstruction -- it looks like a mess.\n", - "\n", - "The source pixels have noisy and unsmooth values, and it is hard to make out if a source is even being reconstructed. \n", - "\n", - "In fact, the linear inversion is (over-)fitting noise in the image data, meaning this system of equations is \n", - "ill-posed. We need to apply some form of smoothing on the source reconstruction to avoid over fitting noise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Regularization Matrix (H)__\n", - "\n", - "Regularization adds a linear regularization term $G_{\\rm L}$ to the $\\chi^2$ we solve for giving us a new merit \n", - "function $G$ (equation 11 WD03):\n", - "\n", - " $G = \\chi^2 + \\lambda \\, G_{\\rm L}$\n", - " \n", - "where $\\lambda$ is the `regularization_coefficient` which describes the magnitude of smoothness that is applied. A \n", - "higher $\\lambda$ will regularize the source more, leading to a smoother source reconstruction.\n", - " \n", - "Different forms for $G_{\\rm L}$ can be defined which regularize the source reconstruction in different ways. The \n", - "`Constant` regularization scheme used in this example applies gradient regularization (equation 14 WD03):\n", - "\n", - " $G_{\\rm L} = \\sum_{\\rm i}^{I} \\sum_{\\rm n=1}^{N} [s_{i} - s_{i, v}]$\n", - "\n", - "This regularization scheme is easier to express in words -- the summation goes to each source pixel,\n", - "determines all source pixels with which it shares a direct vertex (e.g. its neighbors) and penalizes solutions \n", - "where the difference in reconstructed flux of these two neighboring source pixels is large.\n", - "\n", - "The summation does this for all pixels, thus it favours solutions where neighboring source\n", - "pixels reconstruct similar values to one another (e.g. it favours a smooth source reconstruction).\n", - "\n", - "We now define the `regularization matrix`, $H$, which allows us to include this smoothing when we solve for $s$. $H$\n", - "has dimensions `(total_source_pixels, total_source_pixels)`.\n", - "\n", - "This relates to $G_{\\rm L}$ as (equation 13 WD03):\n", - "\n", - " $H_{ik} = \\frac{1}{2} \\frac{\\partial G_{\\rm L}}{\\partial s_{i} \\partial s_{k}}$\n", - "\n", - "$H$ has the `regularization_coefficient` $\\lambda$ folded into it such $\\lambda$'s control on the degree of smoothing\n", - "is accounted for." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", - " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", - " neighbors=mapper.neighbors,\n", - " neighbors_sizes=mapper.neighbors.sizes,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot the regularization matrix and note that:\n", - "\n", - " - non-zero entries indicate that two source-pixels are neighbors and therefore are regularized with one another.\n", - " \n", - " - Zeros indicate the two source pixels do not neighbor one another.\n", - " \n", - "The majority of entries are zero, because the majority of source pixels are not neighbors with one another." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(regularization_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__F + Lamdba H__\n", - "\n", - "$H$ enters the linear algebra system we solve for as follows (WD03 equation (12)):\n", - "\n", - " $s = [F + H]^{-1} D$" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Reconstruction (s)__\n", - "\n", - "We can now solve the linear system above using NumPy linear algebra. \n", - "\n", - "Note that the for loop used above to prevent a LinAlgException is no longer required." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting this source reconstruction we can see that regularization has lead us to reconstruct a smoother source,\n", - "which actually looks like a galaxy! This also implies we are not over-fitting the noise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image Reconstruction__\n", - "\n", - "Using the reconstructed source pixel fluxes we can map the source reconstruction back to the image plane (via\n", - "the `blurred mapping_matrix`) and produce a reconstruction of the image data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_operated_data = (\n", - " al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(\n", - " mapping_matrix=blurred_mapping_matrix, reconstruction=reconstruction\n", - " )\n", - ")\n", - "\n", - "mapped_reconstructed_operated_data = al.Array2D(\n", - " values=mapped_reconstructed_operated_data, mask=mask\n", - ")\n", - "\n", - "aplt.plot_array(array=mapped_reconstructed_operated_data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our lens model and source reconstruction. \n", - "\n", - "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", - "\n", - "The likelihood function for lens modeling consists of five terms:\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + s^{T} H s + \\mathrm{ln} \\, \\left[ \\mathrm{det} (F + H) \\right] - { \\mathrm{ln}} \\, \\left[ \\mathrm{det} (H) \\right] + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "This expression was first derived by Suyu 2006 (https://arxiv.org/abs/astro-ph/0601493) and is given by equation (19).\n", - "It was derived into **PyAutoLens** notation in Dye 2008 (https://arxiv.org/abs/0804.4002) equation (5).\n", - "\n", - "We now explain what each of these terms mean.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", - "\n", - " - `model_data` = `mapped_reconstructed_operated_data` + `lens_light_convolved_image`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", - "\n", - "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", - "for, corresponding to a fit with a lower likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image = convolved_image_2d + mapped_reconstructed_operated_data\n", - "\n", - "residual_map = masked_dataset.data - model_image\n", - "normalized_residual_map = residual_map / masked_dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(chi_squared)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", - "\n", - "aplt.plot_array(array=chi_squared_map, title=\"\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Regularization Term__\n", - "\n", - "The second term, $s^{T} H s$, corresponds to the $\\lambda $G_{\\rm L}$ regularization term we added to our merit \n", - "function above.\n", - "\n", - "This is the term which sums up the difference in flux of all reconstructed source pixels, and reduces the likelihood of \n", - "solutions where there are large differences in flux (e.g. the source is less smooth and more likely to be \n", - "overfitting noise).\n", - "\n", - "We compute it below via matrix multiplication, noting that the `regularization_coefficient`, $\\lambda$, is built into \n", - "the `regularization_matrix` already." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "regularization_term = np.matmul(\n", - " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", - ")\n", - "\n", - "print(regularization_term)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Complexity Terms__\n", - "\n", - "Up to this point, it is unclear why we chose a value of `regularization_coefficient=1.0`. \n", - "\n", - "We cannot rely on the `chi_squared` and `regularization_term` above to optimally choose its value, because increasing \n", - "the `regularization_coefficient` smooths the solution more and therefore:\n", - " \n", - " - Decreases `chi_squared` by fitting the data worse, producing a lower `log_likelihood`.\n", - " \n", - " - Increases the `regularization_term` by penalizing the differences between source pixel fluxes more, again reducing\n", - " the inferred `log_likelihood`.\n", - "\n", - "If we set the regularization coefficient based purely on these two terms, we would set a value of 0.0 and be back where\n", - "we started over-fitting noise!\n", - "\n", - "The terms $\\left[ \\mathrm{det} (F + H) \\right]$ and $ - { \\mathrm{ln}} \\, \\left[ \\mathrm{det} (H) \\right]$ address \n", - "this problem. \n", - "\n", - "They quantify how complex the source reconstruction is, and penalize solutions where *it is more complex*. Reducing \n", - "the `regularization_coefficient` makes the source reconstruction more complex (because a source that is \n", - "smoothed less uses more flexibility to fit the data better).\n", - "\n", - "These two terms therefore counteract the `chi_squared` and `regularization_term`, so as to attribute a higher\n", - "`log_likelihood` to solutions which fit the data with a more smoothed and less complex source (e.g. one with a higher \n", - "`regularization_coefficient`).\n", - "\n", - "In **HowToLens** -> `chapter 4` -> `tutorial_4_bayesian_regularization` we expand on this further and give a more\n", - "detailed description of how these different terms impact the `log_likelihood_function`. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", - "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]\n", - "\n", - "print(log_curvature_reg_matrix_term)\n", - "print(log_regularization_matrix_term)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", - "\n", - "The final term ins the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", - "of the log of every noise-map value squared. \n", - "\n", - "Given the `noise_map` is fixed, this term does not change during the lens modeling process and has no impact on the \n", - "model we infer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "We can now, finally, compute the `log_likelihood` of the lens model, by combining the five terms computed above using\n", - "the likelihood function defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_evidence = float(\n", - " -0.5\n", - " * (\n", - " chi_squared\n", - " + regularization_term\n", - " + log_curvature_reg_matrix_term\n", - " - log_regularization_matrix_term\n", - " + noise_normalization\n", - " )\n", - ")\n", - "\n", - "print(log_evidence)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(\n", - " dataset=masked_dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_border_relocator=True),\n", - ")\n", - "fit_log_evidence = fit.log_evidence\n", - "print(fit_log_evidence)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `Nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the **PyAutoLens** likelihood function, which uses a pixelization, \n", - "regularization scheme and inversion to reconstruct the source galaxy.\n", - "\n", - "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", - "are described above but worth refreshing ourselves on:\n", - "\n", - " - **Sub-gridding**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", - " ray-traced to the source-plane and paired fractionally with each source pixel.\n", - " \n", - " - **Source-plane Interpolation**: Using a RectangularUniform mesh with bilinear interpolation\n", - " to pair each image (sub-)pixel to multiple source-plane pixels with interpolation weights.\n", - " \n", - " - **Source Morphology Pixelization Adaption**: Adapting the pixelization such that is congregates source pixels around\n", - " the source's brightest regions, as opposed to the magnification-based pixelization used here.\n", - " \n", - " - **Luminosity Weighted Regularization**: Using an adaptive regularization coefficient which adapts the level of \n", - " regularization applied to the source based on its luminosity." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Pixelization__\n", + "\n", + "This script provides a step-by-step guide of the **PyAutoLens** `log_likelihood_function` which is used to fit\n", + "`Imaging` data with a pixelization (`RectangularUniform` mesh and `Constant` regularization scheme`).\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can understand the\n", + " likelihood function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + " - To make pixelized reconstructions less of a \"black-box\" to users.\n", + "\n", + "__Contents__\n", + "\n", + "- **Simplifications:** This example uses a `RectangularUniform` mesh, where all rectangular source pixels have the same.\n", + "- **Prerequisites:** The likelihood function of a pixelization builds on that used for standard light profiles and.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Masked Image Grid:** To perform lensing calculations we first must define the 2D image-plane (y,x) coordinates used in.\n", + "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", + "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", + "- **Lens Galaxy:** We set up a lens galaxy with the lens light and mass, which we will use to demonstrate a pixelized.\n", + "- **Source Galaxy Pixelization and Regularization:** The source galaxy is reconstructed using a pixel-grid, in this example a RectangularUniform mesh.\n", + "- **Lens Light:** Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the.\n", + "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the.\n", + "- **Border Relocation:** Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace.\n", + "- **Source Pixel Centre Calculation:** In order to reconstruct the source galaxy using a RectangularUniform mesh, we need to determine the.\n", + "- **Interpolation:** We now combine grids computed above to create an `Interpolator`, which describes how image grid.\n", + "- **Mapper:** We now use the interpolator to create a `Mapper`, which describes the mapping between every image.\n", + "- **Alternative Meshes:** We can briefly consider how this step differs for other mesh types.\n", + "- **Mapping Matrix:** The `mapping_matrix` represents the image-pixel to source-pixel mappings above in a 2D matrix.\n", + "- **Image Reconstruction:** Using the reconstructed source pixel fluxes we can map the source reconstruction back to the image.\n", + "- **Likelihood Function:** We now quantify the goodness-of-fit of our lens model and source reconstruction.\n", + "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", + "- **Regularization Term:** The second term, $s^{T} H s$, corresponds to the $\\lambda $G_{\\rm L}$ regularization term we added.\n", + "- **Complexity Terms:** Up to this point, it is unclear why we chose a value of `regularization_coefficient=1.0`.\n", + "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", + "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the lens model, by combining the five terms.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Simplifications__\n", + "\n", + "This example uses a `RectangularUniform` mesh, where all rectangular source pixels have the same size. Most\n", + "pixelization examples use a `RectangularAdaptDensity` mesh, which adapts the size of source pixels to the\n", + "density of points in the source-plane (e.g. the caustic).\n", + "\n", + "The `RectangularUniform` mesh is used here because it is simpler to explain the likelihood function\n", + "and illustrate the key steps in the calculation. The same principles apply to other mesh types, which this\n", + "example will explain where relevant.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function of a pixelization builds on that used for standard light profiles and\n", + "linear light profiles, therefore you must read the following notebooks before this script:\n", + "\n", + "- `imaging/likelihood_function.ipynb`.\n", + "- `imaging/linear_light_profile/likelihood_function.ipynb`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "In order to perform a likelihood evaluation, we first load a dataset.\n", + "\n", + "This example fits a simulated strong lens which is simulated using a 0.1 arcsecond-per-pixel resolution (this is lower\n", + "resolution than the best quality Hubble Space Telescope imaging and close to that of the Euclid space satellite)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "I will use in-built visualization tools for plotting. \n", + "\n", + "For example, using the `aplt.subplot_imaging_dataset` I can plot the imaging dataset we performed a likelihood evaluation on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "The likelihood is only evaluated using image pixels contained within a 2D mask, which we choose before performing\n", + "lens modeling.\n", + "\n", + "Below, we define a 2D circular mask with a 3.0\" radius." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When we plot the masked imaging, only the circular masked region is shown." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=masked_dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling evaluates a light profile using multiple samples of its intensity per image-pixel.\n", + "\n", + "For simplicity, we disable over sampling in this guide by setting `sub_size=1`. \n", + "\n", + "a full description of over sampling and how to use it is given in `autolens_workspace/*/guides/over_sampling.py`.\n", + "\n", + "This example will explain how sub-gridding changes the linear algebra at the relevent section." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_dataset = masked_dataset.apply_over_sampling(\n", + " over_sample_size_lp=1,\n", + " over_sample_size_pixelization=1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masked Image Grid__\n", + "\n", + "To perform lensing calculations we first must define the 2D image-plane (y,x) coordinates used in the calculation.\n", + "\n", + "For light profiles these are given by `masked_dataset.lp`, which is a uniform grid of (y,x) Cartesian coordinates\n", + "which have had the 3.0\" circular mask applied.\n", + "\n", + "A pixelization uses a separate grid of (y,x) coordinates, called `masked_dataset.grids.pixelization`, which is\n", + "identical to the light profile grid but may of had a different over-sampling scale applied (but in this example\n", + "does not).\n", + "\n", + "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", + "used to construct a pixelization there is a straight forward mapping between the image data and pixelization pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=masked_dataset.grids.pixelization, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", + "set below to 28 x 28. \n", + "\n", + "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", + "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", + "mesh, the same number of pixels must be used in the y and x directions.\n", + "\n", + "__Edge Zeroing__\n", + "\n", + "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", + "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", + "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", + "\n", + "For a rectangular mesh, the source code computes edge pixels internally using the known\n", + "pixels at the edge of the mesh. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "We set up a lens galaxy with the lens light and mass, which we will use to demonstrate a pixelized source\n", + "reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Pixelization and Regularization__\n", + "\n", + "The source galaxy is reconstructed using a pixel-grid, in this example a RectangularUniform mesh, which accounts for \n", + "irregularities and asymmetries in the source's surface brightness. \n", + "\n", + "A constant regularization scheme is applied which applies a smoothness prior on the reconstruction. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.RectangularUniform(shape=mesh_shape),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Light__\n", + "\n", + "Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the `Sersic` \n", + "bulge). \n", + "\n", + "This computes the `image` of each `LightProfile` and adds them together. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To convolve the lens's 2D image with the imaging data's PSF, we need its `blurring_image`. This represents all flux \n", + "values not within the mask, which are close enough to it that their flux blurs into the mask after PSF convolution.\n", + "\n", + "To compute this, a `blurring_mask` and `blurring_grid` are used, corresponding to these pixels near the edge of the \n", + "actual mask whose light blurs into the image:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Light Convolution + Subtraction__\n", + "\n", + "Convolve the 2D lens light images above with the PSF using a `Convolver`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", + " image=image, blurring_image=blurring_image_2d\n", + ")\n", + "\n", + "aplt.plot_array(array=convolved_image_2d, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now subtract this image from the observed image to produce a `lens_subtracted_image`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_subtracted_image = masked_dataset.data - convolved_image_2d\n", + "\n", + "aplt.plot_array(array=lens_subtracted_image, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the image-plane to its (y,x) \n", + "source-plane coordinate $\\beta$ using the summed deflection angles $\\alpha$ of the mass profiles:\n", + "\n", + " $\\beta = \\theta - \\alpha(\\theta)$\n", + "\n", + "The likelihood function of a pixelized source reconstruction ray-traces one grid from the image-plane to the source-plane:\n", + "\n", + " 1) A 2D grid of (y,x) coordinates aligned with the imaging data's image-pixels.\n", + " \n", + "The function below computes the 2D deflection angles of the tracer's lens galaxies and subtracts them from the \n", + "image-plane 2D (y,x) coordinates $\\theta$ of each grid, thus ray-tracing their coordinates to the source plane to \n", + "compute their $\\beta$ values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The source code gets quite complex when handling grids for a pixelization, but it is all handled in\n", + "the `TracerToInversion` objects.\n", + "\n", + "The plots at the bottom of this cell show the traced grids used by the source pixelization, showing\n", + "how the traced image pixels are constructed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(tracer=tracer, dataset=masked_dataset)\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", + "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.pixelization\n", + ")[-1]\n", + "\n", + "\n", + "aplt.plot_grid(grid=traced_grid_pixelization, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Border Relocation__\n", + "\n", + "Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace to far outskirts of\n", + "the source-plane. \n", + "\n", + "We relocate these pixels (for both grids above) to the edge of the source-plane border (defined via the border of the \n", + "image-plane mask). This is detailed in **HowToLens chapter 4 tutorial 5** and figure 2 of https://arxiv.org/abs/1708.07377." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", + "\n", + "border_relocator = BorderRelocator(mask=masked_dataset.mask, sub_size=1)\n", + "\n", + "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", + "\n", + "\n", + "aplt.plot_grid(grid=relocated_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Pixel Centre Calculation__\n", + "\n", + "In order to reconstruct the source galaxy using a RectangularUniform mesh, we need to determine the centres of \n", + "the rectangular mesh's source pixels.\n", + "\n", + "We do this by overlying a rectangular grid on the relocated traced image-plane grid computed above.\n", + "\n", + "This distributes the rectangular mesh so it fully overlaps the region of the source-plane containing the traced \n", + "image-pixels without having edge pixels that extend beyond this region." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoarray.inversion.mesh.mesh.rectangular_adapt_density import overlay_grid_from\n", + "\n", + "mesh_grid = overlay_grid_from(\n", + " shape_native=mesh_shape, grid=al.Grid2DIrregular(relocated_grid)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolation__\n", + "\n", + "We now combine grids computed above to create an `Interpolator`, which describes how image grid pixel maps to\n", + "every rectangular mesh pixel. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "interpolator = pixelization.mesh.interpolator_from(\n", + " source_plane_data_grid=relocated_grid,\n", + " source_plane_mesh_grid=mesh_grid,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For the rectangular mesh, the interpolation scheme is called bilinear interpolation, which means that every image \n", + "pixel maps to the rectangular pixel it lands in and the three neighboring rectangular pixels. \n", + "\n", + "The weight of each mapping is determined by the bilinear interpolation scheme, which is a function of how close the \n", + "image pixel is to the centre of the rectangular pixel it lands in and the three neighboring rectangular pixels.\n", + "\n", + "We can print the mappings and weights, for example of the first image pixel, to confirm this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(interpolator.mappings[0])\n", + "print(interpolator.weights[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapper__\n", + "\n", + "We now use the interpolator to create a `Mapper`, which describes the mapping between every image pixel and every \n", + "rectangular pixel, based on the interpolation scheme above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = al.Mapper(interpolator=interpolator)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting the RectangularUniform mesh shows that the source-plane and been discretized into a grid of rectangular pixels.\n", + "\n", + "Below, we plot the RectangularUniform mesh without the traced image-grid pixels (for clarity) and with them as black \n", + "dots in order to show how each set of image-pixels fall within a RectangularUniform pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Mapper` contains:\n", + "\n", + " 1) `source_plane_data_grid`: the traced grid of (y,x) image-pixel coordinate centres (`relocated_grid`).\n", + " 2) `source_plane_mesh_grid`: The RectangularUniform mesh of traced (y,x) source-pixel coordinates (`mesh_grid`).\n", + "\n", + "We have therefore discretized the source-plane into a rectangular mesh, and can pair every traced image-pixel \n", + "coordinate with the corresponding source pixel it lands in.\n", + "\n", + "These quantities are both in the source-plane, and do not by themselves describe the mapping between the image and \n", + "source planes. The mapping is described by the `pix_indexes_for_sub_slim_index`, which maps every image-pixel index to \n", + "every pixelization pixel index.\n", + "\n", + "`pix_indexes` refers to the pixelization pixel indexes (e.g. rectangular pixel 0, 1, 2 etc.) and `sub_slim_index` \n", + "refers to the index of an image pixel (e.g. image-pixel 0, 1, 2 etc.). \n", + "\n", + "For example, printing the first ten entries of `pix_indexes_for_sub_slim_index` shows the first ten source-pixel\n", + "indexes these image sub-pixels map too." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pix_indexes_for_sub_slim_index = mapper.pix_indexes_for_sub_slim_index\n", + "\n", + "print(pix_indexes_for_sub_slim_index[0:9])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This array can be used to visualize how an input list of image-pixel indexes map to the source-plane.\n", + "\n", + "It also shows that image-pixel indexing begins from the top-left and goes rightwards and downwards, accounting for \n", + "all image-pixels which are not masked." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(\n", + " array=lens_subtracted_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", + ")\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The reverse mappings of source-pixels to image-pixels can also be used.\n", + "\n", + "If we choose the right source-pixel index, we can see that multiple imaging occur whereby image-pixels in different\n", + "regions of the image-plane are grouped into the same source-pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pix_indexes = [[200]]\n", + "\n", + "indexes = mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)\n", + "\n", + "\n", + "aplt.plot_array(\n", + " array=lens_subtracted_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", + ")\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolation__\n", + "\n", + "The right hand plot shows more laying over source pixel 200 than its retangular black lines. Pixels further \n", + "out than the pixel appear to be mapped to this source pixel. \n", + "\n", + "This is because the mesh uses an interpolation mapping scheme whereby each image pixels is paired with four source \n", + "pixels. For a rectangular mesh, this scheme is called bilinear interpolation, and it means that every pixel maps\n", + "not only to the rectangular source pixel it lands in, but also the three neighbouring source pixels. Interpolation is \n", + "key to ensuring that the pixelization can reconstruct smooth source morphologies.\n", + "\n", + "We can confirm that every image pixel maps to four source pixels by printing \n", + "the `pix_sizes_for_sub_slim_index`, which gives the number of mapped source pixels for every image pixel.\n", + "\n", + "We can also confirm that the interpolation introduces weights to each mapping by printing the \n", + "`pix_weights_for_sub_slim_index`, which gives the weight of each mapping for every image pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(mapper.pix_sizes_for_sub_slim_index[0:9])\n", + "print(mapper.pix_weights_for_sub_slim_index[0:9])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Lets quickly think about what happens when we use over sampling in the pixelization (e.g. `sub_size>1`). For\n", + "the `sub_size=1` case above, each image pixel maps to 4 source pixels (due to bilinear interpolation)\n", + "with a weight determined from the bilinear interpolation scheme.\n", + "\n", + "However, the default over sampling for a pixelization is `sub_size=4`, meaning each image pixel is divided\n", + "into a 4x4 grid of sub-pixels (16 sub-pixels in total). Each of these sub-pixels maps to 4 source pixels\n", + "(due to bilinear interpolation), where the weight of each mapping is determined by the bilinear interpolation\n", + "scheme divided by 16 (because there are 16 sub-pixels).\n", + "\n", + "This example therefore used a `sub_size=1` to keep the explanation of the likelihood, visualization of the\n", + "arrays above and understanding of the mapping scheme as simple as possible. You can manually increase the\n", + "`sub_size` above and re-run the notebook to see how this changes the mapping scheme.\n", + "\n", + "__Alternative Meshes__\n", + "\n", + "We can briefly consider how this step differs for other mesh types. Above, we simply overlaid a uniform rectangular\n", + "grid to define the source pixel centres and then mapped image pixels to these source pixels.\n", + "\n", + "The `RectangularAdaptDensity` mesh pretty much works exactly the same, its just that a calculation (which we don't\n", + "describe here) works out how to make a grid of rectangular pixels that adapt to the source-plane density and thus\n", + "vary in size. \n", + "\n", + "There is also a `RectangularAdaptImage` mesh which uses the image of the lensed source to adapt\n", + "the rectangular pixel sizes. This often puts even smaller pixels in the brightest regions of the source,\n", + "even if it lies offset or away from the caustic.\n", + "\n", + "There is also a `Delaunay` mesh which uses a Delaunay triangulation to define an irregular grid of source pixels.\n", + "This is described fully in the `delaunay` example including a likelihood function guide.\n", + "\n", + "__Mapping Matrix__\n", + "\n", + "The `mapping_matrix` represents the image-pixel to source-pixel mappings above in a 2D matrix. \n", + "\n", + "It has dimensions `(total_image_pixels, total_source_pixels)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", + " pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,\n", + " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,\n", + " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,\n", + " pixels=mapper.pixels,\n", + " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", + " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", + " sub_fraction=mapper.over_sampler.sub_fraction,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `mapping_matrix` shows of all image-source pixel mappings.\n", + "\n", + "No row of pixels has more than one non-zero entry. It is not possible for two image pixels to map to the same source \n", + "pixel (meaning that there are no correlated pixels in the mapping matrix)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each column of the `mapping_matrix` can therefore be used to show all image-pixels it maps too. \n", + "\n", + "For example, above, we plotted all image-pixels of source-pixel 200 (as well as 202 and 204). We can extract all\n", + "image-pixel indexes of source pixels 200 using the `mapping_matrix` and use them to plot the image of this\n", + "source-pixel (which corresponds to only values of zeros or ones)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "indexes_source_pix_200 = np.nonzero(mapping_matrix[:, 200])\n", + "\n", + "print(indexes_source_pix_200[0])\n", + "\n", + "array_2d = al.Array2D(values=mapping_matrix[:, 200], mask=masked_dataset.mask)\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Blurred Mapping Matrix ($f$)__\n", + "\n", + "Each source-pixel can therefore be thought of as an image (where all entries of this image are zeros and ones). \n", + "\n", + "To incorporate the imaging data's PSF, we simply blur each one of these source-pixel images with the imaging data's \n", + "Point Spread Function (PSF) via 2D convolution.\n", + "\n", + "This operation does not change the dimensions of the mapping matrix, meaning the `blurred_mapping_matrix` also has\n", + "dimensions `(total_image_pixels, total_source_pixels)`. It turns the values of zeros and ones into \n", + "non-integer values which have been blurred by the PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "blurred_mapping_matrix = masked_dataset.psf.convolved_mapping_matrix_from(\n", + " mapping_matrix=mapping_matrix, mask=masked_dataset.mask\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `blurred_mapping_matrix` shows all image-source pixel mappings including PSF blurring.\n", + "\n", + "Note how, unlike for the `mapping_matrix`, every row of image-pixels now has multiple non-zero entries. It is now \n", + "possible for two image pixels to map to the same source pixel, because they become correlated by PSF convolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " blurred_mapping_matrix,\n", + " aspect=(blurred_mapping_matrix.shape[1] / blurred_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each column of the `blurred_mapping_matrix` shows all image-pixels it maps to after PSF blurring. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "indexes_source_pix_200 = np.nonzero(blurred_mapping_matrix[:, 200])\n", + "\n", + "print(indexes_source_pix_200[0])\n", + "\n", + "array_2d = al.Array2D(values=blurred_mapping_matrix[:, 200], mask=masked_dataset.mask)\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) the `blurred_mapping_matrix` is denoted $f_{ij}$\n", + "where $i$ maps over all $I$ source pixels and $j$ maps over all $J$ image pixels. \n", + "\n", + "For example: \n", + "\n", + " - $f_{0, 2} = 0.3$ indicates that image-pixel $2$ maps to source-pixel $0$ with a weight of $0.3$ after PSF convolution.\n", + " - $f_{4, 8} = 0$ indicates that image-pixel $8$ does not map to source-pixel $4$, even after PSF convolution.\n", + "\n", + "The indexing of the `mapping_matrix` is reversed compared to the notation of WD03 (e.g. image pixels\n", + "are the first entry of `mapping_matrix` whereas for $f$ they are the second index)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Mapping between image pixel 0 and source pixel 2 = {mapping_matrix[0, 2]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Data Vector (D)__\n", + "\n", + "To solve for the source pixel fluxes we now pose the problem as a linear inversion.\n", + "\n", + "This requires us to convert the `blurred_mapping_matrix` and our `data` and `noise map` into matrices of certain dimensions. \n", + "\n", + "The `data_vector`, $D$, is the first matrix and it has dimensions `(total_source_pixels,)`.\n", + "\n", + "In WD03 (https://arxiv.org/abs/astro-ph/0302587) and N15 (https://arxiv.org/abs/1412.7436) the data vector \n", + "is give by: \n", + "\n", + " $\\vec{D}_{i} = \\sum_{\\rm j=1}^{J}f_{ij}(d_{j} - b_{j})/\\sigma_{j}^2 \\, \\, .$\n", + "\n", + "Where:\n", + "\n", + " - $d_{\\rm j}$ are the image-pixel data flux values.\n", + " - $b_{\\rm j}$ are the brightness values of the lens light model (therefore $d_{\\rm j} - b_{\\rm j}$ is the lens light\n", + " subtracted image).\n", + " - $\\sigma{\\rm _j}^2$ are the statistical uncertainties of each image-pixel value.\n", + "\n", + "$i$ maps over all $I$ source pixels and $j$ maps over all $J$ image pixels. \n", + "\n", + "NOTE: WD03 assume the data is already lens subtracted thus $b_{j}$ is omitted (e.g. all values are zero)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_vector = al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from(\n", + " blurred_mapping_matrix=blurred_mapping_matrix,\n", + " image=np.array(lens_subtracted_image),\n", + " noise_map=np.array(masked_dataset.noise_map),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$D$ describes which deconvolved source pixels trace to which image-plane pixels. This ensures the source reconstruction\n", + "fully accounts for the PSF when fitting the data.\n", + "\n", + "We can plot $D$ as a column vector:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Matrix (F)__\n", + "\n", + "The `curvature_matrix` $F$ is the second matrix and it has dimensions `(total_source_pixels, total_source_pixels)`.\n", + "\n", + "In WD03 / N15 (https://arxiv.org/abs/astro-ph/0302587) the curvature matrix is a 2D matrix given by:\n", + "\n", + " ${F}_{ik} = \\sum_{\\rm j=1}^{J}f_{ij}f_{kj}/\\sigma_{j}^2 \\, \\, .$\n", + "\n", + "NOTE: this notation implicitly assumes a summation over $K$, where $k$ runs over all source-pixel indexes $K$.\n", + "\n", + "Note how summation over $J$ runs over $f$ twice, such that every entry of $F$ is the sum of the multiplication\n", + "between all values in every two columns of $f$.\n", + "\n", + "For example, $F_{0,1}$ is the sum of every blurred image pixels values in $f$ of source pixel 0 multiplied by\n", + "every blurred image pixel value of source pixel 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix, noise_map=masked_dataset.noise_map\n", + ")\n", + "\n", + "plt.imshow(curvature_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For $F_{ik}$ to be non-zero, this requires that the images of source pixels $i$ and $k$ share at least one\n", + "image-pixel, which we saw above is only possible due to PSF blurring.\n", + "\n", + "For example, we can see a non-zero entry for $F_{100,101}$ and plotting their images\n", + "show overlap." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_pixel_0 = 0\n", + "source_pixel_1 = 1\n", + "\n", + "print(curvature_matrix[source_pixel_0, source_pixel_1])\n", + "\n", + "array_2d = al.Array2D(\n", + " values=blurred_mapping_matrix[:, source_pixel_0], mask=masked_dataset.mask\n", + ")\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")\n", + "\n", + "array_2d = al.Array2D(\n", + " values=blurred_mapping_matrix[:, source_pixel_1], mask=masked_dataset.mask\n", + ")\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The following chi-squared is minimized when we perform the inversion and reconstruct the source:\n", + "\n", + "$\\chi^2 = \\sum_{\\rm j=1}^{J} \\bigg[ \\frac{(\\sum_{\\rm i=1}^{I} s_{i} f_{ij}) + b_{j} - d_{j}}{\\sigma_{j}} \\bigg]$\n", + "\n", + "Where $s$ is the reconstructed source pixel fluxes in all $I$ source pixels.\n", + "\n", + "The solution for $s$ is therefore given by (equation 5 WD03):\n", + "\n", + " $s = F^{-1} D$\n", + " \n", + "We can compute this using NumPy linear algebra:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# Because we are no using regularizartion (see below) it is common for the curvature matrix to be singular and lead\n", + "# to a LinAlgException. The loop below mitigates this -- you can ignore it as it is not important for understanding\n", + "# the PyAutoLens likelihood function.\n", + "\n", + "for i in range(curvature_matrix.shape[0]):\n", + " curvature_matrix[i, i] += 1e-8\n", + "\n", + "reconstruction = np.linalg.solve(curvature_matrix, data_vector)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot this source reconstruction -- it looks like a mess.\n", + "\n", + "The source pixels have noisy and unsmooth values, and it is hard to make out if a source is even being reconstructed. \n", + "\n", + "In fact, the linear inversion is (over-)fitting noise in the image data, meaning this system of equations is \n", + "ill-posed. We need to apply some form of smoothing on the source reconstruction to avoid over fitting noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Regularization Matrix (H)__\n", + "\n", + "Regularization adds a linear regularization term $G_{\\rm L}$ to the $\\chi^2$ we solve for giving us a new merit \n", + "function $G$ (equation 11 WD03):\n", + "\n", + " $G = \\chi^2 + \\lambda \\, G_{\\rm L}$\n", + " \n", + "where $\\lambda$ is the `regularization_coefficient` which describes the magnitude of smoothness that is applied. A \n", + "higher $\\lambda$ will regularize the source more, leading to a smoother source reconstruction.\n", + " \n", + "Different forms for $G_{\\rm L}$ can be defined which regularize the source reconstruction in different ways. The \n", + "`Constant` regularization scheme used in this example applies gradient regularization (equation 14 WD03):\n", + "\n", + " $G_{\\rm L} = \\sum_{\\rm i}^{I} \\sum_{\\rm n=1}^{N} [s_{i} - s_{i, v}]$\n", + "\n", + "This regularization scheme is easier to express in words -- the summation goes to each source pixel,\n", + "determines all source pixels with which it shares a direct vertex (e.g. its neighbors) and penalizes solutions \n", + "where the difference in reconstructed flux of these two neighboring source pixels is large.\n", + "\n", + "The summation does this for all pixels, thus it favours solutions where neighboring source\n", + "pixels reconstruct similar values to one another (e.g. it favours a smooth source reconstruction).\n", + "\n", + "We now define the `regularization matrix`, $H$, which allows us to include this smoothing when we solve for $s$. $H$\n", + "has dimensions `(total_source_pixels, total_source_pixels)`.\n", + "\n", + "This relates to $G_{\\rm L}$ as (equation 13 WD03):\n", + "\n", + " $H_{ik} = \\frac{1}{2} \\frac{\\partial G_{\\rm L}}{\\partial s_{i} \\partial s_{k}}$\n", + "\n", + "$H$ has the `regularization_coefficient` $\\lambda$ folded into it such $\\lambda$'s control on the degree of smoothing\n", + "is accounted for." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", + " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", + " neighbors=mapper.neighbors,\n", + " neighbors_sizes=mapper.neighbors.sizes,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot the regularization matrix and note that:\n", + "\n", + " - non-zero entries indicate that two source-pixels are neighbors and therefore are regularized with one another.\n", + " \n", + " - Zeros indicate the two source pixels do not neighbor one another.\n", + " \n", + "The majority of entries are zero, because the majority of source pixels are not neighbors with one another." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(regularization_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__F + Lamdba H__\n", + "\n", + "$H$ enters the linear algebra system we solve for as follows (WD03 equation (12)):\n", + "\n", + " $s = [F + H]^{-1} D$" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Reconstruction (s)__\n", + "\n", + "We can now solve the linear system above using NumPy linear algebra. \n", + "\n", + "Note that the for loop used above to prevent a LinAlgException is no longer required." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting this source reconstruction we can see that regularization has lead us to reconstruct a smoother source,\n", + "which actually looks like a galaxy! This also implies we are not over-fitting the noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image Reconstruction__\n", + "\n", + "Using the reconstructed source pixel fluxes we can map the source reconstruction back to the image plane (via\n", + "the `blurred mapping_matrix`) and produce a reconstruction of the image data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_operated_data = (\n", + " al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from(\n", + " mapping_matrix=blurred_mapping_matrix, reconstruction=reconstruction\n", + " )\n", + ")\n", + "\n", + "mapped_reconstructed_operated_data = al.Array2D(\n", + " values=mapped_reconstructed_operated_data, mask=mask\n", + ")\n", + "\n", + "aplt.plot_array(array=mapped_reconstructed_operated_data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our lens model and source reconstruction. \n", + "\n", + "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", + "\n", + "The likelihood function for lens modeling consists of five terms:\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + s^{T} H s + \\mathrm{ln} \\, \\left[ \\mathrm{det} (F + H) \\right] - { \\mathrm{ln}} \\, \\left[ \\mathrm{det} (H) \\right] + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "This expression was first derived by Suyu 2006 (https://arxiv.org/abs/astro-ph/0601493) and is given by equation (19).\n", + "It was derived into **PyAutoLens** notation in Dye 2008 (https://arxiv.org/abs/0804.4002) equation (5).\n", + "\n", + "We now explain what each of these terms mean.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", + "\n", + " - `model_data` = `mapped_reconstructed_operated_data` + `lens_light_convolved_image`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", + "\n", + "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", + "for, corresponding to a fit with a lower likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image = convolved_image_2d + mapped_reconstructed_operated_data\n", + "\n", + "residual_map = masked_dataset.data - model_image\n", + "normalized_residual_map = residual_map / masked_dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(chi_squared)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", + "\n", + "aplt.plot_array(array=chi_squared_map, title=\"\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Regularization Term__\n", + "\n", + "The second term, $s^{T} H s$, corresponds to the $\\lambda $G_{\\rm L}$ regularization term we added to our merit \n", + "function above.\n", + "\n", + "This is the term which sums up the difference in flux of all reconstructed source pixels, and reduces the likelihood of \n", + "solutions where there are large differences in flux (e.g. the source is less smooth and more likely to be \n", + "overfitting noise).\n", + "\n", + "We compute it below via matrix multiplication, noting that the `regularization_coefficient`, $\\lambda$, is built into \n", + "the `regularization_matrix` already." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "regularization_term = np.matmul(\n", + " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", + ")\n", + "\n", + "print(regularization_term)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Complexity Terms__\n", + "\n", + "Up to this point, it is unclear why we chose a value of `regularization_coefficient=1.0`. \n", + "\n", + "We cannot rely on the `chi_squared` and `regularization_term` above to optimally choose its value, because increasing \n", + "the `regularization_coefficient` smooths the solution more and therefore:\n", + " \n", + " - Decreases `chi_squared` by fitting the data worse, producing a lower `log_likelihood`.\n", + " \n", + " - Increases the `regularization_term` by penalizing the differences between source pixel fluxes more, again reducing\n", + " the inferred `log_likelihood`.\n", + "\n", + "If we set the regularization coefficient based purely on these two terms, we would set a value of 0.0 and be back where\n", + "we started over-fitting noise!\n", + "\n", + "The terms $\\left[ \\mathrm{det} (F + H) \\right]$ and $ - { \\mathrm{ln}} \\, \\left[ \\mathrm{det} (H) \\right]$ address \n", + "this problem. \n", + "\n", + "They quantify how complex the source reconstruction is, and penalize solutions where *it is more complex*. Reducing \n", + "the `regularization_coefficient` makes the source reconstruction more complex (because a source that is \n", + "smoothed less uses more flexibility to fit the data better).\n", + "\n", + "These two terms therefore counteract the `chi_squared` and `regularization_term`, so as to attribute a higher\n", + "`log_likelihood` to solutions which fit the data with a more smoothed and less complex source (e.g. one with a higher \n", + "`regularization_coefficient`).\n", + "\n", + "In **HowToLens** -> `chapter 4` -> `tutorial_4_bayesian_regularization` we expand on this further and give a more\n", + "detailed description of how these different terms impact the `log_likelihood_function`. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", + "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]\n", + "\n", + "print(log_curvature_reg_matrix_term)\n", + "print(log_regularization_matrix_term)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", + "\n", + "The final term ins the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", + "of the log of every noise-map value squared. \n", + "\n", + "Given the `noise_map` is fixed, this term does not change during the lens modeling process and has no impact on the \n", + "model we infer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "We can now, finally, compute the `log_likelihood` of the lens model, by combining the five terms computed above using\n", + "the likelihood function defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_evidence = float(\n", + " -0.5\n", + " * (\n", + " chi_squared\n", + " + regularization_term\n", + " + log_curvature_reg_matrix_term\n", + " - log_regularization_matrix_term\n", + " + noise_normalization\n", + " )\n", + ")\n", + "\n", + "print(log_evidence)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(\n", + " dataset=masked_dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_border_relocator=True),\n", + ")\n", + "fit_log_evidence = fit.log_evidence\n", + "print(fit_log_evidence)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `Nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the **PyAutoLens** likelihood function, which uses a pixelization, \n", + "regularization scheme and inversion to reconstruct the source galaxy.\n", + "\n", + "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", + "are described above but worth refreshing ourselves on:\n", + "\n", + " - **Sub-gridding**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", + " ray-traced to the source-plane and paired fractionally with each source pixel.\n", + " \n", + " - **Source-plane Interpolation**: Using a RectangularUniform mesh with bilinear interpolation\n", + " to pair each image (sub-)pixel to multiple source-plane pixels with interpolation weights.\n", + " \n", + " - **Source Morphology Pixelization Adaption**: Adapting the pixelization such that is congregates source pixels around\n", + " the source's brightest regions, as opposed to the magnification-based pixelization used here.\n", + " \n", + " - **Luminosity Weighted Regularization**: Using an adaptive regularization coefficient which adapts the level of \n", + " regularization applied to the source based on its luminosity." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/pixelization/modeling.ipynb b/notebooks/imaging/features/pixelization/modeling.ipynb index 25922fefe..10e91525d 100644 --- a/notebooks/imaging/features/pixelization/modeling.ipynb +++ b/notebooks/imaging/features/pixelization/modeling.ipynb @@ -1,759 +1,796 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Pixelization Modeling\n", - "===============================\n", - "\n", - "A pixelization reconstructs the source's light using a pixel-grid, which is regularized using a prior that forces\n", - "the solution to have a degree of smoothness.\n", - "\n", - "This script fits a source galaxy model which uses a pixelization to reconstruct the source's light.\n", - "\n", - "A Rectangular mesh and constant regularization scheme are used, which are the simplest forms of mesh and regularization\n", - "with provide computationally fast and accurate solutions.\n", - "\n", - "For simplicity, the lens galaxy\u2019s light is omitted from both the simulated data and the model. Including the lens\n", - "galaxy\u2019s light is straightforward and can be done in exactly the same framework.\n", - "\n", - "You may wish to first read the pixelization/fit.py example, which demonstrates how a pixelized source reconstruction\n", - "is applied to a single dataset.\n", - "\n", - "Pixelizations are covered in detail in chapter 4 of the **HowToLens** lectures.\n", - "\n", - "__GPU Run Times__\n", - "\n", - "On consumer laptop GPUs, the run times of pixelization modeling can be a bit prohibitive, for example this example\n", - "takes around 30 minutes to run on a laptop GPU. However, HPC GPUs or very modern laptop GPUs with more\n", - "VRAM can run this example in under 10 minutes. Therefore don't be put off if this example feels slow as\n", - "speed up is possible with better hardware.\n", - "\n", - "__CPU Users__\n", - "\n", - "On CPU, JAX pixelization calculations are not accelerated and are therefore relatively slow. CPU users\n", - "should therefore read the `pixelization/cpu_fast_modeling` example after this one, and copy its set up,\n", - "to get the fastest run times for pixelization modeling on CPU.\n", - "\n", - "On GPU, JAX run times are fast and this example is all that is required, however we do set a few settings\n", - "to make things run extra fast on GPU.\n", - "\n", - "__Contents__\n", - "\n", - "- **GPU Run Times:** On consumer laptop GPUs, the run times of pixelization modeling can be a bit prohibitive, for.\n", - "- **CPU Users:** On CPU, JAX pixelization calculations are not accelerated and are therefore relatively slow.\n", - "- **Advantages & Disadvantages:** Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", - "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Position Likelihood:** We add a penalty term ot the likelihood function, which penalizes models where the brightest.\n", - "- **Brief Description:** Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH`.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Mask Extra Galaxies:** There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "- **Chaining:** Modeling with a pixelization can be made more efficient, robust, and automated using the non-linear.\n", - "- **HowToGalaxy:** A full description of how pixelizations work\u2014which relies heavily on linear algebra, Bayesian.\n", - "\n", - "__Advantages__\n", - "\n", - "Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies. Such structures\n", - "cannot be well approximated by analytic light profiles such as a S\u00e9rsic profile, or even combinations of multiple\n", - "S\u00e9rsic components. pixelizations are therefore required to accurately reconstruct this irregular source-plane light.\n", - "\n", - "Even alternative basis-function approaches, such as shapelets or multi-Gaussian expansions, struggle to accurately\n", - "reconstruct sources with highly complex morphologies or multiple distinct source galaxies.\n", - "\n", - "Pixelized source models are also essential for robustly constraining detailed components of the lens mass\n", - "distribution (e.g. the mass density slope or the presence of dark matter substructure). By fitting all of the lensed\n", - "source light, they reduce degeneracies between the source and lens mass model.\n", - "\n", - "Finally, many science applications aim to study the highly magnified source galaxy itself, in order to learn about\n", - "distant and intrinsically faint galaxies. pixelizations reconstruct the unlensed source emission, enabling detailed\n", - "studies of the source-plane structure.\n", - "\n", - "__Disadvantages__\n", - "\n", - "Pixelized source reconstructions are computationally more expensive than analytic source models. For high-resolution\n", - "imaging data (e.g. Hubble Space Telescope observations), it is common for lens models using pixelizations to require\n", - "hours or even days to fit.\n", - "\n", - "Lens modeling with pixelizations is also conceptually more complex. There are additional failure modes, such as\n", - "solutions where the source is reconstructed in a highly demagnified configuration due to an unphysical lens mass\n", - "model (e.g. too little or too much mass). These issues are discussed in detail later in the workspace.\n", - "\n", - "As a result, learning to successfully fit lens models with pixelizations typically requires more time and experience\n", - "than the simpler modeling approaches introduced elsewhere in the workspace.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algabra solver which allows for positive and negative\n", - "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical. For a pixelizaiton, this often produces negative source pixels which over-fit\n", - "the data, producing unphysical solutions.\n", - "\n", - "All pixelized source reconstructions use a positive-only solver, meaning that every source-pixel is only allowed\n", - "to reconstruct positive flux values. This ensures that the source reconstruction is physical and that we don't\n", - "reconstruct negative flux values that don't exist in the real source galaxy (a common systematic solution in lens\n", - "analysis).\n", - "\n", - "It may be surprising to hear that this is a feature worth pointing out, but it turns out setting up the linear algebra\n", - "to enforce positive reconstructions is difficult to make efficient. A lot of development time went into making this\n", - "possible, where a bespoke fast non-negative linear solver was developed to achieve this.\n", - "\n", - "Other methods in the literature often do not use a positive only solver, and therefore suffer from these\n", - "unphysical solutions, which can degrade the results of lens model in general.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's surface-brightness is reconstructed using a `RectangularAdaptDensity` mesh\n", - " and `Constant` regularization scheme.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple__no_lens_light` via .fits files" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "if not (dataset_path / \"positions.json\").exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "A pixelization uses a separate grid for ray tracing, with its own over sampling scheme, which below we set to a \n", - "uniform grid of values of 2. \n", - "\n", - "The pixelization only reconstructs the source galaxy, therefore the adaptive over sampling used for the lens galaxy's \n", - "light in other examples is not applied to the pixelization. \n", - "\n", - "This example does not model lens light, for examples which combine lens light and a pixelization both over sampling \n", - "schemes should be used, with the lens light adaptive and the pixelization uniform.\n", - "\n", - "Note that the over sampling is input into the `over_sample_size_pixelization` because we are using a `Pixelization`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_over_sampling(\n", - " over_sample_size_pixelization=4,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", - "set below to 28 x 28. \n", - "\n", - "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", - "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", - "mesh, the same number of pixels must be used in the y and x directions.\n", - "\n", - "__Edge Zeroing__\n", - "\n", - "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", - "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", - "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", - "\n", - "For a rectangular mesh, the source code computes edge pixels internally using the known\n", - "pixels at the edge of the mesh. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 8\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In this \n", - "example fits a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source-galaxy's light uses a 20 x 20 `RectangularAdaptDensity` mesh [0 parameters].\n", - "\n", - " - This pixelization is regularized using a `Constant` scheme which smooths every source pixel equally [1 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=6. \n", - "\n", - "It is worth noting the pixelization fits the source using significantly fewer parameters (1 parameter for \n", - "regularization) than fitting the source using light profiles or an MGE (4+ parameters). \n", - "\n", - "The lens model therefore includes a mesh and regularization scheme, which are used together to create the \n", - "pixelization. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.PowerLaw)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", - "regularization = af.Model(al.reg.Constant)\n", - "\n", - "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the source galaxy's has a mesh and regularization scheme, which are combined into a pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\"),\n", - " name=\"pixelization\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=10, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=50000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Position Likelihood__\n", - "\n", - "We add a penalty term ot the likelihood function, which penalizes models where the brightest multiple images of\n", - "the lensed source galaxy do not trace close to one another in the source plane. This removes \"demagnified source\n", - "solutions\" from the source pixelization, which one is likely to infer without this penalty.\n", - "\n", - "A comprehensive description of why we do this is given at the following readthedocs page. I strongly recommend you \n", - "read this page in full if you are not familiar with the positions likelihood penalty and demagnified source \n", - "reconstructions:\n", - "\n", - " https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", - "\n", - "__Brief Description__\n", - "\n", - "Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH` object, which\n", - "includes the positions we loaded above, alongside a `threshold`.\n", - "\n", - "This is because `Inversion`'s suffer a bias whereby they fit unphysical lens models where the source galaxy is \n", - "reconstructed as a demagnified version of the lensed source. \n", - "\n", - "To prevent these solutions biasing the model-fit we specify a `position_threshold` of 0.5\", which requires that a \n", - "mass model traces the four (y,x) coordinates specified by our positions (that correspond to the brightest regions of the \n", - "lensed source) within 0.5\" of one another in the source-plane. If this criteria is not met, a large penalty term is\n", - "added to likelihood that massively reduces the overall likelihood. This penalty is larger if the ``positions``\n", - "trace further from one another.\n", - "\n", - "This ensures the unphysical solutions that bias a pixelization have a lower likelihood that the physical solutions\n", - "we desire. Furthermore, the penalty term reduces as the image-plane multiple image positions trace closer in the \n", - "source-plane, ensuring Nautilus converges towards an accurate mass model. It does this very fast, as \n", - "ray-tracing just a few multiple image positions is computationally cheap. \n", - "\n", - "The threshold of 0.3\" is large. For an accurate lens model we would anticipate the positions trace within < 0.01\" of\n", - "one another. The high threshold ensures only the initial mass models at the start of the fit are penalized.\n", - "\n", - "Position thresholding is described in more detail in the \n", - "script `autolens_workspace/*/guides/modeling/customize`\n", - "\n", - "The arc-second positions of the multiply imaged lensed source galaxy were drawn onto the\n", - "image via the GUI described in the file `autolens_workspace/*/imaging/data_preparation/gui/positions.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n", - "\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data. \n", - "\n", - "The `positions_likelihood_list` is passed to the analysis, which applies the likelihood penalty described above\n", - "for everyone lens mass model.\n", - "\n", - "Pixelized source reconstructions have settings which determine their behavour and run time. Below, we input\n", - "a setting which performs a subset of calculations using mixed precision, which can speed up run times on consumer\n", - "laptop GPUs significantly. \n", - "\n", - "If you are using a high end GPU which can handle the full precision calculations, you can set this to `False` \n", - "for more accurate results without slow down. For modeling which is close to science grade, I recommend setting \n", - "this to `False`, to ensure full accuracy, but for quick model-fits to test out the API and understand \n", - "how pixelizations work, setting this to `True` is a good option." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " settings=al.Settings(use_mixed_precision=True),\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", - "required by a model.\n", - "\n", - "Pixelizations use a lot more VRAM than light profile-only models, with the amount required depending on the size of\n", - "dataset and the number of source pixels in the pixelization's mesh. For 784 source pixels, around 0.05 GB per batched\n", - "likelihood of VRAM is used. \n", - "\n", - "This is why the `batch_size` above is 20, lower than other examples, because reducing the batch size ensures a more \n", - "modest amount of VRAM is used. If you have a GPU with more VRAM, increasing the batch size will lead to faster run times.\n", - "\n", - "Given VRAM use is an important consideration, we print out the estimated VRAM required for this\n", - "model-fit and advise you do this for your own pixelization model-fits.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The run time of a pixelization are fast provided that the GPU VRAM exceeds the amount of memory required to perform\n", - "a likelihood evaluation.\n", - "\n", - "Assuming the use of a 20 x 20 mesh grid above means this is the case, the run times of this model-fit on a GPU\n", - "should take under 10 minutes. If VRAM is exceeded, the run time will be significantly longer (3+ hours). CPU run\n", - "times are also of order hours, but can be sped up using the `numba` library (see the `pixelization/cpu` example).\n", - "\n", - "The run times of pixelizations slow down as the data becomes higher resolution. In this example, data with a pixel\n", - "scale of 0.1\" gives of order 10 minute run times (when VRAM is under control), for a pixel scale of 0.05\" this\n", - "becomes around 30 minutes, and an hour for 0.03\".\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this\n", - "does not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix this):\n", - "\n", - "This confirms that the source galaxy's has a mesh and regularization scheme, which are combined into a pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "The end of this example provides a detailed description of all result options for a pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The example `pixelization/fit` provides a full description of the different calculations that can be performed\n", - "with the result of a pixelization model-fit.\n", - "\n", - "__Source Science (Magnification, Flux and More)__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", - "size of the source.\n", - "\n", - "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", - "including examples using a pixelized source reconstruction. \n", - "\n", - "If you want to study the source galaxy after modeling has reconstructed its unlensed, then check out this example.\n", - "\n", - "__Mask Extra Galaxies__\n", - "\n", - "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source.\n", - "\n", - "If their emission is significant, and close enough to the lens and source, we may simply remove it from the data\n", - "to ensure it does not impact the model-fit. A standard masking approach would be to remove the image pixels containing\n", - "the emission of these galaxies altogether. This is analogous to what the circular masks used throughout the examples\n", - "does.\n", - "\n", - "For fits using a pixelization, masking regions of the image in a way that removes their image pixels entirely from\n", - "the fit. This can produce discontinuities in the pixelixation used to reconstruct the source and produce unexpected\n", - "systematics and unsatisfactory results. In this case, applying the mask in a way where the image pixels are not\n", - "removed from the fit, but their data and noise-map values are scaled such that they contribute negligibly to the fit,\n", - "is a better approach.\n", - "\n", - "We illustrate the API for doing this below, using the `extra_galaxies` dataset which has extra galaxies whose emission\n", - "needs to be removed via scaling in this way. We apply the scaling and show the subplot imaging where the extra\n", - "galaxies mask has scaled the data values to zeros, increasing the noise-map values to large values and in turn made\n", - "the signal to noise of its pixels effectively zero." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_extra_galaxies = al.Mask2D.from_fits(\n", - " file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"),\n", - " pixel_scales=0.1,\n", - " invert=True, # Note that we invert the mask here as `True` means a pixel is scaled.\n", - ")\n", - "\n", - "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native, pixel_scales=0.1, centre=(0.0, 0.0), radius=6.0\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We do not explictly fit this data, for the sake of brevity, however if your data has these nearby galaxies you should\n", - "apply the mask as above before fitting the data.\n", - "\n", - "__Wrap Up__\n", - "\n", - "Pixelizations are the most complex but also the most powerful way to model a galaxy\u2019s light.\n", - "\n", - "Whether you need to use them depends on the science you are doing. If you are only interested in measuring simple\n", - "global quantities (for example, total flux, size, or axis ratio), analytic light profiles such as a S\u00e9rsic, MGE, or\n", - "shapelets are often sufficient. For low-resolution data, pixelizations are also unnecessary, as the complex\n", - "structure of the galaxy is not resolved.\n", - "\n", - "However, modeling galaxies with complex, irregular, or highly structured light distributions requires this level of\n", - "flexibility. Furthermore, if you are interested in studying the detailed morphology of a galaxy itself, there is no\n", - "better approach than using a pixelization.\n", - "\n", - "__Chaining__\n", - "\n", - "Modeling with a pixelization can be made more efficient, robust, and automated using the non-linear chaining feature\n", - "to compose a pipeline that begins by fitting a simpler model using parametric light profiles.\n", - "\n", - "More information on chaining is provided in the\n", - "`autogalaxy_workspace/notebooks/guides/modeling/chaining` folder and in chapter 3 of the **HowToGalaxy** lectures.\n", - "\n", - "__HowToGalaxy__\n", - "\n", - "A full description of how pixelizations work\u2014which relies heavily on linear algebra, Bayesian statistics, and\n", - "2D geometry\u2014is provided in chapter 4 of the **HowToGalaxy** lectures.\n", - "\n", - "__Future Ideas / Contributions__\n", - "\n", - "Here are a list of things I would like to add to this tutorial but haven't found the time. If you are interested\n", - "in having a go at adding them contact me on SLACK! :)\n", - "\n", - "- More diagnostic quantities for reconstructed galaxy light.\n", - "- Gradient calculations of the reconstructed light distribution.\n", - "- Quantifying spatial variations in galaxy structure across the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Pixelization Modeling\n", + "===============================\n", + "\n", + "A pixelization reconstructs the source's light using a pixel-grid, which is regularized using a prior that forces\n", + "the solution to have a degree of smoothness.\n", + "\n", + "This script fits a source galaxy model which uses a pixelization to reconstruct the source's light.\n", + "\n", + "A Rectangular mesh and constant regularization scheme are used, which are the simplest forms of mesh and regularization\n", + "with provide computationally fast and accurate solutions.\n", + "\n", + "For simplicity, the lens galaxy\u2019s light is omitted from both the simulated data and the model. Including the lens\n", + "galaxy\u2019s light is straightforward and can be done in exactly the same framework.\n", + "\n", + "You may wish to first read the pixelization/fit.py example, which demonstrates how a pixelized source reconstruction\n", + "is applied to a single dataset.\n", + "\n", + "Pixelizations are covered in detail in chapter 4 of the **HowToLens** lectures.\n", + "\n", + "__GPU Run Times__\n", + "\n", + "On consumer laptop GPUs, the run times of pixelization modeling can be a bit prohibitive, for example this example\n", + "takes around 30 minutes to run on a laptop GPU. However, HPC GPUs or very modern laptop GPUs with more\n", + "VRAM can run this example in under 10 minutes. Therefore don't be put off if this example feels slow as\n", + "speed up is possible with better hardware.\n", + "\n", + "__CPU Users__\n", + "\n", + "On CPU, JAX pixelization calculations are not accelerated and are therefore relatively slow. CPU users\n", + "should therefore read the `pixelization/cpu_fast_modeling` example after this one, and copy its set up,\n", + "to get the fastest run times for pixelization modeling on CPU.\n", + "\n", + "On GPU, JAX run times are fast and this example is all that is required, however we do set a few settings\n", + "to make things run extra fast on GPU.\n", + "\n", + "__Contents__\n", + "\n", + "- **GPU Run Times:** On consumer laptop GPUs, the run times of pixelization modeling can be a bit prohibitive, for.\n", + "- **CPU Users:** On CPU, JAX pixelization calculations are not accelerated and are therefore relatively slow.\n", + "- **Advantages & Disadvantages:** Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", + "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Position Likelihood:** We add a penalty term ot the likelihood function, which penalizes models where the brightest.\n", + "- **Brief Description:** Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH`.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Mask Extra Galaxies:** There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "- **Chaining:** Modeling with a pixelization can be made more efficient, robust, and automated using the non-linear.\n", + "- **HowToGalaxy:** A full description of how pixelizations work\u2014which relies heavily on linear algebra, Bayesian.\n", + "\n", + "__Advantages__\n", + "\n", + "Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies. Such structures\n", + "cannot be well approximated by analytic light profiles such as a S\u00e9rsic profile, or even combinations of multiple\n", + "S\u00e9rsic components. pixelizations are therefore required to accurately reconstruct this irregular source-plane light.\n", + "\n", + "Even alternative basis-function approaches, such as shapelets or multi-Gaussian expansions, struggle to accurately\n", + "reconstruct sources with highly complex morphologies or multiple distinct source galaxies.\n", + "\n", + "Pixelized source models are also essential for robustly constraining detailed components of the lens mass\n", + "distribution (e.g. the mass density slope or the presence of dark matter substructure). By fitting all of the lensed\n", + "source light, they reduce degeneracies between the source and lens mass model.\n", + "\n", + "Finally, many science applications aim to study the highly magnified source galaxy itself, in order to learn about\n", + "distant and intrinsically faint galaxies. pixelizations reconstruct the unlensed source emission, enabling detailed\n", + "studies of the source-plane structure.\n", + "\n", + "__Disadvantages__\n", + "\n", + "Pixelized source reconstructions are computationally more expensive than analytic source models. For high-resolution\n", + "imaging data (e.g. Hubble Space Telescope observations), it is common for lens models using pixelizations to require\n", + "hours or even days to fit.\n", + "\n", + "Lens modeling with pixelizations is also conceptually more complex. There are additional failure modes, such as\n", + "solutions where the source is reconstructed in a highly demagnified configuration due to an unphysical lens mass\n", + "model (e.g. too little or too much mass). These issues are discussed in detail later in the workspace.\n", + "\n", + "As a result, learning to successfully fit lens models with pixelizations typically requires more time and experience\n", + "than the simpler modeling approaches introduced elsewhere in the workspace.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algabra solver which allows for positive and negative\n", + "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical. For a pixelizaiton, this often produces negative source pixels which over-fit\n", + "the data, producing unphysical solutions.\n", + "\n", + "All pixelized source reconstructions use a positive-only solver, meaning that every source-pixel is only allowed\n", + "to reconstruct positive flux values. This ensures that the source reconstruction is physical and that we don't\n", + "reconstruct negative flux values that don't exist in the real source galaxy (a common systematic solution in lens\n", + "analysis).\n", + "\n", + "It may be surprising to hear that this is a feature worth pointing out, but it turns out setting up the linear algebra\n", + "to enforce positive reconstructions is difficult to make efficient. A lot of development time went into making this\n", + "possible, where a bespoke fast non-negative linear solver was developed to achieve this.\n", + "\n", + "Other methods in the literature often do not use a positive only solver, and therefore suffer from these\n", + "unphysical solutions, which can degrade the results of lens model in general.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's surface-brightness is reconstructed using a `RectangularAdaptDensity` mesh\n", + " and `Constant` regularization scheme.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple__no_lens_light` via .fits files" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "if not (dataset_path / \"positions.json\").exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "A pixelization uses a separate grid for ray tracing, with its own over sampling scheme, which below we set to a \n", + "uniform grid of values of 2. \n", + "\n", + "The pixelization only reconstructs the source galaxy, therefore the adaptive over sampling used for the lens galaxy's \n", + "light in other examples is not applied to the pixelization. \n", + "\n", + "This example does not model lens light, for examples which combine lens light and a pixelization both over sampling \n", + "schemes should be used, with the lens light adaptive and the pixelization uniform.\n", + "\n", + "Note that the over sampling is input into the `over_sample_size_pixelization` because we are using a `Pixelization`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_over_sampling(\n", + " over_sample_size_pixelization=4,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", + "set below to 28 x 28. \n", + "\n", + "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", + "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", + "mesh, the same number of pixels must be used in the y and x directions.\n", + "\n", + "__Edge Zeroing__\n", + "\n", + "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", + "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", + "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", + "\n", + "For a rectangular mesh, the source code computes edge pixels internally using the known\n", + "pixels at the edge of the mesh. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 8\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In this \n", + "example fits a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source-galaxy's light uses a 20 x 20 `RectangularAdaptDensity` mesh [0 parameters].\n", + "\n", + " - This pixelization is regularized using a `Constant` scheme which smooths every source pixel equally [1 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=6. \n", + "\n", + "It is worth noting the pixelization fits the source using significantly fewer parameters (1 parameter for \n", + "regularization) than fitting the source using light profiles or an MGE (4+ parameters). \n", + "\n", + "The lens model therefore includes a mesh and regularization scheme, which are used together to create the \n", + "pixelization. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.PowerLaw)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", + "regularization = af.Model(al.reg.Constant)\n", + "\n", + "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the source galaxy's has a mesh and regularization scheme, which are combined into a pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\"),\n", + " name=\"pixelization\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=10, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=50000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Position Likelihood__\n", + "\n", + "We add a penalty term ot the likelihood function, which penalizes models where the brightest multiple images of\n", + "the lensed source galaxy do not trace close to one another in the source plane. This removes \"demagnified source\n", + "solutions\" from the source pixelization, which one is likely to infer without this penalty.\n", + "\n", + "A comprehensive description of why we do this is given at the following readthedocs page. I strongly recommend you \n", + "read this page in full if you are not familiar with the positions likelihood penalty and demagnified source \n", + "reconstructions:\n", + "\n", + " https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", + "\n", + "__Brief Description__\n", + "\n", + "Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH` object, which\n", + "includes the positions we loaded above, alongside a `threshold`.\n", + "\n", + "This is because `Inversion`'s suffer a bias whereby they fit unphysical lens models where the source galaxy is \n", + "reconstructed as a demagnified version of the lensed source. \n", + "\n", + "To prevent these solutions biasing the model-fit we specify a `position_threshold` of 0.5\", which requires that a \n", + "mass model traces the four (y,x) coordinates specified by our positions (that correspond to the brightest regions of the \n", + "lensed source) within 0.5\" of one another in the source-plane. If this criteria is not met, a large penalty term is\n", + "added to likelihood that massively reduces the overall likelihood. This penalty is larger if the ``positions``\n", + "trace further from one another.\n", + "\n", + "This ensures the unphysical solutions that bias a pixelization have a lower likelihood that the physical solutions\n", + "we desire. Furthermore, the penalty term reduces as the image-plane multiple image positions trace closer in the \n", + "source-plane, ensuring Nautilus converges towards an accurate mass model. It does this very fast, as \n", + "ray-tracing just a few multiple image positions is computationally cheap. \n", + "\n", + "The threshold of 0.3\" is large. For an accurate lens model we would anticipate the positions trace within < 0.01\" of\n", + "one another. The high threshold ensures only the initial mass models at the start of the fit are penalized.\n", + "\n", + "Position thresholding is described in more detail in the \n", + "script `autolens_workspace/*/guides/modeling/customize`\n", + "\n", + "The arc-second positions of the multiply imaged lensed source galaxy were drawn onto the\n", + "image via the GUI described in the file `autolens_workspace/*/imaging/data_preparation/gui/positions.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n", + "\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data. \n", + "\n", + "The `positions_likelihood_list` is passed to the analysis, which applies the likelihood penalty described above\n", + "for everyone lens mass model.\n", + "\n", + "Pixelized source reconstructions have settings which determine their behavour and run time. Below, we input\n", + "a setting which performs a subset of calculations using mixed precision, which can speed up run times on consumer\n", + "laptop GPUs significantly. \n", + "\n", + "If you are using a high end GPU which can handle the full precision calculations, you can set this to `False` \n", + "for more accurate results without slow down. For modeling which is close to science grade, I recommend setting \n", + "this to `False`, to ensure full accuracy, but for quick model-fits to test out the API and understand \n", + "how pixelizations work, setting this to `True` is a good option." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " settings=al.Settings(use_mixed_precision=True),\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", + "required by a model.\n", + "\n", + "Pixelizations use a lot more VRAM than light profile-only models, with the amount required depending on the size of\n", + "dataset and the number of source pixels in the pixelization's mesh. For 784 source pixels, around 0.05 GB per batched\n", + "likelihood of VRAM is used. \n", + "\n", + "This is why the `batch_size` above is 20, lower than other examples, because reducing the batch size ensures a more \n", + "modest amount of VRAM is used. If you have a GPU with more VRAM, increasing the batch size will lead to faster run times.\n", + "\n", + "Given VRAM use is an important consideration, we print out the estimated VRAM required for this\n", + "model-fit and advise you do this for your own pixelization model-fits.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The run time of a pixelization are fast provided that the GPU VRAM exceeds the amount of memory required to perform\n", + "a likelihood evaluation.\n", + "\n", + "Assuming the use of a 20 x 20 mesh grid above means this is the case, the run times of this model-fit on a GPU\n", + "should take under 10 minutes. If VRAM is exceeded, the run time will be significantly longer (3+ hours). CPU run\n", + "times are also of order hours, but can be sped up using the `numba` library (see the `pixelization/cpu` example).\n", + "\n", + "The run times of pixelizations slow down as the data becomes higher resolution. In this example, data with a pixel\n", + "scale of 0.1\" gives of order 10 minute run times (when VRAM is under control), for a pixel scale of 0.05\" this\n", + "becomes around 30 minutes, and an hour for 0.03\".\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this\n", + "does not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix this):\n", + "\n", + "This confirms that the source galaxy's has a mesh and regularization scheme, which are combined into a pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "The end of this example provides a detailed description of all result options for a pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The example `pixelization/fit` provides a full description of the different calculations that can be performed\n", + "with the result of a pixelization model-fit.\n", + "\n", + "__Source Science (Magnification, Flux and More)__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", + "size of the source.\n", + "\n", + "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", + "including examples using a pixelized source reconstruction. \n", + "\n", + "If you want to study the source galaxy after modeling has reconstructed its unlensed, then check out this example.\n", + "\n", + "__Mask Extra Galaxies__\n", + "\n", + "There may be extra galaxies nearby the lens and source galaxies, whose emission blends with the lens and source.\n", + "\n", + "If their emission is significant, and close enough to the lens and source, we may simply remove it from the data\n", + "to ensure it does not impact the model-fit. A standard masking approach would be to remove the image pixels containing\n", + "the emission of these galaxies altogether. This is analogous to what the circular masks used throughout the examples\n", + "does.\n", + "\n", + "For fits using a pixelization, masking regions of the image in a way that removes their image pixels entirely from\n", + "the fit. This can produce discontinuities in the pixelixation used to reconstruct the source and produce unexpected\n", + "systematics and unsatisfactory results. In this case, applying the mask in a way where the image pixels are not\n", + "removed from the fit, but their data and noise-map values are scaled such that they contribute negligibly to the fit,\n", + "is a better approach.\n", + "\n", + "We illustrate the API for doing this below, using the `extra_galaxies` dataset which has extra galaxies whose emission\n", + "needs to be removed via scaling in this way. We apply the scaling and show the subplot imaging where the extra\n", + "galaxies mask has scaled the data values to zeros, increasing the noise-map values to large values and in turn made\n", + "the signal to noise of its pixels effectively zero." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/extra_galaxies/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_extra_galaxies = al.Mask2D.from_fits(\n", + " file_path=Path(dataset_path, \"mask_extra_galaxies.fits\"),\n", + " pixel_scales=0.1,\n", + " invert=True, # Note that we invert the mask here as `True` means a pixel is scaled.\n", + ")\n", + "\n", + "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native, pixel_scales=0.1, centre=(0.0, 0.0), radius=6.0\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We do not explictly fit this data, for the sake of brevity, however if your data has these nearby galaxies you should\n", + "apply the mask as above before fitting the data.\n", + "\n", + "__Wrap Up__\n", + "\n", + "Pixelizations are the most complex but also the most powerful way to model a galaxy\u2019s light.\n", + "\n", + "Whether you need to use them depends on the science you are doing. If you are only interested in measuring simple\n", + "global quantities (for example, total flux, size, or axis ratio), analytic light profiles such as a S\u00e9rsic, MGE, or\n", + "shapelets are often sufficient. For low-resolution data, pixelizations are also unnecessary, as the complex\n", + "structure of the galaxy is not resolved.\n", + "\n", + "However, modeling galaxies with complex, irregular, or highly structured light distributions requires this level of\n", + "flexibility. Furthermore, if you are interested in studying the detailed morphology of a galaxy itself, there is no\n", + "better approach than using a pixelization.\n", + "\n", + "__Chaining__\n", + "\n", + "Modeling with a pixelization can be made more efficient, robust, and automated using the non-linear chaining feature\n", + "to compose a pipeline that begins by fitting a simpler model using parametric light profiles.\n", + "\n", + "More information on chaining is provided in the\n", + "`autogalaxy_workspace/notebooks/guides/modeling/chaining` folder and in chapter 3 of the **HowToGalaxy** lectures.\n", + "\n", + "__HowToGalaxy__\n", + "\n", + "A full description of how pixelizations work\u2014which relies heavily on linear algebra, Bayesian statistics, and\n", + "2D geometry\u2014is provided in chapter 4 of the **HowToGalaxy** lectures.\n", + "\n", + "__Future Ideas / Contributions__\n", + "\n", + "Here are a list of things I would like to add to this tutorial but haven't found the time. If you are interested\n", + "in having a go at adding them contact me on SLACK! :)\n", + "\n", + "- More diagnostic quantities for reconstructed galaxy light.\n", + "- Gradient calculations of the reconstructed light distribution.\n", + "- Quantifying spatial variations in galaxy structure across the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/pixelization/slam.ipynb b/notebooks/imaging/features/pixelization/slam.ipynb index 6547cf1bf..0cea13fe5 100644 --- a/notebooks/imaging/features/pixelization/slam.ipynb +++ b/notebooks/imaging/features/pixelization/slam.ipynb @@ -1,667 +1,704 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: SLaM\n", - "==================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for pixelized source modeling.\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", - "guide before working through this example.\n", - "\n", - "Because the SLaM pipelines are designed around pixelized source modeling, the example `slam_start_here` fully\n", - "describes all design choices and modeling decisions made in this script. This script therefore does not repeat\n", - "that documentation, and `slam_start_here` should be read first.\n", - "\n", - "The differences from `slam_start_here` are:\n", - "\n", - " - The SOURCE PIX PIPELINE 2 uses `AdaptSplit` regularization instead of `Adapt`.\n", - " - The LIGHT LP PIPELINE and MASS TOTAL PIPELINE use `use_jax=True` in their analyses.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`, except `AdaptSplit` regularization is used instead of `Adapt`.\n", - "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "You can still run the script without fully understanding the guide, but reviewing it later will\n", - "make the structure and choices of the SLaM workflow clearer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, except `AdaptSplit` regularization is used instead of `Adapt`.\n", - "\n", - "`AdaptSplit` splits the regularization into two components: one for the source and one for the image, enabling\n", - "more flexible regularization that better adapts to the pixelization mesh." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__LIGHT LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def light_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - "\n", - " lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_result_for_lens: af.Result,\n", - " source_result_for_source: af.Result,\n", - " light_result: af.Result,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_result_for_lens\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_result_for_source.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " use_jax=True,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " bulge = light_result.instance.galaxies.lens.bulge\n", - " disk = light_result.instance.galaxies.lens.disk\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=bulge,\n", - " disk=disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"imaging\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.AdaptSplit,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: SLaM\n", + "==================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for pixelized source modeling.\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", + "guide before working through this example.\n", + "\n", + "Because the SLaM pipelines are designed around pixelized source modeling, the example `slam_start_here` fully\n", + "describes all design choices and modeling decisions made in this script. This script therefore does not repeat\n", + "that documentation, and `slam_start_here` should be read first.\n", + "\n", + "The differences from `slam_start_here` are:\n", + "\n", + " - The SOURCE PIX PIPELINE 2 uses `AdaptSplit` regularization instead of `Adapt`.\n", + " - The LIGHT LP PIPELINE and MASS TOTAL PIPELINE use `use_jax=True` in their analyses.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`, except `AdaptSplit` regularization is used instead of `Adapt`.\n", + "- **LIGHT LP PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "You can still run the script without fully understanding the guide, but reviewing it later will\n", + "make the structure and choices of the SLaM workflow clearer." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, except `AdaptSplit` regularization is used instead of `Adapt`.\n", + "\n", + "`AdaptSplit` splits the regularization into two components: one for the source and one for the image, enabling\n", + "more flexible regularization that better adapts to the pixelization mesh." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__LIGHT LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def light_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + "\n", + " lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_result_for_lens: af.Result,\n", + " source_result_for_source: af.Result,\n", + " light_result: af.Result,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_result_for_lens\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_result_for_source.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " use_jax=True,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " bulge = light_result.instance.galaxies.lens.bulge\n", + " disk = light_result.instance.galaxies.lens.disk\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=bulge,\n", + " disk=disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"imaging\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.AdaptSplit,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/pixelization/source_science.ipynb b/notebooks/imaging/features/pixelization/source_science.ipynb index 8478b5bc3..502bf1a41 100644 --- a/notebooks/imaging/features/pixelization/source_science.ipynb +++ b/notebooks/imaging/features/pixelization/source_science.ipynb @@ -1,769 +1,806 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: Source Science\n", - "============================\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source pixelization, we can compute key quantities such as the magnification, total flux, and\n", - "intrinsic size of the source.\n", - "\n", - "For pixelized source reconstructions, these calculations can be quite involved as they required speciifc code to\n", - "handle irregular mesh pixels and other quantities. We illustrate how to perform these calculations below.\n", - "\n", - "However, this does make the source reconstructions different to share with other people, as it would mean they need\n", - "to understand how to manipulate irregular meshes. The end of this example shows how a .csv source reconstruction file\n", - "is output by a pixelization model-fit, which allows anyone to easy interpolate the source reconstruction on to a uniform grid\n", - "for analysis without the need for PyAutoLens.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model Fit:** Perform the model-fit using the search and analysis.\n", - "- **Interpolated Source:** The simplest way to perform source science calculations on a pixelized source reconstruction is to.\n", - "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", - "- **Zoom:** The interpolation grid above was large in extent (-3.0\" to 3.0\" in both the y and x directions).\n", - "- **Errors:** The interpolated errors on the source reconstruction can also be computed, which will allow you to.\n", - "- **Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", - "- **Masking:** Reconstructions can be imperfect, for example having faint source flux in pixels at the edge of the.\n", - "- **Magnification via Mesh:** The calculations above used an interpolation of the source-plane reconstruction to a 2D grid of.\n", - "- **Reconstruction CSV:** In the results `image` folder there is a .csv file called `source_plane_reconstruction_0.csv` which." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "The code below is identical to the pixelizaiton `modeling` example, crucially creating a model-fit which\n", - "outputs the pixelization source reconstruction to a .csv file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "if not (dataset_path / \"positions.json\").exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "dataset = dataset.apply_over_sampling(\n", - " over_sample_size_pixelization=4,\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", - "\n", - "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", - "regularization = al.reg.Constant(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitImaging(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - ")\n", - "\n", - "inversion = fit.inversion\n", - "\n", - "mapper = inversion.cls_list_from(cls=al.Mapper)[\n", - " 0\n", - "] # Extract the mapper from the inversion\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the fit, confirming that the pixelized source reconstruction provides a good fit to the data.\n", - "\n", - "Note how the pixelized source reconstruction is performed on an irregular adaptive grid of rectangular pixels,\n", - "which is denser in regions of high magnification. This non-uniform distribution of pixels means we need to be care\n", - "when performing source science calculations, especially a quantity like the magnification which depends on area." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "All information about the pixelized source reconstruction is contained in the `Inversion` object, which can be\n", - "accessed via `fit.inversion`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "print(f\"Inversion Object: {inversion}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For example, the reconstructed source pixel flux values are stored in the `reconstruction` attribute of the inversion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = inversion.reconstruction\n", - "\n", - "print(f\"Reconstructed Source Pixel Fluxes: {reconstruction}\")\n", - "\n", - "total_flux = np.sum(reconstruction)\n", - "\n", - "print(f\"Total Source Flux via Pixelization: {total_flux} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In order to perform source science calculations we need to know which flux value corresponds to which pixel in the \n", - "source-plane.\n", - "\n", - "This information is available in the inversion, below we print the (y,x) centre of each source pixel corresponding to \n", - "the `reconstruction` values printed above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = inversion.cls_list_from(cls=al.Mapper)[\n", - " 0\n", - "] # Extract the mapper from the inversion\n", - "\n", - "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", - "\n", - "print(f\"Source Plane Mesh Grid Coordinates: {source_plane_mesh_grid}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The image-plane reconstruction can also be computed from the inversion, which is called \n", - "the `mapped_reconstructed_operated_data` and as seen above is needed to compute the magnification." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_operated_data = inversion.mapped_reconstructed_operated_data\n", - "\n", - "print(f\"Mapped Reconstructed Image: {mapped_reconstructed_operated_data}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolated Source__\n", - "\n", - "The simplest way to perform source science calculations on a pixelized source reconstruction is to interpolate\n", - "its values to a uniform 2D grid of pixels, which can therefore be stored using a `Array2D` object,\n", - "which is basically just a 2D numpy array (see the `Data Structure` section at the top of this example).\n", - "\n", - "We interpolate the rectangular pixelized source reconstruction to a new uniform grid we call the `interpolation_grid`.\n", - "This calculation can be quite slow, so to make this example run fast we use a relatively small grid, but in practice\n", - "you may wish to use a larger grid (e.g. 100x1000 pixels or larger for actual science calculations)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.interpolate import griddata\n", - "\n", - "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", - "\n", - "interpolated_reconstruction = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", - ")\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_ndarray,\n", - " pixel_scales=interpolation_grid.pixel_scales,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By printing the interpolated array, we confirm it is a 2D array and can see the pixel values of the source \n", - "reconstruction.\n", - "\n", - "We also plot the interpolated source reconstruction using an `aplt.plot_array`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(interpolated_reconstruction.native)\n", - "\n", - "aplt.plot_array(array=interpolated_reconstruction, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Flux__\n", - "\n", - "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", - "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", - "\n", - "The total flux of the source reconstruction can now be computed by summing the interpolated array.\n", - "\n", - "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", - "we will assume everything is in electrons per second (`e- s^-1`), which is typical for Hubble Space Telescope imaging data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_source_flux = np.sum(interpolated_reconstruction)\n", - "\n", - "print(f\"Total Source Flux via Interpolated Pixelization: {total_source_flux} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Zoom__\n", - "\n", - "The interpolation grid above was large in extent (-3.0\" to 3.0\" in both the y and x directions), meaning that\n", - "the source was a small flux was a small region of this grid.\n", - "\n", - "By changing the `extent` of the interpolation grid, we can performed the interpolation zoomed in on only the\n", - "regions of the source-plane where the source reconstruction has non-negligible flux. This\n", - "makes the interpolation more accurate, as the interpolation ican use more pixels in the region of interest,\n", - "and also makes visualizing the source reconstruction easier." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extent = (-1.0, 1.0, -1.0, 1.0)\n", - "shape_native = (401, 401)\n", - "\n", - "interpolation_grid_zoom = al.Grid2D.from_extent(\n", - " extent=extent,\n", - " shape_native=shape_native,\n", - ")\n", - "\n", - "interpolated_reconstruction = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid_zoom\n", - ")\n", - "\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", - " interpolation_grid_zoom.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_ndarray,\n", - " pixel_scales=interpolation_grid_zoom.pixel_scales,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors__\n", - "\n", - "The interpolated errors on the source reconstruction can also be computed, which will allow you to perform\n", - "model-fitting of the source reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction_noise_map = inversion.reconstruction_noise_map\n", - "\n", - "interpolated_noise_map = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", - ")\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_noise_map_ndarray = interpolated_noise_map.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_noise_map = al.Array2D.no_mask(\n", - " values=interpolated_noise_map_ndarray, pixel_scales=interpolation_grid.pixel_scales\n", - ")\n", - "\n", - "aplt.plot_array(array=interpolated_noise_map, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Magnification__\n", - "\n", - "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", - "total surface brightness in the source-plane.\n", - "\n", - "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", - "We therefore explicitly mention how area folds into the calculation below.\n", - "\n", - "The interpolated source reconstruction above has different sized pixels in the image-plane and source-plane, so \n", - "we need to explicitly account for area when computing the magnification.\n", - "\n", - "The `pixel_area` attribute of the `Array2D` object gives us the area of each pixel in arcseconds squared, which we\n", - "can use to compute the magnification below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnification = np.sum(\n", - " mapped_reconstructed_operated_data * mapped_reconstructed_operated_data.pixel_area\n", - ") / np.sum(interpolated_reconstruction * interpolated_reconstruction.pixel_area)\n", - "\n", - "print(f\"Magnification via Interpolated Source: {magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masking__\n", - "\n", - "Reconstructions can be imperfect, for example having faint source flux in pixels at the edge of the\n", - "source-plane that through comparison to the data are not a genuine part of the source. This can impact\n", - "the calculation of the source flux and magnification.\n", - "\n", - "If you want to be extra careful, you can use a mask to zero the source-plane pixels that you do not trust and use\n", - "this to remove pixels from source science calculations.\n", - "\n", - "Another approach, which we use below, is we create a source-plane signal-to-noise map and use this to create a mask \n", - "that removes all pixels with a signal-to-noise < 5.0." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "signal_to_noise_map = reconstruction / reconstruction_noise_map\n", - "\n", - "mesh_pixel_mask = signal_to_noise_map < 5.0\n", - "\n", - "reconstruction_masked = reconstruction.copy()\n", - "reconstruction_masked[mesh_pixel_mask] = 0.0\n", - "\n", - "interpolated_reconstruction_masked = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction_masked, xi=interpolation_grid\n", - ")\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_reconstruction_masked_ndarray = interpolated_reconstruction_masked.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction_masked = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_masked_ndarray,\n", - " pixel_scales=interpolation_grid.pixel_scales,\n", - ")\n", - "\n", - "aplt.plot_array(array=interpolated_reconstruction_masked, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Magnification via Mesh__\n", - "\n", - "The calculations above used an interpolation of the source-plane reconstruction to a 2D grid of 1000 x 1000\n", - "pixels.\n", - "\n", - "However, we can use directly the irregular rectangular mesh of the pixelized source reconstruction to compute\n", - "quantities. This is more accurate as it does not introduce interpolation errors, but requires more care as the\n", - "pixels are irregularly spaced and have different areas. \n", - "\n", - "We have already computed the total source flux using the mesh above, but we can also compute the magnification.\n", - "\n", - "Computed the areas of every pixel in the irregular rectangular mesh is a bit involved, therefore the values can be\n", - "accessed from the source code via the `mesh_areas` attribute of the `Mapper` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_areas = mapper.mesh_geometry.areas_for_magnification\n", - "\n", - "magnification = np.sum(\n", - " mapped_reconstructed_operated_data * mapped_reconstructed_operated_data.pixel_area\n", - ") / np.sum(reconstruction * mesh_areas)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction CSV__\n", - "\n", - "In the results `image` folder there is a .csv file called `source_plane_reconstruction_0.csv` which contains the\n", - "y and x coordinates of the pixelization mesh, the reconstruct values and the noise map of these values.\n", - "\n", - "This file is provides all information on the source reconstruction in a format that does not depend autolens\n", - "and therefore be easily loaded to create images of the source or shared collaborations who do not have PyAutoLens\n", - "installed.\n", - "\n", - "We now perform a lens model fit, which will create this .csv file in the modeling output folder.\n", - "\n", - "First, lets load `source_plane_reconstruction_0.csv` as a dictionary, using basic `csv` functionality in Python." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.PowerLaw)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", - "regularization = af.Model(al.reg.Constant)\n", - "\n", - "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"features\"),\n", - " name=\"pixelization\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20,\n", - " iterations_per_quick_update=50000,\n", - ")\n", - "\n", - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n", - "\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " use_jax=True,\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction CSV__\n", - "\n", - "The file ``source_plane_reconstruction_0.csv` provides all information on the source reconstruction in a format that \n", - "does not depend autolens and therefore be easily loaded to create images of the source or shared collaborations who \n", - "do not have PyAutoLens installed.\n", - "\n", - "First, lets load `source_plane_reconstruction_0.csv` as a dictionary, using basic `csv` functionality in Python.\n", - "\n", - "NOTE: If the .csv file does not exist, we create a dictionary with the same format but with dummy values so the rest of\n", - "the script can be run." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import csv\n", - "\n", - "try:\n", - "\n", - " with open(\n", - " search.paths.image_path / \"source_plane_reconstruction_0.csv\", mode=\"r\"\n", - " ) as file:\n", - " reader = csv.reader(file)\n", - " header_list = next(reader) # ['y', 'x', 'reconstruction', 'noise_map']\n", - "\n", - " reconstruction_dict = {header: [] for header in header_list}\n", - "\n", - " for row in reader:\n", - " for key, value in zip(header_list, row):\n", - " reconstruction_dict[key].append(float(value))\n", - "\n", - " # Convert lists to NumPy arrays\n", - " for key in reconstruction_dict:\n", - " reconstruction_dict[key] = np.array(reconstruction_dict[key])\n", - "\n", - "except FileNotFoundError:\n", - "\n", - " print(\"`source_plane_reconstruction_0.csv` not found. Using dummy data instead.\")\n", - "\n", - " x = np.array([-1.0, 0.0, 1.0, -1.0, 0.0, 1.0, -1.0, 0.0, 1.0])\n", - " y = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0, -1.0, -1.0, -1.0])\n", - " reconstruction = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])\n", - " noise_map = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])\n", - "\n", - " reconstruction_dict = {\n", - " \"x\": x,\n", - " \"y\": y,\n", - " \"reconstruction\": reconstruction,\n", - " \"noise_map\": noise_map,\n", - " }\n", - "\n", - "print(reconstruction_dict[\"y\"])\n", - "print(reconstruction_dict[\"x\"])\n", - "print(reconstruction_dict[\"reconstruction\"])\n", - "print(reconstruction_dict[\"noise_map\"])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can now use standard libraries to performed calculations with the reconstruction on the mesh, again avoiding\n", - "the need to use autolens.\n", - "\n", - "For example, we can create a RectangularAdaptDensity mesh using the scipy.spatial library, which is a triangulation\n", - "of the y and x coordinates of the pixelization mesh. This is useful for visualizing the pixelization\n", - "and performing calculations on the mesh." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import scipy\n", - "\n", - "points = np.stack(arrays=(reconstruction_dict[\"x\"], reconstruction_dict[\"y\"]), axis=-1)\n", - "\n", - "mesh = scipy.spatial.Delaunay(points)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Interpolating the result to a uniform grid is also possible using the scipy.interpolate library, which means the result\n", - "can be turned into a uniform 2D image which can be useful to analyse the source with tools which require an uniform grid.\n", - "\n", - "Below, we interpolate the result onto a 201 x 201 grid of pixels with the extent spanning -1.0\" to 1.0\", which\n", - "capture the majority of the source reconstruction without being too high resolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.interpolate import griddata\n", - "\n", - "values = reconstruction_dict[\"reconstruction\"]\n", - "\n", - "interpolation_grid = al.Grid2D.from_extent(\n", - " extent=(-1.0, 1.0, -1.0, 1.0), shape_native=(201, 201)\n", - ")\n", - "\n", - "interpolated_array = griddata(points=points, values=values, xi=interpolation_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: Source Science\n", + "============================\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source pixelization, we can compute key quantities such as the magnification, total flux, and\n", + "intrinsic size of the source.\n", + "\n", + "For pixelized source reconstructions, these calculations can be quite involved as they required speciifc code to\n", + "handle irregular mesh pixels and other quantities. We illustrate how to perform these calculations below.\n", + "\n", + "However, this does make the source reconstructions different to share with other people, as it would mean they need\n", + "to understand how to manipulate irregular meshes. The end of this example shows how a .csv source reconstruction file\n", + "is output by a pixelization model-fit, which allows anyone to easy interpolate the source reconstruction on to a uniform grid\n", + "for analysis without the need for PyAutoLens.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model Fit:** Perform the model-fit using the search and analysis.\n", + "- **Interpolated Source:** The simplest way to perform source science calculations on a pixelized source reconstruction is to.\n", + "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", + "- **Zoom:** The interpolation grid above was large in extent (-3.0\" to 3.0\" in both the y and x directions).\n", + "- **Errors:** The interpolated errors on the source reconstruction can also be computed, which will allow you to.\n", + "- **Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", + "- **Masking:** Reconstructions can be imperfect, for example having faint source flux in pixels at the edge of the.\n", + "- **Magnification via Mesh:** The calculations above used an interpolation of the source-plane reconstruction to a 2D grid of.\n", + "- **Reconstruction CSV:** In the results `image` folder there is a .csv file called `source_plane_reconstruction_0.csv` which." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "The code below is identical to the pixelizaiton `modeling` example, crucially creating a model-fit which\n", + "outputs the pixelization source reconstruction to a .csv file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "if not (dataset_path / \"positions.json\").exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "dataset = dataset.apply_over_sampling(\n", + " over_sample_size_pixelization=4,\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", + "\n", + "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", + "regularization = al.reg.Constant(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitImaging(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + ")\n", + "\n", + "inversion = fit.inversion\n", + "\n", + "mapper = inversion.cls_list_from(cls=al.Mapper)[\n", + " 0\n", + "] # Extract the mapper from the inversion\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the fit, confirming that the pixelized source reconstruction provides a good fit to the data.\n", + "\n", + "Note how the pixelized source reconstruction is performed on an irregular adaptive grid of rectangular pixels,\n", + "which is denser in regions of high magnification. This non-uniform distribution of pixels means we need to be care\n", + "when performing source science calculations, especially a quantity like the magnification which depends on area." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "All information about the pixelized source reconstruction is contained in the `Inversion` object, which can be\n", + "accessed via `fit.inversion`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "print(f\"Inversion Object: {inversion}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For example, the reconstructed source pixel flux values are stored in the `reconstruction` attribute of the inversion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = inversion.reconstruction\n", + "\n", + "print(f\"Reconstructed Source Pixel Fluxes: {reconstruction}\")\n", + "\n", + "total_flux = np.sum(reconstruction)\n", + "\n", + "print(f\"Total Source Flux via Pixelization: {total_flux} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In order to perform source science calculations we need to know which flux value corresponds to which pixel in the \n", + "source-plane.\n", + "\n", + "This information is available in the inversion, below we print the (y,x) centre of each source pixel corresponding to \n", + "the `reconstruction` values printed above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = inversion.cls_list_from(cls=al.Mapper)[\n", + " 0\n", + "] # Extract the mapper from the inversion\n", + "\n", + "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", + "\n", + "print(f\"Source Plane Mesh Grid Coordinates: {source_plane_mesh_grid}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The image-plane reconstruction can also be computed from the inversion, which is called \n", + "the `mapped_reconstructed_operated_data` and as seen above is needed to compute the magnification." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_operated_data = inversion.mapped_reconstructed_operated_data\n", + "\n", + "print(f\"Mapped Reconstructed Image: {mapped_reconstructed_operated_data}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolated Source__\n", + "\n", + "The simplest way to perform source science calculations on a pixelized source reconstruction is to interpolate\n", + "its values to a uniform 2D grid of pixels, which can therefore be stored using a `Array2D` object,\n", + "which is basically just a 2D numpy array (see the `Data Structure` section at the top of this example).\n", + "\n", + "We interpolate the rectangular pixelized source reconstruction to a new uniform grid we call the `interpolation_grid`.\n", + "This calculation can be quite slow, so to make this example run fast we use a relatively small grid, but in practice\n", + "you may wish to use a larger grid (e.g. 100x1000 pixels or larger for actual science calculations)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.interpolate import griddata\n", + "\n", + "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", + "\n", + "interpolated_reconstruction = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", + ")\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_ndarray,\n", + " pixel_scales=interpolation_grid.pixel_scales,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By printing the interpolated array, we confirm it is a 2D array and can see the pixel values of the source \n", + "reconstruction.\n", + "\n", + "We also plot the interpolated source reconstruction using an `aplt.plot_array`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(interpolated_reconstruction.native)\n", + "\n", + "aplt.plot_array(array=interpolated_reconstruction, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Flux__\n", + "\n", + "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", + "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", + "\n", + "The total flux of the source reconstruction can now be computed by summing the interpolated array.\n", + "\n", + "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", + "we will assume everything is in electrons per second (`e- s^-1`), which is typical for Hubble Space Telescope imaging data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_source_flux = np.sum(interpolated_reconstruction)\n", + "\n", + "print(f\"Total Source Flux via Interpolated Pixelization: {total_source_flux} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Zoom__\n", + "\n", + "The interpolation grid above was large in extent (-3.0\" to 3.0\" in both the y and x directions), meaning that\n", + "the source was a small flux was a small region of this grid.\n", + "\n", + "By changing the `extent` of the interpolation grid, we can performed the interpolation zoomed in on only the\n", + "regions of the source-plane where the source reconstruction has non-negligible flux. This\n", + "makes the interpolation more accurate, as the interpolation ican use more pixels in the region of interest,\n", + "and also makes visualizing the source reconstruction easier." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extent = (-1.0, 1.0, -1.0, 1.0)\n", + "shape_native = (401, 401)\n", + "\n", + "interpolation_grid_zoom = al.Grid2D.from_extent(\n", + " extent=extent,\n", + " shape_native=shape_native,\n", + ")\n", + "\n", + "interpolated_reconstruction = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid_zoom\n", + ")\n", + "\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", + " interpolation_grid_zoom.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_ndarray,\n", + " pixel_scales=interpolation_grid_zoom.pixel_scales,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors__\n", + "\n", + "The interpolated errors on the source reconstruction can also be computed, which will allow you to perform\n", + "model-fitting of the source reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction_noise_map = inversion.reconstruction_noise_map\n", + "\n", + "interpolated_noise_map = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", + ")\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_noise_map_ndarray = interpolated_noise_map.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_noise_map = al.Array2D.no_mask(\n", + " values=interpolated_noise_map_ndarray, pixel_scales=interpolation_grid.pixel_scales\n", + ")\n", + "\n", + "aplt.plot_array(array=interpolated_noise_map, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Magnification__\n", + "\n", + "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", + "total surface brightness in the source-plane.\n", + "\n", + "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", + "We therefore explicitly mention how area folds into the calculation below.\n", + "\n", + "The interpolated source reconstruction above has different sized pixels in the image-plane and source-plane, so \n", + "we need to explicitly account for area when computing the magnification.\n", + "\n", + "The `pixel_area` attribute of the `Array2D` object gives us the area of each pixel in arcseconds squared, which we\n", + "can use to compute the magnification below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnification = np.sum(\n", + " mapped_reconstructed_operated_data * mapped_reconstructed_operated_data.pixel_area\n", + ") / np.sum(interpolated_reconstruction * interpolated_reconstruction.pixel_area)\n", + "\n", + "print(f\"Magnification via Interpolated Source: {magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masking__\n", + "\n", + "Reconstructions can be imperfect, for example having faint source flux in pixels at the edge of the\n", + "source-plane that through comparison to the data are not a genuine part of the source. This can impact\n", + "the calculation of the source flux and magnification.\n", + "\n", + "If you want to be extra careful, you can use a mask to zero the source-plane pixels that you do not trust and use\n", + "this to remove pixels from source science calculations.\n", + "\n", + "Another approach, which we use below, is we create a source-plane signal-to-noise map and use this to create a mask \n", + "that removes all pixels with a signal-to-noise < 5.0." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "signal_to_noise_map = reconstruction / reconstruction_noise_map\n", + "\n", + "mesh_pixel_mask = signal_to_noise_map < 5.0\n", + "\n", + "reconstruction_masked = reconstruction.copy()\n", + "reconstruction_masked[mesh_pixel_mask] = 0.0\n", + "\n", + "interpolated_reconstruction_masked = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction_masked, xi=interpolation_grid\n", + ")\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_reconstruction_masked_ndarray = interpolated_reconstruction_masked.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction_masked = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_masked_ndarray,\n", + " pixel_scales=interpolation_grid.pixel_scales,\n", + ")\n", + "\n", + "aplt.plot_array(array=interpolated_reconstruction_masked, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Magnification via Mesh__\n", + "\n", + "The calculations above used an interpolation of the source-plane reconstruction to a 2D grid of 1000 x 1000\n", + "pixels.\n", + "\n", + "However, we can use directly the irregular rectangular mesh of the pixelized source reconstruction to compute\n", + "quantities. This is more accurate as it does not introduce interpolation errors, but requires more care as the\n", + "pixels are irregularly spaced and have different areas. \n", + "\n", + "We have already computed the total source flux using the mesh above, but we can also compute the magnification.\n", + "\n", + "Computed the areas of every pixel in the irregular rectangular mesh is a bit involved, therefore the values can be\n", + "accessed from the source code via the `mesh_areas` attribute of the `Mapper` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_areas = mapper.mesh_geometry.areas_for_magnification\n", + "\n", + "magnification = np.sum(\n", + " mapped_reconstructed_operated_data * mapped_reconstructed_operated_data.pixel_area\n", + ") / np.sum(reconstruction * mesh_areas)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction CSV__\n", + "\n", + "In the results `image` folder there is a .csv file called `source_plane_reconstruction_0.csv` which contains the\n", + "y and x coordinates of the pixelization mesh, the reconstruct values and the noise map of these values.\n", + "\n", + "This file is provides all information on the source reconstruction in a format that does not depend autolens\n", + "and therefore be easily loaded to create images of the source or shared collaborations who do not have PyAutoLens\n", + "installed.\n", + "\n", + "We now perform a lens model fit, which will create this .csv file in the modeling output folder.\n", + "\n", + "First, lets load `source_plane_reconstruction_0.csv` as a dictionary, using basic `csv` functionality in Python." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.PowerLaw)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", + "regularization = af.Model(al.reg.Constant)\n", + "\n", + "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"features\"),\n", + " name=\"pixelization\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20,\n", + " iterations_per_quick_update=50000,\n", + ")\n", + "\n", + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n", + "\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " use_jax=True,\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction CSV__\n", + "\n", + "The file ``source_plane_reconstruction_0.csv` provides all information on the source reconstruction in a format that \n", + "does not depend autolens and therefore be easily loaded to create images of the source or shared collaborations who \n", + "do not have PyAutoLens installed.\n", + "\n", + "First, lets load `source_plane_reconstruction_0.csv` as a dictionary, using basic `csv` functionality in Python.\n", + "\n", + "NOTE: If the .csv file does not exist, we create a dictionary with the same format but with dummy values so the rest of\n", + "the script can be run." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import csv\n", + "\n", + "try:\n", + "\n", + " with open(\n", + " search.paths.image_path / \"source_plane_reconstruction_0.csv\", mode=\"r\"\n", + " ) as file:\n", + " reader = csv.reader(file)\n", + " header_list = next(reader) # ['y', 'x', 'reconstruction', 'noise_map']\n", + "\n", + " reconstruction_dict = {header: [] for header in header_list}\n", + "\n", + " for row in reader:\n", + " for key, value in zip(header_list, row):\n", + " reconstruction_dict[key].append(float(value))\n", + "\n", + " # Convert lists to NumPy arrays\n", + " for key in reconstruction_dict:\n", + " reconstruction_dict[key] = np.array(reconstruction_dict[key])\n", + "\n", + "except FileNotFoundError:\n", + "\n", + " print(\"`source_plane_reconstruction_0.csv` not found. Using dummy data instead.\")\n", + "\n", + " x = np.array([-1.0, 0.0, 1.0, -1.0, 0.0, 1.0, -1.0, 0.0, 1.0])\n", + " y = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0, -1.0, -1.0, -1.0])\n", + " reconstruction = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])\n", + " noise_map = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])\n", + "\n", + " reconstruction_dict = {\n", + " \"x\": x,\n", + " \"y\": y,\n", + " \"reconstruction\": reconstruction,\n", + " \"noise_map\": noise_map,\n", + " }\n", + "\n", + "print(reconstruction_dict[\"y\"])\n", + "print(reconstruction_dict[\"x\"])\n", + "print(reconstruction_dict[\"reconstruction\"])\n", + "print(reconstruction_dict[\"noise_map\"])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can now use standard libraries to performed calculations with the reconstruction on the mesh, again avoiding\n", + "the need to use autolens.\n", + "\n", + "For example, we can create a RectangularAdaptDensity mesh using the scipy.spatial library, which is a triangulation\n", + "of the y and x coordinates of the pixelization mesh. This is useful for visualizing the pixelization\n", + "and performing calculations on the mesh." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import scipy\n", + "\n", + "points = np.stack(arrays=(reconstruction_dict[\"x\"], reconstruction_dict[\"y\"]), axis=-1)\n", + "\n", + "mesh = scipy.spatial.Delaunay(points)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Interpolating the result to a uniform grid is also possible using the scipy.interpolate library, which means the result\n", + "can be turned into a uniform 2D image which can be useful to analyse the source with tools which require an uniform grid.\n", + "\n", + "Below, we interpolate the result onto a 201 x 201 grid of pixels with the extent spanning -1.0\" to 1.0\", which\n", + "capture the majority of the source reconstruction without being too high resolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.interpolate import griddata\n", + "\n", + "values = reconstruction_dict[\"reconstruction\"]\n", + "\n", + "interpolation_grid = al.Grid2D.from_extent(\n", + " extent=(-1.0, 1.0, -1.0, 1.0), shape_native=(201, 201)\n", + ")\n", + "\n", + "interpolated_array = griddata(points=points, values=values, xi=interpolation_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/scaling_relation/fit.ipynb b/notebooks/imaging/features/scaling_relation/fit.ipynb index ded3760da..46f507bf6 100644 --- a/notebooks/imaging/features/scaling_relation/fit.ipynb +++ b/notebooks/imaging/features/scaling_relation/fit.ipynb @@ -1,527 +1,564 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Scaling Relation Fit\n", - "==============================\n", - "\n", - "A strong lens system often has many foreground galaxies near the line of sight to the source, in addition to the\n", - "primary lens. As the number of foreground galaxies grows, modelling each one individually with its own free\n", - "`einstein_radius` parameter rapidly becomes intractable.\n", - "\n", - "A common solution is to split foreground galaxies into two tiers:\n", - "\n", - " - **Individually-modelled extras** \u2014 the brighter, closer companions that contribute non-trivially to the lensing\n", - " on their own. Each gets its own free Einstein radius.\n", - " - **Scaling-relation extras** \u2014 the long tail of fainter companions whose Einstein radii are tied together via a\n", - " shared two-parameter relation\n", - " einstein_radius = scaling_factor * luminosity ** scaling_exponent\n", - " so adding more galaxies to this tier does not grow the model.\n", - "\n", - "This script illustrates the API for performing a fit to a strong lens with both tiers active, via the standard\n", - "`Tracer` and `FitImaging` objects, without invoking a non-linear search. It is intended to make the per-galaxy\n", - "deflection composition concrete before the reader moves on to `modeling.py` (search-based).\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Adaptive over-sampling at every galaxy centre.\n", - "- **Centres + Luminosities:** Load extra-galaxy centres (JSON) and scaling-tier centres + luminosities (CSV).\n", - "- **MGE Basis:** Build a `Basis` of linear Gaussians for the source.\n", - "- **Galaxies:** Concrete composition \u2014 lens, individually-modelled extras, scaling-tier extras, source.\n", - "- **Tracer:** Build the `Tracer` and fit the dataset.\n", - "- **Scaling Relation Tour:** Per-galaxy deflections sum into the tracer's total deflection. The scaling-tier\n", - " galaxies' Einstein radii come from `scaling_factor * luminosity ** scaling_exponent`.\n", - "- **Intensities:** The solved-for linear light profile `intensity` values for each MGE Gaussian.\n", - "- **Wrap Up:** Summary and next steps.\n", - "\n", - "__Prerequisites__\n", - "\n", - "This script focuses on the API specific to a mixed individually-modelled + scaling-tier extras population. For\n", - "background on the underlying single-plane fit API and the MGE source parameterization, you should read first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/fit.py` \u2014 the standard single-plane fit.\n", - " - `autolens_workspace/scripts/imaging/features/scaling_relation/modeling.py` \u2014 the search-based version of this\n", - " script, which composes the same model via `af.Model` with free `scaling_factor` and `scaling_exponent` priors.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py` \u2014 the MGE source `Basis` API.\n", - "\n", - "All non-linear parameters below are set to the simulator's true values, so the fit visibly recovers the lens\n", - "configuration without a search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `extra_and_scaling_galaxies` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_and_scaling_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/scaling_relation/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 6.0\" circular mask, large enough to include all four extra galaxies (the most distant sits at radius ~5\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 6.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres + Luminosities__\n", - "\n", - "Load the centres of the individually-modelled extras and the scaling-tier extras. The scaling-tier galaxies need\n", - "both centres AND a measured luminosity each, so they're loaded from a CSV via `al.galaxy_table_from_csv`.\n", - "\n", - "In a real analysis, the scaling-tier luminosities come from a prior light-only fit \u2014 see\n", - "`scripts/group/features/scaling_relation/modeling_for_luminosities.py` for the standalone version of that fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "individual_extras_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "scaling_table = al.galaxy_table_from_csv(\n", - " file_path=dataset_path / \"scaling_galaxies.csv\"\n", - ")\n", - "scaling_extras_centres = scaling_table.centres\n", - "scaling_extras_luminosities = scaling_table.luminosities\n", - "\n", - "print(f\"Individually-modelled extras centres: {list(individual_extras_centres)}\")\n", - "print(f\"Scaling-tier extras centres: {list(scaling_extras_centres)}\")\n", - "print(f\"Scaling-tier extras luminosities: {scaling_extras_luminosities}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling at every galaxy centre, so each light profile is evaluated accurately at its peak." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_galaxy_centres = (\n", - " [(0.0, 0.0)]\n", - " + [tuple(c) for c in individual_extras_centres]\n", - " + [tuple(c) for c in scaling_extras_centres]\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_galaxy_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MGE Basis__\n", - "\n", - "A `Basis` of 30 linear Gaussians for the source galaxy. The `intensity` of each Gaussian is solved for via linear\n", - "algebra at fit time." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 30\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "source_bulge = build_source_basis(centre=(0.0, 0.1))\n", - "\n", - "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "We compose four populations:\n", - "\n", - " - `lens` (z=0.5): `SersicSph` light + `IsothermalSph` mass at the origin. The simulator's true Einstein radius\n", - " is 1.6\".\n", - " - `individual_extras` (z=0.5): two close companions, each modelled with its own `SersicSph` light +\n", - " `IsothermalSph` mass. Simulator-true Einstein radii: 0.4\" and 0.5\".\n", - " - `scaling_extras` (z=0.5): two further-out, fainter companions whose Einstein radii are derived from the\n", - " scaling relation. With `scaling_factor=0.3` and `scaling_exponent=1.0` and per-galaxy luminosity=0.45, each\n", - " acquires `einstein_radius = 0.3 * 0.45 ** 1.0 = 0.135` \u2014 matches the simulator.\n", - " - `source` (z=1.0): the MGE basis above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=1.5, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6),\n", - ")\n", - "\n", - "individual_extras_truth = [\n", - " dict(intensity=0.9, effective_radius=0.6, sersic_index=2.5, einstein_radius=0.4),\n", - " dict(intensity=0.8, effective_radius=0.6, sersic_index=2.5, einstein_radius=0.5),\n", - "]\n", - "\n", - "individual_extras = []\n", - "for centre, truth in zip(individual_extras_centres, individual_extras_truth):\n", - " individual_extras.append(\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=tuple(centre),\n", - " intensity=truth[\"intensity\"],\n", - " effective_radius=truth[\"effective_radius\"],\n", - " sersic_index=truth[\"sersic_index\"],\n", - " ),\n", - " mass=al.mp.IsothermalSph(\n", - " centre=tuple(centre), einstein_radius=truth[\"einstein_radius\"]\n", - " ),\n", - " )\n", - " )\n", - "\n", - "scaling_factor = 0.3\n", - "scaling_exponent = 1.0\n", - "\n", - "scaling_extras = []\n", - "scaling_extras_einstein_radii = []\n", - "for centre, luminosity in zip(scaling_extras_centres, scaling_extras_luminosities):\n", - " einstein_radius = scaling_factor * luminosity**scaling_exponent\n", - " scaling_extras_einstein_radii.append(einstein_radius)\n", - " scaling_extras.append(\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=tuple(centre),\n", - " intensity=luminosity,\n", - " effective_radius=0.5,\n", - " sersic_index=2.5,\n", - " ),\n", - " mass=al.mp.IsothermalSph(\n", - " centre=tuple(centre), einstein_radius=einstein_radius\n", - " ),\n", - " )\n", - " )\n", - "\n", - "source = al.Galaxy(redshift=1.0, bulge=source_bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "The `Tracer` performs the ray-tracing. It queries every mass profile attached to every galaxy in the lens plane\n", - "and sums their deflections. For our mixed population, this means the lens's `IsothermalSph`, each individually-\n", - "modelled extra's `IsothermalSph`, and each scaling-tier extra's `IsothermalSph` all contribute independent\n", - "deflections that sum before mapping image-plane coordinates onto the source-plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens] + individual_extras + scaling_extras + [source])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "Pass the `Tracer` to a `FitImaging` to fit the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling Relation Tour__\n", - "\n", - "The scaling-tier galaxies' Einstein radii are NOT free parameters in the model \u2014 they're computed from a shared\n", - "two-parameter relation and per-galaxy luminosity. With `scaling_factor=0.3` and `scaling_exponent=1.0`, we have:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for centre, luminosity, er in zip(\n", - " scaling_extras_centres, scaling_extras_luminosities, scaling_extras_einstein_radii\n", - "):\n", - " print(\n", - " f\" scaling galaxy @ {tuple(centre)}: luminosity = {luminosity:.3f}, \"\n", - " f\"einstein_radius = {scaling_factor:.3f} * {luminosity:.3f} ** {scaling_exponent:.1f} = {er:.4f}\"\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The lens-plane total deflection is the SUM of every mass profile's contribution. We verify this by computing\n", - "each one explicitly and confirming the sum equals what the `Tracer` returns." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = dataset.grid\n", - "\n", - "alpha_lens = lens.mass.deflections_yx_2d_from(grid=grid)\n", - "alpha_individual = [g.mass.deflections_yx_2d_from(grid=grid) for g in individual_extras]\n", - "alpha_scaling = [g.mass.deflections_yx_2d_from(grid=grid) for g in scaling_extras]\n", - "\n", - "print(f\"\\nalpha_lens (first coord): {alpha_lens[0]}\")\n", - "print(f\"alpha_individual_extra_0 (first coord): {alpha_individual[0][0]}\")\n", - "print(f\"alpha_individual_extra_1 (first coord): {alpha_individual[1][0]}\")\n", - "print(f\"alpha_scaling_extra_0 (first coord): {alpha_scaling[0][0]}\")\n", - "print(f\"alpha_scaling_extra_1 (first coord): {alpha_scaling[1][0]}\")\n", - "\n", - "alpha_total_summed = alpha_lens + sum(alpha_individual) + sum(alpha_scaling)\n", - "\n", - "traced_grids = tracer.traced_grid_2d_list_from(grid=grid)\n", - "alpha_total_tracer = grid - traced_grids[1]\n", - "\n", - "print(f\"\\nalpha_total (summed by hand, first 3): {alpha_total_summed[:3]}\")\n", - "print(f\"alpha_total (from tracer, first 3): {alpha_total_tracer[:3]}\")\n", - "\n", - "assert np.allclose(np.asarray(alpha_total_summed), np.asarray(alpha_total_tracer))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The scaling-tier deflections per galaxy are visibly smaller than the lens and individually-modelled extras\n", - "because their Einstein radii are an order of magnitude smaller. The shared relation lets us include them in the\n", - "model with zero additional free parameters \u2014 adding more scaling galaxies would not grow the parameter space.\n", - "\n", - "__Intensities__\n", - "\n", - "After the fit, every linear Gaussian in the source MGE basis has been assigned an `intensity` via linear algebra." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"\\nFirst Gaussian intensity, source = \"\n", - " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", - ")\n", - "\n", - "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "subplot_basis_image(basis=tracer_fitted.galaxies[-1].bulge, grid=plot_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script demonstrated the mixed-strategy API for handling many foreground galaxies \u2014 a small number of\n", - "individually-modelled extras for the bright, close companions, and an arbitrary number of scaling-tier extras for\n", - "the long tail of fainter ones. The scaling relation collapses what would otherwise be N free `einstein_radius`\n", - "parameters into 2 shared parameters (`scaling_factor` and `scaling_exponent`), keeping the model dimensionality\n", - "manageable as galaxy count grows.\n", - "\n", - "In a real modeling workflow:\n", - "\n", - " - `modeling.py` runs the search-based version, where `scaling_factor` and `scaling_exponent` are free `af.Model`\n", - " parameters with `UniformPrior`s. The luminosities still come from a prior light-only fit.\n", - " - For group-scale lenses with multiple main lens galaxies, see\n", - " `autolens_workspace/scripts/group/features/scaling_relation/` \u2014 the three-tier API generalises this script to\n", - " `lens_dict + extra_galaxies + scaling_galaxies` collections.\n", - "\n", - "The key takeaway is that scaling relations let the lens model stay tractable even when 10s or 100s of foreground\n", - "galaxies sit on it \u2014 the lens-plane deflection is still a simple sum of per-galaxy contributions, but the\n", - "contributions from the scaling tier are parameterized through luminosity rather than per-galaxy free parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Scaling Relation Fit\n", + "==============================\n", + "\n", + "A strong lens system often has many foreground galaxies near the line of sight to the source, in addition to the\n", + "primary lens. As the number of foreground galaxies grows, modelling each one individually with its own free\n", + "`einstein_radius` parameter rapidly becomes intractable.\n", + "\n", + "A common solution is to split foreground galaxies into two tiers:\n", + "\n", + " - **Individually-modelled extras** \u2014 the brighter, closer companions that contribute non-trivially to the lensing\n", + " on their own. Each gets its own free Einstein radius.\n", + " - **Scaling-relation extras** \u2014 the long tail of fainter companions whose Einstein radii are tied together via a\n", + " shared two-parameter relation\n", + " einstein_radius = scaling_factor * luminosity ** scaling_exponent\n", + " so adding more galaxies to this tier does not grow the model.\n", + "\n", + "This script illustrates the API for performing a fit to a strong lens with both tiers active, via the standard\n", + "`Tracer` and `FitImaging` objects, without invoking a non-linear search. It is intended to make the per-galaxy\n", + "deflection composition concrete before the reader moves on to `modeling.py` (search-based).\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Adaptive over-sampling at every galaxy centre.\n", + "- **Centres + Luminosities:** Load extra-galaxy centres (JSON) and scaling-tier centres + luminosities (CSV).\n", + "- **MGE Basis:** Build a `Basis` of linear Gaussians for the source.\n", + "- **Galaxies:** Concrete composition \u2014 lens, individually-modelled extras, scaling-tier extras, source.\n", + "- **Tracer:** Build the `Tracer` and fit the dataset.\n", + "- **Scaling Relation Tour:** Per-galaxy deflections sum into the tracer's total deflection. The scaling-tier\n", + " galaxies' Einstein radii come from `scaling_factor * luminosity ** scaling_exponent`.\n", + "- **Intensities:** The solved-for linear light profile `intensity` values for each MGE Gaussian.\n", + "- **Wrap Up:** Summary and next steps.\n", + "\n", + "__Prerequisites__\n", + "\n", + "This script focuses on the API specific to a mixed individually-modelled + scaling-tier extras population. For\n", + "background on the underlying single-plane fit API and the MGE source parameterization, you should read first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/fit.py` \u2014 the standard single-plane fit.\n", + " - `autolens_workspace/scripts/imaging/features/scaling_relation/modeling.py` \u2014 the search-based version of this\n", + " script, which composes the same model via `af.Model` with free `scaling_factor` and `scaling_exponent` priors.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/fit.py` \u2014 the MGE source `Basis` API.\n", + "\n", + "All non-linear parameters below are set to the simulator's true values, so the fit visibly recovers the lens\n", + "configuration without a search." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autogalaxy.profiles.plot.basis_plots import subplot_image as subplot_basis_image" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `extra_and_scaling_galaxies` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_and_scaling_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/scaling_relation/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 6.0\" circular mask, large enough to include all four extra galaxies (the most distant sits at radius ~5\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 6.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres + Luminosities__\n", + "\n", + "Load the centres of the individually-modelled extras and the scaling-tier extras. The scaling-tier galaxies need\n", + "both centres AND a measured luminosity each, so they're loaded from a CSV via `al.galaxy_table_from_csv`.\n", + "\n", + "In a real analysis, the scaling-tier luminosities come from a prior light-only fit \u2014 see\n", + "`scripts/group/features/scaling_relation/modeling_for_luminosities.py` for the standalone version of that fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "individual_extras_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "scaling_table = al.galaxy_table_from_csv(\n", + " file_path=dataset_path / \"scaling_galaxies.csv\"\n", + ")\n", + "scaling_extras_centres = scaling_table.centres\n", + "scaling_extras_luminosities = scaling_table.luminosities\n", + "\n", + "print(f\"Individually-modelled extras centres: {list(individual_extras_centres)}\")\n", + "print(f\"Scaling-tier extras centres: {list(scaling_extras_centres)}\")\n", + "print(f\"Scaling-tier extras luminosities: {scaling_extras_luminosities}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling at every galaxy centre, so each light profile is evaluated accurately at its peak." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_galaxy_centres = (\n", + " [(0.0, 0.0)]\n", + " + [tuple(c) for c in individual_extras_centres]\n", + " + [tuple(c) for c in scaling_extras_centres]\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_galaxy_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MGE Basis__\n", + "\n", + "A `Basis` of 30 linear Gaussians for the source galaxy. The `intensity` of each Gaussian is solved for via linear\n", + "algebra at fit time." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 30\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "source_bulge = build_source_basis(centre=(0.0, 0.1))\n", + "\n", + "plot_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "We compose four populations:\n", + "\n", + " - `lens` (z=0.5): `SersicSph` light + `IsothermalSph` mass at the origin. The simulator's true Einstein radius\n", + " is 1.6\".\n", + " - `individual_extras` (z=0.5): two close companions, each modelled with its own `SersicSph` light +\n", + " `IsothermalSph` mass. Simulator-true Einstein radii: 0.4\" and 0.5\".\n", + " - `scaling_extras` (z=0.5): two further-out, fainter companions whose Einstein radii are derived from the\n", + " scaling relation. With `scaling_factor=0.3` and `scaling_exponent=1.0` and per-galaxy luminosity=0.45, each\n", + " acquires `einstein_radius = 0.3 * 0.45 ** 1.0 = 0.135` \u2014 matches the simulator.\n", + " - `source` (z=1.0): the MGE basis above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=1.5, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6),\n", + ")\n", + "\n", + "individual_extras_truth = [\n", + " dict(intensity=0.9, effective_radius=0.6, sersic_index=2.5, einstein_radius=0.4),\n", + " dict(intensity=0.8, effective_radius=0.6, sersic_index=2.5, einstein_radius=0.5),\n", + "]\n", + "\n", + "individual_extras = []\n", + "for centre, truth in zip(individual_extras_centres, individual_extras_truth):\n", + " individual_extras.append(\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=tuple(centre),\n", + " intensity=truth[\"intensity\"],\n", + " effective_radius=truth[\"effective_radius\"],\n", + " sersic_index=truth[\"sersic_index\"],\n", + " ),\n", + " mass=al.mp.IsothermalSph(\n", + " centre=tuple(centre), einstein_radius=truth[\"einstein_radius\"]\n", + " ),\n", + " )\n", + " )\n", + "\n", + "scaling_factor = 0.3\n", + "scaling_exponent = 1.0\n", + "\n", + "scaling_extras = []\n", + "scaling_extras_einstein_radii = []\n", + "for centre, luminosity in zip(scaling_extras_centres, scaling_extras_luminosities):\n", + " einstein_radius = scaling_factor * luminosity**scaling_exponent\n", + " scaling_extras_einstein_radii.append(einstein_radius)\n", + " scaling_extras.append(\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=tuple(centre),\n", + " intensity=luminosity,\n", + " effective_radius=0.5,\n", + " sersic_index=2.5,\n", + " ),\n", + " mass=al.mp.IsothermalSph(\n", + " centre=tuple(centre), einstein_radius=einstein_radius\n", + " ),\n", + " )\n", + " )\n", + "\n", + "source = al.Galaxy(redshift=1.0, bulge=source_bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "The `Tracer` performs the ray-tracing. It queries every mass profile attached to every galaxy in the lens plane\n", + "and sums their deflections. For our mixed population, this means the lens's `IsothermalSph`, each individually-\n", + "modelled extra's `IsothermalSph`, and each scaling-tier extra's `IsothermalSph` all contribute independent\n", + "deflections that sum before mapping image-plane coordinates onto the source-plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens] + individual_extras + scaling_extras + [source])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "Pass the `Tracer` to a `FitImaging` to fit the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling Relation Tour__\n", + "\n", + "The scaling-tier galaxies' Einstein radii are NOT free parameters in the model \u2014 they're computed from a shared\n", + "two-parameter relation and per-galaxy luminosity. With `scaling_factor=0.3` and `scaling_exponent=1.0`, we have:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for centre, luminosity, er in zip(\n", + " scaling_extras_centres, scaling_extras_luminosities, scaling_extras_einstein_radii\n", + "):\n", + " print(\n", + " f\" scaling galaxy @ {tuple(centre)}: luminosity = {luminosity:.3f}, \"\n", + " f\"einstein_radius = {scaling_factor:.3f} * {luminosity:.3f} ** {scaling_exponent:.1f} = {er:.4f}\"\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The lens-plane total deflection is the SUM of every mass profile's contribution. We verify this by computing\n", + "each one explicitly and confirming the sum equals what the `Tracer` returns." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = dataset.grid\n", + "\n", + "alpha_lens = lens.mass.deflections_yx_2d_from(grid=grid)\n", + "alpha_individual = [g.mass.deflections_yx_2d_from(grid=grid) for g in individual_extras]\n", + "alpha_scaling = [g.mass.deflections_yx_2d_from(grid=grid) for g in scaling_extras]\n", + "\n", + "print(f\"\\nalpha_lens (first coord): {alpha_lens[0]}\")\n", + "print(f\"alpha_individual_extra_0 (first coord): {alpha_individual[0][0]}\")\n", + "print(f\"alpha_individual_extra_1 (first coord): {alpha_individual[1][0]}\")\n", + "print(f\"alpha_scaling_extra_0 (first coord): {alpha_scaling[0][0]}\")\n", + "print(f\"alpha_scaling_extra_1 (first coord): {alpha_scaling[1][0]}\")\n", + "\n", + "alpha_total_summed = alpha_lens + sum(alpha_individual) + sum(alpha_scaling)\n", + "\n", + "traced_grids = tracer.traced_grid_2d_list_from(grid=grid)\n", + "alpha_total_tracer = grid - traced_grids[1]\n", + "\n", + "print(f\"\\nalpha_total (summed by hand, first 3): {alpha_total_summed[:3]}\")\n", + "print(f\"alpha_total (from tracer, first 3): {alpha_total_tracer[:3]}\")\n", + "\n", + "assert np.allclose(np.asarray(alpha_total_summed), np.asarray(alpha_total_tracer))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The scaling-tier deflections per galaxy are visibly smaller than the lens and individually-modelled extras\n", + "because their Einstein radii are an order of magnitude smaller. The shared relation lets us include them in the\n", + "model with zero additional free parameters \u2014 adding more scaling galaxies would not grow the parameter space.\n", + "\n", + "__Intensities__\n", + "\n", + "After the fit, every linear Gaussian in the source MGE basis has been assigned an `intensity` via linear algebra." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"\\nFirst Gaussian intensity, source = \"\n", + " f\"{fit.linear_light_profile_intensity_dict[source_bulge.profile_list[0]]}\"\n", + ")\n", + "\n", + "tracer_fitted = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "subplot_basis_image(basis=tracer_fitted.galaxies[-1].bulge, grid=plot_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script demonstrated the mixed-strategy API for handling many foreground galaxies \u2014 a small number of\n", + "individually-modelled extras for the bright, close companions, and an arbitrary number of scaling-tier extras for\n", + "the long tail of fainter ones. The scaling relation collapses what would otherwise be N free `einstein_radius`\n", + "parameters into 2 shared parameters (`scaling_factor` and `scaling_exponent`), keeping the model dimensionality\n", + "manageable as galaxy count grows.\n", + "\n", + "In a real modeling workflow:\n", + "\n", + " - `modeling.py` runs the search-based version, where `scaling_factor` and `scaling_exponent` are free `af.Model`\n", + " parameters with `UniformPrior`s. The luminosities still come from a prior light-only fit.\n", + " - For group-scale lenses with multiple main lens galaxies, see\n", + " `autolens_workspace/scripts/group/features/scaling_relation/` \u2014 the three-tier API generalises this script to\n", + " `lens_dict + extra_galaxies + scaling_galaxies` collections.\n", + "\n", + "The key takeaway is that scaling relations let the lens model stay tractable even when 10s or 100s of foreground\n", + "galaxies sit on it \u2014 the lens-plane deflection is still a simple sum of per-galaxy contributions, but the\n", + "contributions from the scaling tier are parameterized through luminosity rather than per-galaxy free parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/scaling_relation/likelihood_function.ipynb b/notebooks/imaging/features/scaling_relation/likelihood_function.ipynb index eeae16a70..4adfb93b5 100644 --- a/notebooks/imaging/features/scaling_relation/likelihood_function.ipynb +++ b/notebooks/imaging/features/scaling_relation/likelihood_function.ipynb @@ -1,432 +1,469 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Scaling Relation__\n", - "\n", - "This script describes the additional steps required to compute the `log_likelihood` for a strong lens whose\n", - "foreground galaxy population is split between two tiers \u2014 individually-modelled extras (each with its own free\n", - "`einstein_radius`) and scaling-tier extras (whose Einstein radii are derived from a shared two-parameter\n", - "relation `einstein_radius = scaling_factor * luminosity ** scaling_exponent`).\n", - "\n", - "This script does NOT repeat the steps shared with single-plane lensing (mask, image-plane grid, PSF convolution,\n", - "chi-squared, noise normalization, linear-algebra solver for MGE source intensities). It documents only the part\n", - "of the likelihood function which is specific to a scaling-relation tier: the lens-plane deflection composition.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function below builds directly on the standard imaging and MGE likelihood functions. You should\n", - "read these notebooks first:\n", - "\n", - " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 the canonical single-plane log-likelihood\n", - " walkthrough.\n", - " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a `Basis`\n", - " of linear Gaussians is solved for via linear algebra.\n", - " - `autolens_workspace/scripts/imaging/features/scaling_relation/modeling.py` \u2014 the search-based version of the\n", - " composition demonstrated here.\n", - "\n", - "Sections covered in those scripts (e.g. \"Chi Squared\", \"Noise Normalization Term\", \"Mapping Matrix\") are not\n", - "repeated here.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script (see above).\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Centres + Luminosities:** Load extras and scaling-tier galaxy data from JSON + CSV.\n", - "- **Galaxies:** Lens + individually-modelled extras + scaling-tier extras + MGE source.\n", - "- **Per-Galaxy Deflection:** Per-tier per-galaxy contributions, plus the scaling-relation evaluation.\n", - "- **Manual Ray-Tracing:** Hand-compute the source-plane grid and confirm it matches `Tracer`.\n", - "- **Source-Plane Image:** Source MGE evaluated at the ray-traced grid.\n", - "- **Fit Check:** `FitImaging.log_likelihood`.\n", - "- **Wrap Up.**\n", - "\n", - "__What Changes For A Scaling Relation__\n", - "\n", - "For a single-component lens with one mass profile, the lens-plane deflection is just one profile evaluated at\n", - "each image-plane coordinate:\n", - "\n", - " alpha_lens(theta) = alpha_total(theta ; Isothermal parameters)\n", - "\n", - "For a lens with a population of individually-modelled extras, the lens-plane deflection becomes a sum:\n", - "\n", - " alpha_lens(theta) = alpha_main(theta) + sum_i alpha_extra_individual_i(theta)\n", - "\n", - "For a lens with BOTH tiers active (this script), the deflection sum extends to the scaling-tier extras, but each\n", - "scaling-tier galaxy's Einstein radius is NOT a free parameter \u2014 it is derived from a shared two-parameter\n", - "relation and the galaxy's own luminosity:\n", - "\n", - " alpha_lens(theta) = alpha_main(theta)\n", - " + sum_i alpha_extra_individual_i(theta)\n", - " + sum_j alpha_extra_scaling_j(theta)\n", - "\n", - " where alpha_extra_scaling_j is the deflection of a mass profile whose\n", - " einstein_radius_j = scaling_factor * luminosity_j ** scaling_exponent.\n", - "\n", - "The model gains exactly 2 free parameters (`scaling_factor`, `scaling_exponent`) regardless of how many galaxies\n", - "sit on the scaling-tier. Every other step of the likelihood (PSF convolution, chi-squared, noise normalization,\n", - "MGE linear-algebra solver) is unchanged." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the extra_and_scaling_galaxies dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_and_scaling_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/features/scaling_relation/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "mask_radius = 6.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres + Luminosities__\n", - "\n", - "Load the individually-modelled extras' centres from JSON and the scaling-tier extras' centres + luminosities\n", - "from CSV." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "individual_extras_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "scaling_table = al.galaxy_table_from_csv(\n", - " file_path=dataset_path / \"scaling_galaxies.csv\"\n", - ")\n", - "scaling_extras_centres = scaling_table.centres\n", - "scaling_extras_luminosities = scaling_table.luminosities" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "The three populations that participate in the ray-tracing:\n", - "\n", - " - `lens` (z=0.5): `IsothermalSph` mass at the origin with `einstein_radius=1.6` (simulator truth).\n", - " - `individual_extras` (z=0.5): two `IsothermalSph` masses with simulator-true Einstein radii 0.4 and 0.5.\n", - " - `scaling_extras` (z=0.5): two `IsothermalSph` masses with Einstein radii derived from the scaling relation\n", - " `einstein_radius = 0.3 * luminosity ** 1.0` (simulator truth).\n", - " - `source` (z=1.0): an MGE light component (a simple basis of 10 linear Gaussians)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 10\n", - "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", - "\n", - "\n", - "def build_source_basis(centre):\n", - " gaussian_list = [\n", - " al.lp_linear.Gaussian(\n", - " centre=centre,\n", - " ell_comps=(0.0, 0.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " for i in range(total_gaussians)\n", - " ]\n", - " return al.lp_basis.Basis(profile_list=gaussian_list)\n", - "\n", - "\n", - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6),\n", - ")\n", - "\n", - "individual_extras_einstein_radii = [0.4, 0.5]\n", - "individual_extras = [\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=tuple(centre), einstein_radius=er),\n", - " )\n", - " for centre, er in zip(individual_extras_centres, individual_extras_einstein_radii)\n", - "]\n", - "\n", - "scaling_factor = 0.3\n", - "scaling_exponent = 1.0\n", - "\n", - "scaling_extras = []\n", - "for centre, luminosity in zip(scaling_extras_centres, scaling_extras_luminosities):\n", - " einstein_radius = scaling_factor * luminosity**scaling_exponent\n", - " scaling_extras.append(\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(\n", - " centre=tuple(centre), einstein_radius=einstein_radius\n", - " ),\n", - " )\n", - " )\n", - "\n", - "source = al.Galaxy(redshift=1.0, bulge=build_source_basis(centre=(0.0, 0.1)))\n", - "\n", - "tracer = al.Tracer(galaxies=[lens] + individual_extras + scaling_extras + [source])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Galaxy Deflection__\n", - "\n", - "The `Tracer.traced_grid_2d_list_from(...)` call performs the standard image-plane \u2192 source-plane ray-tracing.\n", - "Internally it queries every mass profile in the lens plane and sums their deflections.\n", - "\n", - "To make the decomposition concrete we re-compute the same source-plane grid by hand. Each profile exposes its\n", - "own `deflections_yx_2d_from`; the SUM of the main lens deflection plus every per-galaxy contribution from BOTH\n", - "tiers is what the tracer applies internally:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_grid = dataset.grid\n", - "\n", - "alpha_lens = lens.mass.deflections_yx_2d_from(grid=masked_grid)\n", - "alpha_individual = [\n", - " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in individual_extras\n", - "]\n", - "alpha_scaling = [\n", - " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in scaling_extras\n", - "]\n", - "\n", - "alpha_total = alpha_lens + sum(alpha_individual) + sum(alpha_scaling)\n", - "\n", - "print(f\"alpha_lens (first coord): {alpha_lens[0]}\")\n", - "print(f\"alpha_individual (tier sum) (first coord): \" f\"{sum(alpha_individual)[0]}\")\n", - "print(f\"alpha_scaling (tier sum) (first coord): \" f\"{sum(alpha_scaling)[0]}\")\n", - "print(f\"alpha_total (across all) (first coord): {alpha_total[0]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The scaling-tier contributions are computed from the scaling relation:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for centre, luminosity in zip(scaling_extras_centres, scaling_extras_luminosities):\n", - " er = scaling_factor * luminosity**scaling_exponent\n", - " print(\n", - " f\" scaling galaxy @ {tuple(centre)}: \"\n", - " f\"einstein_radius = {scaling_factor:.2f} * {luminosity:.3f} ** {scaling_exponent:.1f} = {er:.4f}\"\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Manual Ray-Tracing__\n", - "\n", - "The source-plane grid is the image-plane grid minus the total deflection. We compute it by hand and compare to\n", - "the `Tracer`-produced grid; they should be identical to within floating-point precision." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_source_manual = masked_grid - alpha_total\n", - "\n", - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=masked_grid)\n", - "grid_source_tracer = traced_grid_list[1]\n", - "\n", - "print(f\"\\nsource-plane grid (first coord, manual): {grid_source_manual[0]}\")\n", - "print(f\"source-plane grid (first coord, tracer): {grid_source_tracer[0]}\")\n", - "\n", - "assert np.allclose(np.asarray(grid_source_manual), np.asarray(grid_source_tracer))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source-Plane Image__\n", - "\n", - "The source galaxy's MGE basis is evaluated at the ray-traced source-plane grid, producing image-plane pixel\n", - "values per Gaussian with a placeholder `intensity=1.0`. The true `intensity` of each Gaussian is solved for at\n", - "the linear-algebra step (see the MGE likelihood prerequisite)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image_unconvolved = tracer.image_2d_from(grid=masked_grid)\n", - "\n", - "aplt.plot_array(\n", - " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "What `image_2d_from` does internally for our two-tier extras population:\n", - "\n", - " 1. Computes `alpha_lens(theta) = alpha_main + sum_i alpha_extra_individual_i + sum_j alpha_extra_scaling_j`,\n", - " where each `alpha_extra_scaling_j` is the deflection of a profile whose `einstein_radius` was derived from\n", - " `scaling_factor * luminosity_j ** scaling_exponent`.\n", - " 2. Ray-traces the image-plane grid to obtain `grid_source = grid - alpha_lens`.\n", - " 3. Evaluates the source MGE at `grid_source`, producing its image-plane contribution.\n", - "\n", - "For a single-lens system there is just one mass-profile contributing to step 1; for our mixed-strategy lens\n", - "there are `1 + len(individual_extras) + len(scaling_extras)` contributions, but the model only gains\n", - "`len(individual_extras)` free `einstein_radius` parameters plus 2 shared scaling parameters.\n", - "\n", - "__Likelihood__\n", - "\n", - "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model image\n", - "above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", - "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`. The MGE source's per-Gaussian\n", - "intensities are solved for via the linear-algebra step documented in\n", - "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py`.\n", - "\n", - "We delegate the remaining steps to `FitImaging`, which assembles the full `log_likelihood`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n", - "\n", - "print(f\"\\nLog likelihood of the manual scaling-relation fit: {fit.log_likelihood}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The scaling-relation `log_likelihood` differs from a population-of-individually-modelled-extras case in exactly\n", - "one place: some galaxies' `einstein_radius` values are not free parameters \u2014 they're derived from a shared\n", - "two-parameter relation plus a per-galaxy luminosity. Every other step (ray-tracing, source-plane evaluation, PSF\n", - "convolution, chi-squared, noise normalization, linear algebra) is shared with the standard imaging likelihood\n", - "and documented in the prerequisite scripts.\n", - "\n", - "This is what lets the model dimensionality stay tractable as foreground galaxy count grows: 100 scaling-tier\n", - "galaxies cost the same 2 shared parameters as 2 do." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Scaling Relation__\n", + "\n", + "This script describes the additional steps required to compute the `log_likelihood` for a strong lens whose\n", + "foreground galaxy population is split between two tiers \u2014 individually-modelled extras (each with its own free\n", + "`einstein_radius`) and scaling-tier extras (whose Einstein radii are derived from a shared two-parameter\n", + "relation `einstein_radius = scaling_factor * luminosity ** scaling_exponent`).\n", + "\n", + "This script does NOT repeat the steps shared with single-plane lensing (mask, image-plane grid, PSF convolution,\n", + "chi-squared, noise normalization, linear-algebra solver for MGE source intensities). It documents only the part\n", + "of the likelihood function which is specific to a scaling-relation tier: the lens-plane deflection composition.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function below builds directly on the standard imaging and MGE likelihood functions. You should\n", + "read these notebooks first:\n", + "\n", + " - `autolens_workspace/scripts/imaging/likelihood_function.py` \u2014 the canonical single-plane log-likelihood\n", + " walkthrough.\n", + " - `autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py` \u2014 how a `Basis`\n", + " of linear Gaussians is solved for via linear algebra.\n", + " - `autolens_workspace/scripts/imaging/features/scaling_relation/modeling.py` \u2014 the search-based version of the\n", + " composition demonstrated here.\n", + "\n", + "Sections covered in those scripts (e.g. \"Chi Squared\", \"Noise Normalization Term\", \"Mapping Matrix\") are not\n", + "repeated here.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script (see above).\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Centres + Luminosities:** Load extras and scaling-tier galaxy data from JSON + CSV.\n", + "- **Galaxies:** Lens + individually-modelled extras + scaling-tier extras + MGE source.\n", + "- **Per-Galaxy Deflection:** Per-tier per-galaxy contributions, plus the scaling-relation evaluation.\n", + "- **Manual Ray-Tracing:** Hand-compute the source-plane grid and confirm it matches `Tracer`.\n", + "- **Source-Plane Image:** Source MGE evaluated at the ray-traced grid.\n", + "- **Fit Check:** `FitImaging.log_likelihood`.\n", + "- **Wrap Up.**\n", + "\n", + "__What Changes For A Scaling Relation__\n", + "\n", + "For a single-component lens with one mass profile, the lens-plane deflection is just one profile evaluated at\n", + "each image-plane coordinate:\n", + "\n", + " alpha_lens(theta) = alpha_total(theta ; Isothermal parameters)\n", + "\n", + "For a lens with a population of individually-modelled extras, the lens-plane deflection becomes a sum:\n", + "\n", + " alpha_lens(theta) = alpha_main(theta) + sum_i alpha_extra_individual_i(theta)\n", + "\n", + "For a lens with BOTH tiers active (this script), the deflection sum extends to the scaling-tier extras, but each\n", + "scaling-tier galaxy's Einstein radius is NOT a free parameter \u2014 it is derived from a shared two-parameter\n", + "relation and the galaxy's own luminosity:\n", + "\n", + " alpha_lens(theta) = alpha_main(theta)\n", + " + sum_i alpha_extra_individual_i(theta)\n", + " + sum_j alpha_extra_scaling_j(theta)\n", + "\n", + " where alpha_extra_scaling_j is the deflection of a mass profile whose\n", + " einstein_radius_j = scaling_factor * luminosity_j ** scaling_exponent.\n", + "\n", + "The model gains exactly 2 free parameters (`scaling_factor`, `scaling_exponent`) regardless of how many galaxies\n", + "sit on the scaling-tier. Every other step of the likelihood (PSF convolution, chi-squared, noise normalization,\n", + "MGE linear-algebra solver) is unchanged." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the extra_and_scaling_galaxies dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_and_scaling_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/features/scaling_relation/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask_radius = 6.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres + Luminosities__\n", + "\n", + "Load the individually-modelled extras' centres from JSON and the scaling-tier extras' centres + luminosities\n", + "from CSV." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "individual_extras_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "scaling_table = al.galaxy_table_from_csv(\n", + " file_path=dataset_path / \"scaling_galaxies.csv\"\n", + ")\n", + "scaling_extras_centres = scaling_table.centres\n", + "scaling_extras_luminosities = scaling_table.luminosities" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "The three populations that participate in the ray-tracing:\n", + "\n", + " - `lens` (z=0.5): `IsothermalSph` mass at the origin with `einstein_radius=1.6` (simulator truth).\n", + " - `individual_extras` (z=0.5): two `IsothermalSph` masses with simulator-true Einstein radii 0.4 and 0.5.\n", + " - `scaling_extras` (z=0.5): two `IsothermalSph` masses with Einstein radii derived from the scaling relation\n", + " `einstein_radius = 0.3 * luminosity ** 1.0` (simulator truth).\n", + " - `source` (z=1.0): an MGE light component (a simple basis of 10 linear Gaussians)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 10\n", + "log10_sigma_list = np.linspace(-2, np.log10(0.5), total_gaussians)\n", + "\n", + "\n", + "def build_source_basis(centre):\n", + " gaussian_list = [\n", + " al.lp_linear.Gaussian(\n", + " centre=centre,\n", + " ell_comps=(0.0, 0.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " for i in range(total_gaussians)\n", + " ]\n", + " return al.lp_basis.Basis(profile_list=gaussian_list)\n", + "\n", + "\n", + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6),\n", + ")\n", + "\n", + "individual_extras_einstein_radii = [0.4, 0.5]\n", + "individual_extras = [\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=tuple(centre), einstein_radius=er),\n", + " )\n", + " for centre, er in zip(individual_extras_centres, individual_extras_einstein_radii)\n", + "]\n", + "\n", + "scaling_factor = 0.3\n", + "scaling_exponent = 1.0\n", + "\n", + "scaling_extras = []\n", + "for centre, luminosity in zip(scaling_extras_centres, scaling_extras_luminosities):\n", + " einstein_radius = scaling_factor * luminosity**scaling_exponent\n", + " scaling_extras.append(\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(\n", + " centre=tuple(centre), einstein_radius=einstein_radius\n", + " ),\n", + " )\n", + " )\n", + "\n", + "source = al.Galaxy(redshift=1.0, bulge=build_source_basis(centre=(0.0, 0.1)))\n", + "\n", + "tracer = al.Tracer(galaxies=[lens] + individual_extras + scaling_extras + [source])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Galaxy Deflection__\n", + "\n", + "The `Tracer.traced_grid_2d_list_from(...)` call performs the standard image-plane \u2192 source-plane ray-tracing.\n", + "Internally it queries every mass profile in the lens plane and sums their deflections.\n", + "\n", + "To make the decomposition concrete we re-compute the same source-plane grid by hand. Each profile exposes its\n", + "own `deflections_yx_2d_from`; the SUM of the main lens deflection plus every per-galaxy contribution from BOTH\n", + "tiers is what the tracer applies internally:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_grid = dataset.grid\n", + "\n", + "alpha_lens = lens.mass.deflections_yx_2d_from(grid=masked_grid)\n", + "alpha_individual = [\n", + " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in individual_extras\n", + "]\n", + "alpha_scaling = [\n", + " g.mass.deflections_yx_2d_from(grid=masked_grid) for g in scaling_extras\n", + "]\n", + "\n", + "alpha_total = alpha_lens + sum(alpha_individual) + sum(alpha_scaling)\n", + "\n", + "print(f\"alpha_lens (first coord): {alpha_lens[0]}\")\n", + "print(f\"alpha_individual (tier sum) (first coord): \" f\"{sum(alpha_individual)[0]}\")\n", + "print(f\"alpha_scaling (tier sum) (first coord): \" f\"{sum(alpha_scaling)[0]}\")\n", + "print(f\"alpha_total (across all) (first coord): {alpha_total[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The scaling-tier contributions are computed from the scaling relation:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for centre, luminosity in zip(scaling_extras_centres, scaling_extras_luminosities):\n", + " er = scaling_factor * luminosity**scaling_exponent\n", + " print(\n", + " f\" scaling galaxy @ {tuple(centre)}: \"\n", + " f\"einstein_radius = {scaling_factor:.2f} * {luminosity:.3f} ** {scaling_exponent:.1f} = {er:.4f}\"\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Manual Ray-Tracing__\n", + "\n", + "The source-plane grid is the image-plane grid minus the total deflection. We compute it by hand and compare to\n", + "the `Tracer`-produced grid; they should be identical to within floating-point precision." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_source_manual = masked_grid - alpha_total\n", + "\n", + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=masked_grid)\n", + "grid_source_tracer = traced_grid_list[1]\n", + "\n", + "print(f\"\\nsource-plane grid (first coord, manual): {grid_source_manual[0]}\")\n", + "print(f\"source-plane grid (first coord, tracer): {grid_source_tracer[0]}\")\n", + "\n", + "assert np.allclose(np.asarray(grid_source_manual), np.asarray(grid_source_tracer))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source-Plane Image__\n", + "\n", + "The source galaxy's MGE basis is evaluated at the ray-traced source-plane grid, producing image-plane pixel\n", + "values per Gaussian with a placeholder `intensity=1.0`. The true `intensity` of each Gaussian is solved for at\n", + "the linear-algebra step (see the MGE likelihood prerequisite)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image_unconvolved = tracer.image_2d_from(grid=masked_grid)\n", + "\n", + "aplt.plot_array(\n", + " array=model_image_unconvolved, title=\"Model image before PSF convolution\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "What `image_2d_from` does internally for our two-tier extras population:\n", + "\n", + " 1. Computes `alpha_lens(theta) = alpha_main + sum_i alpha_extra_individual_i + sum_j alpha_extra_scaling_j`,\n", + " where each `alpha_extra_scaling_j` is the deflection of a profile whose `einstein_radius` was derived from\n", + " `scaling_factor * luminosity_j ** scaling_exponent`.\n", + " 2. Ray-traces the image-plane grid to obtain `grid_source = grid - alpha_lens`.\n", + " 3. Evaluates the source MGE at `grid_source`, producing its image-plane contribution.\n", + "\n", + "For a single-lens system there is just one mass-profile contributing to step 1; for our mixed-strategy lens\n", + "there are `1 + len(individual_extras) + len(scaling_extras)` contributions, but the model only gains\n", + "`len(individual_extras)` free `einstein_radius` parameters plus 2 shared scaling parameters.\n", + "\n", + "__Likelihood__\n", + "\n", + "PSF convolution and chi-squared / noise normalization are unchanged from the single-plane case. The model image\n", + "above is convolved with the PSF and compared to the data via the standard imaging chi-squared expression\n", + "documented in `autolens_workspace/scripts/imaging/likelihood_function.py`. The MGE source's per-Gaussian\n", + "intensities are solved for via the linear-algebra step documented in\n", + "`autolens_workspace/scripts/imaging/features/multi_gaussian_expansion/likelihood_function.py`.\n", + "\n", + "We delegate the remaining steps to `FitImaging`, which assembles the full `log_likelihood`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n", + "\n", + "print(f\"\\nLog likelihood of the manual scaling-relation fit: {fit.log_likelihood}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The scaling-relation `log_likelihood` differs from a population-of-individually-modelled-extras case in exactly\n", + "one place: some galaxies' `einstein_radius` values are not free parameters \u2014 they're derived from a shared\n", + "two-parameter relation plus a per-galaxy luminosity. Every other step (ray-tracing, source-plane evaluation, PSF\n", + "convolution, chi-squared, noise normalization, linear algebra) is shared with the standard imaging likelihood\n", + "and documented in the prerequisite scripts.\n", + "\n", + "This is what lets the model dimensionality stay tractable as foreground galaxy count grows: 100 scaling-tier\n", + "galaxies cost the same 2 shared parameters as 2 do." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/scaling_relation/modeling.ipynb b/notebooks/imaging/features/scaling_relation/modeling.ipynb index 39d684031..b8f0443b8 100644 --- a/notebooks/imaging/features/scaling_relation/modeling.ipynb +++ b/notebooks/imaging/features/scaling_relation/modeling.ipynb @@ -1,596 +1,633 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Scaling Relations\n", - "====================================\n", - "\n", - "Strong lenses often have many galaxies surrounding the lens galaxy whose mass contributes to the ray-tracing of the\n", - "source. The `extra_galaxies` example shows how to add each of these galaxies into the model with their own light and\n", - "mass profiles, fixing their centres to the observed centres of light. That works well when only a handful of extra\n", - "galaxies are present, but rapidly becomes unwieldy as the number grows. With 10 extra galaxies modelled with\n", - "`IsothermalSph` mass profiles, the lens model gains 10 additional `einstein_radius` free parameters; with 30, the\n", - "parameter space is so large that the non-linear search struggles to converge and the data lacks the information content\n", - "to constrain every galaxy individually.\n", - "\n", - "A common solution is to model the lensing contribution of these galaxies via a **scaling relation**. An easier-to-measure\n", - "property of each galaxy (typically luminosity, but it could also be stellar mass or velocity dispersion) is related to\n", - "its mass profile via a small number of shared free parameters:\n", - "\n", - " einstein_radius = scaling_factor * (luminosity ** scaling_exponent)\n", - "\n", - "The free parameters are now `scaling_factor` and `scaling_exponent` only \u2014 two parameters total, regardless of how many\n", - "galaxies sit on the relation. The luminosities act as priors on the masses, ensuring each galaxy's contribution stays\n", - "physically reasonable.\n", - "\n", - "This example demonstrates the **mixed-strategy** pattern: a single `extra_galaxies` collection that contains BOTH\n", - "galaxies modelled individually (each with its own free Einstein radius) AND galaxies on a shared scaling relation. This\n", - "is the typical real-world configuration: the brighter / closer companions get individual mass parameters because they\n", - "contribute non-trivially to the lensing on their own, while the long tail of fainter companions sit on the relation.\n", - "\n", - "The dataset used here is `dataset/imaging/extra_and_scaling_galaxies`, simulated by the paired script\n", - "`scripts/imaging/features/scaling_relation/simulator.py`. It contains a galaxy-scale lens at the origin, two close\n", - "companions (the \"individual\" tier here), and two fainter further-out companions (the \"scaling-relation\" tier). All four\n", - "companions live in the same `extra_galaxies` collection in this imaging-context example \u2014 the terminology\n", - "`scaling_galaxies` for a separate top-level collection is reserved for the group-scale example.\n", - "\n", - "__Contents__\n", - "\n", - "- **Two Strategies, One Collection:** Why mix individual + relational extras in the same `extra_galaxies` collection.\n", - "- **Centres:** Two JSON files load the centres of the individually-modelled extras and the scaling-relation extras.\n", - "- **Luminosities:** The scaling-relation tier needs a measured luminosity per galaxy.\n", - "- **Where do luminosities come from?:** The `modeling_for_luminosities.py` example and the SLAM `source_lp[0]` step.\n", - "- **Redshifts:** All foreground galaxies are at the same redshift as the lens galaxy.\n", - "- **Group vs Imaging:** Where the group-scale variant lives.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask.\n", - "- **Lens & Source:** MGE bulge + Isothermal mass + ExternalShear lens; MGE source.\n", - "- **Individually-Modelled Extras:** Bounded `UniformPrior` on `einstein_radius` per galaxy.\n", - "- **Scaling-Relation Extras:** Shared `scaling_factor` and `scaling_exponent` priors.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Over Sampling:** Adaptive over-sampling at every galaxy centre.\n", - "- **Search and Analysis:** Configure the non-linear search and run the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Two Strategies, One Collection__\n", - "\n", - "There is no architectural distinction in PyAutoLens between \"individual\" and \"scaling-relation\" extra galaxies \u2014 both\n", - "sit in the same `extra_galaxies = af.Collection([...])`. The distinction is purely in how the per-galaxy `Galaxy` model\n", - "is built:\n", - "\n", - " - For an individually-modelled galaxy: `mass.einstein_radius = af.UniformPrior(...)` \u2014 one free parameter per galaxy.\n", - " - For a scaling-relation galaxy: `mass.einstein_radius = scaling_factor * luminosity ** scaling_exponent` \u2014 zero new\n", - " free parameters per galaxy, because the relation parameters are shared across the whole tier.\n", - "\n", - "The two strategies coexist freely in the same collection. This script builds a Python list, pushes the individually-\n", - "modelled galaxies onto it first, then the relational galaxies, and wraps the whole thing in a single `af.Collection`.\n", - "\n", - "__Where do luminosities come from?__\n", - "\n", - "In a real analysis the luminosities used by the scaling relation are not known a priori \u2014 they have to be measured\n", - "from the data itself. Two production patterns:\n", - "\n", - " - **Standalone light-only fit.** Run a single non-linear search whose model is just MGE bulges for every galaxy\n", - " (no mass, no source). After the fit, compute total luminosity per galaxy from the bulge gaussian intensities:\n", - " `total_luminosity = sum(2 * pi * sigma**2 / axis_ratio * intensity) / pixel_scale**2`. Feed those numbers into the\n", - " scaling-relation model below. See `scripts/group/features/scaling_relation/modeling_for_luminosities.py` for a\n", - " worked example.\n", - "\n", - " - **As the first stage of a SLaM pipeline.** The Source-Light-Mass (SLaM) pipelines define a `source_lp[0]` stage\n", - " whose only job is to fit a light-only MGE model to the lens, extras and scaling galaxies. The next stage chains\n", - " from that result to compute luminosities and bound / scale the mass models. See `scripts/group/slam.py`,\n", - " `scripts/group/features/pixelization/slam.py`, and the other group `slam.py` variants for production examples.\n", - "\n", - "This tutorial loads the luminosities from a `scaling_galaxies.csv` written by the simulator (see\n", - "`al.galaxy_table_from_csv` further down). In a real analysis the same CSV would be the *output* of one of the patterns\n", - "above \u2014 `modeling_for_luminosities.py` already writes its result in this format, and the SLAM `source_lp[0]` stage can\n", - "similarly emit one.\n", - "\n", - "__Redshifts__\n", - "\n", - "In this example all foreground galaxies are at the same redshift as the lens galaxy, meaning multi-plane lensing is not\n", - "used. To enable multi-plane lensing, define per-galaxy redshifts and pass them when constructing each\n", - "`af.Model(al.Galaxy, ...)`.\n", - "\n", - "__Group vs Imaging__\n", - "\n", - "This is the **imaging-context** example: there is a single main lens galaxy and all companions live in a single\n", - "`extra_galaxies` collection. For the group-scale variant \u2014 multiple \"main\" lens galaxies AND a top-level\n", - "`scaling_galaxies` collection separate from `extra_galaxies` \u2014 see\n", - "`autolens_workspace/scripts/group/features/scaling_relation/modeling.py`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "We use the `dataset/imaging/extra_and_scaling_galaxies` dataset, which contains:\n", - "\n", - " - a galaxy-scale lens at the origin\n", - " - two close companions (the \"individual\" tier here)\n", - " - two fainter further-out companions (the \"scaling-relation\" tier)\n", - "\n", - "The simulator at `scripts/imaging/features/scaling_relation/simulator.py` writes two centre JSON files\n", - "(`extra_galaxies_centres.json` and `scaling_galaxies_centres.json`), one per modeling strategy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_and_scaling_galaxies\"\n", - "dataset_path = Path(\"dataset\", \"imaging\", dataset_name)\n", - "\n", - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/scaling_relation/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "A 6.0\" circular mask, large enough to enclose the lens, the close companions, and the further-out scaling galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 6.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres__\n", - "\n", - "The individually-modelled tier loads its centres from a JSON file (a list of (y, x) tuples) \u2014 the centres are the only\n", - "input the modeling script needs for that tier." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "individual_extras_centres = al.from_json(\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", - ")\n", - "\n", - "print(f\"Individually-modelled extras: {individual_extras_centres}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centres + Luminosities (scaling-relation tier)__\n", - "\n", - "The scaling-relation tier needs both centres AND a measured luminosity per galaxy. There are two equally-supported\n", - "ways to provide them in PyAutoLens \u2014 both shown below so you can pick whichever fits your workflow.\n", - "\n", - "**Option A \u2014 CSV via `al.galaxy_table_from_csv` (recommended for non-trivial galaxy counts).** The simulator writes a\n", - "`scaling_galaxies.csv` with columns `y, x, luminosity` (and optional `redshift`) alongside the centre JSONs. We load it\n", - "in one call which returns a typed `GalaxyTable` with `.centres` (a `Grid2DIrregular`), `.luminosities`, and (optionally)\n", - "`.redshifts`. This scales naturally to populations of tens or hundreds of galaxies \u2014 the source of truth lives in a\n", - "single editable file.\n", - "\n", - "In a real analysis, a prior light-only fit produces this CSV \u2014 see\n", - "`scripts/group/features/scaling_relation/modeling_for_luminosities.py` for the standalone version of that fit, or the\n", - "SLAM `source_lp[0]` stage in `scripts/group/slam.py` for the chained-pipeline equivalent.\n", - "\n", - "**Option B \u2014 JSON centres + hardcoded luminosity list (the original API, fine for short, fixed-length tutorials).**\n", - "Load the centres from `scaling_galaxies_centres.json` with `al.from_json` (the same loader used for the\n", - "individually-modelled tier above) and define the luminosities as a Python list. Concise and obvious for small\n", - "populations; awkward once you have more than a handful.\n", - "\n", - "We use Option A by default below. The Option B equivalent is shown commented out \u2014 uncomment it (and comment out\n", - "Option A) to switch." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Option A: CSV (recommended)\n", - "relational_extras_table = al.galaxy_table_from_csv(\n", - " file_path=dataset_path / \"scaling_galaxies.csv\"\n", - ")\n", - "relational_extras_centres = relational_extras_table.centres\n", - "relational_extras_luminosity_list = relational_extras_table.luminosities\n", - "\n", - "# Option B: JSON centres + hardcoded luminosities (uncomment to use instead)\n", - "# relational_extras_centres = al.from_json(\n", - "# file_path=dataset_path / \"scaling_galaxies_centres.json\"\n", - "# )\n", - "# relational_extras_luminosity_list = [0.45, 0.45]\n", - "# assert len(relational_extras_luminosity_list) == len(list(relational_extras_centres)), (\n", - "# \"Number of luminosities must match number of scaling-relation extra galaxy centres.\"\n", - "# )\n", - "\n", - "print(f\"Scaling-relation extras: {relational_extras_centres}\")\n", - "print(f\"Scaling-relation luminosities: {relational_extras_luminosity_list}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens__\n", - "\n", - "Standard MGE bulge + `Isothermal` mass + `ExternalShear`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", - ")\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Individually-Modelled Extras__\n", - "\n", - "The first tier inside `extra_galaxies`. Each galaxy gets:\n", - "\n", - " - an MGE bulge with `centre_fixed` (the light is fit but the centre is pinned)\n", - " - an `IsothermalSph` mass with bounded uniform-prior `einstein_radius`\n", - "\n", - "Each adds 1 free Einstein-radius parameter to the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_list = []\n", - "\n", - "for centre in individual_extras_centres:\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=10, centre_fixed=tuple(centre)\n", - " )\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = tuple(centre)\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=1.5)\n", - "\n", - " extra_galaxies_list.append(\n", - " af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling-Relation Extras__\n", - "\n", - "The second tier inside `extra_galaxies`. The two relation priors are defined ONCE outside the loop, so every galaxy\n", - "in this tier shares them. Adding more galaxies to this tier does not add free parameters.\n", - "\n", - "For each galaxy:\n", - "\n", - " - an MGE bulge with `centre_fixed`\n", - " - an `Isothermal` mass with `einstein_radius = scaling_factor * luminosity ** scaling_exponent`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "scaling_exponent = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", - "\n", - "for relational_centre, relational_luminosity in zip(\n", - " relational_extras_centres, relational_extras_luminosity_list\n", - "):\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=10,\n", - " centre_fixed=tuple(relational_centre),\n", - " )\n", - "\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = tuple(relational_centre)\n", - " mass.einstein_radius = scaling_factor * relational_luminosity**scaling_exponent\n", - "\n", - " extra_galaxies_list.append(\n", - " af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", - " )\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "Two top-level components: `galaxies` (lens + source) and `extra_galaxies` (the mixed individual + relational tier).\n", - "Keeping all extras in one collection matches the `features/extra_galaxies` naming convention while still letting us\n", - "mix the two strategies internally." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source=source),\n", - " extra_galaxies=extra_galaxies,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `model.info` attribute prints the composed model. Notice that the first two extras have independent\n", - "`einstein_radius` priors, while the last two share `scaling_factor` and `scaling_exponent` \u2014 the relation in action." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling at every galaxy centre \u2014 lens, individually-modelled extras, and scaling-relation extras alike." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "all_centres = (\n", - " [(0.0, 0.0)] + list(individual_extras_centres) + list(relational_extras_centres)\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_centres,\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\") / \"features\",\n", - " name=\"scaling_relation\",\n", - " unique_tag=dataset_name,\n", - " n_live=200,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The mixed-strategy model adds a small per-galaxy likelihood overhead but keeps the parameter space compact: only 2\n", - "extra parameters from the individually-modelled tier (one Einstein radius each) plus 2 shared parameters from the\n", - "scaling-relation tier, no matter how many galaxies sit on it.\n", - "\n", - "GPU log-likelihood evaluation is < 0.005 s per call; CPU is < 0.05 s. Expected end-to-end run time is ~15 minutes on\n", - "GPU, ~30 minutes on CPU.\n", - "\n", - "__Model Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This example showed how to mix two strategies for `extra_galaxies` modeling \u2014 individually-modelled and on a shared\n", - "scaling relation \u2014 within a single `extra_galaxies` collection. The same pattern works with any mass profile and any\n", - "measured property (swap the `Isothermal` for `PowerLaw`, or the luminosity for stellar mass, and the structure is\n", - "unchanged).\n", - "\n", - "For the production-style luminosity-fitting workflow that produces the `relational_extras_luminosity_list` used here,\n", - "see:\n", - "\n", - " - `autolens_workspace/scripts/group/features/scaling_relation/modeling_for_luminosities.py` \u2014 a standalone light-only\n", - " fit that produces per-galaxy total luminosities.\n", - " - `autolens_workspace/scripts/group/slam.py` and `autolens_workspace/scripts/group/features/pixelization/slam.py` \u2014\n", - " the SLAM `source_lp[0]` stage that does the same job inside a chained pipeline.\n", - "\n", - "For the group-scale variant \u2014 multiple \"main\" lens galaxies AND a top-level `scaling_galaxies` collection separate from\n", - "`extra_galaxies` \u2014 see `autolens_workspace/scripts/group/features/scaling_relation/modeling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Scaling Relations\n", + "====================================\n", + "\n", + "Strong lenses often have many galaxies surrounding the lens galaxy whose mass contributes to the ray-tracing of the\n", + "source. The `extra_galaxies` example shows how to add each of these galaxies into the model with their own light and\n", + "mass profiles, fixing their centres to the observed centres of light. That works well when only a handful of extra\n", + "galaxies are present, but rapidly becomes unwieldy as the number grows. With 10 extra galaxies modelled with\n", + "`IsothermalSph` mass profiles, the lens model gains 10 additional `einstein_radius` free parameters; with 30, the\n", + "parameter space is so large that the non-linear search struggles to converge and the data lacks the information content\n", + "to constrain every galaxy individually.\n", + "\n", + "A common solution is to model the lensing contribution of these galaxies via a **scaling relation**. An easier-to-measure\n", + "property of each galaxy (typically luminosity, but it could also be stellar mass or velocity dispersion) is related to\n", + "its mass profile via a small number of shared free parameters:\n", + "\n", + " einstein_radius = scaling_factor * (luminosity ** scaling_exponent)\n", + "\n", + "The free parameters are now `scaling_factor` and `scaling_exponent` only \u2014 two parameters total, regardless of how many\n", + "galaxies sit on the relation. The luminosities act as priors on the masses, ensuring each galaxy's contribution stays\n", + "physically reasonable.\n", + "\n", + "This example demonstrates the **mixed-strategy** pattern: a single `extra_galaxies` collection that contains BOTH\n", + "galaxies modelled individually (each with its own free Einstein radius) AND galaxies on a shared scaling relation. This\n", + "is the typical real-world configuration: the brighter / closer companions get individual mass parameters because they\n", + "contribute non-trivially to the lensing on their own, while the long tail of fainter companions sit on the relation.\n", + "\n", + "The dataset used here is `dataset/imaging/extra_and_scaling_galaxies`, simulated by the paired script\n", + "`scripts/imaging/features/scaling_relation/simulator.py`. It contains a galaxy-scale lens at the origin, two close\n", + "companions (the \"individual\" tier here), and two fainter further-out companions (the \"scaling-relation\" tier). All four\n", + "companions live in the same `extra_galaxies` collection in this imaging-context example \u2014 the terminology\n", + "`scaling_galaxies` for a separate top-level collection is reserved for the group-scale example.\n", + "\n", + "__Contents__\n", + "\n", + "- **Two Strategies, One Collection:** Why mix individual + relational extras in the same `extra_galaxies` collection.\n", + "- **Centres:** Two JSON files load the centres of the individually-modelled extras and the scaling-relation extras.\n", + "- **Luminosities:** The scaling-relation tier needs a measured luminosity per galaxy.\n", + "- **Where do luminosities come from?:** The `modeling_for_luminosities.py` example and the SLAM `source_lp[0]` step.\n", + "- **Redshifts:** All foreground galaxies are at the same redshift as the lens galaxy.\n", + "- **Group vs Imaging:** Where the group-scale variant lives.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask.\n", + "- **Lens & Source:** MGE bulge + Isothermal mass + ExternalShear lens; MGE source.\n", + "- **Individually-Modelled Extras:** Bounded `UniformPrior` on `einstein_radius` per galaxy.\n", + "- **Scaling-Relation Extras:** Shared `scaling_factor` and `scaling_exponent` priors.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Over Sampling:** Adaptive over-sampling at every galaxy centre.\n", + "- **Search and Analysis:** Configure the non-linear search and run the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Two Strategies, One Collection__\n", + "\n", + "There is no architectural distinction in PyAutoLens between \"individual\" and \"scaling-relation\" extra galaxies \u2014 both\n", + "sit in the same `extra_galaxies = af.Collection([...])`. The distinction is purely in how the per-galaxy `Galaxy` model\n", + "is built:\n", + "\n", + " - For an individually-modelled galaxy: `mass.einstein_radius = af.UniformPrior(...)` \u2014 one free parameter per galaxy.\n", + " - For a scaling-relation galaxy: `mass.einstein_radius = scaling_factor * luminosity ** scaling_exponent` \u2014 zero new\n", + " free parameters per galaxy, because the relation parameters are shared across the whole tier.\n", + "\n", + "The two strategies coexist freely in the same collection. This script builds a Python list, pushes the individually-\n", + "modelled galaxies onto it first, then the relational galaxies, and wraps the whole thing in a single `af.Collection`.\n", + "\n", + "__Where do luminosities come from?__\n", + "\n", + "In a real analysis the luminosities used by the scaling relation are not known a priori \u2014 they have to be measured\n", + "from the data itself. Two production patterns:\n", + "\n", + " - **Standalone light-only fit.** Run a single non-linear search whose model is just MGE bulges for every galaxy\n", + " (no mass, no source). After the fit, compute total luminosity per galaxy from the bulge gaussian intensities:\n", + " `total_luminosity = sum(2 * pi * sigma**2 / axis_ratio * intensity) / pixel_scale**2`. Feed those numbers into the\n", + " scaling-relation model below. See `scripts/group/features/scaling_relation/modeling_for_luminosities.py` for a\n", + " worked example.\n", + "\n", + " - **As the first stage of a SLaM pipeline.** The Source-Light-Mass (SLaM) pipelines define a `source_lp[0]` stage\n", + " whose only job is to fit a light-only MGE model to the lens, extras and scaling galaxies. The next stage chains\n", + " from that result to compute luminosities and bound / scale the mass models. See `scripts/group/slam.py`,\n", + " `scripts/group/features/pixelization/slam.py`, and the other group `slam.py` variants for production examples.\n", + "\n", + "This tutorial loads the luminosities from a `scaling_galaxies.csv` written by the simulator (see\n", + "`al.galaxy_table_from_csv` further down). In a real analysis the same CSV would be the *output* of one of the patterns\n", + "above \u2014 `modeling_for_luminosities.py` already writes its result in this format, and the SLAM `source_lp[0]` stage can\n", + "similarly emit one.\n", + "\n", + "__Redshifts__\n", + "\n", + "In this example all foreground galaxies are at the same redshift as the lens galaxy, meaning multi-plane lensing is not\n", + "used. To enable multi-plane lensing, define per-galaxy redshifts and pass them when constructing each\n", + "`af.Model(al.Galaxy, ...)`.\n", + "\n", + "__Group vs Imaging__\n", + "\n", + "This is the **imaging-context** example: there is a single main lens galaxy and all companions live in a single\n", + "`extra_galaxies` collection. For the group-scale variant \u2014 multiple \"main\" lens galaxies AND a top-level\n", + "`scaling_galaxies` collection separate from `extra_galaxies` \u2014 see\n", + "`autolens_workspace/scripts/group/features/scaling_relation/modeling.py`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "We use the `dataset/imaging/extra_and_scaling_galaxies` dataset, which contains:\n", + "\n", + " - a galaxy-scale lens at the origin\n", + " - two close companions (the \"individual\" tier here)\n", + " - two fainter further-out companions (the \"scaling-relation\" tier)\n", + "\n", + "The simulator at `scripts/imaging/features/scaling_relation/simulator.py` writes two centre JSON files\n", + "(`extra_galaxies_centres.json` and `scaling_galaxies_centres.json`), one per modeling strategy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_and_scaling_galaxies\"\n", + "dataset_path = Path(\"dataset\", \"imaging\", dataset_name)\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/scaling_relation/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "A 6.0\" circular mask, large enough to enclose the lens, the close companions, and the further-out scaling galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 6.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres__\n", + "\n", + "The individually-modelled tier loads its centres from a JSON file (a list of (y, x) tuples) \u2014 the centres are the only\n", + "input the modeling script needs for that tier." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "individual_extras_centres = al.from_json(\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\"\n", + ")\n", + "\n", + "print(f\"Individually-modelled extras: {individual_extras_centres}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centres + Luminosities (scaling-relation tier)__\n", + "\n", + "The scaling-relation tier needs both centres AND a measured luminosity per galaxy. There are two equally-supported\n", + "ways to provide them in PyAutoLens \u2014 both shown below so you can pick whichever fits your workflow.\n", + "\n", + "**Option A \u2014 CSV via `al.galaxy_table_from_csv` (recommended for non-trivial galaxy counts).** The simulator writes a\n", + "`scaling_galaxies.csv` with columns `y, x, luminosity` (and optional `redshift`) alongside the centre JSONs. We load it\n", + "in one call which returns a typed `GalaxyTable` with `.centres` (a `Grid2DIrregular`), `.luminosities`, and (optionally)\n", + "`.redshifts`. This scales naturally to populations of tens or hundreds of galaxies \u2014 the source of truth lives in a\n", + "single editable file.\n", + "\n", + "In a real analysis, a prior light-only fit produces this CSV \u2014 see\n", + "`scripts/group/features/scaling_relation/modeling_for_luminosities.py` for the standalone version of that fit, or the\n", + "SLAM `source_lp[0]` stage in `scripts/group/slam.py` for the chained-pipeline equivalent.\n", + "\n", + "**Option B \u2014 JSON centres + hardcoded luminosity list (the original API, fine for short, fixed-length tutorials).**\n", + "Load the centres from `scaling_galaxies_centres.json` with `al.from_json` (the same loader used for the\n", + "individually-modelled tier above) and define the luminosities as a Python list. Concise and obvious for small\n", + "populations; awkward once you have more than a handful.\n", + "\n", + "We use Option A by default below. The Option B equivalent is shown commented out \u2014 uncomment it (and comment out\n", + "Option A) to switch." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Option A: CSV (recommended)\n", + "relational_extras_table = al.galaxy_table_from_csv(\n", + " file_path=dataset_path / \"scaling_galaxies.csv\"\n", + ")\n", + "relational_extras_centres = relational_extras_table.centres\n", + "relational_extras_luminosity_list = relational_extras_table.luminosities\n", + "\n", + "# Option B: JSON centres + hardcoded luminosities (uncomment to use instead)\n", + "# relational_extras_centres = al.from_json(\n", + "# file_path=dataset_path / \"scaling_galaxies_centres.json\"\n", + "# )\n", + "# relational_extras_luminosity_list = [0.45, 0.45]\n", + "# assert len(relational_extras_luminosity_list) == len(list(relational_extras_centres)), (\n", + "# \"Number of luminosities must match number of scaling-relation extra galaxy centres.\"\n", + "# )\n", + "\n", + "print(f\"Scaling-relation extras: {relational_extras_centres}\")\n", + "print(f\"Scaling-relation luminosities: {relational_extras_luminosity_list}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens__\n", + "\n", + "Standard MGE bulge + `Isothermal` mass + `ExternalShear`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=True\n", + ")\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Individually-Modelled Extras__\n", + "\n", + "The first tier inside `extra_galaxies`. Each galaxy gets:\n", + "\n", + " - an MGE bulge with `centre_fixed` (the light is fit but the centre is pinned)\n", + " - an `IsothermalSph` mass with bounded uniform-prior `einstein_radius`\n", + "\n", + "Each adds 1 free Einstein-radius parameter to the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_list = []\n", + "\n", + "for centre in individual_extras_centres:\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=10, centre_fixed=tuple(centre)\n", + " )\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = tuple(centre)\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=1.5)\n", + "\n", + " extra_galaxies_list.append(\n", + " af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling-Relation Extras__\n", + "\n", + "The second tier inside `extra_galaxies`. The two relation priors are defined ONCE outside the loop, so every galaxy\n", + "in this tier shares them. Adding more galaxies to this tier does not add free parameters.\n", + "\n", + "For each galaxy:\n", + "\n", + " - an MGE bulge with `centre_fixed`\n", + " - an `Isothermal` mass with `einstein_radius = scaling_factor * luminosity ** scaling_exponent`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scaling_factor = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "scaling_exponent = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)\n", + "\n", + "for relational_centre, relational_luminosity in zip(\n", + " relational_extras_centres, relational_extras_luminosity_list\n", + "):\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=10,\n", + " centre_fixed=tuple(relational_centre),\n", + " )\n", + "\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = tuple(relational_centre)\n", + " mass.einstein_radius = scaling_factor * relational_luminosity**scaling_exponent\n", + "\n", + " extra_galaxies_list.append(\n", + " af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass)\n", + " )\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "Two top-level components: `galaxies` (lens + source) and `extra_galaxies` (the mixed individual + relational tier).\n", + "Keeping all extras in one collection matches the `features/extra_galaxies` naming convention while still letting us\n", + "mix the two strategies internally." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source=source),\n", + " extra_galaxies=extra_galaxies,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `model.info` attribute prints the composed model. Notice that the first two extras have independent\n", + "`einstein_radius` priors, while the last two share `scaling_factor` and `scaling_exponent` \u2014 the relation in action." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling at every galaxy centre \u2014 lens, individually-modelled extras, and scaling-relation extras alike." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "all_centres = (\n", + " [(0.0, 0.0)] + list(individual_extras_centres) + list(relational_extras_centres)\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_centres,\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\") / \"features\",\n", + " name=\"scaling_relation\",\n", + " unique_tag=dataset_name,\n", + " n_live=200,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The mixed-strategy model adds a small per-galaxy likelihood overhead but keeps the parameter space compact: only 2\n", + "extra parameters from the individually-modelled tier (one Einstein radius each) plus 2 shared parameters from the\n", + "scaling-relation tier, no matter how many galaxies sit on it.\n", + "\n", + "GPU log-likelihood evaluation is < 0.005 s per call; CPU is < 0.05 s. Expected end-to-end run time is ~15 minutes on\n", + "GPU, ~30 minutes on CPU.\n", + "\n", + "__Model Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This example showed how to mix two strategies for `extra_galaxies` modeling \u2014 individually-modelled and on a shared\n", + "scaling relation \u2014 within a single `extra_galaxies` collection. The same pattern works with any mass profile and any\n", + "measured property (swap the `Isothermal` for `PowerLaw`, or the luminosity for stellar mass, and the structure is\n", + "unchanged).\n", + "\n", + "For the production-style luminosity-fitting workflow that produces the `relational_extras_luminosity_list` used here,\n", + "see:\n", + "\n", + " - `autolens_workspace/scripts/group/features/scaling_relation/modeling_for_luminosities.py` \u2014 a standalone light-only\n", + " fit that produces per-galaxy total luminosities.\n", + " - `autolens_workspace/scripts/group/slam.py` and `autolens_workspace/scripts/group/features/pixelization/slam.py` \u2014\n", + " the SLAM `source_lp[0]` stage that does the same job inside a chained pipeline.\n", + "\n", + "For the group-scale variant \u2014 multiple \"main\" lens galaxies AND a top-level `scaling_galaxies` collection separate from\n", + "`extra_galaxies` \u2014 see `autolens_workspace/scripts/group/features/scaling_relation/modeling.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/scaling_relation/simulator.ipynb b/notebooks/imaging/features/scaling_relation/simulator.ipynb index f0b7511a0..7f15af839 100644 --- a/notebooks/imaging/features/scaling_relation/simulator.ipynb +++ b/notebooks/imaging/features/scaling_relation/simulator.ipynb @@ -1,458 +1,495 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Extra and Scaling Galaxies\n", - "=====================================\n", - "\n", - "This script simulates a galaxy-scale strong lens with **two populations of foreground extra galaxies** in front of the\n", - "lensed source:\n", - "\n", - " - Two **individually-modelled** extras close to the lens, each bright enough to warrant its own free Einstein radius\n", - " in the lens model.\n", - " - Two **scaling-relation** extras further out / fainter, whose Einstein radii are tied together via a shared\n", - " luminosity-mass relation in the modeling stage.\n", - "\n", - "Both populations are dumped to separate JSON centre files (`extra_galaxies_centres.json` and\n", - "`scaling_galaxies_centres.json`) so the modeling script can load each independently and apply the appropriate\n", - "strategy. They both still live under the umbrella of \"extra galaxies\" in imaging-context terminology.\n", - "\n", - "This dataset is consumed by `scripts/imaging/features/scaling_relation/modeling.py` and\n", - "`scripts/imaging/features/scaling_relation/modeling_for_luminosities.py` (if added in a future task)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"extra_and_scaling_galaxies\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "A galaxy-scale field of view: 130x130 pixels at 0.1\"/pixel = 13\" wide. Big enough to enclose the lens, two close\n", - "companions, two further-out companions, and the lensed source; small enough to remain a galaxy-scale tutorial." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(130, 130),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Centres__\n", - "\n", - "Two centre lists, one per modeling strategy. Both populations are foreground galaxies near the main lens \u2014 the split\n", - "is purely about how they're modelled downstream." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_centres = [(3.5, 2.5), (-2.0, -3.5)]\n", - "scaling_galaxies_centres = [(5.0, -1.0), (-1.0, 5.0)]\n", - "\n", - "all_galaxy_centres = [(0.0, 0.0)] + extra_galaxies_centres + scaling_galaxies_centres" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Adaptive over-sampling at every galaxy centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=all_galaxy_centres,\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__PSF + Simulator__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")\n", - "\n", - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "A standard galaxy-scale primary lens: spherical Sersic light + Isothermal mass at the origin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(0.0, 0.0), intensity=0.7, effective_radius=1.5, sersic_index=3.0\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Individually-Modelled Extras__\n", - "\n", - "Two close, brighter companions. Each gets its own Sersic light + Isothermal mass with a non-trivial Einstein radius \u2014\n", - "in the modeling script the corresponding tier gives each a free `einstein_radius` parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.6, sersic_index=2.5\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.4),\n", - ")\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-2.0, -3.5), intensity=0.8, effective_radius=0.6, sersic_index=2.5\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-2.0, -3.5), einstein_radius=0.5),\n", - ")\n", - "\n", - "individual_extras = [extra_galaxy_0, extra_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Scaling-Relation Extras__\n", - "\n", - "Two further-out, fainter companions whose true Einstein radii are consistent with\n", - "``einstein_radius = 0.3 * luminosity ** 1.0`` (luminosities ~0.45 -> Einstein radii ~0.135). In the modeling script\n", - "they share two scaling-relation priors regardless of how many are added here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "scaling_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(5.0, -1.0), intensity=0.45, effective_radius=0.5, sersic_index=2.5\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(5.0, -1.0), einstein_radius=0.135),\n", - ")\n", - "\n", - "scaling_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.SersicSph(\n", - " centre=(-1.0, 5.0), intensity=0.45, effective_radius=0.5, sersic_index=2.5\n", - " ),\n", - " mass=al.mp.IsothermalSph(centre=(-1.0, 5.0), einstein_radius=0.135),\n", - ")\n", - "\n", - "relational_extras = [scaling_galaxy_0, scaling_galaxy_1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=3.0,\n", - " effective_radius=0.2,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Tracer order: lens, individual extras, relational extras, source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy] + individual_extras + relational_extras + [source_galaxy]\n", - ")\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Centre JSON Files__\n", - "\n", - "Two JSON files, one per population, matching the names the modeling script loads." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", - " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(scaling_galaxies_centres),\n", - " file_path=Path(dataset_path, \"scaling_galaxies_centres.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Population CSVs__\n", - "\n", - "The modeling script loads luminosities (and centres) for the scaling-relation tier from a CSV\n", - "written here. The simulator knows the truth values of the per-galaxy luminosities so we write\n", - "them out alongside the centre JSONs.\n", - "\n", - "The CSV schema is `y, x, luminosity, redshift?` -- see `al.galaxy_table_from_csv` /\n", - "`al.galaxy_table_to_csv` (`autogalaxy/galaxy/galaxy_table.py`). Centre JSONs above are kept for\n", - "backward compatibility; new consumers should prefer the CSV." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_luminosities = [0.9, 0.8]\n", - "scaling_galaxies_luminosities = [0.45, 0.45]\n", - "\n", - "al.galaxy_table_to_csv(\n", - " centres=extra_galaxies_centres,\n", - " luminosities=extra_galaxies_luminosities,\n", - " file_path=Path(dataset_path, \"extra_galaxies.csv\"),\n", - ")\n", - "\n", - "al.galaxy_table_to_csv(\n", - " centres=scaling_galaxies_centres,\n", - " luminosities=scaling_galaxies_luminosities,\n", - " file_path=Path(dataset_path, \"scaling_galaxies.csv\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Extra and Scaling Galaxies\n", + "=====================================\n", + "\n", + "This script simulates a galaxy-scale strong lens with **two populations of foreground extra galaxies** in front of the\n", + "lensed source:\n", + "\n", + " - Two **individually-modelled** extras close to the lens, each bright enough to warrant its own free Einstein radius\n", + " in the lens model.\n", + " - Two **scaling-relation** extras further out / fainter, whose Einstein radii are tied together via a shared\n", + " luminosity-mass relation in the modeling stage.\n", + "\n", + "Both populations are dumped to separate JSON centre files (`extra_galaxies_centres.json` and\n", + "`scaling_galaxies_centres.json`) so the modeling script can load each independently and apply the appropriate\n", + "strategy. They both still live under the umbrella of \"extra galaxies\" in imaging-context terminology.\n", + "\n", + "This dataset is consumed by `scripts/imaging/features/scaling_relation/modeling.py` and\n", + "`scripts/imaging/features/scaling_relation/modeling_for_luminosities.py` (if added in a future task)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"extra_and_scaling_galaxies\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "A galaxy-scale field of view: 130x130 pixels at 0.1\"/pixel = 13\" wide. Big enough to enclose the lens, two close\n", + "companions, two further-out companions, and the lensed source; small enough to remain a galaxy-scale tutorial." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(130, 130),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Centres__\n", + "\n", + "Two centre lists, one per modeling strategy. Both populations are foreground galaxies near the main lens \u2014 the split\n", + "is purely about how they're modelled downstream." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_centres = [(3.5, 2.5), (-2.0, -3.5)]\n", + "scaling_galaxies_centres = [(5.0, -1.0), (-1.0, 5.0)]\n", + "\n", + "all_galaxy_centres = [(0.0, 0.0)] + extra_galaxies_centres + scaling_galaxies_centres" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Adaptive over-sampling at every galaxy centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=all_galaxy_centres,\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__PSF + Simulator__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")\n", + "\n", + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "A standard galaxy-scale primary lens: spherical Sersic light + Isothermal mass at the origin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(0.0, 0.0), intensity=0.7, effective_radius=1.5, sersic_index=3.0\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Individually-Modelled Extras__\n", + "\n", + "Two close, brighter companions. Each gets its own Sersic light + Isothermal mass with a non-trivial Einstein radius \u2014\n", + "in the modeling script the corresponding tier gives each a free `einstein_radius` parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(3.5, 2.5), intensity=0.9, effective_radius=0.6, sersic_index=2.5\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(3.5, 2.5), einstein_radius=0.4),\n", + ")\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-2.0, -3.5), intensity=0.8, effective_radius=0.6, sersic_index=2.5\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-2.0, -3.5), einstein_radius=0.5),\n", + ")\n", + "\n", + "individual_extras = [extra_galaxy_0, extra_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Scaling-Relation Extras__\n", + "\n", + "Two further-out, fainter companions whose true Einstein radii are consistent with\n", + "``einstein_radius = 0.3 * luminosity ** 1.0`` (luminosities ~0.45 -> Einstein radii ~0.135). In the modeling script\n", + "they share two scaling-relation priors regardless of how many are added here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "scaling_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(5.0, -1.0), intensity=0.45, effective_radius=0.5, sersic_index=2.5\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(5.0, -1.0), einstein_radius=0.135),\n", + ")\n", + "\n", + "scaling_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.SersicSph(\n", + " centre=(-1.0, 5.0), intensity=0.45, effective_radius=0.5, sersic_index=2.5\n", + " ),\n", + " mass=al.mp.IsothermalSph(centre=(-1.0, 5.0), einstein_radius=0.135),\n", + ")\n", + "\n", + "relational_extras = [scaling_galaxy_0, scaling_galaxy_1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=3.0,\n", + " effective_radius=0.2,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Tracer order: lens, individual extras, relational extras, source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy] + individual_extras + relational_extras + [source_galaxy]\n", + ")\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Centre JSON Files__\n", + "\n", + "Two JSON files, one per population, matching the names the modeling script loads." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(extra_galaxies_centres),\n", + " file_path=Path(dataset_path, \"extra_galaxies_centres.json\"),\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(scaling_galaxies_centres),\n", + " file_path=Path(dataset_path, \"scaling_galaxies_centres.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Population CSVs__\n", + "\n", + "The modeling script loads luminosities (and centres) for the scaling-relation tier from a CSV\n", + "written here. The simulator knows the truth values of the per-galaxy luminosities so we write\n", + "them out alongside the centre JSONs.\n", + "\n", + "The CSV schema is `y, x, luminosity, redshift?` -- see `al.galaxy_table_from_csv` /\n", + "`al.galaxy_table_to_csv` (`autogalaxy/galaxy/galaxy_table.py`). Centre JSONs above are kept for\n", + "backward compatibility; new consumers should prefer the CSV." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_luminosities = [0.9, 0.8]\n", + "scaling_galaxies_luminosities = [0.45, 0.45]\n", + "\n", + "al.galaxy_table_to_csv(\n", + " centres=extra_galaxies_centres,\n", + " luminosities=extra_galaxies_luminosities,\n", + " file_path=Path(dataset_path, \"extra_galaxies.csv\"),\n", + ")\n", + "\n", + "al.galaxy_table_to_csv(\n", + " centres=scaling_galaxies_centres,\n", + " luminosities=scaling_galaxies_luminosities,\n", + " file_path=Path(dataset_path, \"scaling_galaxies.csv\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/features/simulator_manual_signal_to_noise_ratio.ipynb b/notebooks/imaging/features/simulator_manual_signal_to_noise_ratio.ipynb index 39c38f7cc..95f73e0ba 100644 --- a/notebooks/imaging/features/simulator_manual_signal_to_noise_ratio.ipynb +++ b/notebooks/imaging/features/simulator_manual_signal_to_noise_ratio.ipynb @@ -1,403 +1,440 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Manual Signal to Noise Ratio\n", - "=======================================\n", - "\n", - "When simulating `Imaging` of a strong lens, one is often not concerned with the actual units of the light (e.g.\n", - "electrons per second, counts, etc.) but instead simple wants the data to correspond to a certain signal to noise\n", - "value.\n", - "\n", - "This can be difficult to achieve when specifying the `intensity` of the input light profiles, especially given the\n", - "unknown contribution of the mass model's magnification.\n", - "\n", - "This script illustrates the `lp_snr` light profiles, which when used to simulate a dataset via a tracer, set the\n", - "signal to noise of each light profile to an input value. This uses the `exposure_time` and `background_sky_level`\n", - "of the `SimulatorImaging` object to choose the `intensity` of each light profile such that the input signal to\n", - "noise is used.\n", - "\n", - "For normal light profiles, the `intensity` is defined in units of electrons per second, meaning that the\n", - "`exposure_time` and `background_sky_level` are used to convert this to counts when adding noise. When the `lp_snr`\n", - "profiles are used, the `exposure_time` and `background_sky_level` are instead used to set its S/N, meaning their input\n", - "values do not set the S/N.\n", - "\n", - "However, the ratio of `exposure_time` and `background_sky_level` does set how much noise is due to Poisson count\n", - "statistics in the CCD imaging detector relative to the background sky. If one doubles the `exposure_time`, the\n", - "Poisson count component will contribute more compared to the background sky component. For detailed scientific\n", - "analysis, one should therefore make sure their values are chosen to produce images with realistic noise properties.\n", - "\n", - "The use of the `light_snr` profiles changes the meaning of `exposure_time` and `background_sky_level`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's bulge is an `Sersic` with a S/N of 50.0.\n", - " - The lens galaxy's disk is an `Exponential` with a S/N of 20.0.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is two `Sersic`;s with S/N of 20.0 and 10.0.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_label = \"misc\"\n", - "dataset_name = \"manual_signal_to_noise_ratio\"\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "\n", - "the `lens_galaxy` uses light profile signal-to-noise objects (`lp_snr`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp_snr.Sersic(\n", - " signal_to_noise_ratio=50.0,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " disk=al.lp_snr.Exponential(\n", - " signal_to_noise_ratio=20.0,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", - " effective_radius=1.6,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_snr.Sersic(\n", - " signal_to_noise_ratio=20.0,\n", - " centre=(0.25, 0.15),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=120.0),\n", - " effective_radius=0.7,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "source_galaxy_1 = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_snr.Sersic(\n", - " signal_to_noise_ratio=10.0,\n", - " centre=(0.7, -0.5),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", - " effective_radius=1.6,\n", - " sersic_index=3.0,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/misc/manual_signal_to_noise_ratio`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Manual Signal to Noise Ratio\n", + "=======================================\n", + "\n", + "When simulating `Imaging` of a strong lens, one is often not concerned with the actual units of the light (e.g.\n", + "electrons per second, counts, etc.) but instead simple wants the data to correspond to a certain signal to noise\n", + "value.\n", + "\n", + "This can be difficult to achieve when specifying the `intensity` of the input light profiles, especially given the\n", + "unknown contribution of the mass model's magnification.\n", + "\n", + "This script illustrates the `lp_snr` light profiles, which when used to simulate a dataset via a tracer, set the\n", + "signal to noise of each light profile to an input value. This uses the `exposure_time` and `background_sky_level`\n", + "of the `SimulatorImaging` object to choose the `intensity` of each light profile such that the input signal to\n", + "noise is used.\n", + "\n", + "For normal light profiles, the `intensity` is defined in units of electrons per second, meaning that the\n", + "`exposure_time` and `background_sky_level` are used to convert this to counts when adding noise. When the `lp_snr`\n", + "profiles are used, the `exposure_time` and `background_sky_level` are instead used to set its S/N, meaning their input\n", + "values do not set the S/N.\n", + "\n", + "However, the ratio of `exposure_time` and `background_sky_level` does set how much noise is due to Poisson count\n", + "statistics in the CCD imaging detector relative to the background sky. If one doubles the `exposure_time`, the\n", + "Poisson count component will contribute more compared to the background sky component. For detailed scientific\n", + "analysis, one should therefore make sure their values are chosen to produce images with realistic noise properties.\n", + "\n", + "The use of the `light_snr` profiles changes the meaning of `exposure_time` and `background_sky_level`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Ray Tracing:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's bulge is an `Sersic` with a S/N of 50.0.\n", + " - The lens galaxy's disk is an `Exponential` with a S/N of 20.0.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is two `Sersic`;s with S/N of 20.0 and 10.0.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_label = \"misc\"\n", + "dataset_name = \"manual_signal_to_noise_ratio\"\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "\n", + "the `lens_galaxy` uses light profile signal-to-noise objects (`lp_snr`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp_snr.Sersic(\n", + " signal_to_noise_ratio=50.0,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " disk=al.lp_snr.Exponential(\n", + " signal_to_noise_ratio=20.0,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=30.0),\n", + " effective_radius=1.6,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_snr.Sersic(\n", + " signal_to_noise_ratio=20.0,\n", + " centre=(0.25, 0.15),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=120.0),\n", + " effective_radius=0.7,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "source_galaxy_1 = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_snr.Sersic(\n", + " signal_to_noise_ratio=10.0,\n", + " centre=(0.7, -0.5),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", + " effective_radius=1.6,\n", + " sersic_index=3.0,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Imaging` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/misc/manual_signal_to_noise_ratio`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/fit.ipynb b/notebooks/imaging/fit.ipynb index 69b726dfd..05c26046c 100644 --- a/notebooks/imaging/fit.ipynb +++ b/notebooks/imaging/fit.ipynb @@ -1,756 +1,802 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fits\n", - "====\n", - "\n", - "This guide shows how to fit data using the `FitImaging` object, including visualizing and interpreting its results.\n", - "\n", - "References\n", - "----------\n", - "\n", - "This example uses functionality described fully in other examples in the `guides` package:\n", - "\n", - "- `guides/plot`: Using the plotting API (`aplt.plot_array`, `aplt.subplot_fit_imaging`, etc.) to visualize figures.\n", - "- `guides/units`: The source code unit conventions (e.g. arc seconds for distances and how to convert to physical units).\n", - "- `guides/data_structures`: The bespoke data structures used to store 1D and 2d arrays.\n", - "\n", - "__Contents__\n", - "\n", - "- **Loading Data:** We we begin by loading the strong lens dataset `simple__no_lens_light` from .fits files, which is.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", - "- **Bad Fit:** A bad lens model will show features in the residual-map and chi-squared map.\n", - "- **Fit Quantities:** The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", - "- **Figures of Merit:** There are single valued floats which quantify the goodness of fit.\n", - "- **Plane Quantities:** The `FitImaging` object has specific quantities which break down each image of each plane.\n", - "- **Unmasked Quantities:** All of the quantities above are computed using the mask which was used to fit the data.\n", - "- **Pixel Counting:** An alternative way to quantify residuals like the lens light residuals is pixel counting.\n", - "- **Outputting Results:** You may wish to output certain results to .fits files for later inspection.\n", - "\n", - "__JAX__\n", - "\n", - "This script constructs a `FitImaging` directly from a tracer and dataset\n", - "(no Analysis / no non-linear search). The fit itself is JAX-friendly:\n", - "any quantities it computes (`fit.model_image`, `fit.residual_map`,\n", - "`fit.log_likelihood`, etc.) work on either backend and return arrays\n", - "backed by `numpy.ndarray` on the default path or `jax.Array` if you\n", - "constructed the upstream objects with `xp=jnp`.\n", - "\n", - "For the standard analysis-driven modeling path \u2014 where `AnalysisImaging`\n", - "auto-enables `use_jax=True` and the search driver handles the JIT\n", - "internally \u2014 see `start_here.py` / `modeling.py`. For the advanced path\n", - "where you wrap your own `@jax.jit` around `FitImaging` construction, see\n", - "`likelihood_function.py`'s `__JAX__` section and the `lens_calc.py` guide." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Data__\n", - "\n", - "We we begin by loading the strong lens dataset `simple__no_lens_light` from .fits files, which is the dataset \n", - "we will use to demonstrate fitting.\n", - "\n", - "This dataset was simulated using the `imaging/simulator` example, read through that to have a better\n", - "understanding of how the data this exam fits was generated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `aplt.subplot_imaging_dataset` contains a subplot which plots all the key properties of the dataset simultaneously.\n", - "\n", - "This includes the observed image data, RMS noise map, Point Spread Function and other information." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Noise Scaling__\n", - "\n", - "Before masking, we must deal with any extra galaxies in the data: nearby galaxies (or foreground stars, or\n", - "data-reduction artefacts) whose emission is not associated with the strong lens but blends into the field. If\n", - "their light is left in the data it will contaminate the fit and bias the inferred lens model. It is too easy to\n", - "skip straight to fitting without checking for these, so we make this step an explicit part of the workflow.\n", - "\n", - "To prevent extra galaxies from impacting the fit, we do not mask them entirely from the fit, which would be\n", - "analogous to making the circular mask smaller or using a more refined mask. When pixels are masked and removed\n", - "entirely from the fit, their coordinates are not used when performing ray-tracing and the light of the lens and\n", - "source galaxies in these pixels not evaluated.\n", - "\n", - "Instead, the pixels are kept in the fit, but their data values are scaled to zero and their noise-map values\n", - "are increased to very large values. This means that during the fit, these pixels contribute negligibly to the\n", - "likelihood, and therefore do not impact the lens model.\n", - "\n", - "This approach is used because for certain types of modeling approaches, like a pixelized source reconstruction,\n", - "masking regions of the image in a way that removes their image pixels entirely from the fit can produce\n", - "discontinuities in the pixelization. This can lead to unexpected systematics and unsatisfactory results.\n", - "\n", - "The dataset includes a faint extra galaxy, and a `mask_extra_galaxies.fits` covering it is shipped with the\n", - "dataset (created by the simulator). If you are fitting your own data with an extra galaxy, you must either:\n", - "\n", - " - Create a `mask_extra_galaxies.fits` for it using the data-preparation tools (the GUI\n", - " `autolens_workspace/*/imaging/data_preparation/gui/mask_extra_galaxies.py`, or the manual\n", - " `autolens_workspace/*/imaging/data_preparation/examples/optional/mask_extra_galaxies.py`), then load it\n", - " as below; or\n", - " - Shrink the circular mask below so the extra galaxy lies outside it and is removed from the fit entirely.\n", - "\n", - "After scaling, the extra galaxy's pixels have their data set to zero and noise-map increased, making their\n", - "signal-to-noise effectively zero." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_extra_galaxies = al.Mask2D.from_fits(\n", - " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", - " pixel_scales=dataset.pixel_scales,\n", - " invert=True, # `True` means a pixel is scaled.\n", - ")\n", - "\n", - "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We now mask the data, so that regions where there is no signal (e.g. the edges) are omitted from the fit.\n", - "\n", - "We use a ``Mask2D`` object, which for this example is a 3.0\" circular mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now combine the imaging dataset with the mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now plot the image with the mask applied, where the image automatically zooms around the mask to make the lensed \n", - "source appear bigger." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=dataset.data, title=\"Image Data With Mask Applied\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The mask is also used to compute a `Grid2D`, where the (y,x) arc-second coordinates are only computed in unmasked \n", - "pixels within the masks' circle. \n", - "\n", - "As shown in the previous overview example, this grid will be used to perform lensing calculations when fitting the\n", - "data below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=dataset.grid, title=\"Grid2D of Masked Dataset\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fitting__\n", - "\n", - "Following the previous overview example, we can make a tracer from a collection of light profiles, mass profiles\n", - "and galaxies.\n", - "\n", - "The combination of light and mass profiels below is the same as those used to generate the simulated \n", - "dataset we loaded above.\n", - "\n", - "It therefore produces a tracer whose image looks exactly like the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Because the tracer's light and mass profiles are the same used to make the dataset, its image is nearly the same as the\n", - "observed image.\n", - "\n", - "However, the tracer's image does appear different to the data, in that its ring appears a bit thinner. This is\n", - "because its image has not been blurred with the telescope optics PSF, which the data has.\n", - "\n", - "[For those not familiar with Astronomy data, the PSF describes how the observed emission of the galaxy is blurred by\n", - "the telescope optics when it is observed. It mimicks this blurring effect via a 2D convolution operation]." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now use a `FitImaging` object to fit this tracer to the dataset. \n", - "\n", - "The fit creates a `model_image` which we fit the data with, which includes performing the step of blurring the tracer`s \n", - "image with the imaging dataset's PSF. We can see this by comparing the tracer`s image (which isn't PSF convolved) and \n", - "the fit`s model image (which is)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.plot_array(array=fit.model_data, title=\"Model Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit does a lot more than just blur the tracer's image with the PSF, it also creates the following:\n", - "\n", - " - The `residual_map`: The `model_image` subtracted from the observed dataset`s `data`.\n", - " - The `normalized_residual_map`: The `residual_map `divided by the observed dataset's `noise_map`.\n", - " - The `chi_squared_map`: The `normalized_residual_map` squared.\n", - "\n", - "For a good lens model where the model image and tracer are representative of the strong lens system the\n", - "residuals, normalized residuals and chi-squareds are minimized:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", - "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", - "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi Squared Map\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A subplot can be plotted which contains all of the above quantities, as well as other information contained in the\n", - "tracer such as the source-plane image, a zoom in of the source-plane and a normalized residual map where the colorbar\n", - "goes from 1.0 sigma to -1.0 sigma, to highlight regions where the fit is poor." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit also provides us with a ``log_likelihood``, a single value quantifying how good the tracer fitted the dataset.\n", - "\n", - "Lens modeling, describe in the next overview example, effectively tries to maximize this log likelihood value." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Bad Fit__\n", - "\n", - "A bad lens model will show features in the residual-map and chi-squared map.\n", - "\n", - "We can produce such an image by creating a tracer with different lens and source galaxies. In the example below, we \n", - "change the centre of the source galaxy from (0.0, 0.0) to (0.05, 0.05), which leads to residuals appearing\n", - "in the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.1, 0.1),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.1, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A new fit using this plane shows residuals, normalized residuals and chi-squared which are non-zero. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We also note that its likelihood decreases." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit Quantities__\n", - "\n", - "The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", - "\n", - "There is a `model_image`, which is the image-plane image of the tracer we inspected in the previous tutorial\n", - "blurred with the imaging data's PSF. \n", - "\n", - "This is the image that is fitted to the data in order to compute the log likelihood and therefore quantify the \n", - "goodness-of-fit.\n", - "\n", - "If you are unclear on what `slim` means, refer to the section `Data Structure` at the top of this example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_data.slim)\n", - "\n", - "# The native property provides quantities in 2D NumPy Arrays.\n", - "# print(fit.model_data.native)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There are numerous ndarrays showing the goodness of fit: \n", - "\n", - " - `residual_map`: Residuals = (Data - Model_Data).\n", - " - `normalized_residual_map`: Normalized_Residual = (Data - Model_Data) / Noise\n", - " - `chi_squared_map`: Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.residual_map.slim)\n", - "print(fit.normalized_residual_map.slim)\n", - "print(fit.chi_squared_map.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Figures of Merit__\n", - "\n", - "There are single valued floats which quantify the goodness of fit:\n", - "\n", - " - `chi_squared`: The sum of the `chi_squared_map`.\n", - "\n", - " - `noise_normalization`: The normalizing noise term in the likelihood function \n", - " where [Noise_Term] = sum(log(2*pi*[Noise]**2.0)).\n", - "\n", - " - `log_likelihood`: The log likelihood value of the fit where [LogLikelihood] = -0.5*[Chi_Squared_Term + Noise_Term]." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.chi_squared)\n", - "print(fit.noise_normalization)\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Plane Quantities__\n", - "\n", - "The `FitImaging` object has specific quantities which break down each image of each plane:\n", - "\n", - " - `model_images_of_planes_list`: Model-images of each individual plane, which in this example is a model image of the \n", - " lens galaxy and model image of the lensed source galaxy. Both images are convolved with the imaging's PSF.\n", - "\n", - " - `subtracted_images_of_planes_list`: Subtracted images of each individual plane, which are the data's image with\n", - " all other plane's model-images subtracted. For example, the first subtracted image has the source galaxy's model image\n", - " subtracted and therefore is of only the lens galaxy's emission. The second subtracted image is of the lensed source,\n", - " with the lens galaxy's light removed.\n", - "\n", - "For multi-plane lens systems these lists will be extended to provide information on every individual plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_images_of_planes_list[0].slim)\n", - "print(fit.model_images_of_planes_list[1].slim)\n", - "\n", - "print(fit.subtracted_images_of_planes_list[0].slim)\n", - "print(fit.subtracted_images_of_planes_list[1].slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Unmasked Quantities__\n", - "\n", - "All of the quantities above are computed using the mask which was used to fit the data.\n", - "\n", - "The `FitImaging` can also compute the unmasked blurred image of each plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.unmasked_blurred_image.native)\n", - "print(fit.unmasked_blurred_image_of_planes_list[0].native)\n", - "print(fit.unmasked_blurred_image_of_planes_list[1].native)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We can use the `Mask2D` object to mask regions of one of the fit's maps and estimate quantities of it.\n", - "\n", - "Below, we estimate the average absolute normalized residuals within a 1.0\" circular mask, which would inform us of\n", - "how accurate the lens light subtraction of a model fit is and if it leaves any significant residuals" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(\n", - " shape_native=fit.dataset.shape_native,\n", - " pixel_scales=fit.dataset.pixel_scales,\n", - " radius=1.0,\n", - ")\n", - "\n", - "normalized_residuals = fit.normalized_residual_map.apply_mask(mask=mask)\n", - "\n", - "print(np.mean(np.abs(normalized_residuals.slim)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixel Counting__\n", - "\n", - "An alternative way to quantify residuals like the lens light residuals is pixel counting. For example, we could sum\n", - "up the number of pixels whose chi-squared values are above 10 which indicates a poor fit to the data.\n", - "\n", - "Whereas computing the mean above the average level of residuals, pixel counting informs us how spatially large the\n", - "residuals extend. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(\n", - " shape_native=fit.dataset.shape_native,\n", - " pixel_scales=fit.dataset.pixel_scales,\n", - " radius=1.0,\n", - ")\n", - "\n", - "chi_squared_map = fit.chi_squared_map.apply_mask(mask=mask)\n", - "\n", - "print(np.sum(chi_squared_map > 10.0))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Outputting Results__\n", - "\n", - "You may wish to output certain results to .fits files for later inspection. \n", - "\n", - "For example, one could output the lens light subtracted image of the lensed source galaxy to a .fits file such that\n", - "we could fit this source-only image again with an independent pipeline." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_subtracted_image = fit.subtracted_images_of_planes_list[1]\n", - "aplt.fits_array(\n", - " array=lens_subtracted_image,\n", - " file_path=dataset_path / \"lens_subtracted_data.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fits\n", + "====\n", + "\n", + "This guide shows how to fit data using the `FitImaging` object, including visualizing and interpreting its results.\n", + "\n", + "References\n", + "----------\n", + "\n", + "This example uses functionality described fully in other examples in the `guides` package:\n", + "\n", + "- `guides/plot`: Using the plotting API (`aplt.plot_array`, `aplt.subplot_fit_imaging`, etc.) to visualize figures.\n", + "- `guides/units`: The source code unit conventions (e.g. arc seconds for distances and how to convert to physical units).\n", + "- `guides/data_structures`: The bespoke data structures used to store 1D and 2d arrays.\n", + "\n", + "__Contents__\n", + "\n", + "- **Loading Data:** We we begin by loading the strong lens dataset `simple__no_lens_light` from .fits files, which is.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", + "- **Bad Fit:** A bad lens model will show features in the residual-map and chi-squared map.\n", + "- **Fit Quantities:** The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", + "- **Figures of Merit:** There are single valued floats which quantify the goodness of fit.\n", + "- **Plane Quantities:** The `FitImaging` object has specific quantities which break down each image of each plane.\n", + "- **Unmasked Quantities:** All of the quantities above are computed using the mask which was used to fit the data.\n", + "- **Pixel Counting:** An alternative way to quantify residuals like the lens light residuals is pixel counting.\n", + "- **Outputting Results:** You may wish to output certain results to .fits files for later inspection.\n", + "\n", + "__JAX__\n", + "\n", + "This script constructs a `FitImaging` directly from a tracer and dataset\n", + "(no Analysis / no non-linear search). The fit itself is JAX-friendly:\n", + "any quantities it computes (`fit.model_image`, `fit.residual_map`,\n", + "`fit.log_likelihood`, etc.) work on either backend and return arrays\n", + "backed by `numpy.ndarray` on the default path or `jax.Array` if you\n", + "constructed the upstream objects with `xp=jnp`.\n", + "\n", + "For the standard analysis-driven modeling path \u2014 where `AnalysisImaging`\n", + "auto-enables `use_jax=True` and the search driver handles the JIT\n", + "internally \u2014 see `start_here.py` / `modeling.py`. For the advanced path\n", + "where you wrap your own `@jax.jit` around `FitImaging` construction, see\n", + "`likelihood_function.py`'s `__JAX__` section and the `lens_calc.py` guide." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Data__\n", + "\n", + "We we begin by loading the strong lens dataset `simple__no_lens_light` from .fits files, which is the dataset \n", + "we will use to demonstrate fitting.\n", + "\n", + "This dataset was simulated using the `imaging/simulator` example, read through that to have a better\n", + "understanding of how the data this exam fits was generated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "# PSF convolution runs at the image resolution (sub size 1), which is the fastest\n", + "# option and accurate for well-sampled PSFs. Supplying a PSF at a multiple of the\n", + "# image resolution and raising this value improves blurring fidelity for\n", + "# undersampled PSFs (e.g. HST / Euclid VIS) at extra compute cost \u2014 see\n", + "# `guides/advanced/over_sampling.py` and the simulator's `__Oversampled PSF__` section.\n", + "psf_convolve_over_sample_size = 1\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " convolve_over_sample_size_lp=psf_convolve_over_sample_size,\n", + " convolve_over_sample_size_pixelization=psf_convolve_over_sample_size,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `aplt.subplot_imaging_dataset` contains a subplot which plots all the key properties of the dataset simultaneously.\n", + "\n", + "This includes the observed image data, RMS noise map, Point Spread Function and other information." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Noise Scaling__\n", + "\n", + "Before masking, we must deal with any extra galaxies in the data: nearby galaxies (or foreground stars, or\n", + "data-reduction artefacts) whose emission is not associated with the strong lens but blends into the field. If\n", + "their light is left in the data it will contaminate the fit and bias the inferred lens model. It is too easy to\n", + "skip straight to fitting without checking for these, so we make this step an explicit part of the workflow.\n", + "\n", + "To prevent extra galaxies from impacting the fit, we do not mask them entirely from the fit, which would be\n", + "analogous to making the circular mask smaller or using a more refined mask. When pixels are masked and removed\n", + "entirely from the fit, their coordinates are not used when performing ray-tracing and the light of the lens and\n", + "source galaxies in these pixels not evaluated.\n", + "\n", + "Instead, the pixels are kept in the fit, but their data values are scaled to zero and their noise-map values\n", + "are increased to very large values. This means that during the fit, these pixels contribute negligibly to the\n", + "likelihood, and therefore do not impact the lens model.\n", + "\n", + "This approach is used because for certain types of modeling approaches, like a pixelized source reconstruction,\n", + "masking regions of the image in a way that removes their image pixels entirely from the fit can produce\n", + "discontinuities in the pixelization. This can lead to unexpected systematics and unsatisfactory results.\n", + "\n", + "The dataset includes a faint extra galaxy, and a `mask_extra_galaxies.fits` covering it is shipped with the\n", + "dataset (created by the simulator). If you are fitting your own data with an extra galaxy, you must either:\n", + "\n", + " - Create a `mask_extra_galaxies.fits` for it using the data-preparation tools (the GUI\n", + " `autolens_workspace/*/imaging/data_preparation/gui/mask_extra_galaxies.py`, or the manual\n", + " `autolens_workspace/*/imaging/data_preparation/examples/optional/mask_extra_galaxies.py`), then load it\n", + " as below; or\n", + " - Shrink the circular mask below so the extra galaxy lies outside it and is removed from the fit entirely.\n", + "\n", + "After scaling, the extra galaxy's pixels have their data set to zero and noise-map increased, making their\n", + "signal-to-noise effectively zero." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_extra_galaxies = al.Mask2D.from_fits(\n", + " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", + " pixel_scales=dataset.pixel_scales,\n", + " invert=True, # `True` means a pixel is scaled.\n", + ")\n", + "\n", + "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We now mask the data, so that regions where there is no signal (e.g. the edges) are omitted from the fit.\n", + "\n", + "We use a ``Mask2D`` object, which for this example is a 3.0\" circular mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now combine the imaging dataset with the mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now plot the image with the mask applied, where the image automatically zooms around the mask to make the lensed \n", + "source appear bigger." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=dataset.data, title=\"Image Data With Mask Applied\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The mask is also used to compute a `Grid2D`, where the (y,x) arc-second coordinates are only computed in unmasked \n", + "pixels within the masks' circle. \n", + "\n", + "As shown in the previous overview example, this grid will be used to perform lensing calculations when fitting the\n", + "data below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=dataset.grid, title=\"Grid2D of Masked Dataset\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fitting__\n", + "\n", + "Following the previous overview example, we can make a tracer from a collection of light profiles, mass profiles\n", + "and galaxies.\n", + "\n", + "The combination of light and mass profiels below is the same as those used to generate the simulated \n", + "dataset we loaded above.\n", + "\n", + "It therefore produces a tracer whose image looks exactly like the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Because the tracer's light and mass profiles are the same used to make the dataset, its image is nearly the same as the\n", + "observed image.\n", + "\n", + "However, the tracer's image does appear different to the data, in that its ring appears a bit thinner. This is\n", + "because its image has not been blurred with the telescope optics PSF, which the data has.\n", + "\n", + "[For those not familiar with Astronomy data, the PSF describes how the observed emission of the galaxy is blurred by\n", + "the telescope optics when it is observed. It mimicks this blurring effect via a 2D convolution operation]." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now use a `FitImaging` object to fit this tracer to the dataset. \n", + "\n", + "The fit creates a `model_image` which we fit the data with, which includes performing the step of blurring the tracer`s \n", + "image with the imaging dataset's PSF. We can see this by comparing the tracer`s image (which isn't PSF convolved) and \n", + "the fit`s model image (which is)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.plot_array(array=fit.model_data, title=\"Model Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit does a lot more than just blur the tracer's image with the PSF, it also creates the following:\n", + "\n", + " - The `residual_map`: The `model_image` subtracted from the observed dataset`s `data`.\n", + " - The `normalized_residual_map`: The `residual_map `divided by the observed dataset's `noise_map`.\n", + " - The `chi_squared_map`: The `normalized_residual_map` squared.\n", + "\n", + "For a good lens model where the model image and tracer are representative of the strong lens system the\n", + "residuals, normalized residuals and chi-squareds are minimized:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=fit.residual_map, title=\"Residual Map\")\n", + "aplt.plot_array(array=fit.normalized_residual_map, title=\"Normalized Residual Map\")\n", + "aplt.plot_array(array=fit.chi_squared_map, title=\"Chi Squared Map\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A subplot can be plotted which contains all of the above quantities, as well as other information contained in the\n", + "tracer such as the source-plane image, a zoom in of the source-plane and a normalized residual map where the colorbar\n", + "goes from 1.0 sigma to -1.0 sigma, to highlight regions where the fit is poor." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit also provides us with a ``log_likelihood``, a single value quantifying how good the tracer fitted the dataset.\n", + "\n", + "Lens modeling, describe in the next overview example, effectively tries to maximize this log likelihood value." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Bad Fit__\n", + "\n", + "A bad lens model will show features in the residual-map and chi-squared map.\n", + "\n", + "We can produce such an image by creating a tracer with different lens and source galaxies. In the example below, we \n", + "change the centre of the source galaxy from (0.0, 0.0) to (0.05, 0.05), which leads to residuals appearing\n", + "in the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.1, 0.1),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.1, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A new fit using this plane shows residuals, normalized residuals and chi-squared which are non-zero. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We also note that its likelihood decreases." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit Quantities__\n", + "\n", + "The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", + "\n", + "There is a `model_image`, which is the image-plane image of the tracer we inspected in the previous tutorial\n", + "blurred with the imaging data's PSF. \n", + "\n", + "This is the image that is fitted to the data in order to compute the log likelihood and therefore quantify the \n", + "goodness-of-fit.\n", + "\n", + "If you are unclear on what `slim` means, refer to the section `Data Structure` at the top of this example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_data.slim)\n", + "\n", + "# The native property provides quantities in 2D NumPy Arrays.\n", + "# print(fit.model_data.native)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There are numerous ndarrays showing the goodness of fit: \n", + "\n", + " - `residual_map`: Residuals = (Data - Model_Data).\n", + " - `normalized_residual_map`: Normalized_Residual = (Data - Model_Data) / Noise\n", + " - `chi_squared_map`: Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.residual_map.slim)\n", + "print(fit.normalized_residual_map.slim)\n", + "print(fit.chi_squared_map.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Figures of Merit__\n", + "\n", + "There are single valued floats which quantify the goodness of fit:\n", + "\n", + " - `chi_squared`: The sum of the `chi_squared_map`.\n", + "\n", + " - `noise_normalization`: The normalizing noise term in the likelihood function \n", + " where [Noise_Term] = sum(log(2*pi*[Noise]**2.0)).\n", + "\n", + " - `log_likelihood`: The log likelihood value of the fit where [LogLikelihood] = -0.5*[Chi_Squared_Term + Noise_Term]." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.chi_squared)\n", + "print(fit.noise_normalization)\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Plane Quantities__\n", + "\n", + "The `FitImaging` object has specific quantities which break down each image of each plane:\n", + "\n", + " - `model_images_of_planes_list`: Model-images of each individual plane, which in this example is a model image of the \n", + " lens galaxy and model image of the lensed source galaxy. Both images are convolved with the imaging's PSF.\n", + "\n", + " - `subtracted_images_of_planes_list`: Subtracted images of each individual plane, which are the data's image with\n", + " all other plane's model-images subtracted. For example, the first subtracted image has the source galaxy's model image\n", + " subtracted and therefore is of only the lens galaxy's emission. The second subtracted image is of the lensed source,\n", + " with the lens galaxy's light removed.\n", + "\n", + "For multi-plane lens systems these lists will be extended to provide information on every individual plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_images_of_planes_list[0].slim)\n", + "print(fit.model_images_of_planes_list[1].slim)\n", + "\n", + "print(fit.subtracted_images_of_planes_list[0].slim)\n", + "print(fit.subtracted_images_of_planes_list[1].slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Unmasked Quantities__\n", + "\n", + "All of the quantities above are computed using the mask which was used to fit the data.\n", + "\n", + "The `FitImaging` can also compute the unmasked blurred image of each plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.unmasked_blurred_image.native)\n", + "print(fit.unmasked_blurred_image_of_planes_list[0].native)\n", + "print(fit.unmasked_blurred_image_of_planes_list[1].native)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We can use the `Mask2D` object to mask regions of one of the fit's maps and estimate quantities of it.\n", + "\n", + "Below, we estimate the average absolute normalized residuals within a 1.0\" circular mask, which would inform us of\n", + "how accurate the lens light subtraction of a model fit is and if it leaves any significant residuals" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=fit.dataset.shape_native,\n", + " pixel_scales=fit.dataset.pixel_scales,\n", + " radius=1.0,\n", + ")\n", + "\n", + "normalized_residuals = fit.normalized_residual_map.apply_mask(mask=mask)\n", + "\n", + "print(np.mean(np.abs(normalized_residuals.slim)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixel Counting__\n", + "\n", + "An alternative way to quantify residuals like the lens light residuals is pixel counting. For example, we could sum\n", + "up the number of pixels whose chi-squared values are above 10 which indicates a poor fit to the data.\n", + "\n", + "Whereas computing the mean above the average level of residuals, pixel counting informs us how spatially large the\n", + "residuals extend. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=fit.dataset.shape_native,\n", + " pixel_scales=fit.dataset.pixel_scales,\n", + " radius=1.0,\n", + ")\n", + "\n", + "chi_squared_map = fit.chi_squared_map.apply_mask(mask=mask)\n", + "\n", + "print(np.sum(chi_squared_map > 10.0))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Outputting Results__\n", + "\n", + "You may wish to output certain results to .fits files for later inspection. \n", + "\n", + "For example, one could output the lens light subtracted image of the lensed source galaxy to a .fits file such that\n", + "we could fit this source-only image again with an independent pipeline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_subtracted_image = fit.subtracted_images_of_planes_list[1]\n", + "aplt.fits_array(\n", + " array=lens_subtracted_image,\n", + " file_path=dataset_path / \"lens_subtracted_data.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/likelihood_function.ipynb b/notebooks/imaging/likelihood_function.ipynb index 4d86ee960..197322be4 100644 --- a/notebooks/imaging/likelihood_function.ipynb +++ b/notebooks/imaging/likelihood_function.ipynb @@ -1,934 +1,980 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Inversion (Parametric)__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data with\n", - "a lens light profile and source light profile (e.g. an elliptical Sersic lens and source).\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers using, so that readers can understand the likelihood\n", - " function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", - "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", - "packages are called when the likelihood is evaluated.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Masked Image Grid:** To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", - "- **Lens Galaxy Mass:** We next define the mass profiles which represents the lens galaxy's mass, which will be used to.\n", - "- **Lens Galaxy:** We now combine the light and mass profiles into a single `Galaxy` object for the lens galaxy.\n", - "- **Source Galaxy Light Profile:** The source galaxy is fitted using another analytic light profile, in this example another.\n", - "- **Lens Light:** Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the an.\n", - "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the.\n", - "- **Source Image:** We pass the traced grid and blurring grid of coordinates to the source galaxy to evaluate its 2D.\n", - "- **Convolution:** Convolve the 2D image of the lens and source above with the PSF in real-space (as opposed to via an.\n", - "- **Likelihood Function:** We now quantify the goodness-of-fit of our lens and source model.\n", - "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", - "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", - "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the lens model, by combining the two terms.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autoarray as aa\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "In order to perform a likelihood evaluation, we first load a dataset.\n", - "\n", - "This example fits a simulated galaxy where the imaging resolution is 0.1 arcsecond-per-pixel resolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This guide uses in-built visualization tools for plotting. \n", - "\n", - "For example, using the `aplt.subplot_imaging_dataset` the imaging dataset we perform a likelihood evaluation on is plotted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Noise Scaling__\n", - "\n", - "Before masking, we must deal with any extra galaxies in the data: nearby galaxies (or foreground stars, or\n", - "data-reduction artefacts) whose emission is not associated with the strong lens but blends into the field. If\n", - "their light is left in the data it will contaminate the likelihood evaluation and bias the inferred lens model.\n", - "It is too easy to skip straight to modeling without checking for these, so we make this step explicit.\n", - "\n", - "To prevent extra galaxies from impacting the fit, we do not mask them entirely from the fit. Instead, the pixels\n", - "are kept in the fit but their data values are scaled to zero and their noise-map values increased to very large\n", - "values, so they contribute negligibly to the likelihood. This is preferable to removing the pixels entirely\n", - "(e.g. for a pixelized source reconstruction, removing pixels can produce discontinuities in the pixelization).\n", - "\n", - "The `simple` dataset includes a faint extra galaxy, and a `mask_extra_galaxies.fits` covering it is shipped with\n", - "the dataset (created by the simulator). If you are modeling your own data with an extra galaxy, you must either\n", - "create such a mask using the data-preparation tools\n", - "(`autolens_workspace/*/imaging/data_preparation/gui/mask_extra_galaxies.py`, or the manual\n", - "`data_preparation/examples/optional/mask_extra_galaxies.py`), or shrink the circular mask below so the extra\n", - "galaxy lies outside it and is removed from the fit entirely." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_extra_galaxies = al.Mask2D.from_fits(\n", - " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", - " pixel_scales=dataset.pixel_scales,\n", - " invert=True, # `True` means a pixel is scaled.\n", - ")\n", - "\n", - "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "The likelihood is only evaluated using image pixels contained within a 2D mask, which we choose before performing\n", - "lens modeling.\n", - "\n", - "Below, we define a 2D circular mask with a 3.0\" radius." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "masked_dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When we plot the masked imaging, only the circular masked region is shown." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=masked_dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling evaluates a light profile using multiple samples of its intensity per image-pixel.\n", - "\n", - "For simplicity, we disable over sampling in this guide by setting `sub_size=1`. \n", - "\n", - "a full description of over sampling and how to use it is given in `autolens_workspace/*/guides/over_sampling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masked Image Grid__\n", - "\n", - "To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", - "\n", - "These are given by `masked_dataset.grids.lp`, which we can plot and see is a uniform grid of (y,x) Cartesian \n", - "coordinates which have had the 3.0\" circular mask applied.\n", - "\n", - "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", - "used to perform ray-tracing and evaluate a light profile the intensity of the profile at the centre of each \n", - "image-pixel is computed, making it straight forward to compute the light profile's image to the image data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=masked_dataset.grids.lp, title=\"\")\n", - "\n", - "print(\n", - " f\"(y,x) coordinates of first ten unmasked image-pixels {masked_dataset.grid[0:9]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To perform lensing calculations we convert this 2D (y,x) grid of coordinates to elliptical coordinates:\n", - "\n", - " $\\eta = \\sqrt{(x - x_c)^2 + (y - y_c)^2/q^2}$\n", - "\n", - "Where:\n", - "\n", - " - $y$ and $x$ are the (y,x) arc-second coordinates of each unmasked image-pixel, given by `masked_dataset.grids.lp`.\n", - " - $y_c$ and $x_c$ are the (y,x) arc-second `centre` of the light or mass profile used to perform lensing calculations.\n", - " - $q$ is the axis-ratio of the elliptical light or mass profile (`axis_ratio=1.0` for spherical profiles).\n", - " - The elliptical coordinates is rotated by position angle $\\phi$, defined counter-clockwise from the positive \n", - " x-axis.\n", - "\n", - "$q$ and $\\phi$ are not used to parameterize a light profile but expresses these as \"elliptical components\", \n", - "or `ell_comps` for short:\n", - "\n", - "$\\epsilon_{1} =\\frac{1-q}{1+q} \\sin 2\\phi, \\,\\,$\n", - "$\\epsilon_{2} =\\frac{1-q}{1+q} \\cos 2\\phi.$\n", - "\n", - "Note that `Ell` is used as shorthand for elliptical and `Sph` for spherical." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "profile = al.EllProfile(centre=(0.1, 0.2), ell_comps=(0.1, 0.2))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Transform `masked_dataset.grids.lp` to the centre of profile and rotate it using its angle `phi`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "transformed_grid = profile.transformed_to_reference_frame_grid_from(\n", - " grid=masked_dataset.grids.lp\n", - ")\n", - "\n", - "aplt.plot_grid(grid=transformed_grid, title=\"\")\n", - "print(\n", - " f\"transformed coordinates of first ten unmasked image-pixels {transformed_grid[0:9]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Using these transformed (y',x') values we compute the elliptical coordinates $\\eta = \\sqrt{(x')^2 + (y')^2/q^2}$" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "elliptical_radii = profile.elliptical_radii_grid_from(grid=transformed_grid)\n", - "\n", - "print(\n", - " f\"elliptical coordinates of first ten unmasked image-pixels {elliptical_radii[0:9]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy Light (Setup)__\n", - "\n", - "To perform a likelihood evaluation we now compose our lens model.\n", - "\n", - "We first define the light profiles which represents the lens galaxy's light, which will be used to fit the lens \n", - "light.\n", - "\n", - "A light profile is defined by its intensity $I (\\eta_{\\rm l}) $, for example the Sersic profile:\n", - "\n", - "$I_{\\rm Ser} (\\eta_{\\rm l}) = I \\exp \\bigg\\{ -k \\bigg[ \\bigg( \\frac{\\eta}{R} \\bigg)^{\\frac{1}{n}} - 1 \\bigg] \\bigg\\}$\n", - "\n", - "Where:\n", - "\n", - " - $\\eta$ are the elliptical coordinates (see above) or the masked image-grid.\n", - " - $I$ is the `intensity`, which controls the overall brightness of the Sersic profile.\n", - " - $n$ is the ``sersic_index``, which via $k$ controls the steepness of the inner profile.\n", - " - $R$ is the `effective_radius`, which defines the arc-second radius of a circle containing half the light.\n", - "\n", - "In this example, we assume our lens is composed of one light profile, an elliptical Sersic which represent the \n", - "bulge of the lens. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=4.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Using the masked 2D grid defined above, we can calculate and plot images of each light profile component.\n", - "\n", - "(The transformation to elliptical coordinates above are built into the `image_2d_from` function and performed \n", - "implicitly)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_2d_bulge = bulge.image_2d_from(grid=masked_dataset.grid)\n", - "\n", - "aplt.plot_array(array=bulge.image_2d_from(grid=masked_dataset.grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy Mass__\n", - "\n", - "We next define the mass profiles which represents the lens galaxy's mass, which will be used to ray-trace the \n", - "image-plane 2D grid of (y,x) coordinates to the source-plane so that the source model can be evaluated.\n", - "\n", - "In this example, we assume our lens is composed of an elliptical isothermal mass distribution and external shear.\n", - "\n", - "A mass profile is defined by its convergence $\\kappa (\\eta)$, which is related to\n", - "the surface density of the mass distribution as\n", - "\n", - "$\\kappa(\\eta)=\\frac{\\Sigma(\\eta)}{\\Sigma_\\mathrm{crit}},$\n", - "\n", - "where\n", - "\n", - "$\\Sigma_\\mathrm{crit}=\\frac{{\\rm c}^2}{4{\\rm \\pi} {\\rm G}}\\frac{D_{\\rm s}}{D_{\\rm l} D_{\\rm ls}},$\n", - "\n", - "and\n", - "\n", - " - `c` is the speed of light.\n", - " - $D_{\\rm l}$, $D_{\\rm s}$, and $D_{\\rm ls}$ are respectively the angular diameter distances to the lens, to the \n", - " source, and from the lens to the source.\n", - "\n", - "For readers less familiar with lensing, we can think of $\\kappa(\\eta)$ as a convenient and\n", - "dimensionless way to describe how light is gravitationally lensed after assuming a cosmology.\n", - "\n", - "For the for the isothermal profile:\n", - "\n", - "$\\kappa(\\eta) = \\frac{1.0}{1 + q} \\bigg( \\frac{\\theta_{\\rm E}}{\\eta} \\bigg)$\n", - "\n", - "Where:\n", - "\n", - " - $\\theta_{\\rm E}$ is the `einstein_radius` (which is rescaled compared to other einstein radius\n", - " definitions)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "aplt.plot_array(\n", - " array=mass.convergence_2d_from(grid=masked_dataset.grid), title=\"Convergence\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From each mass profile we can compute its deflection angles, which describe how due to gravitational lensing\n", - "image-pixels are ray-traced to the source plane.\n", - "\n", - "The deflection angles are computed by integrating $\\kappa$: \n", - "\n", - "$\\vec{{\\alpha}}_{\\rm x,y} (\\vec{x}) = \\frac{1}{\\pi} \\int \\frac{\\vec{x} - \\vec{x'}}{\\left | \\vec{x} - \\vec{x'} \\right |^2} \\kappa(\\vec{x'}) d\\vec{x'} \\, ,$" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_yx_2d = mass.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "\n", - "deflections = mass.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=masked_dataset.grid.mask)\n", - "aplt.plot_array(array=deflections_y, title=\"Deflections Y\")\n", - "deflections = mass.deflections_yx_2d_from(grid=masked_dataset.grid)\n", - "deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=masked_dataset.grid.mask)\n", - "aplt.plot_array(array=deflections_x, title=\"Deflections X\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "We now combine the light and mass profiles into a single `Galaxy` object for the lens galaxy.\n", - "\n", - "When computing quantities for the light and mass profiles from this object, it computes each individual quantity and \n", - "adds them together. \n", - "\n", - "For example, for the `bulge`, when it computes their 2D images it computes each individually and then adds\n", - "them together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Light Profile__\n", - "\n", - "The source galaxy is fitted using another analytic light profile, in this example another elliptical Sersic." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Light__\n", - "\n", - "Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the an MGE \n", - "bulge). \n", - "\n", - "This computes the `lens_image_2d` of each `LightProfile` and adds them together. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To convolve the lens's 2D image with the imaging data's PSF, we need its `blurring_image`. This represents all flux \n", - "values not within the mask, which are close enough to it that their flux blurs into the mask after PSF convolution.\n", - "\n", - "To compute this, a `blurring_mask` and `blurring_grid` are used, corresponding to these pixels near the edge of the \n", - "actual mask whose light blurs into the image:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the image-plane to its (y,x) \n", - "source-plane coordinate $\\beta$ using the summed deflection angles $\\alpha$ of the mass profiles:\n", - "\n", - " $\\beta = \\theta - \\alpha(\\theta)$\n", - "\n", - "The likelihood function of a source light profile ray-traces two grids from the image-plane to the source-plane:\n", - "\n", - " 1) A 2D grid of (y,x) coordinates aligned with the imaging data's image-pixels.\n", - " \n", - " 2) The 2D blurring grid (used for the lens light above) which accounts for pixels at the edge of the mask whose\n", - " light blurs into the mask.\n", - " \n", - "The function below computes the 2D deflection angles of the tracer's lens galaxies and subtracts them from the \n", - "image-plane 2D (y,x) coordinates $\\theta$ of each grid, thus ray-tracing their coordinates to the source plane to \n", - "compute their $\\beta$ values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", - "\n", - "\n", - "aplt.plot_grid(grid=traced_grid, title=\"\")\n", - "\n", - "traced_blurring_grid = tracer.traced_grid_2d_list_from(\n", - " grid=masked_dataset.grids.blurring\n", - ")[-1]\n", - "\n", - "\n", - "aplt.plot_grid(grid=traced_blurring_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Image__\n", - "\n", - "We pass the traced grid and blurring grid of coordinates to the source galaxy to evaluate its 2D image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_image_2d = source_galaxy.image_2d_from(grid=traced_grid)\n", - "\n", - "\n", - "source_blurring_image_2d = source_galaxy.image_2d_from(grid=traced_blurring_grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens + Source Light Addition__\n", - "\n", - "We add the lens and source galaxy images and blurring together, to create an overall image of the strong lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = lens_image_2d + source_image_2d\n", - "\n", - "aplt.plot_array(array=image, title=\"\")\n", - "\n", - "blurring_image_2d = lens_blurring_image_2d + source_blurring_image_2d\n", - "\n", - "aplt.plot_array(array=blurring_image_2d, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Convolution__\n", - "\n", - "Convolve the 2D image of the lens and source above with the PSF in real-space (as opposed to via an FFT) using \n", - "a `Kernal2D`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", - " image=image, blurring_image=blurring_image_2d\n", - ")\n", - "\n", - "aplt.plot_array(array=convolved_image_2d, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our lens and source model.\n", - "\n", - "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", - "\n", - "The likelihood function for parametric lens modeling consists of two terms:\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "We now explain what each of these terms mean.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", - "\n", - " - `model_data` = `convolved_image_2d`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", - "\n", - "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", - "for, corresponding to a fit with a lower likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_image = convolved_image_2d\n", - "\n", - "residual_map = masked_dataset.data - model_image\n", - "normalized_residual_map = residual_map / masked_dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(chi_squared)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", - "\n", - "aplt.plot_array(array=chi_squared_map, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", - "\n", - "The final term ins the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", - "of the log of every noise-map value squared. \n", - "\n", - "Given the `noise_map` is fixed, this term does not change during the lens modeling process and has no impact on the \n", - "model we infer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "We can now, finally, compute the `log_likelihood` of the lens model, by combining the two terms computed above using\n", - "the likelihood function defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(figure_of_merit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This 11 step process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)\n", - "fit_figure_of_merit = fit.figure_of_merit\n", - "print(fit_figure_of_merit)\n", - "\n", - "aplt.subplot_fit_imaging(fit=fit)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `Nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the parametric likelihood function, which uses analytic\n", - "light profiles to fit the lens and source light.\n", - "\n", - "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", - "are described in additional notebooks found in this package. In brief, these describe:\n", - "\n", - " - **Sub-gridding**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually\n", - " ray-traced to the source-plane and used to evaluate the light profile more accurately.\n", - "\n", - "__JAX__\n", - "\n", - "The step-by-step likelihood you've just walked through can be JAX-\n", - "accelerated by wrapping the whole construction in `@jax.jit`. The\n", - "pattern:\n", - "\n", - "```python\n", - "import jax\n", - "import jax.numpy as jnp\n", - "from autolens.jax import register_tracer_classes\n", - "\n", - "# One-time setup: register Tracer + Galaxy + profile classes as JAX\n", - "# pytrees so the tracer can cross the @jax.jit boundary as an argument.\n", - "register_tracer_classes(tracer)\n", - "\n", - "@jax.jit\n", - "def my_log_likelihood(instance):\n", - " tracer = al.Tracer(galaxies=instance.galaxies)\n", - " fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", - " return fit.log_likelihood\n", - "```\n", - "\n", - "To validate the JAX path matches the NumPy chi-squared you just\n", - "computed, use `Fitness._vmap` (the production validation pattern \u2014\n", - "single `jax.jit(fn)(concrete)` hides un-threaded `xp` sites that\n", - "`vmap(jit(call))` exposes):\n", - "\n", - "```python\n", - "from autofit.non_linear.fitness import Fitness\n", - "\n", - "fitness = Fitness(\n", - " model=model,\n", - " analysis=al.AnalysisImaging(dataset=dataset),\n", - " fom_is_log_likelihood=True,\n", - ")\n", - "log_l_jax = fitness._vmap(jnp.array([instance_parameters]))[0]\n", - "assert np.isclose(log_l_jax, log_l_numpy_from_walkthrough)\n", - "```\n", - "\n", - "For the canonical Analysis-driven modeling path (where you write zero\n", - "JAX code), see `start_here.py` / `modeling.py`. For JIT-ing library\n", - "methods directly (`tracer.image_2d_from`, `LensCalc.magnification_2d_via_hessian_from`,\n", - "etc.) without going through `FitImaging`, see\n", - "`scripts/guides/lens_calc.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Inversion (Parametric)__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Imaging` data with\n", + "a lens light profile and source light profile (e.g. an elliptical Sersic lens and source).\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers using, so that readers can understand the likelihood\n", + " function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", + "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", + "packages are called when the likelihood is evaluated.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Masked Image Grid:** To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", + "- **Lens Galaxy Mass:** We next define the mass profiles which represents the lens galaxy's mass, which will be used to.\n", + "- **Lens Galaxy:** We now combine the light and mass profiles into a single `Galaxy` object for the lens galaxy.\n", + "- **Source Galaxy Light Profile:** The source galaxy is fitted using another analytic light profile, in this example another.\n", + "- **Lens Light:** Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the an.\n", + "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the.\n", + "- **Source Image:** We pass the traced grid and blurring grid of coordinates to the source galaxy to evaluate its 2D.\n", + "- **Convolution:** Convolve the 2D image of the lens and source above with the PSF in real-space (as opposed to via an.\n", + "- **Likelihood Function:** We now quantify the goodness-of-fit of our lens and source model.\n", + "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", + "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", + "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the lens model, by combining the two terms.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autoarray as aa\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "In order to perform a likelihood evaluation, we first load a dataset.\n", + "\n", + "This example fits a simulated galaxy where the imaging resolution is 0.1 arcsecond-per-pixel resolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", \"imaging\", \"simple\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "# PSF convolution runs at the image resolution (sub size 1), which is the fastest\n", + "# option and accurate for well-sampled PSFs. Supplying a PSF at a multiple of the\n", + "# image resolution and raising this value improves blurring fidelity for\n", + "# undersampled PSFs (e.g. HST / Euclid VIS) at extra compute cost \u2014 see\n", + "# `guides/advanced/over_sampling.py` and the simulator's `__Oversampled PSF__` section.\n", + "psf_convolve_over_sample_size = 1\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " convolve_over_sample_size_lp=psf_convolve_over_sample_size,\n", + " convolve_over_sample_size_pixelization=psf_convolve_over_sample_size,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This guide uses in-built visualization tools for plotting. \n", + "\n", + "For example, using the `aplt.subplot_imaging_dataset` the imaging dataset we perform a likelihood evaluation on is plotted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Noise Scaling__\n", + "\n", + "Before masking, we must deal with any extra galaxies in the data: nearby galaxies (or foreground stars, or\n", + "data-reduction artefacts) whose emission is not associated with the strong lens but blends into the field. If\n", + "their light is left in the data it will contaminate the likelihood evaluation and bias the inferred lens model.\n", + "It is too easy to skip straight to modeling without checking for these, so we make this step explicit.\n", + "\n", + "To prevent extra galaxies from impacting the fit, we do not mask them entirely from the fit. Instead, the pixels\n", + "are kept in the fit but their data values are scaled to zero and their noise-map values increased to very large\n", + "values, so they contribute negligibly to the likelihood. This is preferable to removing the pixels entirely\n", + "(e.g. for a pixelized source reconstruction, removing pixels can produce discontinuities in the pixelization).\n", + "\n", + "The `simple` dataset includes a faint extra galaxy, and a `mask_extra_galaxies.fits` covering it is shipped with\n", + "the dataset (created by the simulator). If you are modeling your own data with an extra galaxy, you must either\n", + "create such a mask using the data-preparation tools\n", + "(`autolens_workspace/*/imaging/data_preparation/gui/mask_extra_galaxies.py`, or the manual\n", + "`data_preparation/examples/optional/mask_extra_galaxies.py`), or shrink the circular mask below so the extra\n", + "galaxy lies outside it and is removed from the fit entirely." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_extra_galaxies = al.Mask2D.from_fits(\n", + " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", + " pixel_scales=dataset.pixel_scales,\n", + " invert=True, # `True` means a pixel is scaled.\n", + ")\n", + "\n", + "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "The likelihood is only evaluated using image pixels contained within a 2D mask, which we choose before performing\n", + "lens modeling.\n", + "\n", + "Below, we define a 2D circular mask with a 3.0\" radius." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "masked_dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When we plot the masked imaging, only the circular masked region is shown." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=masked_dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling evaluates a light profile using multiple samples of its intensity per image-pixel.\n", + "\n", + "For simplicity, we disable over sampling in this guide by setting `sub_size=1`. \n", + "\n", + "a full description of over sampling and how to use it is given in `autolens_workspace/*/guides/over_sampling.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "masked_dataset = masked_dataset.apply_over_sampling(over_sample_size_lp=1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masked Image Grid__\n", + "\n", + "To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", + "\n", + "These are given by `masked_dataset.grids.lp`, which we can plot and see is a uniform grid of (y,x) Cartesian \n", + "coordinates which have had the 3.0\" circular mask applied.\n", + "\n", + "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", + "used to perform ray-tracing and evaluate a light profile the intensity of the profile at the centre of each \n", + "image-pixel is computed, making it straight forward to compute the light profile's image to the image data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=masked_dataset.grids.lp, title=\"\")\n", + "\n", + "print(\n", + " f\"(y,x) coordinates of first ten unmasked image-pixels {masked_dataset.grid[0:9]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To perform lensing calculations we convert this 2D (y,x) grid of coordinates to elliptical coordinates:\n", + "\n", + " $\\eta = \\sqrt{(x - x_c)^2 + (y - y_c)^2/q^2}$\n", + "\n", + "Where:\n", + "\n", + " - $y$ and $x$ are the (y,x) arc-second coordinates of each unmasked image-pixel, given by `masked_dataset.grids.lp`.\n", + " - $y_c$ and $x_c$ are the (y,x) arc-second `centre` of the light or mass profile used to perform lensing calculations.\n", + " - $q$ is the axis-ratio of the elliptical light or mass profile (`axis_ratio=1.0` for spherical profiles).\n", + " - The elliptical coordinates is rotated by position angle $\\phi$, defined counter-clockwise from the positive \n", + " x-axis.\n", + "\n", + "$q$ and $\\phi$ are not used to parameterize a light profile but expresses these as \"elliptical components\", \n", + "or `ell_comps` for short:\n", + "\n", + "$\\epsilon_{1} =\\frac{1-q}{1+q} \\sin 2\\phi, \\,\\,$\n", + "$\\epsilon_{2} =\\frac{1-q}{1+q} \\cos 2\\phi.$\n", + "\n", + "Note that `Ell` is used as shorthand for elliptical and `Sph` for spherical." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "profile = al.EllProfile(centre=(0.1, 0.2), ell_comps=(0.1, 0.2))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Transform `masked_dataset.grids.lp` to the centre of profile and rotate it using its angle `phi`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "transformed_grid = profile.transformed_to_reference_frame_grid_from(\n", + " grid=masked_dataset.grids.lp\n", + ")\n", + "\n", + "aplt.plot_grid(grid=transformed_grid, title=\"\")\n", + "print(\n", + " f\"transformed coordinates of first ten unmasked image-pixels {transformed_grid[0:9]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Using these transformed (y',x') values we compute the elliptical coordinates $\\eta = \\sqrt{(x')^2 + (y')^2/q^2}$" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "elliptical_radii = profile.elliptical_radii_grid_from(grid=transformed_grid)\n", + "\n", + "print(\n", + " f\"elliptical coordinates of first ten unmasked image-pixels {elliptical_radii[0:9]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy Light (Setup)__\n", + "\n", + "To perform a likelihood evaluation we now compose our lens model.\n", + "\n", + "We first define the light profiles which represents the lens galaxy's light, which will be used to fit the lens \n", + "light.\n", + "\n", + "A light profile is defined by its intensity $I (\\eta_{\\rm l}) $, for example the Sersic profile:\n", + "\n", + "$I_{\\rm Ser} (\\eta_{\\rm l}) = I \\exp \\bigg\\{ -k \\bigg[ \\bigg( \\frac{\\eta}{R} \\bigg)^{\\frac{1}{n}} - 1 \\bigg] \\bigg\\}$\n", + "\n", + "Where:\n", + "\n", + " - $\\eta$ are the elliptical coordinates (see above) or the masked image-grid.\n", + " - $I$ is the `intensity`, which controls the overall brightness of the Sersic profile.\n", + " - $n$ is the ``sersic_index``, which via $k$ controls the steepness of the inner profile.\n", + " - $R$ is the `effective_radius`, which defines the arc-second radius of a circle containing half the light.\n", + "\n", + "In this example, we assume our lens is composed of one light profile, an elliptical Sersic which represent the \n", + "bulge of the lens. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=4.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Using the masked 2D grid defined above, we can calculate and plot images of each light profile component.\n", + "\n", + "(The transformation to elliptical coordinates above are built into the `image_2d_from` function and performed \n", + "implicitly)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_2d_bulge = bulge.image_2d_from(grid=masked_dataset.grid)\n", + "\n", + "aplt.plot_array(array=bulge.image_2d_from(grid=masked_dataset.grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy Mass__\n", + "\n", + "We next define the mass profiles which represents the lens galaxy's mass, which will be used to ray-trace the \n", + "image-plane 2D grid of (y,x) coordinates to the source-plane so that the source model can be evaluated.\n", + "\n", + "In this example, we assume our lens is composed of an elliptical isothermal mass distribution and external shear.\n", + "\n", + "A mass profile is defined by its convergence $\\kappa (\\eta)$, which is related to\n", + "the surface density of the mass distribution as\n", + "\n", + "$\\kappa(\\eta)=\\frac{\\Sigma(\\eta)}{\\Sigma_\\mathrm{crit}},$\n", + "\n", + "where\n", + "\n", + "$\\Sigma_\\mathrm{crit}=\\frac{{\\rm c}^2}{4{\\rm \\pi} {\\rm G}}\\frac{D_{\\rm s}}{D_{\\rm l} D_{\\rm ls}},$\n", + "\n", + "and\n", + "\n", + " - `c` is the speed of light.\n", + " - $D_{\\rm l}$, $D_{\\rm s}$, and $D_{\\rm ls}$ are respectively the angular diameter distances to the lens, to the \n", + " source, and from the lens to the source.\n", + "\n", + "For readers less familiar with lensing, we can think of $\\kappa(\\eta)$ as a convenient and\n", + "dimensionless way to describe how light is gravitationally lensed after assuming a cosmology.\n", + "\n", + "For the for the isothermal profile:\n", + "\n", + "$\\kappa(\\eta) = \\frac{1.0}{1 + q} \\bigg( \\frac{\\theta_{\\rm E}}{\\eta} \\bigg)$\n", + "\n", + "Where:\n", + "\n", + " - $\\theta_{\\rm E}$ is the `einstein_radius` (which is rescaled compared to other einstein radius\n", + " definitions)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "aplt.plot_array(\n", + " array=mass.convergence_2d_from(grid=masked_dataset.grid), title=\"Convergence\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From each mass profile we can compute its deflection angles, which describe how due to gravitational lensing\n", + "image-pixels are ray-traced to the source plane.\n", + "\n", + "The deflection angles are computed by integrating $\\kappa$: \n", + "\n", + "$\\vec{{\\alpha}}_{\\rm x,y} (\\vec{x}) = \\frac{1}{\\pi} \\int \\frac{\\vec{x} - \\vec{x'}}{\\left | \\vec{x} - \\vec{x'} \\right |^2} \\kappa(\\vec{x'}) d\\vec{x'} \\, ,$" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_yx_2d = mass.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "\n", + "deflections = mass.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=masked_dataset.grid.mask)\n", + "aplt.plot_array(array=deflections_y, title=\"Deflections Y\")\n", + "deflections = mass.deflections_yx_2d_from(grid=masked_dataset.grid)\n", + "deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=masked_dataset.grid.mask)\n", + "aplt.plot_array(array=deflections_x, title=\"Deflections X\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "We now combine the light and mass profiles into a single `Galaxy` object for the lens galaxy.\n", + "\n", + "When computing quantities for the light and mass profiles from this object, it computes each individual quantity and \n", + "adds them together. \n", + "\n", + "For example, for the `bulge`, when it computes their 2D images it computes each individually and then adds\n", + "them together." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(redshift=0.5, bulge=bulge, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Light Profile__\n", + "\n", + "The source galaxy is fitted using another analytic light profile, in this example another elliptical Sersic." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Light__\n", + "\n", + "Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the an MGE \n", + "bulge). \n", + "\n", + "This computes the `lens_image_2d` of each `LightProfile` and adds them together. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To convolve the lens's 2D image with the imaging data's PSF, we need its `blurring_image`. This represents all flux \n", + "values not within the mask, which are close enough to it that their flux blurs into the mask after PSF convolution.\n", + "\n", + "To compute this, a `blurring_mask` and `blurring_grid` are used, corresponding to these pixels near the edge of the \n", + "actual mask whose light blurs into the image:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_blurring_image_2d = lens_galaxy.image_2d_from(grid=masked_dataset.grids.blurring)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the image-plane to its (y,x) \n", + "source-plane coordinate $\\beta$ using the summed deflection angles $\\alpha$ of the mass profiles:\n", + "\n", + " $\\beta = \\theta - \\alpha(\\theta)$\n", + "\n", + "The likelihood function of a source light profile ray-traces two grids from the image-plane to the source-plane:\n", + "\n", + " 1) A 2D grid of (y,x) coordinates aligned with the imaging data's image-pixels.\n", + " \n", + " 2) The 2D blurring grid (used for the lens light above) which accounts for pixels at the edge of the mask whose\n", + " light blurs into the mask.\n", + " \n", + "The function below computes the 2D deflection angles of the tracer's lens galaxies and subtracts them from the \n", + "image-plane 2D (y,x) coordinates $\\theta$ of each grid, thus ray-tracing their coordinates to the source plane to \n", + "compute their $\\beta$ values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=masked_dataset.grid)[-1]\n", + "\n", + "\n", + "aplt.plot_grid(grid=traced_grid, title=\"\")\n", + "\n", + "traced_blurring_grid = tracer.traced_grid_2d_list_from(\n", + " grid=masked_dataset.grids.blurring\n", + ")[-1]\n", + "\n", + "\n", + "aplt.plot_grid(grid=traced_blurring_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Image__\n", + "\n", + "We pass the traced grid and blurring grid of coordinates to the source galaxy to evaluate its 2D image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_image_2d = source_galaxy.image_2d_from(grid=traced_grid)\n", + "\n", + "\n", + "source_blurring_image_2d = source_galaxy.image_2d_from(grid=traced_blurring_grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens + Source Light Addition__\n", + "\n", + "We add the lens and source galaxy images and blurring together, to create an overall image of the strong lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = lens_image_2d + source_image_2d\n", + "\n", + "aplt.plot_array(array=image, title=\"\")\n", + "\n", + "blurring_image_2d = lens_blurring_image_2d + source_blurring_image_2d\n", + "\n", + "aplt.plot_array(array=blurring_image_2d, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Convolution__\n", + "\n", + "Convolve the 2D image of the lens and source above with the PSF in real-space (as opposed to via an FFT) using \n", + "a `Kernal2D`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "convolved_image_2d = masked_dataset.psf.convolved_image_from(\n", + " image=image, blurring_image=blurring_image_2d\n", + ")\n", + "\n", + "aplt.plot_array(array=convolved_image_2d, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our lens and source model.\n", + "\n", + "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", + "\n", + "The likelihood function for parametric lens modeling consists of two terms:\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "We now explain what each of these terms mean.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", + "\n", + " - `model_data` = `convolved_image_2d`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", + "\n", + "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", + "for, corresponding to a fit with a lower likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_image = convolved_image_2d\n", + "\n", + "residual_map = masked_dataset.data - model_image\n", + "normalized_residual_map = residual_map / masked_dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(chi_squared)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_map = al.Array2D(values=chi_squared_map, mask=mask)\n", + "\n", + "aplt.plot_array(array=chi_squared_map, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", + "\n", + "The final term ins the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", + "of the log of every noise-map value squared. \n", + "\n", + "Given the `noise_map` is fixed, this term does not change during the lens modeling process and has no impact on the \n", + "model we infer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(np.sum(np.log(2 * np.pi * masked_dataset.noise_map**2.0)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "We can now, finally, compute the `log_likelihood` of the lens model, by combining the two terms computed above using\n", + "the likelihood function defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(figure_of_merit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This 11 step process to perform a likelihood function evaluation is what is performed in the `FitImaging` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)\n", + "fit_figure_of_merit = fit.figure_of_merit\n", + "print(fit_figure_of_merit)\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `Nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the parametric likelihood function, which uses analytic\n", + "light profiles to fit the lens and source light.\n", + "\n", + "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", + "are described in additional notebooks found in this package. In brief, these describe:\n", + "\n", + " - **Sub-gridding**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually\n", + " ray-traced to the source-plane and used to evaluate the light profile more accurately.\n", + "\n", + "__JAX__\n", + "\n", + "The step-by-step likelihood you've just walked through can be JAX-\n", + "accelerated by wrapping the whole construction in `@jax.jit`. The\n", + "pattern:\n", + "\n", + "```python\n", + "import jax\n", + "import jax.numpy as jnp\n", + "from autolens.jax import register_tracer_classes\n", + "\n", + "# One-time setup: register Tracer + Galaxy + profile classes as JAX\n", + "# pytrees so the tracer can cross the @jax.jit boundary as an argument.\n", + "register_tracer_classes(tracer)\n", + "\n", + "@jax.jit\n", + "def my_log_likelihood(instance):\n", + " tracer = al.Tracer(galaxies=instance.galaxies)\n", + " fit = al.FitImaging(dataset=dataset, tracer=tracer)\n", + " return fit.log_likelihood\n", + "```\n", + "\n", + "To validate the JAX path matches the NumPy chi-squared you just\n", + "computed, use `Fitness._vmap` (the production validation pattern \u2014\n", + "single `jax.jit(fn)(concrete)` hides un-threaded `xp` sites that\n", + "`vmap(jit(call))` exposes):\n", + "\n", + "```python\n", + "from autofit.non_linear.fitness import Fitness\n", + "\n", + "fitness = Fitness(\n", + " model=model,\n", + " analysis=al.AnalysisImaging(dataset=dataset),\n", + " fom_is_log_likelihood=True,\n", + ")\n", + "log_l_jax = fitness._vmap(jnp.array([instance_parameters]))[0]\n", + "assert np.isclose(log_l_jax, log_l_numpy_from_walkthrough)\n", + "```\n", + "\n", + "For the canonical Analysis-driven modeling path (where you write zero\n", + "JAX code), see `start_here.py` / `modeling.py`. For JIT-ing library\n", + "methods directly (`tracer.image_2d_from`, `LensCalc.magnification_2d_via_hessian_from`,\n", + "etc.) without going through `FitImaging`, see\n", + "`scripts/guides/lens_calc.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/modeling.ipynb b/notebooks/imaging/modeling.ipynb index 3f61f29ef..1f54c96f3 100644 --- a/notebooks/imaging/modeling.ipynb +++ b/notebooks/imaging/modeling.ipynb @@ -1,978 +1,1024 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Iamging: Modeling\n", - "=================\n", - "\n", - "This script is the starting point for lens modeling of CCD imaging data (E.g. Hubble Space Telescope, Euclid) with\n", - "**PyAutoLens** and it provides an overview of the lens modeling API.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Plotters:** Overview of plotting tools used for visualization.\n", - "- **Simulation:** Overview of how the simulated dataset was generated.\n", - "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Extra Galaxies Noise Scaling:** Scale the noise of nearby contaminating galaxies so they do not impact the fit.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Model Composition:** Compose the lens model using the Model and Collection API.\n", - "- **Coordinates:** Coordinate system assumptions for the model-fit.\n", - "- **Improved Lens Model:** The previous model used S\u00e9rsic light profiles for the lens and source galaxies.\n", - "- **Linear Light Profiles:** The MGE model below uses a **linear light profile** for the bulge and disk via the ``lp_linear``.\n", - "- **Concise API:** The MGE model composition API is quite long and technical, so we simply load the MGE models for the.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Unique Identifier:** In the path above, the `unique_identifier` appears as a collection of characters, where this.\n", - "- **Iterations Per Update:** Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and.\n", - "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", - "- **VRAM Use:** When running with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM.\n", - "- **Run Times:** Profiling the expected run time of the model-fit.\n", - "- **Output Folder Layout:** Description of the structure of the `output` folder where results are written.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Features:** This script gives a concise overview of the basic modeling API, fitting one the simplest lens.\n", - "- **HowToLens:** This `start_here.py` script, and the features examples above, do not explain many details of how.\n", - "- **Modeling Customization:** The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a Multi Gaussian Expansion bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a Multi Gaussian Expansion.\n", - "\n", - "This lens model is simple and computationally fast to fit, and therefore acts as a good starting point for new\n", - "users.\n", - "\n", - "__Plotters__\n", - "\n", - "To produce images of the data plotting function objects are used, which are high-level wrappers of matplotlib\n", - "code which produce high quality visualization of strong lenses.\n", - "\n", - "The plotting function API is described in the script `autolens_workspace/*/guides/plot`.\n", - "\n", - "__Simulation__\n", - "\n", - "This script fits a simulated `Imaging` dataset of a strong lens, which is produced in the\n", - "script `autolens_workspace/*/imaging/simulator.py`\n", - "\n", - "__Data Preparation__\n", - "\n", - "The `Imaging` dataset fitted in this example confirms to a number of standard that make it suitable to be fitted in\n", - "**PyAutoLens**.\n", - "\n", - "If you are intending to fit your own strong lens data, you will need to ensure it conforms to these standards, which are\n", - "described in the script `autolens_workspace/*/imaging/data_preparation/start_here.ipynb`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens dataset `simple` via .fits files, which is a data format used by astronomers to store images.\n", - "\n", - "The `pixel_scales` define the arc-second to pixel conversion factor of the image, which for the dataset we are using \n", - "is 0.1\" / pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use an `aplt.subplot_imaging_dataset` the plot the data, including: \n", - "\n", - " - `data`: The image of the strong lens.\n", - " - `noise_map`: The noise-map of the image, which quantifies the noise in every pixel as their RMS values.\n", - " - `psf`: The point spread function of the image, which describes the blurring of the image by the telescope optics.\n", - " - `signal_to_noise_map`: Quantifies the signal-to-noise in every pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Noise Scaling__\n", - "\n", - "Before masking, we must deal with any extra galaxies in the data: nearby galaxies (or foreground stars, or\n", - "data-reduction artefacts) whose emission is not associated with the strong lens but blends into the field. If\n", - "their light is left in the data it will contaminate the model-fit and bias the inferred lens model. It is too\n", - "easy to skip straight to modeling without checking for these, so we make this step an explicit part of the\n", - "workflow.\n", - "\n", - "To prevent extra galaxies from impacting the model-fit, we do not mask them entirely from the fit, which\n", - "would be analogous to making the circular mask smaller or using a more refined mask. When pixels are masked and\n", - "removed entirely from the fit, their coordinates are not used when performing ray-tracing and the light of the\n", - "lens and source galaxies in these pixels not evaluated.\n", - "\n", - "Instead, the pixels are kept in the fit, but their data values are scaled to zero and their noise-map values\n", - "are increased to very large values. This means that during the model-fit, these pixels contribute negligibly to\n", - "the likelihood of the fit, and therefore do not impact the lens model.\n", - "\n", - "This approach is used because for certain types of modeling approaches, like a pixelized source reconstruction,\n", - "masking regions of the image in a way that removes their image pixels entirely from the fit can produce\n", - "discontinuities in the pixelization. This can lead to unexpected systematics and unsatisfactory results.\n", - "\n", - "In this case, applying the mask in a way where the image pixels are not removed from the fit, but their data and\n", - "noise-map values are scaled such that they contribute negligibly to the fit, is a better approach.\n", - "\n", - "The `simple` dataset includes a faint extra galaxy, and a `mask_extra_galaxies.fits` covering it is shipped with\n", - "the dataset (created by the simulator). If you are modeling your own data with an extra galaxy, you must either:\n", - "\n", - " - Create a `mask_extra_galaxies.fits` for it using the data-preparation tools (the GUI\n", - " `autolens_workspace/*/imaging/data_preparation/gui/mask_extra_galaxies.py`, or the manual\n", - " `autolens_workspace/*/imaging/data_preparation/examples/optional/mask_extra_galaxies.py`), then load it\n", - " as below; or\n", - " - Shrink the circular mask below so the extra galaxy lies outside it and is removed from the fit entirely.\n", - "\n", - "We then plot the dataset, where the extra galaxy's pixels now have their data scaled to zero and noise-map\n", - "values increased, making their signal-to-noise effectively zero." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_extra_galaxies = al.Mask2D.from_fits(\n", - " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", - " pixel_scales=dataset.pixel_scales,\n", - " invert=True, # `True` means a pixel is scaled.\n", - ")\n", - "\n", - "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "The model-fit requires a 2D mask defining the regions of the image we fit the lens model to the data.\n", - "\n", - "We create a 3.0 arcsecond circular mask and apply it to the `Imaging` object that the lens model fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we plot the masked data, the mask removes the exterior regions of the image where there is no emission from the \n", - "lens and lensed source galaxies.\n", - "\n", - "The mask used to fit the data can be customized, as described in \n", - "the script `autolens_workspace/*/guides/modeling/customize`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated \n", - "on a higher resolution grid than the image data to ensure the calculation is accurate. \n", - "\n", - "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of \n", - "over sampling to ensure the lensed images are evaluated accurately.\n", - "\n", - "For a new user, the details of over-sampling are not important, therefore just be aware that calculations either:\n", - "\n", - " (i) use adaptive over sampling for the foregorund lens's light, which ensures high accuracy across. \n", - " (ii) use cored light profiles for the background source galaxy, where the core ensures low levels of over-sampling \n", - " produce numerically accurate but fast to compute results.\n", - "\n", - "\n", - "Once you are more experienced, you should read up on over-sampling in more detail via \n", - "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The imaging subplot updates the bottom two panels to reflect the update to over sampling, which now uses a higher\n", - "values in the centre.\n", - "\n", - "Whilst you may not yet understand the details of over-sampling, you can at least track it visually in the plots\n", - "and later learnt more about it in the `over_sampling.ipynb` guide." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "In this example we compose a lens model where:\n", - "\n", - " - The lens galaxy's light is a `Sersic` bulge [7 parameters].\n", - " \n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - " \n", - " - The source galaxy's light is a `SersicCore` bulge [7 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=21.\n", - "\n", - "__Model Composition__\n", - "\n", - "The API below for composing a lens model uses the `Model` and `Collection` objects, which are imported from \n", - "**PyAutoLens**'s parent project **PyAutoFit** \n", - "\n", - "The API is fairly self explanatory and is straight forward to extend, for example adding more light profiles\n", - "to the lens and source or using a different mass profile.\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", - "\n", - "__Coordinates__\n", - "\n", - "The model fitting default settings assume that the lens galaxy centre is near the coordinates (0.0\", 0.0\"). \n", - "\n", - "If for your dataset the lens is not centred at (0.0\", 0.0\"), we recommend that you either: \n", - "\n", - " - Reduce your data so that the centre is (`autolens_workspace/*/data_preparation`). \n", - " - Manually override the lens model priors (`autolens_workspace/*/guides/modeling/customize`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lp.Sersic)\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "The `info` below may not display optimally on your computer screen, for example the whitespace between parameter\n", - "names on the left and parameter priors on the right may lead them to appear across multiple lines. This is a\n", - "common issue in Jupyter notebooks.\n", - "\n", - "The`info_whitespace_length` parameter in the file `config/general.yaml` in the [output] section can be changed to \n", - "increase or decrease the amount of whitespace (The Jupyter notebook kernel will need to be reset for this change to \n", - "appear in a notebook)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Improved Lens Model__\n", - "\n", - "The previous model used S\u00e9rsic light profiles for the lens and source galaxies. This makes the model API concise, \n", - "readable, and easy to follow.\n", - "\n", - "However, single S\u00e9rsic profiles perform poorly for most strong lenses. Symmetric profiles (e.g. elliptical S\u00e9rsics) \n", - "typically leave significant residuals because they cannot capture the irregular and asymmetric morphology of real \n", - "galaxies (e.g. isophotal twists, radially varying ellipticity).\n", - "\n", - "This example therefore uses a lens model that combines two features, described in detail elsewhere (but a brief \n", - "overview is provided below):\n", - "\n", - "- **Linear light profiles** (see ``autolens_workspace/*/imaging/features/linear_light_profiles``)\n", - "- **Multi-Gaussian Expansion (MGE) light profiles** (see ``autolens_workspace/*/imaging/features/multi_gaussian_expansion``)\n", - "\n", - "These features avoid wasted effort trying to fit S\u00e9rsic profiles to complex data, which is likely to fail unless the \n", - "lens is extremely simple. This does mean the model composition is more complex and as a user its a steeper learning\n", - "curve to understand the API, but its worth it for the improved accuracy and speed of lens modeling.\n", - "\n", - "__Multi-Gaussian Expansion (MGE)__\n", - "\n", - "A Multi-Gaussian Expansion (MGE) decomposes the lens and source light into ~50\u2013100 Gaussians with varying ellipticities \n", - "and sizes. An MGE captures irregular features far more effectively than S\u00e9rsic profiles, leading to more accurate lens m\n", - "odels.\n", - "\n", - "Remarkably, modeling with MGEs is also significantly faster than using S\u00e9rsics: they remain efficient in JAX (on CPU \n", - "or GPU), require fewer non-linear parameters despite their flexibility, and yield simpler parameter spaces that\n", - "sample in far fewer iterations. \n", - "\n", - "__Linear Light Profiles__\n", - "\n", - "The MGE model below uses a **linear light profile** for the bulge and disk via the ``lp_linear`` API, instead of the \n", - "standard ``lp`` light profiles used above.\n", - "\n", - "A linear light profile solves for the *intensity* of each component via a linear inversion, rather than treating it as \n", - "a free parameter. This reduces the dimensionality of the non-linear parameter space: a model with ~80 Gaussians\n", - "does not introduce ~80 additional free parameters.\n", - "\n", - "Linear light profiles therefore improve speed and accuracy, and they are used by default in all modeling example.\n", - "\n", - "__Concise API__\n", - "\n", - "The MGE model composition API is quite long and technical, so we simply load the MGE models for the lens and source \n", - "below via a utility function `mge_model_from` which hides the API to make the code in this introduction example ready \n", - "to read. We then use the PyAutoLens Model API to compose the over lens model.\n", - "\n", - "The full MGE composition API is given in the `features/multi_gaussian_expansion` package." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing the model info confirms the model has Gaussians for both the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The lens model is fitted to the data using a non-linear search. \n", - "\n", - "All examples in the autolens workspace use the nested sampling algorithm \n", - "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has revealed gives the most \n", - "accurate and efficient modeling results.\n", - "\n", - "Nautilus has one main setting that trades-off accuracy and computational run-time, the number of `live_points`. \n", - "A higher number of live points gives a more accurate result, but increases the run-time. A lower value give \n", - "less reliable lens modeling (e.g. the fit may infer a local maxima), but is faster. \n", - "\n", - "The suitable value depends on the model complexity whereby models with more parameters require more live points. \n", - "The default value of 200 is sufficient for the vast majority of common lens models. Lower values often given reliable\n", - "results though, and speed up the run-times. In this example, given the model is quite simple (N=21 parameters), we \n", - "reduce the number of live points to 100 to speed up the run-time.\n", - "\n", - "__Unique Identifier__\n", - "\n", - "In the path above, the `unique_identifier` appears as a collection of characters, where this identifier is generated \n", - "based on the model, search and dataset that are used in the fit.\n", - " \n", - "An identical combination of model and search generates the same identifier, meaning that rerunning the script will use \n", - "the existing results to resume the model-fit. In contrast, if you change the model or search, a new unique identifier \n", - "will be generated, ensuring that the model-fit results are output into a separate folder.\n", - "\n", - "We additionally want the unique identifier to be specific to the dataset fitted, so that if we fit different datasets\n", - "with the same model and search results are output to a different folder. We achieve this below by passing \n", - "the `dataset_name` to the search's `unique_tag`.\n", - "\n", - "__Iterations Per Update__\n", - "\n", - "Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and its best fit\n", - "image to the Jupyter Notebook display and to hard-disk.\n", - "\n", - "This process takes around ~10 seconds, so we don't want it to happen too often so as to slow down the overall\n", - "fit, but we also want it to happen frequently enough that we can track the progress.\n", - "\n", - "The value of 10000 below means this output happens every few minutes on GPU and every ~10 minutes on CPU, a good balance.\n", - "\n", - "__Live Visual Update__\n", - "\n", - "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", - "live display surface:\n", - "\n", - "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", - " watch the fit converge without leaving your terminal.\n", - "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", - " refreshes in place every `iterations_per_quick_update`.\n", - "\n", - "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", - "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"imaging\"), # The path where results and output are stored.\n", - " name=\"modeling\", # The name of the fit and folder results are output to.\n", - " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", - " n_live=100, # The number of Nautilus \"live\" points, increase for more complex models.\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=10000, # Every N iterations the max likelihood model is visualized in the Jupter Notebook and output to hard-disk.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We next create an `AnalysisImaging` object, which can be given many inputs customizing how the lens model is \n", - "fitted to the data (in this example they are omitted for simplicity).\n", - "\n", - "Internally, this object defines the `log_likelihood_function` used by the non-linear search to fit the model to \n", - "the `Imaging` dataset. \n", - "\n", - "It is not vital that you as a user understand the details of how the `log_likelihood_function` fits a lens model to \n", - "data, but interested readers can find a step-by-step guide of the likelihood \n", - "function at ``autolens_workspace/*/imaging/log_likelihood_function`\n", - "\n", - "__JAX__\n", - "\n", - "Analysis uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU\n", - "support, your fits will run much faster (around 10 minutes instead of an hour). If only a CPU is available,\n", - "JAX will still provide a speed up via multithreading, with fits taking around 20-30 minutes.\n", - "\n", - "If you don\u2019t have a GPU locally, consider Google Colab which provides free GPUs, so your modeling runs are much faster." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM Use__\n", - "\n", - "When running with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM. If insufficient VRAM is \n", - "available, the analysis will fail with an out-of-memory error, typically during JIT compilation or the first \n", - "likelihood call.\n", - "\n", - "Two factors dictate the VRAM usage of an analysis:\n", - "\n", - "- The number of arrays and other data structures JAX must store in VRAM to fit the model\n", - " to the data in the likelihood function. This is dictated by the model complexity and dataset size.\n", - "\n", - "- The `batch_size` sets how many likelihood evaluations are performed simultaneously.\n", - " Increasing the batch size increases VRAM usage but can reduce overall run time,\n", - " while decreasing it lowers VRAM usage at the cost of slower execution.\n", - "\n", - "Before running an analysis, users should check that the estimated VRAM usage for the\n", - "chosen batch size is comfortably below their GPU\u2019s total VRAM.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits.\n", - "\n", - "For a MGE model with the low resolution dataset fitted in this example VRAM use is relatively low (~0.027GB) For other \n", - "models (e.g. pixelized sources) and higher resolution datasets it can be much higher (> 1GB going beyond 10GB)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "Lens modeling can be a computationally expensive process. When fitting complex models to high resolution datasets \n", - "run times can be of order hours, days, weeks or even months.\n", - "\n", - "Run times are dictated by two factors:\n", - "\n", - " - The log likelihood evaluation time: the time it takes for a single `instance` of the lens model to be fitted to \n", - " the dataset such that a log likelihood is returned.\n", - " \n", - " - The number of iterations (e.g. log likelihood evaluations) performed by the non-linear search: more complex lens\n", - " models require more iterations to converge to a solution.\n", - " \n", - "For this analysis, the log likelihood evaluation time is < 0.001 seconds on GPU, < 0.01 seconds on CPU, which is \n", - "extremely fast for lens modeling. \n", - "\n", - "To estimate the expected overall run time of the model-fit we multiply the log likelihood evaluation time by an \n", - "estimate of the number of iterations the non-linear search will perform, which is around 10000 to 30000 for this model.\n", - "\n", - "GPU run times are around 10 minutes, CPU run times are around 30 minutes.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We can now begin the model-fit by passing the model and analysis object to the search, which performs the \n", - "Nautilus non-linear search in order to find which models fit the data with the highest likelihood.\n", - "\n", - "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce \n", - "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output Folder Layout__\n", - "\n", - "Now the fit is running you should checkout the `autolens_workspace/output` folder. This is where results are\n", - "written to hard-disk in human-readable formats \u2014 `.json`, `.csv`, `.fits`, `.png` and plain text.\n", - "\n", - "As the fit progresses, results are written on the fly using the highest likelihood model found by the\n", - "non-linear search so far. This means you can inspect the model-fit as it runs, without waiting for the\n", - "non-linear search to terminate.\n", - "\n", - "Each completed fit lives at a path like::\n", - "\n", - " output/imaging//modeling//\n", - " files/ <- JSON + CSV: loadable Python objects\n", - " tracer.json <- max log likelihood Tracer\n", - " model.json <- fitted af.Collection model\n", - " samples.csv <- full Nautilus samples\n", - " samples_summary.json <- max log likelihood parameter values + errors\n", - " samples_info.json <- metadata about the samples\n", - " search.json <- non-linear search configuration\n", - " settings.json <- search settings\n", - " cosmology.json <- cosmology used for the fit\n", - " covariance.csv <- parameter covariance matrix\n", - " image/ <- FITS + PNG: imaging products\n", - " dataset.fits <- data, noise-map and PSF\n", - " fit.fits <- model image, residuals, chi-squared map\n", - " tracer.fits <- tracer image-plane images per galaxy\n", - " source_plane_images.fits <- source plane reconstructions\n", - " model_galaxy_images.fits <- per-galaxy model images\n", - " galaxy_images.fits <- per-galaxy images\n", - " dataset.png, fit.png, tracer.png <- visualisations\n", - " model.info <- human-readable model summary\n", - " model.results <- human-readable fit summary\n", - " search.summary <- search run summary\n", - " search_internal/ <- internal files used to resume / visualise the search\n", - " metadata <- run metadata\n", - "\n", - "The `` is a 32-character identifier derived from the model, search and dataset, so re-running the\n", - "same configuration resumes from the existing fit automatically.\n", - "\n", - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", - "\n", - "[Above, we discussed that the `info_whitespace_length` parameter in the config files could b changed to make \n", - "the `model.info` attribute display optimally on your computer. This attribute also controls the whitespace of the\n", - "`result.info` attribute.]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Result` object also contains:\n", - "\n", - " - The model corresponding to the maximum log likelihood solution in parameter space.\n", - " - The corresponding maximum log likelihood `Tracer` and `FitImaging` objects.\n", - " \n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It also contains information on the posterior as estimated by the non-linear search (in this example `Nautilus`). \n", - "\n", - "Below, we make a corner plot of the \"Probability Density Function\" of every parameter in the model-fit.\n", - "\n", - "The plot is labeled with short hand parameter names (e.g. `sersic_index` is mapped to the short hand \n", - "parameter `n`). These mappings ate specified in the `config/notation.yaml` file and can be customized by users.\n", - "\n", - "The superscripts of labels correspond to the name each component was given in the model (e.g. for the `Isothermal`\n", - "mass its name `mass` defined when making the `Model` above is used)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading From Output Folder__\n", - "\n", - "Everything the `Result` object contains has also been written to hard-disk, inside the fit's output folder. Each\n", - "file loads back into a full Python object with a single line \u2014 much faster and simpler than re-running the fit.\n", - "\n", - "For example, the maximum log likelihood `Tracer` is saved as a `.json` file and the tracer image-plane images as\n", - "a `.fits` file:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoconf.dictable import from_json\n", - "\n", - "result_path = search.paths.output_path # Points at the fit's unique output folder.\n", - "\n", - "if (result_path / \"files\" / \"tracer.json\").exists():\n", - " tracer = from_json(file_path=result_path / \"files\" / \"tracer.json\")\n", - "\n", - " tracer_fits = al.Array2D.from_fits(\n", - " file_path=result_path / \"image\" / \"tracer.fits\", hdu=0, pixel_scales=0.1\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The output folder also contains `model.json`, `samples.csv`, `dataset.fits`, `fit.fits` and more. A full walkthrough\n", - "of loading results from the output folder \u2014 covering both single-fit (`from_json`) and multi-fit (aggregator)\n", - "workflows \u2014 is given in:\n", - "\n", - " `autolens_workspace/*/guides/results/start_here.py`\n", - "\n", - "__Source Science (Magnification, Flux and More)__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", - "size of the source.\n", - "\n", - "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", - "including examples using a pixelized source reconstruction. \n", - "\n", - "If you want to study the source galaxy after modeling has reconstructed its unlensed, then check out this example.\n", - "\n", - "__Features__\n", - "\n", - "This script gives a concise overview of the basic modeling API, fitting one the simplest lens models possible.\n", - "\n", - "Lets now consider what features you should read about to improve your lens modeling, especially if you are aiming\n", - "to fit more complex models to your data.\n", - "\n", - "The examples in the `autolens_workspace/*/imaging/features` package illustrate other lens modeling features. \n", - "\n", - "We recommend you checkout the following features, because they make lens modeling in general more reliable and \n", - "efficient (you will therefore benefit from using these features irrespective of the quality of your data and \n", - "scientific topic of study).\n", - "\n", - "We recommend you now checkout the following features:\n", - "\n", - "- ``linear_light_profiles``: The model light profiles use linear algebra to solve for their intensity, reducing model complexity.\n", - "- ``multi_gaussian_expansion``: The lens (or source) light is modeled as ~25-100 Gaussian basis functions.\n", - "- ``pixelization``: The source is reconstructed using an adaptive rectangular or Delaunay mesh\n", - "- ``no_lens_light``: The foreground lens's light is not present in the data and thus omitted from the model.\n", - "\n", - "The files `autolens_workspace/*/guides/modeling/searches` and `autolens_workspace/*/guides/modeling/customize`\n", - "provide guides on how to customize many other aspects of the model-fit. Check them out to see if anything\n", - "sounds useful, but for most users you can get by without using these forms of customization!\n", - " \n", - "__Data Preparation__\n", - "\n", - "If you are looking to fit your own CCD imaging data of a strong lens, checkout \n", - "the `autolens_workspace/*/imaging/data_preparation/start_here.ipynb` script for an overview of how data should be \n", - "prepared before being modeled.\n", - "\n", - "__HowToLens__\n", - "\n", - "This `start_here.py` script, and the features examples above, do not explain many details of how lens modeling is \n", - "performed, for example:\n", - "\n", - " - How does PyAutoLens perform ray-tracing and lensing calculations in order to fit a lens model?\n", - " - How is a lens model fitted to data? What quantifies the goodness of fit (e.g. how is a log likelihood computed?).\n", - " - How does Nautilus find the highest likelihood lens models? What exactly is a \"non-linear search\"?\n", - "\n", - "You do not need to be able to answer these questions in order to fit lens models with PyAutoLens and do science.\n", - "However, having a deeper understanding of how it all works is both interesting and will benefit you as a scientist\n", - "\n", - "This deeper insight is offered by the **HowToLens** Jupyter notebook lectures, which live in their own repository:\n", - "https://github.com/PyAutoLabs/HowToLens.\n", - "\n", - "I recommend that you check them out if you are interested in more details!\n", - "\n", - "__Modeling Customization__\n", - "\n", - "The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative non-linear searches,\n", - "other than Nautilus, that can be used to fit lens models. \n", - "\n", - "They also provide details on how to customize the model-fit, for example the priors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Iamging: Modeling\n", + "=================\n", + "\n", + "This script is the starting point for lens modeling of CCD imaging data (E.g. Hubble Space Telescope, Euclid) with\n", + "**PyAutoLens** and it provides an overview of the lens modeling API.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Plotters:** Overview of plotting tools used for visualization.\n", + "- **Simulation:** Overview of how the simulated dataset was generated.\n", + "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Extra Galaxies Noise Scaling:** Scale the noise of nearby contaminating galaxies so they do not impact the fit.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Model Composition:** Compose the lens model using the Model and Collection API.\n", + "- **Coordinates:** Coordinate system assumptions for the model-fit.\n", + "- **Improved Lens Model:** The previous model used S\u00e9rsic light profiles for the lens and source galaxies.\n", + "- **Linear Light Profiles:** The MGE model below uses a **linear light profile** for the bulge and disk via the ``lp_linear``.\n", + "- **Concise API:** The MGE model composition API is quite long and technical, so we simply load the MGE models for the.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Unique Identifier:** In the path above, the `unique_identifier` appears as a collection of characters, where this.\n", + "- **Iterations Per Update:** Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and.\n", + "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", + "- **VRAM Use:** When running with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM.\n", + "- **Run Times:** Profiling the expected run time of the model-fit.\n", + "- **Output Folder Layout:** Description of the structure of the `output` folder where results are written.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Features:** This script gives a concise overview of the basic modeling API, fitting one the simplest lens.\n", + "- **HowToLens:** This `start_here.py` script, and the features examples above, do not explain many details of how.\n", + "- **Modeling Customization:** The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a Multi Gaussian Expansion bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a Multi Gaussian Expansion.\n", + "\n", + "This lens model is simple and computationally fast to fit, and therefore acts as a good starting point for new\n", + "users.\n", + "\n", + "__Plotters__\n", + "\n", + "To produce images of the data plotting function objects are used, which are high-level wrappers of matplotlib\n", + "code which produce high quality visualization of strong lenses.\n", + "\n", + "The plotting function API is described in the script `autolens_workspace/*/guides/plot`.\n", + "\n", + "__Simulation__\n", + "\n", + "This script fits a simulated `Imaging` dataset of a strong lens, which is produced in the\n", + "script `autolens_workspace/*/imaging/simulator.py`\n", + "\n", + "__Data Preparation__\n", + "\n", + "The `Imaging` dataset fitted in this example confirms to a number of standard that make it suitable to be fitted in\n", + "**PyAutoLens**.\n", + "\n", + "If you are intending to fit your own strong lens data, you will need to ensure it conforms to these standards, which are\n", + "described in the script `autolens_workspace/*/imaging/data_preparation/start_here.ipynb`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens dataset `simple` via .fits files, which is a data format used by astronomers to store images.\n", + "\n", + "The `pixel_scales` define the arc-second to pixel conversion factor of the image, which for the dataset we are using \n", + "is 0.1\" / pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "# PSF convolution runs at the image resolution (sub size 1), which is the fastest\n", + "# option and accurate for well-sampled PSFs. Supplying a PSF at a multiple of the\n", + "# image resolution and raising this value improves blurring fidelity for\n", + "# undersampled PSFs (e.g. HST / Euclid VIS) at extra compute cost \u2014 see\n", + "# `guides/advanced/over_sampling.py` and the simulator's `__Oversampled PSF__` section.\n", + "psf_convolve_over_sample_size = 1\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " convolve_over_sample_size_lp=psf_convolve_over_sample_size,\n", + " convolve_over_sample_size_pixelization=psf_convolve_over_sample_size,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use an `aplt.subplot_imaging_dataset` the plot the data, including: \n", + "\n", + " - `data`: The image of the strong lens.\n", + " - `noise_map`: The noise-map of the image, which quantifies the noise in every pixel as their RMS values.\n", + " - `psf`: The point spread function of the image, which describes the blurring of the image by the telescope optics.\n", + " - `signal_to_noise_map`: Quantifies the signal-to-noise in every pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Noise Scaling__\n", + "\n", + "Before masking, we must deal with any extra galaxies in the data: nearby galaxies (or foreground stars, or\n", + "data-reduction artefacts) whose emission is not associated with the strong lens but blends into the field. If\n", + "their light is left in the data it will contaminate the model-fit and bias the inferred lens model. It is too\n", + "easy to skip straight to modeling without checking for these, so we make this step an explicit part of the\n", + "workflow.\n", + "\n", + "To prevent extra galaxies from impacting the model-fit, we do not mask them entirely from the fit, which\n", + "would be analogous to making the circular mask smaller or using a more refined mask. When pixels are masked and\n", + "removed entirely from the fit, their coordinates are not used when performing ray-tracing and the light of the\n", + "lens and source galaxies in these pixels not evaluated.\n", + "\n", + "Instead, the pixels are kept in the fit, but their data values are scaled to zero and their noise-map values\n", + "are increased to very large values. This means that during the model-fit, these pixels contribute negligibly to\n", + "the likelihood of the fit, and therefore do not impact the lens model.\n", + "\n", + "This approach is used because for certain types of modeling approaches, like a pixelized source reconstruction,\n", + "masking regions of the image in a way that removes their image pixels entirely from the fit can produce\n", + "discontinuities in the pixelization. This can lead to unexpected systematics and unsatisfactory results.\n", + "\n", + "In this case, applying the mask in a way where the image pixels are not removed from the fit, but their data and\n", + "noise-map values are scaled such that they contribute negligibly to the fit, is a better approach.\n", + "\n", + "The `simple` dataset includes a faint extra galaxy, and a `mask_extra_galaxies.fits` covering it is shipped with\n", + "the dataset (created by the simulator). If you are modeling your own data with an extra galaxy, you must either:\n", + "\n", + " - Create a `mask_extra_galaxies.fits` for it using the data-preparation tools (the GUI\n", + " `autolens_workspace/*/imaging/data_preparation/gui/mask_extra_galaxies.py`, or the manual\n", + " `autolens_workspace/*/imaging/data_preparation/examples/optional/mask_extra_galaxies.py`), then load it\n", + " as below; or\n", + " - Shrink the circular mask below so the extra galaxy lies outside it and is removed from the fit entirely.\n", + "\n", + "We then plot the dataset, where the extra galaxy's pixels now have their data scaled to zero and noise-map\n", + "values increased, making their signal-to-noise effectively zero." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_extra_galaxies = al.Mask2D.from_fits(\n", + " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", + " pixel_scales=dataset.pixel_scales,\n", + " invert=True, # `True` means a pixel is scaled.\n", + ")\n", + "\n", + "dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "The model-fit requires a 2D mask defining the regions of the image we fit the lens model to the data.\n", + "\n", + "We create a 3.0 arcsecond circular mask and apply it to the `Imaging` object that the lens model fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If we plot the masked data, the mask removes the exterior regions of the image where there is no emission from the \n", + "lens and lensed source galaxies.\n", + "\n", + "The mask used to fit the data can be customized, as described in \n", + "the script `autolens_workspace/*/guides/modeling/customize`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated \n", + "on a higher resolution grid than the image data to ensure the calculation is accurate. \n", + "\n", + "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of \n", + "over sampling to ensure the lensed images are evaluated accurately.\n", + "\n", + "For a new user, the details of over-sampling are not important, therefore just be aware that calculations either:\n", + "\n", + " (i) use adaptive over sampling for the foregorund lens's light, which ensures high accuracy across. \n", + " (ii) use cored light profiles for the background source galaxy, where the core ensures low levels of over-sampling \n", + " produce numerically accurate but fast to compute results.\n", + "\n", + "\n", + "Once you are more experienced, you should read up on over-sampling in more detail via \n", + "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The imaging subplot updates the bottom two panels to reflect the update to over sampling, which now uses a higher\n", + "values in the centre.\n", + "\n", + "Whilst you may not yet understand the details of over-sampling, you can at least track it visually in the plots\n", + "and later learnt more about it in the `over_sampling.ipynb` guide." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "In this example we compose a lens model where:\n", + "\n", + " - The lens galaxy's light is a `Sersic` bulge [7 parameters].\n", + " \n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + " \n", + " - The source galaxy's light is a `SersicCore` bulge [7 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=21.\n", + "\n", + "__Model Composition__\n", + "\n", + "The API below for composing a lens model uses the `Model` and `Collection` objects, which are imported from \n", + "**PyAutoLens**'s parent project **PyAutoFit** \n", + "\n", + "The API is fairly self explanatory and is straight forward to extend, for example adding more light profiles\n", + "to the lens and source or using a different mass profile.\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", + "\n", + "__Coordinates__\n", + "\n", + "The model fitting default settings assume that the lens galaxy centre is near the coordinates (0.0\", 0.0\"). \n", + "\n", + "If for your dataset the lens is not centred at (0.0\", 0.0\"), we recommend that you either: \n", + "\n", + " - Reduce your data so that the centre is (`autolens_workspace/*/data_preparation`). \n", + " - Manually override the lens model priors (`autolens_workspace/*/guides/modeling/customize`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lp.Sersic)\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "The `info` below may not display optimally on your computer screen, for example the whitespace between parameter\n", + "names on the left and parameter priors on the right may lead them to appear across multiple lines. This is a\n", + "common issue in Jupyter notebooks.\n", + "\n", + "The`info_whitespace_length` parameter in the file `config/general.yaml` in the [output] section can be changed to \n", + "increase or decrease the amount of whitespace (The Jupyter notebook kernel will need to be reset for this change to \n", + "appear in a notebook)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Improved Lens Model__\n", + "\n", + "The previous model used S\u00e9rsic light profiles for the lens and source galaxies. This makes the model API concise, \n", + "readable, and easy to follow.\n", + "\n", + "However, single S\u00e9rsic profiles perform poorly for most strong lenses. Symmetric profiles (e.g. elliptical S\u00e9rsics) \n", + "typically leave significant residuals because they cannot capture the irregular and asymmetric morphology of real \n", + "galaxies (e.g. isophotal twists, radially varying ellipticity).\n", + "\n", + "This example therefore uses a lens model that combines two features, described in detail elsewhere (but a brief \n", + "overview is provided below):\n", + "\n", + "- **Linear light profiles** (see ``autolens_workspace/*/imaging/features/linear_light_profiles``)\n", + "- **Multi-Gaussian Expansion (MGE) light profiles** (see ``autolens_workspace/*/imaging/features/multi_gaussian_expansion``)\n", + "\n", + "These features avoid wasted effort trying to fit S\u00e9rsic profiles to complex data, which is likely to fail unless the \n", + "lens is extremely simple. This does mean the model composition is more complex and as a user its a steeper learning\n", + "curve to understand the API, but its worth it for the improved accuracy and speed of lens modeling.\n", + "\n", + "__Multi-Gaussian Expansion (MGE)__\n", + "\n", + "A Multi-Gaussian Expansion (MGE) decomposes the lens and source light into ~50\u2013100 Gaussians with varying ellipticities \n", + "and sizes. An MGE captures irregular features far more effectively than S\u00e9rsic profiles, leading to more accurate lens m\n", + "odels.\n", + "\n", + "Remarkably, modeling with MGEs is also significantly faster than using S\u00e9rsics: they remain efficient in JAX (on CPU \n", + "or GPU), require fewer non-linear parameters despite their flexibility, and yield simpler parameter spaces that\n", + "sample in far fewer iterations. \n", + "\n", + "__Linear Light Profiles__\n", + "\n", + "The MGE model below uses a **linear light profile** for the bulge and disk via the ``lp_linear`` API, instead of the \n", + "standard ``lp`` light profiles used above.\n", + "\n", + "A linear light profile solves for the *intensity* of each component via a linear inversion, rather than treating it as \n", + "a free parameter. This reduces the dimensionality of the non-linear parameter space: a model with ~80 Gaussians\n", + "does not introduce ~80 additional free parameters.\n", + "\n", + "Linear light profiles therefore improve speed and accuracy, and they are used by default in all modeling example.\n", + "\n", + "__Concise API__\n", + "\n", + "The MGE model composition API is quite long and technical, so we simply load the MGE models for the lens and source \n", + "below via a utility function `mge_model_from` which hides the API to make the code in this introduction example ready \n", + "to read. We then use the PyAutoLens Model API to compose the over lens model.\n", + "\n", + "The full MGE composition API is given in the `features/multi_gaussian_expansion` package." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing the model info confirms the model has Gaussians for both the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The lens model is fitted to the data using a non-linear search. \n", + "\n", + "All examples in the autolens workspace use the nested sampling algorithm \n", + "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has revealed gives the most \n", + "accurate and efficient modeling results.\n", + "\n", + "Nautilus has one main setting that trades-off accuracy and computational run-time, the number of `live_points`. \n", + "A higher number of live points gives a more accurate result, but increases the run-time. A lower value give \n", + "less reliable lens modeling (e.g. the fit may infer a local maxima), but is faster. \n", + "\n", + "The suitable value depends on the model complexity whereby models with more parameters require more live points. \n", + "The default value of 200 is sufficient for the vast majority of common lens models. Lower values often given reliable\n", + "results though, and speed up the run-times. In this example, given the model is quite simple (N=21 parameters), we \n", + "reduce the number of live points to 100 to speed up the run-time.\n", + "\n", + "__Unique Identifier__\n", + "\n", + "In the path above, the `unique_identifier` appears as a collection of characters, where this identifier is generated \n", + "based on the model, search and dataset that are used in the fit.\n", + " \n", + "An identical combination of model and search generates the same identifier, meaning that rerunning the script will use \n", + "the existing results to resume the model-fit. In contrast, if you change the model or search, a new unique identifier \n", + "will be generated, ensuring that the model-fit results are output into a separate folder.\n", + "\n", + "We additionally want the unique identifier to be specific to the dataset fitted, so that if we fit different datasets\n", + "with the same model and search results are output to a different folder. We achieve this below by passing \n", + "the `dataset_name` to the search's `unique_tag`.\n", + "\n", + "__Iterations Per Update__\n", + "\n", + "Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and its best fit\n", + "image to the Jupyter Notebook display and to hard-disk.\n", + "\n", + "This process takes around ~10 seconds, so we don't want it to happen too often so as to slow down the overall\n", + "fit, but we also want it to happen frequently enough that we can track the progress.\n", + "\n", + "The value of 10000 below means this output happens every few minutes on GPU and every ~10 minutes on CPU, a good balance.\n", + "\n", + "__Live Visual Update__\n", + "\n", + "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", + "live display surface:\n", + "\n", + "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", + " watch the fit converge without leaving your terminal.\n", + "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", + " refreshes in place every `iterations_per_quick_update`.\n", + "\n", + "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", + "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"imaging\"), # The path where results and output are stored.\n", + " name=\"modeling\", # The name of the fit and folder results are output to.\n", + " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", + " n_live=100, # The number of Nautilus \"live\" points, increase for more complex models.\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=10000, # Every N iterations the max likelihood model is visualized in the Jupter Notebook and output to hard-disk.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We next create an `AnalysisImaging` object, which can be given many inputs customizing how the lens model is \n", + "fitted to the data (in this example they are omitted for simplicity).\n", + "\n", + "Internally, this object defines the `log_likelihood_function` used by the non-linear search to fit the model to \n", + "the `Imaging` dataset. \n", + "\n", + "It is not vital that you as a user understand the details of how the `log_likelihood_function` fits a lens model to \n", + "data, but interested readers can find a step-by-step guide of the likelihood \n", + "function at ``autolens_workspace/*/imaging/log_likelihood_function`\n", + "\n", + "__JAX__\n", + "\n", + "Analysis uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU\n", + "support, your fits will run much faster (around 10 minutes instead of an hour). If only a CPU is available,\n", + "JAX will still provide a speed up via multithreading, with fits taking around 20-30 minutes.\n", + "\n", + "If you don\u2019t have a GPU locally, consider Google Colab which provides free GPUs, so your modeling runs are much faster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM Use__\n", + "\n", + "When running with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM. If insufficient VRAM is \n", + "available, the analysis will fail with an out-of-memory error, typically during JIT compilation or the first \n", + "likelihood call.\n", + "\n", + "Two factors dictate the VRAM usage of an analysis:\n", + "\n", + "- The number of arrays and other data structures JAX must store in VRAM to fit the model\n", + " to the data in the likelihood function. This is dictated by the model complexity and dataset size.\n", + "\n", + "- The `batch_size` sets how many likelihood evaluations are performed simultaneously.\n", + " Increasing the batch size increases VRAM usage but can reduce overall run time,\n", + " while decreasing it lowers VRAM usage at the cost of slower execution.\n", + "\n", + "Before running an analysis, users should check that the estimated VRAM usage for the\n", + "chosen batch size is comfortably below their GPU\u2019s total VRAM.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits.\n", + "\n", + "For a MGE model with the low resolution dataset fitted in this example VRAM use is relatively low (~0.027GB) For other \n", + "models (e.g. pixelized sources) and higher resolution datasets it can be much higher (> 1GB going beyond 10GB)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "Lens modeling can be a computationally expensive process. When fitting complex models to high resolution datasets \n", + "run times can be of order hours, days, weeks or even months.\n", + "\n", + "Run times are dictated by two factors:\n", + "\n", + " - The log likelihood evaluation time: the time it takes for a single `instance` of the lens model to be fitted to \n", + " the dataset such that a log likelihood is returned.\n", + " \n", + " - The number of iterations (e.g. log likelihood evaluations) performed by the non-linear search: more complex lens\n", + " models require more iterations to converge to a solution.\n", + " \n", + "For this analysis, the log likelihood evaluation time is < 0.001 seconds on GPU, < 0.01 seconds on CPU, which is \n", + "extremely fast for lens modeling. \n", + "\n", + "To estimate the expected overall run time of the model-fit we multiply the log likelihood evaluation time by an \n", + "estimate of the number of iterations the non-linear search will perform, which is around 10000 to 30000 for this model.\n", + "\n", + "GPU run times are around 10 minutes, CPU run times are around 30 minutes.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We can now begin the model-fit by passing the model and analysis object to the search, which performs the \n", + "Nautilus non-linear search in order to find which models fit the data with the highest likelihood.\n", + "\n", + "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce \n", + "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output Folder Layout__\n", + "\n", + "Now the fit is running you should checkout the `autolens_workspace/output` folder. This is where results are\n", + "written to hard-disk in human-readable formats \u2014 `.json`, `.csv`, `.fits`, `.png` and plain text.\n", + "\n", + "As the fit progresses, results are written on the fly using the highest likelihood model found by the\n", + "non-linear search so far. This means you can inspect the model-fit as it runs, without waiting for the\n", + "non-linear search to terminate.\n", + "\n", + "Each completed fit lives at a path like::\n", + "\n", + " output/imaging//modeling//\n", + " files/ <- JSON + CSV: loadable Python objects\n", + " tracer.json <- max log likelihood Tracer\n", + " model.json <- fitted af.Collection model\n", + " samples.csv <- full Nautilus samples\n", + " samples_summary.json <- max log likelihood parameter values + errors\n", + " samples_info.json <- metadata about the samples\n", + " search.json <- non-linear search configuration\n", + " settings.json <- search settings\n", + " cosmology.json <- cosmology used for the fit\n", + " covariance.csv <- parameter covariance matrix\n", + " image/ <- FITS + PNG: imaging products\n", + " dataset.fits <- data, noise-map and PSF\n", + " fit.fits <- model image, residuals, chi-squared map\n", + " tracer.fits <- tracer image-plane images per galaxy\n", + " source_plane_images.fits <- source plane reconstructions\n", + " model_galaxy_images.fits <- per-galaxy model images\n", + " galaxy_images.fits <- per-galaxy images\n", + " dataset.png, fit.png, tracer.png <- visualisations\n", + " model.info <- human-readable model summary\n", + " model.results <- human-readable fit summary\n", + " search.summary <- search run summary\n", + " search_internal/ <- internal files used to resume / visualise the search\n", + " metadata <- run metadata\n", + "\n", + "The `` is a 32-character identifier derived from the model, search and dataset, so re-running the\n", + "same configuration resumes from the existing fit automatically.\n", + "\n", + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", + "\n", + "[Above, we discussed that the `info_whitespace_length` parameter in the config files could b changed to make \n", + "the `model.info` attribute display optimally on your computer. This attribute also controls the whitespace of the\n", + "`result.info` attribute.]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Result` object also contains:\n", + "\n", + " - The model corresponding to the maximum log likelihood solution in parameter space.\n", + " - The corresponding maximum log likelihood `Tracer` and `FitImaging` objects.\n", + " \n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It also contains information on the posterior as estimated by the non-linear search (in this example `Nautilus`). \n", + "\n", + "Below, we make a corner plot of the \"Probability Density Function\" of every parameter in the model-fit.\n", + "\n", + "The plot is labeled with short hand parameter names (e.g. `sersic_index` is mapped to the short hand \n", + "parameter `n`). These mappings ate specified in the `config/notation.yaml` file and can be customized by users.\n", + "\n", + "The superscripts of labels correspond to the name each component was given in the model (e.g. for the `Isothermal`\n", + "mass its name `mass` defined when making the `Model` above is used)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading From Output Folder__\n", + "\n", + "Everything the `Result` object contains has also been written to hard-disk, inside the fit's output folder. Each\n", + "file loads back into a full Python object with a single line \u2014 much faster and simpler than re-running the fit.\n", + "\n", + "For example, the maximum log likelihood `Tracer` is saved as a `.json` file and the tracer image-plane images as\n", + "a `.fits` file:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoconf.dictable import from_json\n", + "\n", + "result_path = search.paths.output_path # Points at the fit's unique output folder.\n", + "\n", + "if (result_path / \"files\" / \"tracer.json\").exists():\n", + " tracer = from_json(file_path=result_path / \"files\" / \"tracer.json\")\n", + "\n", + " tracer_fits = al.Array2D.from_fits(\n", + " file_path=result_path / \"image\" / \"tracer.fits\", hdu=0, pixel_scales=0.1\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The output folder also contains `model.json`, `samples.csv`, `dataset.fits`, `fit.fits` and more. A full walkthrough\n", + "of loading results from the output folder \u2014 covering both single-fit (`from_json`) and multi-fit (aggregator)\n", + "workflows \u2014 is given in:\n", + "\n", + " `autolens_workspace/*/guides/results/start_here.py`\n", + "\n", + "__Source Science (Magnification, Flux and More)__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", + "size of the source.\n", + "\n", + "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", + "including examples using a pixelized source reconstruction. \n", + "\n", + "If you want to study the source galaxy after modeling has reconstructed its unlensed, then check out this example.\n", + "\n", + "__Features__\n", + "\n", + "This script gives a concise overview of the basic modeling API, fitting one the simplest lens models possible.\n", + "\n", + "Lets now consider what features you should read about to improve your lens modeling, especially if you are aiming\n", + "to fit more complex models to your data.\n", + "\n", + "The examples in the `autolens_workspace/*/imaging/features` package illustrate other lens modeling features. \n", + "\n", + "We recommend you checkout the following features, because they make lens modeling in general more reliable and \n", + "efficient (you will therefore benefit from using these features irrespective of the quality of your data and \n", + "scientific topic of study).\n", + "\n", + "We recommend you now checkout the following features:\n", + "\n", + "- ``linear_light_profiles``: The model light profiles use linear algebra to solve for their intensity, reducing model complexity.\n", + "- ``multi_gaussian_expansion``: The lens (or source) light is modeled as ~25-100 Gaussian basis functions.\n", + "- ``pixelization``: The source is reconstructed using an adaptive rectangular or Delaunay mesh\n", + "- ``no_lens_light``: The foreground lens's light is not present in the data and thus omitted from the model.\n", + "\n", + "The files `autolens_workspace/*/guides/modeling/searches` and `autolens_workspace/*/guides/modeling/customize`\n", + "provide guides on how to customize many other aspects of the model-fit. Check them out to see if anything\n", + "sounds useful, but for most users you can get by without using these forms of customization!\n", + " \n", + "__Data Preparation__\n", + "\n", + "If you are looking to fit your own CCD imaging data of a strong lens, checkout \n", + "the `autolens_workspace/*/imaging/data_preparation/start_here.ipynb` script for an overview of how data should be \n", + "prepared before being modeled.\n", + "\n", + "__HowToLens__\n", + "\n", + "This `start_here.py` script, and the features examples above, do not explain many details of how lens modeling is \n", + "performed, for example:\n", + "\n", + " - How does PyAutoLens perform ray-tracing and lensing calculations in order to fit a lens model?\n", + " - How is a lens model fitted to data? What quantifies the goodness of fit (e.g. how is a log likelihood computed?).\n", + " - How does Nautilus find the highest likelihood lens models? What exactly is a \"non-linear search\"?\n", + "\n", + "You do not need to be able to answer these questions in order to fit lens models with PyAutoLens and do science.\n", + "However, having a deeper understanding of how it all works is both interesting and will benefit you as a scientist\n", + "\n", + "This deeper insight is offered by the **HowToLens** Jupyter notebook lectures, which live in their own repository:\n", + "https://github.com/PyAutoLabs/HowToLens.\n", + "\n", + "I recommend that you check them out if you are interested in more details!\n", + "\n", + "__Modeling Customization__\n", + "\n", + "The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative non-linear searches,\n", + "other than Nautilus, that can be used to fit lens models. \n", + "\n", + "They also provide details on how to customize the model-fit, for example the priors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/simulator.ipynb b/notebooks/imaging/simulator.ipynb index 9b07db886..854a7d324 100644 --- a/notebooks/imaging/simulator.ipynb +++ b/notebooks/imaging/simulator.ipynb @@ -1,634 +1,742 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Start Here\n", - "=====================\n", - "\n", - "This script is the starting point for simulating galaxy-galaxy strong lenses as CCD imaging data (E.g. Hubble Space\n", - "Telescope, Euclid) and it provides an overview of the lens simulation API.\n", - "\n", - "After reading this script, the `examples` folder provide examples for simulating more complex lenses in different ways.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Plotters:** Overview of plotting tools used for visualization.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Grid:** Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Ray Tracing:** We now define the lens galaxy's light (elliptical Sersic + Exponential), mass (SIE+Shear) and.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** In the same folder as the .fits files, we also output subplots of the simulated dataset in .png.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "- **Multiple Images:** Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x).\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's light profile is a `Sersic`.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a `Sersic`.\n", - " - A faint extra galaxy is included offset from the lens, whose emission must be removed via noise scaling\n", - " (a `mask_extra_galaxies.fits` is written for this purpose).\n", - "\n", - "__Plotters__\n", - "\n", - "To output images of the simulated data, plotting function objects are used, which are high-level wrappers of matplotlib\n", - "code which produce high quality visualization of strong lenses.\n", - "\n", - "The plotting function API is described in the `autolens_workspace/*/guides/plot` script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", - "\n", - " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/image.fits`.\n", - " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.fits`.\n", - " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/psf.fits`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"imaging\"\n", - "dataset_name = \"simple\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output. \n", - "\n", - "In this example, this is: `/autolens_workspace/dataset/imaging/simple`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and therefore simulated \n", - "on, via the inputs:\n", - "\n", - " - `shape_native`: The (y_pixels, x_pixels) 2D shape of the grid defining the shape of the data that is simulated.\n", - " - `pixel_scales`: The arc-second to pixel conversion factor of the grid and data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxy Centre__\n", - "\n", - "This `simple` dataset deliberately includes a faint extra galaxy offset from the main lens, so that the modeling\n", - "examples can demonstrate the `__Extra Galaxies Noise Scaling__` step end-to-end. Its centre is defined here so it\n", - "can be reused for over-sampling, the galaxy itself and the `mask_extra_galaxies.fits` written further down.\n", - "\n", - "It is placed inside the 3.0\" modeling mask but clear of the lensed source arcs (Einstein radius ~1.6\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_centre = (2.2, 1.6)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated \n", - "on a higher resolution grid than the image data to ensure the calculation is accurate. \n", - "\n", - "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of \n", - "over sampling to ensure the lensed images are evaluated accurately.\n", - "\n", - "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated \n", - "on a higher resolution grid than the image data to ensure the calculation is accurate. \n", - "\n", - "An adaptive oversampling scheme is implemented, evaluating the central regions at (0.0\", 0.0\") of the light profile at a \n", - "resolution of 32x32, transitioning to 8x8 in intermediate areas, and 2x2 in the outskirts. This ensures precise and \n", - "accurate image simulation while focusing computational resources on the bright regions that demand higher oversampling.\n", - "\n", - "An adaptive oversampling grid cannot be defined for the lensed source because its light appears in different regions of \n", - "the image plane for each dataset. For this reason, most workspace examples utilize cored light profiles for the \n", - "source galaxy. Cored light profiles change gradually in their central regions, allowing accurate evaluation without \n", - "requiring oversampling.\n", - "\n", - "Once you are more experienced, you should read up on over-sampling in more detail via \n", - "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0), extra_galaxy_centre],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "All CCD imaging data (e.g. Hubble Space Telescope, Euclid) are blurred by the telescope optics when they are imaged.\n", - "\n", - "The Point Spread Function (PSF) describes the blurring of the image by the telescope optics, in the form of a\n", - "two dimensional convolution kernel. The lens modeling scripts use this PSF when fitting the data, to account for\n", - "this blurring of the image.\n", - "\n", - "In this example, use a simple 2D Gaussian PSF, which is convolved with the image of the lens and source galaxies \n", - "when simulating the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the `Imaging` dataset we first create a simulator, which includes:\n", - "\n", - " - The exposure time of the simulated dataset, increasing this will increase the signal-to-noise of the simulated data.\n", - " - The PSF of the simulated dataset, which is convolved with the image of the lens and source galaxies.\n", - " - The background sky level of the simulated dataset, which is added to the image of the lens and source galaxies and\n", - " leads to a higher level of Poisson noise.\n", - " - Whether the simulated dataset includes Poisson noise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "We now define the lens galaxy's light (elliptical Sersic + Exponential), mass (SIE+Shear) and source galaxy light\n", - "(elliptical Sersic) for this simulated lens.\n", - "\n", - "The following should be noted about the parameters below:\n", - "\n", - " - The native units of light and mass profiles distance parameters (e.g. centres, effective_radius) are arc-seconds. \n", - " - The intensity of the light profiles is in units of electrons per second per arc-second squared.\n", - " - The ellipticity of light and mass profiles are defined using the `ell_comps` parameter, however we below use\n", - " the convert module to input the `axis-ratio` (semi-major axis / semi-minor axis = b/a) and positive \n", - " angle (degrees defined counter clockwise from the positive x-axis).\n", - " - The external shear is defined using the (gamma_1, gamma_2) convention.\n", - " - The input redshifts are used to determine which galaxy is the lens (e.g. lower redshift) and which is the \n", - " source (e.g. higher redshift).\n", - " - The source uses a cored Sersic with a radius half the pixel-scale, ensuring that over-sampling is not necessary." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxy__\n", - "\n", - "We include a single faint extra galaxy offset from the main lens, representing a nearby object whose emission is not\n", - "associated with the strong lens but blends into the field. Its light contaminates the model-fit and must be removed,\n", - "which the modeling examples demonstrate via the `__Extra Galaxies Noise Scaling__` step (loading the\n", - "`mask_extra_galaxies.fits` written below and calling `dataset.apply_noise_scaling`).\n", - "\n", - "We give the extra galaxy a light profile only (no mass), so the lensed source arcs are unchanged and the dataset\n", - "remains a clean galaxy-scale lens for all other examples that load it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " light=al.lp.ExponentialSph(\n", - " centre=extra_galaxy_centre, intensity=1.0, effective_radius=0.3\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now pass these galaxies to a `Tracer`, which performs the ray-tracing calculations they describe and returns\n", - "the image of the strong lens system they produce." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, extra_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot the `Tracer``s image, which is the image we'll next simulate as CCD imaging data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By passing the `Tracer` and grid to the simulator, we create the simulated CCD imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now plot the simulated `Imaging` dataset before outputting it to fits.\n", - "\n", - "Note how unlike the `Tracer` image above, the simulated `Imaging` dataset includes the blurring effects of the \n", - "telescope's PSF and also has noise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files.\n", - "\n", - "If you are unfamiliar with .fits files, this is the standard file format of astronomical data and you can open \n", - "them using the software ds9 (https://sites.google.com/cfa.harvard.edu/saoimageds9/home)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask Extra Galaxies__\n", - "\n", - "Build and output a `mask_extra_galaxies.fits` covering the extra galaxy, so the modeling examples\n", - "(`imaging/modeling.py`, `imaging/fit.py`, `imaging/likelihood_function.py`) can load it directly and apply\n", - "noise scaling without a separate data-preparation step.\n", - "\n", - "The circle is sized to ~3x the galaxy's `effective_radius`, which comfortably covers its light extent. The\n", - "geometry is derived from the same `extra_galaxy_centre` defined above, so it stays in sync with any future tweak." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_extra_galaxies = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " centre=extra_galaxy_centre,\n", - " radius=3.0 * 0.3,\n", - " invert=True, # `True` inside the circle, i.e. the region whose noise is scaled.\n", - ")\n", - "\n", - "aplt.fits_array(\n", - " array=mask_extra_galaxies,\n", - " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "In the same folder as the .fits files, we also output subplots of the simulated dataset in .png format, as well as \n", - "other images which summarise the dataset.\n", - "\n", - "Having .png files like this is useful, as they can be opened quickly and easily by the user to check the dataset.\n", - "\n", - "For a faster run time, this visualization uses a regular grid which does not perferm the iterative ray-tracing." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multiple Images__\n", - "\n", - "Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x) \n", - "coordinates of multiple images of a source galaxy to positions which are far apart from one another \n", - "in the source plane are penalized in the lens model's overall likelihood.\n", - "\n", - "This speeds up lens modeling, helps the non-linear search avoid local maxima and is vital for inferred \n", - "accurate solutions when using pixelized source reconstructions.\n", - "\n", - "For real data, the multiple image positions are determined by eye from the data, for example\n", - "using a Graphical User Interface (GUI) to mark them with mouse clicks. For simulated data, we can save\n", - "ourselves time by using the `PointSolver` to determine the multiple image positions automatically and\n", - "output to a .json file.\n", - "\n", - "If you have not looked in the `point_source` package, the point solver is the core tool used to find\n", - "multiple image positions for point source lens modeling (e.g. lensed quasars)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " file_path=dataset_path / \"positions.json\",\n", - " obj=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/simple`.\n", - "\n", - "__JAX Variant__\n", - "\n", - "For an order-of-magnitude speedup on large or repeated simulations\n", - "(parameter sweeps, mock-data studies, batch figure generation), construct\n", - "the simulator with `use_jax=True` and wrap your call in `@jax.jit`. The\n", - "simulator handles pytree registration internally \u2014 you write nothing\n", - "JAX-specific beyond the decorator.\n", - "\n", - "```python\n", - "import jax\n", - "\n", - "simulator_jax = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - " use_jax=True,\n", - ")\n", - "\n", - "@jax.jit\n", - "def simulate(tracer):\n", - " return simulator_jax.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "dataset_jax = simulate(tracer) # Imaging with jax.Array data\n", - "```\n", - "\n", - "The `dataset_jax.data.array` is a `jax.Array`; `aplt.fits_imaging` and the\n", - "plotters call `numpy.asarray()` internally, so saving / plotting works\n", - "without manual conversion.\n", - "\n", - "Note: eager `simulator_jax.via_tracer_from(tracer, grid)` (no `@jax.jit`)\n", - "already runs on JAX and is sufficient for one-off simulations. The\n", - "`@jax.jit` wrap is only beneficial when you call the function many times.\n", - "\n", - "See `scripts/guides/lens_calc.py` for the advanced \"JIT-it-yourself\"\n", - "pattern that wraps individual library methods like `tracer.image_2d_from`\n", - "directly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Start Here\n", + "=====================\n", + "\n", + "This script is the starting point for simulating galaxy-galaxy strong lenses as CCD imaging data (E.g. Hubble Space\n", + "Telescope, Euclid) and it provides an overview of the lens simulation API.\n", + "\n", + "After reading this script, the `examples` folder provide examples for simulating more complex lenses in different ways.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Plotters:** Overview of plotting tools used for visualization.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Grid:** Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Ray Tracing:** We now define the lens galaxy's light (elliptical Sersic + Exponential), mass (SIE+Shear) and.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** In the same folder as the .fits files, we also output subplots of the simulated dataset in .png.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "- **Multiple Images:** Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x).\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's light profile is a `Sersic`.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a `Sersic`.\n", + " - A faint extra galaxy is included offset from the lens, whose emission must be removed via noise scaling\n", + " (a `mask_extra_galaxies.fits` is written for this purpose).\n", + "\n", + "__Plotters__\n", + "\n", + "To output images of the simulated data, plotting function objects are used, which are high-level wrappers of matplotlib\n", + "code which produce high quality visualization of strong lenses.\n", + "\n", + "The plotting function API is described in the `autolens_workspace/*/guides/plot` script." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", + "\n", + " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/image.fits`.\n", + " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.fits`.\n", + " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/psf.fits`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"imaging\"\n", + "dataset_name = \"simple\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output. \n", + "\n", + "In this example, this is: `/autolens_workspace/dataset/imaging/simple`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2d grid of (y,x) coordinates that the lens and source galaxy images are evaluated and therefore simulated \n", + "on, via the inputs:\n", + "\n", + " - `shape_native`: The (y_pixels, x_pixels) 2D shape of the grid defining the shape of the data that is simulated.\n", + " - `pixel_scales`: The arc-second to pixel conversion factor of the grid and data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxy Centre__\n", + "\n", + "This `simple` dataset deliberately includes a faint extra galaxy offset from the main lens, so that the modeling\n", + "examples can demonstrate the `__Extra Galaxies Noise Scaling__` step end-to-end. Its centre is defined here so it\n", + "can be reused for over-sampling, the galaxy itself and the `mask_extra_galaxies.fits` written further down.\n", + "\n", + "It is placed inside the 3.0\" modeling mask but clear of the lensed source arcs (Einstein radius ~1.6\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_centre = (2.2, 1.6)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated \n", + "on a higher resolution grid than the image data to ensure the calculation is accurate. \n", + "\n", + "For lensing calculations, the high magnification regions of a lensed source galaxy require especially high levels of \n", + "over sampling to ensure the lensed images are evaluated accurately.\n", + "\n", + "Over sampling is a numerical technique where the images of light profiles and galaxies are evaluated \n", + "on a higher resolution grid than the image data to ensure the calculation is accurate. \n", + "\n", + "An adaptive oversampling scheme is implemented, evaluating the central regions at (0.0\", 0.0\") of the light profile at a \n", + "resolution of 32x32, transitioning to 8x8 in intermediate areas, and 2x2 in the outskirts. This ensures precise and \n", + "accurate image simulation while focusing computational resources on the bright regions that demand higher oversampling.\n", + "\n", + "An adaptive oversampling grid cannot be defined for the lensed source because its light appears in different regions of \n", + "the image plane for each dataset. For this reason, most workspace examples utilize cored light profiles for the \n", + "source galaxy. Cored light profiles change gradually in their central regions, allowing accurate evaluation without \n", + "requiring oversampling.\n", + "\n", + "Once you are more experienced, you should read up on over-sampling in more detail via \n", + "the `autolens_workspace/*/guides/over_sampling.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0), extra_galaxy_centre],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "All CCD imaging data (e.g. Hubble Space Telescope, Euclid) are blurred by the telescope optics when they are imaged.\n", + "\n", + "The Point Spread Function (PSF) describes the blurring of the image by the telescope optics, in the form of a\n", + "two dimensional convolution kernel. The lens modeling scripts use this PSF when fitting the data, to account for\n", + "this blurring of the image.\n", + "\n", + "In this example, use a simple 2D Gaussian PSF, which is convolved with the image of the lens and source galaxies \n", + "when simulating the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# PSF convolution runs at the image resolution (sub size 1), which is the fastest\n", + "# option and accurate for well-sampled PSFs. Supplying a PSF at a multiple of the\n", + "# image resolution and raising this value improves blurring fidelity for\n", + "# undersampled PSFs (e.g. HST / Euclid VIS) at extra compute cost \u2014 see\n", + "# `guides/advanced/over_sampling.py` and the simulator's `__Oversampled PSF__` section.\n", + "psf_convolve_over_sample_size = 1\n", + "\n", + "psf = al.Convolver.from_gaussian(\n", + " convolve_over_sample_size=psf_convolve_over_sample_size,\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the `Imaging` dataset we first create a simulator, which includes:\n", + "\n", + " - The exposure time of the simulated dataset, increasing this will increase the signal-to-noise of the simulated data.\n", + " - The PSF of the simulated dataset, which is convolved with the image of the lens and source galaxies.\n", + " - The background sky level of the simulated dataset, which is added to the image of the lens and source galaxies and\n", + " leads to a higher level of Poisson noise.\n", + " - Whether the simulated dataset includes Poisson noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "We now define the lens galaxy's light (elliptical Sersic + Exponential), mass (SIE+Shear) and source galaxy light\n", + "(elliptical Sersic) for this simulated lens.\n", + "\n", + "The following should be noted about the parameters below:\n", + "\n", + " - The native units of light and mass profiles distance parameters (e.g. centres, effective_radius) are arc-seconds. \n", + " - The intensity of the light profiles is in units of electrons per second per arc-second squared.\n", + " - The ellipticity of light and mass profiles are defined using the `ell_comps` parameter, however we below use\n", + " the convert module to input the `axis-ratio` (semi-major axis / semi-minor axis = b/a) and positive \n", + " angle (degrees defined counter clockwise from the positive x-axis).\n", + " - The external shear is defined using the (gamma_1, gamma_2) convention.\n", + " - The input redshifts are used to determine which galaxy is the lens (e.g. lower redshift) and which is the \n", + " source (e.g. higher redshift).\n", + " - The source uses a cored Sersic with a radius half the pixel-scale, ensuring that over-sampling is not necessary." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxy__\n", + "\n", + "We include a single faint extra galaxy offset from the main lens, representing a nearby object whose emission is not\n", + "associated with the strong lens but blends into the field. Its light contaminates the model-fit and must be removed,\n", + "which the modeling examples demonstrate via the `__Extra Galaxies Noise Scaling__` step (loading the\n", + "`mask_extra_galaxies.fits` written below and calling `dataset.apply_noise_scaling`).\n", + "\n", + "We give the extra galaxy a light profile only (no mass), so the lensed source arcs are unchanged and the dataset\n", + "remains a clean galaxy-scale lens for all other examples that load it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " light=al.lp.ExponentialSph(\n", + " centre=extra_galaxy_centre, intensity=1.0, effective_radius=0.3\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now pass these galaxies to a `Tracer`, which performs the ray-tracing calculations they describe and returns\n", + "the image of the strong lens system they produce." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, extra_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot the `Tracer``s image, which is the image we'll next simulate as CCD imaging data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By passing the `Tracer` and grid to the simulator, we create the simulated CCD imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now plot the simulated `Imaging` dataset before outputting it to fits.\n", + "\n", + "Note how unlike the `Tracer` image above, the simulated `Imaging` dataset includes the blurring effects of the \n", + "telescope's PSF and also has noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files.\n", + "\n", + "If you are unfamiliar with .fits files, this is the standard file format of astronomical data and you can open \n", + "them using the software ds9 (https://sites.google.com/cfa.harvard.edu/saoimageds9/home)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask Extra Galaxies__\n", + "\n", + "Build and output a `mask_extra_galaxies.fits` covering the extra galaxy, so the modeling examples\n", + "(`imaging/modeling.py`, `imaging/fit.py`, `imaging/likelihood_function.py`) can load it directly and apply\n", + "noise scaling without a separate data-preparation step.\n", + "\n", + "The circle is sized to ~3x the galaxy's `effective_radius`, which comfortably covers its light extent. The\n", + "geometry is derived from the same `extra_galaxy_centre` defined above, so it stays in sync with any future tweak." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_extra_galaxies = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " centre=extra_galaxy_centre,\n", + " radius=3.0 * 0.3,\n", + " invert=True, # `True` inside the circle, i.e. the region whose noise is scaled.\n", + ")\n", + "\n", + "aplt.fits_array(\n", + " array=mask_extra_galaxies,\n", + " file_path=dataset_path / \"mask_extra_galaxies.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "In the same folder as the .fits files, we also output subplots of the simulated dataset in .png format, as well as \n", + "other images which summarise the dataset.\n", + "\n", + "Having .png files like this is useful, as they can be opened quickly and easily by the user to check the dataset.\n", + "\n", + "For a faster run time, this visualization uses a regular grid which does not perferm the iterative ray-tracing." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multiple Images__\n", + "\n", + "Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x) \n", + "coordinates of multiple images of a source galaxy to positions which are far apart from one another \n", + "in the source plane are penalized in the lens model's overall likelihood.\n", + "\n", + "This speeds up lens modeling, helps the non-linear search avoid local maxima and is vital for inferred \n", + "accurate solutions when using pixelized source reconstructions.\n", + "\n", + "For real data, the multiple image positions are determined by eye from the data, for example\n", + "using a Graphical User Interface (GUI) to mark them with mouse clicks. For simulated data, we can save\n", + "ourselves time by using the `PointSolver` to determine the multiple image positions automatically and\n", + "output to a .json file.\n", + "\n", + "If you have not looked in the `point_source` package, the point solver is the core tool used to find\n", + "multiple image positions for point source lens modeling (e.g. lensed quasars)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " file_path=dataset_path / \"positions.json\",\n", + " obj=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/simple`.\n", + "\n", + "__JAX Variant__\n", + "\n", + "For an order-of-magnitude speedup on large or repeated simulations\n", + "(parameter sweeps, mock-data studies, batch figure generation), construct\n", + "the simulator with `use_jax=True` and wrap your call in `@jax.jit`. The\n", + "simulator handles pytree registration internally \u2014 you write nothing\n", + "JAX-specific beyond the decorator.\n", + "\n", + "```python\n", + "import jax\n", + "\n", + "simulator_jax = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + " use_jax=True,\n", + ")\n", + "\n", + "@jax.jit\n", + "def simulate(tracer):\n", + " return simulator_jax.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "dataset_jax = simulate(tracer) # Imaging with jax.Array data\n", + "```\n", + "\n", + "The `dataset_jax.data.array` is a `jax.Array`; `aplt.fits_imaging` and the\n", + "plotters call `numpy.asarray()` internally, so saving / plotting works\n", + "without manual conversion.\n", + "\n", + "Note: eager `simulator_jax.via_tracer_from(tracer, grid)` (no `@jax.jit`)\n", + "already runs on JAX and is sufficient for one-off simulations. The\n", + "`@jax.jit` wrap is only beneficial when you call the function many times.\n", + "\n", + "See `scripts/guides/lens_calc.py` for the advanced \"JIT-it-yourself\"\n", + "pattern that wraps individual library methods like `tracer.image_2d_from`\n", + "directly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Oversampled PSF__\n", + "\n", + "The simulation above evaluates the lensed image on an over-sampled grid, but the PSF convolution itself is\n", + "performed at the resolution of the image pixels. For most simulations this is accurate enough. However, when the\n", + "PSF is undersampled by the detector (its width is comparable to the pixel scale, as for HST or Euclid VIS imaging)\n", + "or when you want to model the blurring with maximum fidelity, the convolution itself can also be performed at a\n", + "higher resolution.\n", + "\n", + "To do this, supply the PSF at a multiple of the image resolution and set `convolve_over_sample_size`. For example,\n", + "with `convolve_over_sample_size=2` the PSF kernel below has pixels half the size of the image pixels (note the\n", + "`pixel_scales` and the larger `shape_native` covering the same physical area). The simulator then evaluates the\n", + "lensed image on the over-sampled grid, convolves at the fine resolution and bins the result back to the image\n", + "resolution.\n", + "\n", + "Two requirements to be aware of:\n", + "\n", + " - Every entry of the grid's `over_sample_size` must be divisible by `convolve_over_sample_size` (the k x s\n", + " coupling: adaptive evaluation is partially binned to the convolution resolution before blurring, so the\n", + " adaptive radial schemes above compose with an oversampled PSF; a non-divisible combination raises a clear\n", + " error).\n", + " - When you later fit data simulated this way, pass the same fine-resolution PSF and the matching\n", + " `convolve_over_sample_size_lp` / `convolve_over_sample_size_pixelization` to the `Imaging` object.\n", + "\n", + "```python\n", + "grid_fine = al.Grid2D.uniform(\n", + " shape_native=grid.shape_native,\n", + " pixel_scales=grid.pixel_scales,\n", + " over_sample_size=2,\n", + ")\n", + "\n", + "psf_fine = al.Convolver.from_gaussian(\n", + " shape_native=(21, 21), # twice the pixels of an 11x11 image-resolution kernel...\n", + " pixel_scales=grid.pixel_scales[0] / 2, # ...at half the pixel scale, so the same physical extent.\n", + " sigma=0.1,\n", + " normalize=True,\n", + " convolve_over_sample_size=2,\n", + ")\n", + "\n", + "simulator_fine = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf_fine,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")\n", + "\n", + "dataset_fine = simulator_fine.via_tracer_from(tracer=tracer, grid=grid_fine)\n", + "```\n", + "\n", + "The numerical test scripts in `autolens_workspace_test/scripts/imaging/convolution_over_sampled.py` verify this\n", + "machinery against brute-force reference calculations for every supported model surface (standard, linear and\n", + "operated light profiles and pixelized sources).\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/simulator_sample.ipynb b/notebooks/imaging/simulator_sample.ipynb index 9223349b1..31671da48 100644 --- a/notebooks/imaging/simulator_sample.ipynb +++ b/notebooks/imaging/simulator_sample.ipynb @@ -1,346 +1,383 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Sample\n", - "=================\n", - "\n", - "This script illustrates how to simulate a sample of `Imaging` datasets of 'galaxy-scale' strong lenses, which can\n", - "easily be used to simulate hundreds or thousands of strong lenses.\n", - "\n", - "To simulate the sample of lenses, each lens and source galaxies set up using the `Model` object which is also used in\n", - "the `modeling` scripts. This means that the parameters of each simulated strong lens are drawn from the distributions\n", - "defined via priors, which can be customized to simulate a wider range of strong lenses.\n", - "\n", - "This script simulate a sample of `Imaging` datasets of 'galaxy-scale' strong lenses, whose light and mass profiles are\n", - "the same as those used in the `start_here` script.\n", - "\n", - "It is used in `autolens_workspace/notebooks/advanced/graphical` to illustrate how a hierarchical model can\n", - "be fitted to a large sample of strong lenses in order to infer the glboal properties of the lens sample.\n", - "\n", - "This script uses the signal-to-noise based light profiles described in the\n", - "script `imaging/features/simulator_/manual_signal_to_noise_ratio.ipynb`, to make it straight forward to ensure the lens\n", - "and source galaxies are visible in each image.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Sample Model Distributions:** To simulate a sample, we draw random instances of lens and source galaxies where the parameters of.\n", - "- **Sample Instances:** Within a for loop, we will now generate instances of the lens and source galaxies using the.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates a sample of `Imaging` data of 'galaxy-scale' strong lenses where:\n", - "\n", - " - The lens galaxies light profiles are `Sersic`'s.\n", - " - The lens galaxies total mass distributions are `Isothermal` models with `ExternalShear`'s.\n", - " - The source galaxies light profiles are `Sersic`'s.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import numpy as np\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"samples\"\n", - "dataset_type = \"imaging\"\n", - "dataset_sample_name = \"simple\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output/" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", dataset_type, dataset_label, dataset_sample_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the `Imaging` dataset we first create a simulator, which defines the exposure time, background sky,\n", - "noise levels and psf of the dataset that is simulated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sample Truth Distributions__\n", - "\n", - "To simulate a sample, we draw random instances of lens and source galaxies. Each parameter\n", - "is sampled directly from a numpy ``Generator`` and used to construct concrete light/mass\n", - "profile instances \u2014 there is no ``af.Model`` involved here because we are generating\n", - "*truths* for synthetic data, not fitting a model.\n", - "\n", - "The bulges use ``al.lp_snr.Sersic`` so each lens/source hits a target signal-to-noise\n", - "ratio in the data \u2014 SNR is a property of the data, not a fitted parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "rng = np.random.default_rng()\n", - "\n", - "\n", - "def _clipped_ell_comp() -> float:\n", - " return float(np.clip(rng.normal(0.0, 0.2), -1.0, 1.0))\n", - "\n", - "\n", - "def _random_lens_and_source() -> tuple[al.Galaxy, al.Galaxy]:\n", - " lens_bulge = al.lp_snr.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(_clipped_ell_comp(), _clipped_ell_comp()),\n", - " effective_radius=float(rng.uniform(1.0, 5.0)),\n", - " sersic_index=float(np.clip(rng.normal(4.0, 0.5), 0.8, 5.0)),\n", - " signal_to_noise_ratio=float(rng.uniform(20.0, 60.0)),\n", - " )\n", - " mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=(_clipped_ell_comp(), _clipped_ell_comp()),\n", - " einstein_radius=float(rng.uniform(0.2, 1.8)),\n", - " )\n", - " shear = al.mp.ExternalShear(\n", - " gamma_1=float(rng.normal(0.0, 0.05)),\n", - " gamma_2=float(rng.normal(0.0, 0.05)),\n", - " )\n", - " lens = al.Galaxy(redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear)\n", - "\n", - " source_bulge = al.lp_snr.Sersic(\n", - " centre=(float(rng.normal(0.0, 0.3)), float(rng.normal(0.0, 0.3))),\n", - " ell_comps=(_clipped_ell_comp(), _clipped_ell_comp()),\n", - " effective_radius=float(rng.uniform(0.01, 3.0)),\n", - " sersic_index=float(np.clip(rng.normal(2.0, 0.5), 0.8, 5.0)),\n", - " signal_to_noise_ratio=float(rng.uniform(10.0, 30.0)),\n", - " )\n", - " source = al.Galaxy(redshift=1.0, bulge=source_bulge)\n", - "\n", - " return lens, source\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sample Instances__\n", - "\n", - "Within a for loop, we will now generate instances of the lens and source galaxies.\n", - "This loop will run for `total_datasets` iterations, which sets the number of lenses\n", - "that are simulated.\n", - "\n", - "Each iteration of the for loop will then create a tracer and use this to simulate the\n", - "imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_datasets = 3\n", - "\n", - "for sample_index in range(total_datasets):\n", - " dataset_sample_path = Path(dataset_path, f\"dataset_{sample_index}\")\n", - "\n", - " lens_galaxy, source_galaxy = _random_lens_and_source()\n", - "\n", - " \"\"\"\n", - " __Ray Tracing__\n", - "\n", - " Use the sample's lens and source galaxies to setup a tracer, which will generate the image for the \n", - " simulated `Imaging` dataset.\n", - "\n", - " The steps below are expanded on in other `imaging/simulator` scripts, so check them out if anything below is unclear.\n", - " \"\"\"\n", - " tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", - "\n", - " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " \"\"\"\n", - " __Output__\n", - "\n", - " Output the simulated dataset to the dataset path as .fits files.\n", - "\n", - " This uses the updated `dataset_path_sample` which outputs this sample lens to a unique folder.\n", - " \"\"\"\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=Path(dataset_sample_path, \"data.fits\"),\n", - " psf_path=Path(dataset_sample_path, \"psf.fits\"),\n", - " noise_map_path=Path(dataset_sample_path, \"noise_map.fits\"),\n", - " overwrite=True,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Visualize__\n", - "\n", - " Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files.\n", - " \"\"\"\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - " aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - " aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_sample_path, output_format=\"png\"\n", - " )\n", - " aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_sample_path, output_format=\"png\"\n", - " )\n", - "\n", - " \"\"\"\n", - " __Tracer json__\n", - "\n", - " Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - " are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - " This can be loaded via the method `tracer = al.from_json()`.\n", - " \"\"\"\n", - " al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_sample_path, \"tracer.json\"),\n", - " )\n", - "\n", - " \"\"\"\n", - " The dataset can be viewed in the \n", - " folder `autolens_workspace/imaging/sample/light_sersic__mass_sie__source_sersic_{sample_index]`.\n", - " \"\"\"\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Sample\n", + "=================\n", + "\n", + "This script illustrates how to simulate a sample of `Imaging` datasets of 'galaxy-scale' strong lenses, which can\n", + "easily be used to simulate hundreds or thousands of strong lenses.\n", + "\n", + "To simulate the sample of lenses, each lens and source galaxies set up using the `Model` object which is also used in\n", + "the `modeling` scripts. This means that the parameters of each simulated strong lens are drawn from the distributions\n", + "defined via priors, which can be customized to simulate a wider range of strong lenses.\n", + "\n", + "This script simulate a sample of `Imaging` datasets of 'galaxy-scale' strong lenses, whose light and mass profiles are\n", + "the same as those used in the `start_here` script.\n", + "\n", + "It is used in `autolens_workspace/notebooks/advanced/graphical` to illustrate how a hierarchical model can\n", + "be fitted to a large sample of strong lenses in order to infer the glboal properties of the lens sample.\n", + "\n", + "This script uses the signal-to-noise based light profiles described in the\n", + "script `imaging/features/simulator_/manual_signal_to_noise_ratio.ipynb`, to make it straight forward to ensure the lens\n", + "and source galaxies are visible in each image.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Sample Model Distributions:** To simulate a sample, we draw random instances of lens and source galaxies where the parameters of.\n", + "- **Sample Instances:** Within a for loop, we will now generate instances of the lens and source galaxies using the.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates a sample of `Imaging` data of 'galaxy-scale' strong lenses where:\n", + "\n", + " - The lens galaxies light profiles are `Sersic`'s.\n", + " - The lens galaxies total mass distributions are `Isothermal` models with `ExternalShear`'s.\n", + " - The source galaxies light profiles are `Sersic`'s.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import numpy as np\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"samples\"\n", + "dataset_type = \"imaging\"\n", + "dataset_sample_name = \"simple\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output/" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", dataset_type, dataset_label, dataset_sample_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the `Imaging` dataset we first create a simulator, which defines the exposure time, background sky,\n", + "noise levels and psf of the dataset that is simulated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sample Truth Distributions__\n", + "\n", + "To simulate a sample, we draw random instances of lens and source galaxies. Each parameter\n", + "is sampled directly from a numpy ``Generator`` and used to construct concrete light/mass\n", + "profile instances \u2014 there is no ``af.Model`` involved here because we are generating\n", + "*truths* for synthetic data, not fitting a model.\n", + "\n", + "The bulges use ``al.lp_snr.Sersic`` so each lens/source hits a target signal-to-noise\n", + "ratio in the data \u2014 SNR is a property of the data, not a fitted parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "rng = np.random.default_rng()\n", + "\n", + "\n", + "def _clipped_ell_comp() -> float:\n", + " return float(np.clip(rng.normal(0.0, 0.2), -1.0, 1.0))\n", + "\n", + "\n", + "def _random_lens_and_source() -> tuple[al.Galaxy, al.Galaxy]:\n", + " lens_bulge = al.lp_snr.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(_clipped_ell_comp(), _clipped_ell_comp()),\n", + " effective_radius=float(rng.uniform(1.0, 5.0)),\n", + " sersic_index=float(np.clip(rng.normal(4.0, 0.5), 0.8, 5.0)),\n", + " signal_to_noise_ratio=float(rng.uniform(20.0, 60.0)),\n", + " )\n", + " mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=(_clipped_ell_comp(), _clipped_ell_comp()),\n", + " einstein_radius=float(rng.uniform(0.2, 1.8)),\n", + " )\n", + " shear = al.mp.ExternalShear(\n", + " gamma_1=float(rng.normal(0.0, 0.05)),\n", + " gamma_2=float(rng.normal(0.0, 0.05)),\n", + " )\n", + " lens = al.Galaxy(redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear)\n", + "\n", + " source_bulge = al.lp_snr.Sersic(\n", + " centre=(float(rng.normal(0.0, 0.3)), float(rng.normal(0.0, 0.3))),\n", + " ell_comps=(_clipped_ell_comp(), _clipped_ell_comp()),\n", + " effective_radius=float(rng.uniform(0.01, 3.0)),\n", + " sersic_index=float(np.clip(rng.normal(2.0, 0.5), 0.8, 5.0)),\n", + " signal_to_noise_ratio=float(rng.uniform(10.0, 30.0)),\n", + " )\n", + " source = al.Galaxy(redshift=1.0, bulge=source_bulge)\n", + "\n", + " return lens, source\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sample Instances__\n", + "\n", + "Within a for loop, we will now generate instances of the lens and source galaxies.\n", + "This loop will run for `total_datasets` iterations, which sets the number of lenses\n", + "that are simulated.\n", + "\n", + "Each iteration of the for loop will then create a tracer and use this to simulate the\n", + "imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_datasets = 3\n", + "\n", + "for sample_index in range(total_datasets):\n", + " dataset_sample_path = Path(dataset_path, f\"dataset_{sample_index}\")\n", + "\n", + " lens_galaxy, source_galaxy = _random_lens_and_source()\n", + "\n", + " \"\"\"\n", + " __Ray Tracing__\n", + "\n", + " Use the sample's lens and source galaxies to setup a tracer, which will generate the image for the \n", + " simulated `Imaging` dataset.\n", + "\n", + " The steps below are expanded on in other `imaging/simulator` scripts, so check them out if anything below is unclear.\n", + " \"\"\"\n", + " tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", + "\n", + " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " \"\"\"\n", + " __Output__\n", + "\n", + " Output the simulated dataset to the dataset path as .fits files.\n", + "\n", + " This uses the updated `dataset_path_sample` which outputs this sample lens to a unique folder.\n", + " \"\"\"\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=Path(dataset_sample_path, \"data.fits\"),\n", + " psf_path=Path(dataset_sample_path, \"psf.fits\"),\n", + " noise_map_path=Path(dataset_sample_path, \"noise_map.fits\"),\n", + " overwrite=True,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Visualize__\n", + "\n", + " Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files.\n", + " \"\"\"\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + " aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + " aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_sample_path, output_format=\"png\"\n", + " )\n", + " aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_sample_path, output_format=\"png\"\n", + " )\n", + "\n", + " \"\"\"\n", + " __Tracer json__\n", + "\n", + " Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + " are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + " This can be loaded via the method `tracer = al.from_json()`.\n", + " \"\"\"\n", + " al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_sample_path, \"tracer.json\"),\n", + " )\n", + "\n", + " \"\"\"\n", + " The dataset can be viewed in the \n", + " folder `autolens_workspace/imaging/sample/light_sersic__mass_sie__source_sersic_{sample_index]`.\n", + " \"\"\"\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/source_science.ipynb b/notebooks/imaging/source_science.ipynb index 7f3f804dc..ee48e7c32 100644 --- a/notebooks/imaging/source_science.ipynb +++ b/notebooks/imaging/source_science.ipynb @@ -1,437 +1,474 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Source Science\n", - "==============\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", - "size of the source.\n", - "\n", - "This example shows how to perform these calculations using Sersic parametric sources on imaging data, which\n", - "is conceptually the simplest case for source science calculations and a good introduction to the topic.\n", - "\n", - "__Contents__\n", - "\n", - "- **Simulated Dataset:** We load and plot the `simple__no_lens_light` example dataset, which is simulated imaging of a.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Source Values:** Source science calculations for real lenses are performed using the best-fitting model inferred.\n", - "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", - "- **Source Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", - "- **Tracer:** Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to.\n", - "- **Parametric Source Models:** If your lens modeling uses a parametric source model (e.g." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulated Dataset__\n", - "\n", - "We load and plot the `simple__no_lens_light` example dataset, which is simulated imaging of a strong lens\n", - "that we will use to demonstrate source science caluculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We apply a 3.0 arcsecond circular mask and apply it to the `Imaging` object.\n", - "\n", - "Source science calculations are typically performed on masked datasets to ensure only the lensed source is used\n", - "in the calculations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Values__\n", - "\n", - "Source science calculations for real lenses are performed using the best-fitting model inferred from a dataset, \n", - "and this example demonstrates how to use this below.\n", - "\n", - "However, we for simplicity, we demonstrate these calculations using the Sersic source model used to simulate the dataset, \n", - "which we refer to as the \"true\" source model. When analysing real strong lenses, a true underlying model is not known, \n", - "but for simulated datasets it is.\n", - "\n", - "This allows us to illustrate the calculations in a way that does not depend on the specific details of the data or \n", - "on assumptions about how the lens model is inferred.\n", - "\n", - "The `tracer` below corresponds to the same tracer used to simulate the `simple__no_lens_light` dataset, and therefore \n", - "represents the true source model. We also include the 2D grid of (y,x) coordinates which simulate the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the image of the tracer, we confirm it looks identical to the simulated dataset but does not have\n", - "CCD imaging features such as noise or blurring from a PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Flux__\n", - "\n", - "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", - "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", - "\n", - "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", - "we compute the image of the light profile, and then sum the image. \n", - "\n", - "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", - "we will assume everything is in electrons per second (`e- s^-1`), which is typical for Hubble Space Telescope imaging data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Source Galaxy's Intensity {source_galaxy.bulge.intensity} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The total flux, in units of `e- s^-1` , is computed by summing the image of the light profile over all pixels.\n", - "\n", - "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", - "and high enough resolution that it captures all the light from the light profile.\n", - "\n", - "Note that we are using the source galaxy's true light profile, which corresponds to its emission in the source-plane.\n", - "For real datasets, we have to infer this via lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels\n", - "\n", - "print(f\"Total Source Flux: {total_flux} e- s^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Below, we will compare how this true source flux compares to the inferred source fluxes we compute using different\n", - "source modeling techniques (e.g. parametric and pixelized source models). Converting the flux to magnitudes or\n", - "other quantities used for tasks like SED fitting is described in the `autolens_workspace/*/guides/units/flux` example.\n", - "\n", - "__Source Magnification__\n", - "\n", - "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", - "total surface brightness in the source-plane.\n", - "\n", - "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", - "We therefore explicitly mention how area folds into the calculation below.\n", - "\n", - "To ensure the magnification is stable and that we resolve all source emission in both the image-plane and source-plane \n", - "we use a very high resolution grid, higher than we used to compute the total flux above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We repeat our calculation of the source's total flux in the source-plane using this higher resolution grid, note\n", - "that we do not take the area into account, the reason for this is explained below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_source_plane_flux = np.sum(image) # in units e- s^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now need the total flux of the lensed source in the image-plane, that is how much flux we measure after\n", - "gravitational lensing.\n", - "\n", - "To calculation this, we first ray-trace the grid above from the image-plane to the source-plane using the tracer\n", - "and then pass it to the source galaxy's light profile to compute the lensed image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "source_plane_grid = traced_grid_list[1]\n", - "\n", - "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(\n", - " lensed_source_image\n", - ") # in units e- s^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now take the ratio of the total image-plane flux to source-plane flux to estimate the magnification.\n", - "\n", - "Because both fluxes were computed on grids with the same total area and area per pixel, we do not need to\n", - "explicitly account for area in this calculation. This is because the area terms cancel out when taking the ratio.\n", - "Were the grid areas different, we would need to include area terms in the calculation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Magnification: {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to compute\n", - "source science calculations for real datasets.\n", - "\n", - "The code below shows how using a tracer, composed of any combination of lens and source galaxies, we can\n", - "compute the source flux and magnification. It reproduces the calculations above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "image_plane_grid = traced_grid_list[0]\n", - "source_plane_grid = traced_grid_list[1]\n", - "\n", - "lensed_source_image = tracer.planes[1].image_2d_from(grid=source_plane_grid)\n", - "source_plane_image = tracer.planes[1].image_2d_from(grid=image_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(lensed_source_image)\n", - "total_source_plane_flux = np.sum(source_plane_image)\n", - "\n", - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Plane Total Flux via Tracer: {total_source_plane_flux} e- s^-1\")\n", - "print(f\"Source Magnification via Tracer: {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parametric Source Models__\n", - "\n", - "If your lens modeling uses a parametric source model (e.g. Sersic, Multi Gaussian Expansion), the only object\n", - "you need to perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", - "\n", - "Alternatively, as done above, you can manually set up a tracer using the lens and source galaxies inferred\n", - "by lens modeling.\n", - "\n", - "Therefore, you may now wish to go to your results, extract the `max_log_likelihood_tracer`, and use it to compute\n", - "the source flux and magnification as shown above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Source Science\n", + "==============\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", + "size of the source.\n", + "\n", + "This example shows how to perform these calculations using Sersic parametric sources on imaging data, which\n", + "is conceptually the simplest case for source science calculations and a good introduction to the topic.\n", + "\n", + "__Contents__\n", + "\n", + "- **Simulated Dataset:** We load and plot the `simple__no_lens_light` example dataset, which is simulated imaging of a.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Source Values:** Source science calculations for real lenses are performed using the best-fitting model inferred.\n", + "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", + "- **Source Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", + "- **Tracer:** Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to.\n", + "- **Parametric Source Models:** If your lens modeling uses a parametric source model (e.g." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulated Dataset__\n", + "\n", + "We load and plot the `simple__no_lens_light` example dataset, which is simulated imaging of a strong lens\n", + "that we will use to demonstrate source science caluculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We apply a 3.0 arcsecond circular mask and apply it to the `Imaging` object.\n", + "\n", + "Source science calculations are typically performed on masked datasets to ensure only the lensed source is used\n", + "in the calculations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Values__\n", + "\n", + "Source science calculations for real lenses are performed using the best-fitting model inferred from a dataset, \n", + "and this example demonstrates how to use this below.\n", + "\n", + "However, we for simplicity, we demonstrate these calculations using the Sersic source model used to simulate the dataset, \n", + "which we refer to as the \"true\" source model. When analysing real strong lenses, a true underlying model is not known, \n", + "but for simulated datasets it is.\n", + "\n", + "This allows us to illustrate the calculations in a way that does not depend on the specific details of the data or \n", + "on assumptions about how the lens model is inferred.\n", + "\n", + "The `tracer` below corresponds to the same tracer used to simulate the `simple__no_lens_light` dataset, and therefore \n", + "represents the true source model. We also include the 2D grid of (y,x) coordinates which simulate the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the image of the tracer, we confirm it looks identical to the simulated dataset but does not have\n", + "CCD imaging features such as noise or blurring from a PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Flux__\n", + "\n", + "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", + "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", + "\n", + "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", + "we compute the image of the light profile, and then sum the image. \n", + "\n", + "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", + "we will assume everything is in electrons per second (`e- s^-1`), which is typical for Hubble Space Telescope imaging data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Source Galaxy's Intensity {source_galaxy.bulge.intensity} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The total flux, in units of `e- s^-1` , is computed by summing the image of the light profile over all pixels.\n", + "\n", + "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", + "and high enough resolution that it captures all the light from the light profile.\n", + "\n", + "Note that we are using the source galaxy's true light profile, which corresponds to its emission in the source-plane.\n", + "For real datasets, we have to infer this via lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_flux = np.sum(image) # in units e- s^-1 as summed over pixels\n", + "\n", + "print(f\"Total Source Flux: {total_flux} e- s^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below, we will compare how this true source flux compares to the inferred source fluxes we compute using different\n", + "source modeling techniques (e.g. parametric and pixelized source models). Converting the flux to magnitudes or\n", + "other quantities used for tasks like SED fitting is described in the `autolens_workspace/*/guides/units/flux` example.\n", + "\n", + "__Source Magnification__\n", + "\n", + "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", + "total surface brightness in the source-plane.\n", + "\n", + "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", + "We therefore explicitly mention how area folds into the calculation below.\n", + "\n", + "To ensure the magnification is stable and that we resolve all source emission in both the image-plane and source-plane \n", + "we use a very high resolution grid, higher than we used to compute the total flux above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We repeat our calculation of the source's total flux in the source-plane using this higher resolution grid, note\n", + "that we do not take the area into account, the reason for this is explained below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_source_plane_flux = np.sum(image) # in units e- s^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now need the total flux of the lensed source in the image-plane, that is how much flux we measure after\n", + "gravitational lensing.\n", + "\n", + "To calculation this, we first ray-trace the grid above from the image-plane to the source-plane using the tracer\n", + "and then pass it to the source galaxy's light profile to compute the lensed image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "source_plane_grid = traced_grid_list[1]\n", + "\n", + "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(\n", + " lensed_source_image\n", + ") # in units e- s^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now take the ratio of the total image-plane flux to source-plane flux to estimate the magnification.\n", + "\n", + "Because both fluxes were computed on grids with the same total area and area per pixel, we do not need to\n", + "explicitly account for area in this calculation. This is because the area terms cancel out when taking the ratio.\n", + "Were the grid areas different, we would need to include area terms in the calculation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Magnification: {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to compute\n", + "source science calculations for real datasets.\n", + "\n", + "The code below shows how using a tracer, composed of any combination of lens and source galaxies, we can\n", + "compute the source flux and magnification. It reproduces the calculations above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "image_plane_grid = traced_grid_list[0]\n", + "source_plane_grid = traced_grid_list[1]\n", + "\n", + "lensed_source_image = tracer.planes[1].image_2d_from(grid=source_plane_grid)\n", + "source_plane_image = tracer.planes[1].image_2d_from(grid=image_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(lensed_source_image)\n", + "total_source_plane_flux = np.sum(source_plane_image)\n", + "\n", + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Plane Total Flux via Tracer: {total_source_plane_flux} e- s^-1\")\n", + "print(f\"Source Magnification via Tracer: {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parametric Source Models__\n", + "\n", + "If your lens modeling uses a parametric source model (e.g. Sersic, Multi Gaussian Expansion), the only object\n", + "you need to perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", + "\n", + "Alternatively, as done above, you can manually set up a tracer using the lens and source galaxies inferred\n", + "by lens modeling.\n", + "\n", + "Therefore, you may now wish to go to your results, extract the `max_log_likelihood_tracer`, and use it to compute\n", + "the source flux and magnification as shown above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/imaging/start_here.ipynb b/notebooks/imaging/start_here.ipynb index 0f92dd39f..afcad06f6 100644 --- a/notebooks/imaging/start_here.ipynb +++ b/notebooks/imaging/start_here.ipynb @@ -142,7 +142,16 @@ "dataset_name = \"cosmos_web_ring\"\n", "dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", "\n", + "# PSF convolution runs at the image resolution (sub size 1), which is the fastest\n", + "# option and accurate for well-sampled PSFs. Supplying a PSF at a multiple of the\n", + "# image resolution and raising this value improves blurring fidelity for\n", + "# undersampled PSFs (e.g. HST / Euclid VIS) at extra compute cost \u2014 see\n", + "# `guides/advanced/over_sampling.py` and the simulator's `__Oversampled PSF__` section.\n", + "psf_convolve_over_sample_size = 1\n", + "\n", "dataset = al.Imaging.from_fits(\n", + " convolve_over_sample_size_lp=psf_convolve_over_sample_size,\n", + " convolve_over_sample_size_pixelization=psf_convolve_over_sample_size,\n", " data_path=dataset_path / \"data.fits\",\n", " psf_path=dataset_path / \"psf.fits\",\n", " noise_map_path=dataset_path / \"noise_map.fits\",\n", diff --git a/notebooks/interferometer/casa_reduction.ipynb b/notebooks/interferometer/casa_reduction.ipynb index 78bccd847..3bd4c8bdc 100644 --- a/notebooks/interferometer/casa_reduction.ipynb +++ b/notebooks/interferometer/casa_reduction.ipynb @@ -1,594 +1,631 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Interferometer: CASA Reduction\n", - "==============================\n", - "\n", - "This script is a practical walkthrough of how to reduce an ALMA / JVLA interferometer\n", - "observation with **CASA** and export it into the three FITS files that **PyAutoLens**\n", - "expects for lens modeling:\n", - "\n", - "- `data.fits` \u2014 complex visibilities, shape `(n_vis,)`\n", - "- `noise_map.fits` \u2014 complex per-visibility RMS (sigma), shape `(n_vis,)`\n", - "- `uv_wavelengths.fits` \u2014 `(u, v)` baselines in units of *wavelengths*, shape `(n_vis, 2)`\n", - "\n", - "The script is deliberately a hybrid: the sections that must run **inside the CASA\n", - "environment** (`split`, `statwt`, `uvcontsub`, `tb.open`, etc.) are shown as CASA\n", - "snippets you copy into a CASA session. The plain Python helpers at the bottom use\n", - "`numpy` and `astropy` to convert CASA-exported columns into the shapes PyAutoLens\n", - "requires, and can be run either inside CASA (which has a Python interpreter) or\n", - "in a normal Python environment once the FITS files have been produced.\n", - "\n", - "The script is **a starting point**, not a finished tool. Every real dataset has\n", - "quirks (single-field vs. mosaic, spectral line vs. continuum, different numbers\n", - "of spectral windows, flagged antennas, polarisation products, ...). You will\n", - "almost certainly need to tweak it. If you get stuck, please contact us on the\n", - "PyAutoLens **SLACK** \u2014 interferometry support is actively evolving, and your\n", - "feedback directly shapes these examples.\n", - "\n", - "__Contents__\n", - "\n", - "- **PyAutoLens Requirements:** What the `al.Interferometer` object expects as input.\n", - "- **CASA in Five Steps:** The canonical CASA reduction pipeline for PyAutoLens.\n", - "- **Step 1 \u2014 Split by Field / SPW:** Isolate the lens target and each spectral window.\n", - "- **Step 2 \u2014 Channel Averaging:** Reduce visibility count by averaging channels.\n", - "- **Step 3 \u2014 Continuum Subtraction:** Optional `uvcontsub` for line observations.\n", - "- **Step 4 \u2014 Rescale Sigmas:** `statwt` makes the SIGMA column reflect real scatter.\n", - "- **Step 5 \u2014 Export to FITS:** Python helpers to read the .ms and write .fits.\n", - "- **Combining Spectral Windows:** Concatenate per-SPW arrays into a single dataset.\n", - "- **Polarisations:** Averaging XX/YY (or RR/LL) into a single complex visibility.\n", - "- **Building the Interferometer Object:** Load the FITS files into `al.Interferometer`.\n", - "- **Troubleshooting:** Common pitfalls and what to check.\n", - "- **SLACK:** Where to ask for help.\n", - "\n", - "__PyAutoLens Requirements__\n", - "\n", - "`al.Interferometer.from_fits(...)` loads three FITS files. **On disk**, they are\n", - "stored as plain real-valued arrays with a single `n_vis` axis, after polarisations\n", - "and channels have been collapsed into one long visibility list:\n", - "\n", - "- `data.fits` : real array, shape `(n_vis, 2)` \u2014 col 0 real, col 1 imag\n", - "- `noise_map.fits` : real array, shape `(n_vis, 2)` \u2014 col 0 sigma_real, col 1 sigma_imag\n", - "- `uv_wavelengths.fits` : real array, shape `(n_vis, 2)` \u2014 cols are (u, v) in wavelengths\n", - "\n", - "PyAutoLens converts the first two into complex arrays internally when the\n", - "`Interferometer` object is built.\n", - "\n", - "Critically:\n", - "\n", - "- The `uv` coordinates must be in **wavelengths**, not metres. The CASA `UVW`\n", - " column stores metres, so you must multiply by `frequency / c` per channel.\n", - "- You must flatten *all* polarisation and channel axes into the single `n_vis`\n", - " axis before writing the FITS. The standard order is:\n", - " **(1) average polarisations \u2192 (2) reshape channels into n_vis \u2192 (3) concatenate SPWs**.\n", - " If you skip step (1), the polarisations end up interleaved with the channels\n", - " and the noise-map will be wrong.\n", - "\n", - "__CASA in Five Steps__\n", - "\n", - "A typical ALMA calibrated Measurement Set `my_obs.ms.split.cal` becomes PyAutoLens\n", - "input via this pipeline:\n", - "\n", - " 1. split out the field and (optionally) per-SPW \u2192 field/spw isolated .ms\n", - " 2. channel-average to reduce n_vis \u2192 chanaveraged .ms\n", - " 3. (line only) uvcontsub to remove continuum \u2192 .ms.contsub\n", - " 4. statwt to rescale SIGMA \u2192 .ms.statwt\n", - " 5. export DATA / UVW / CHAN_FREQ / SIGMA to FITS \u2192 data/noise/uv .fits\n", - "\n", - "The rest of this script walks through each of these steps in order. All `split(...)`,\n", - "`statwt(...)`, `uvcontsub(...)`, `tb.open(...)` calls must be executed from within\n", - "a running CASA session (either interactively or via `casa -c my_script.py`).\n", - "\n", - "__Step 1 \u2014 Split by Field / SPW__\n", - "\n", - "ALMA `.ms` datasets typically contain multiple fields (calibrators + science\n", - "target) and multiple spectral windows (SPWs). Start by splitting out just the\n", - "science field. You can also split per-SPW at this stage, which makes later\n", - "channel-averaging simpler because different SPWs can have different numbers\n", - "of channels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# CASA:\n", - "# split(\n", - "# vis=\"my_obs.ms.split.cal\",\n", - "# outputvis=\"my_obs_field_SPT-0418_spw0.ms\",\n", - "# keepmms=True,\n", - "# field=\"SPT-0418\",\n", - "# spw=\"0\",\n", - "# datacolumn=\"data\", # use \"corrected\" if CORRECTED_DATA exists\n", - "# keepflags=False,\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Repeat the `split` for each SPW you want to include (e.g. spw=\"1\", \"2\", \"3\").\n", - "If your `.ms` only has a single field and you want *all* SPWs in one file, you\n", - "can omit the `spw` argument.\n", - "\n", - "__Step 2 \u2014 Channel Averaging__\n", - "\n", - "ALMA observations often have thousands of channels per SPW, producing huge\n", - "visibility counts. With the JAX-native `TransformerNUFFT` (backed by `nufftax`),\n", - "the light-profile modeling path scales efficiently to ALMA-class datasets with\n", - "many millions of visibilities, so channel averaging is no longer required purely\n", - "for performance. You may still average channels (typically by setting `width`)\n", - "to reduce the dataset size, simplify the analysis, or maximize signal-to-noise\n", - "for continuum-style lens modeling.\n", - "\n", - "Note that `width` must not exceed the number of channels in the SPW \u2014 check with\n", - "the `get_num_chan` helper below if unsure." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# CASA:\n", - "# split(\n", - "# vis=\"my_obs_field_SPT-0418_spw0.ms\",\n", - "# outputvis=\"my_obs_field_SPT-0418_spw0_chanavg.ms\",\n", - "# keepmms=True,\n", - "# width=128, # average 128 channels together\n", - "# datacolumn=\"data\",\n", - "# keepflags=False,\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Step 3 \u2014 Continuum Subtraction (line observations only)__\n", - "\n", - "If you're modeling a **spectral line** (e.g. CO(9-8) from a high-z lensed galaxy),\n", - "you want to remove the underlying continuum first using `uvcontsub`. Specify\n", - "the line-free channel ranges via `fitspw` and the channels to keep via `spw`.\n", - "Skip this step for pure continuum lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# CASA:\n", - "# uvcontsub(\n", - "# vis=\"my_obs_field_SPT-0418_spw0.ms\",\n", - "# outputvis=\"my_obs_field_SPT-0418_spw0.ms.contsub\",\n", - "# fitspw=\"0:10~200;800~1000\", # line-free channel ranges for the fit\n", - "# spw=\"0\", # channels to keep in the output\n", - "# fitorder=1,\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If you ran `uvcontsub`, use the resulting `.ms.contsub` file as input to Step 4\n", - "and Step 5 instead of the original.\n", - "\n", - "__Step 4 \u2014 Rescale Sigmas with statwt__\n", - "\n", - "The `SIGMA` / `WEIGHT` columns that come out of ALMA calibration are often\n", - "set to nominal values rather than reflecting the true scatter in the visibilities.\n", - "`statwt` measures the scatter per-baseline and rescales SIGMA accordingly. This\n", - "is **strongly recommended** before exporting the noise-map \u2014 without it, your\n", - "per-visibility error bars will be wrong and the likelihood will be biased.\n", - "\n", - "`statwt` modifies the `.ms` in place, so copy it first if you want to keep the\n", - "original SIGMAs for comparison." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# CASA (from the shell / inside a CASA session):\n", - "# import os\n", - "# os.system(\"cp -r my_obs_field_SPT-0418_spw0_chanavg.ms \"\n", - "# \"my_obs_field_SPT-0418_spw0_chanavg.ms.statwt\")\n", - "# statwt(\n", - "# vis=\"my_obs_field_SPT-0418_spw0_chanavg.ms.statwt\",\n", - "# datacolumn=\"data\",\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Step 5 \u2014 Export DATA / UVW / CHAN_FREQ / SIGMA to FITS__\n", - "\n", - "With the reduced `.ms` in hand we now read the relevant columns out with the\n", - "CASA `tb` tool and save them as FITS files. The helpers below can be saved\n", - "as a separate `.py` file and executed inside CASA via\n", - "\n", - " casa -c export_to_fits.py\n", - "\n", - "or pasted directly into an interactive CASA session. They use the global `tb`\n", - "object CASA injects, together with numpy and astropy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "import os\n", - "import numpy as np\n", - "\n", - "try:\n", - " from astropy import units, constants\n", - " from astropy.io import fits\n", - "\n", - " astropy_is_imported = True\n", - "except ImportError:\n", - " astropy_is_imported = False\n", - "\n", - "\n", - "def getcol_wrapper(ms, table, colname):\n", - " \"\"\"\n", - " Open a CASA Measurement Set sub-table and return a squeezed column.\n", - "\n", - " Parameters\n", - " ----------\n", - " ms : str\n", - " Path to the .ms directory.\n", - " table : str\n", - " Sub-table name (e.g. \"SPECTRAL_WINDOW\", \"DATA_DESCRIPTION\"). Empty\n", - " string for the main table.\n", - " colname : str\n", - " Column name (e.g. \"DATA\", \"UVW\", \"CHAN_FREQ\", \"SIGMA\", \"NUM_CHAN\").\n", - "\n", - " Returns\n", - " -------\n", - " np.ndarray\n", - " The requested column, with trivial dimensions removed.\n", - " \"\"\"\n", - " if not Path(ms).is_dir():\n", - " raise IOError(f\"{ms} does not exist\")\n", - "\n", - " tb.open(f\"{ms}/{table}\" if table else ms) # noqa: F821 \u2014 `tb` is the CASA tool\n", - " col = np.squeeze(tb.getcol(colname)) # noqa: F821\n", - " tb.close() # noqa: F821\n", - " return col\n", - "\n", - "\n", - "def get_num_chan(ms):\n", - " \"\"\"Number of channels per SPW (shape `(n_spw,)` or scalar).\"\"\"\n", - " return getcol_wrapper(ms=ms, table=\"SPECTRAL_WINDOW\", colname=\"NUM_CHAN\")\n", - "\n", - "\n", - "def get_spw_ids(ms):\n", - " \"\"\"SPW ids present in the main table (shape `(n_spw,)` or scalar).\"\"\"\n", - " return getcol_wrapper(ms=ms, table=\"DATA_DESCRIPTION\", colname=\"SPECTRAL_WINDOW_ID\")\n", - "\n", - "\n", - "def get_frequencies(ms):\n", - " \"\"\"Channel frequencies in Hz (shape `(n_chan,)` or `(n_chan, n_spw)`).\"\"\"\n", - " return getcol_wrapper(ms=ms, table=\"SPECTRAL_WINDOW\", colname=\"CHAN_FREQ\")\n", - "\n", - "\n", - "def _write_array(filename, data):\n", - " \"\"\"Write a numpy array as FITS if astropy is available, else `.npy`.\"\"\"\n", - " if astropy_is_imported:\n", - " fits.writeto(filename=filename + \".fits\", data=data, overwrite=True)\n", - " else:\n", - " with open(filename + \".numpy\", \"wb\") as f:\n", - " np.save(f, data)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visibilities__\n", - "\n", - "The `DATA` column has complex dtype and shape `(n_pol, n_chan, n_vis)` for a\n", - "single-SPW split. We stack the real and imaginary parts along a new trailing\n", - "axis so downstream code can load them as a real-valued FITS image and\n", - "recombine into a complex array via `vis[..., 0] + 1j * vis[..., 1]`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def get_visibilities(ms):\n", - " data = getcol_wrapper(ms=ms, table=\"\", colname=\"DATA\")\n", - " return np.stack(arrays=(data.real, data.imag), axis=-1)\n", - "\n", - "\n", - "def export_visibilities(ms, filename):\n", - " if Path(filename + \".fits\").is_file() or Path(filename + \".numpy\").is_file():\n", - " print(f\"{filename} already exists \u2014 skipping\")\n", - " return\n", - " visibilities = get_visibilities(ms=ms)\n", - " print(\"shape (visibilities):\", visibilities.shape)\n", - " _write_array(filename=filename, data=visibilities)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__UV Wavelengths__\n", - "\n", - "`UVW` in CASA is stored in **metres**, with shape `(3, n_vis)`. PyAutoLens\n", - "needs `(u, v)` in **wavelengths**, so we multiply by `frequency / c` for each\n", - "channel. This produces a `(2, n_chan, n_vis)` array of u, v wavelengths, which\n", - "is later reshaped to `(n_vis_total, 2)` once you flatten over channels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def convert_array_to_wavelengths(array, frequency):\n", - " if astropy_is_imported:\n", - " return (\n", - " ((array * units.m) * (frequency * units.Hz) / constants.c).decompose().value\n", - " )\n", - " return array * frequency / 299792458.0\n", - "\n", - "\n", - "def get_uv_wavelengths(ms):\n", - " uvw = getcol_wrapper(ms=ms, table=\"\", colname=\"UVW\")\n", - " chan_freq = get_frequencies(ms=ms)\n", - "\n", - " if np.shape(chan_freq):\n", - " n_chan = np.shape(chan_freq)[0]\n", - " u_wavelengths = np.zeros((n_chan, uvw.shape[1]))\n", - " v_wavelengths = np.zeros((n_chan, uvw.shape[1]))\n", - " for i in range(n_chan):\n", - " u_wavelengths[i] = convert_array_to_wavelengths(uvw[0], chan_freq[i])\n", - " v_wavelengths[i] = convert_array_to_wavelengths(uvw[1], chan_freq[i])\n", - " else:\n", - " u_wavelengths = convert_array_to_wavelengths(uvw[0], chan_freq)\n", - " v_wavelengths = convert_array_to_wavelengths(uvw[1], chan_freq)\n", - "\n", - " return np.stack(arrays=(u_wavelengths, v_wavelengths), axis=-1)\n", - "\n", - "\n", - "def export_uv_wavelengths(ms, filename):\n", - " if Path(filename + \".fits\").is_file() or Path(filename + \".numpy\").is_file():\n", - " print(f\"{filename} already exists \u2014 skipping\")\n", - " return\n", - " uv_wavelengths = get_uv_wavelengths(ms=ms)\n", - " print(\"shape (uv_wavelengths):\", uv_wavelengths.shape)\n", - " _write_array(filename=filename, data=uv_wavelengths)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sigma (Noise Map)__\n", - "\n", - "The `SIGMA` column in CASA has shape `(n_pol, n_vis)` \u2014 one value per polarisation\n", - "per visibility \u2014 and CASA assigns the *same* sigma to the real and imaginary\n", - "components. We broadcast it across the channel axis and duplicate the last\n", - "axis so the final shape matches the visibilities: `(n_pol, n_chan, n_vis, 2)`.\n", - "\n", - "For this to be meaningful you **must** have run `statwt` first, otherwise the\n", - "sigmas are nominal placeholders rather than real scatter estimates." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def get_sigma(ms):\n", - " sigma = getcol_wrapper(ms=ms, table=\"\", colname=\"SIGMA\")\n", - " chan_freq = np.atleast_1d(get_frequencies(ms=ms))\n", - " n_chan = chan_freq.shape[0]\n", - "\n", - " # Re-introduce the polarisation axis if np.squeeze removed it (n_pol == 1).\n", - " if sigma.ndim == 1:\n", - " sigma = sigma[np.newaxis, :]\n", - "\n", - " sigma = np.tile(sigma[:, np.newaxis, :], (1, n_chan, 1))\n", - " return np.stack(arrays=(sigma, sigma), axis=-1)\n", - "\n", - "\n", - "def export_sigma(ms, filename):\n", - " if Path(filename + \".fits\").is_file() or Path(filename + \".numpy\").is_file():\n", - " print(f\"{filename} already exists \u2014 skipping\")\n", - " return\n", - " sigma = get_sigma(ms=ms)\n", - " print(\"shape (sigma):\", sigma.shape)\n", - " _write_array(filename=filename, data=sigma)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Polarisations (do this FIRST)__\n", - "\n", - "Raw ALMA data has two polarisations (XX, YY) on axis 0 of `DATA` and `SIGMA`.\n", - "For lens modeling the standard approach is to average the two into a single\n", - "complex visibility with combined sigma, since the source emission is not\n", - "polarised for the vast majority of lensed galaxies. **Do this before\n", - "flattening channels or concatenating SPWs**, otherwise the polarisations end\n", - "up interleaved with the channel axis and the noise-map no longer corresponds\n", - "to the visibility it was measured from.\n", - "\n", - " # vis: (n_pol, n_chan, n_vis, 2) \u2014 real/imag on last axis\n", - " vis_avg = 0.5 * (vis[0] + vis[1]) # -> (n_chan, n_vis, 2)\n", - "\n", - " # sig: (n_pol, n_chan, n_vis, 2) \u2014 sigma_real/sigma_imag on last axis\n", - " sig_avg = 0.5 * np.sqrt(sig[0] ** 2 + sig[1] ** 2) # -> (n_chan, n_vis, 2)\n", - "\n", - "If your science *does* care about polarisation, keep each pol as a separate\n", - "dataset and fit them jointly \u2014 contact us on Slack for the current state of\n", - "joint-polarisation modeling.\n", - "\n", - "__Combining Spectral Windows__\n", - "\n", - "Once you have per-SPW pol-averaged arrays, flatten channels into the n_vis\n", - "axis and concatenate across SPWs. All three arrays (visibilities, sigma,\n", - "uv-wavelengths) use the same ordering so they stay aligned.\n", - "\n", - " # Per-SPW (pol-averaged) shapes:\n", - " # vis_avg : (n_chan_spw, n_vis_spw, 2)\n", - " # sig_avg : (n_chan_spw, n_vis_spw, 2)\n", - " # uv : (n_chan_spw, n_vis_spw, 2)\n", - "\n", - " vis_all = np.concatenate(\n", - " [v.reshape(-1, 2) for v in per_spw_vis], axis=0\n", - " )\n", - " sig_all = np.concatenate(\n", - " [s.reshape(-1, 2) for s in per_spw_sig], axis=0\n", - " )\n", - " uv_all = np.concatenate(\n", - " [u.reshape(-1, 2) for u in per_spw_uv], axis=0\n", - " )\n", - "\n", - "After this step all three arrays have shape `(n_vis_total, 2)` \u2014 exactly\n", - "what `al.Interferometer.from_fits` expects.\n", - "\n", - "__Building the Interferometer Object__\n", - "\n", - "Write the three concatenated arrays to FITS (using `astropy.io.fits.writeto`\n", - "or `autoconf.fitsable.output_to_fits`) and load them via the canonical\n", - "workspace pattern used in `start_here.py`:\n", - "\n", - " import autolens as al\n", - "\n", - " # Small, fast real-space mask \u2014 keep shape_native and radius modest while\n", - " # you are iterating on the reduction. Increase later for production fits.\n", - " real_space_mask = al.Mask2D.circular(\n", - " shape_native=(64, 64),\n", - " pixel_scales=0.05,\n", - " radius=1.5,\n", - " )\n", - "\n", - " dataset = al.Interferometer.from_fits(\n", - " data_path=\"data.fits\",\n", - " noise_map_path=\"noise_map.fits\",\n", - " uv_wavelengths_path=\"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT, # NUFFT scales to large n_vis\n", - " )\n", - "\n", - "The 64x64 / 0.05\" mask above keeps the Fourier transform and any dirty-image\n", - "sanity plots fast while you iterate \u2014 a first-pass check of your reduction\n", - "should finish in seconds, not minutes. Bump `shape_native` up (e.g. 256x256)\n", - "and drop `pixel_scales` (e.g. 0.02\") for the real modeling run. Choosing\n", - "these numbers sensibly for your instrument is covered in\n", - "`scripts/interferometer/data_preparation.py`.\n", - "\n", - "__Troubleshooting__\n", - "\n", - "- **`uv_wavelengths` values look way too big/small** \u2014 you forgot the\n", - " metres \u2192 wavelengths conversion, or you used the wrong frequency column.\n", - "- **Dirty image is offset / upside-down** \u2014 the `(u, v)` sign convention or\n", - " the `(y, x)` axis order may disagree with your real-space mask. Plot the\n", - " dirty image with `aplt.subplot_interferometer_dirty_images` as a sanity\n", - " check.\n", - "- **Sigmas produce a flat chi-squared map** \u2014 you probably forgot `statwt`.\n", - "- **n_vis is enormous** \u2014 ensure you are using `TransformerNUFFT` (the default\n", - " recommendation, backed by JAX-native `nufftax`). It scales to many millions\n", - " of visibilities. Channel averaging via a larger `width` in `split` is still\n", - " useful when you want to reduce dataset size for other reasons.\n", - "- **Only one polarisation present** \u2014 if CASA has flagged one pol the\n", - " averaging code above will produce NaNs. Check with `get_visibilities`\n", - " and branch accordingly.\n", - "\n", - "__SLACK__\n", - "\n", - "These scripts are a starting point, not a polished pipeline. If your\n", - "observation has mosaicking, heterogeneous SPWs, unusual correlators, or you\n", - "just can't get the numbers to line up \u2014 ping us on the PyAutoLens Slack.\n", - "Direct user feedback is actively shaping this workflow and we're happy to\n", - "help debug real datasets.\n", - "\n", - "__Running as a Script__\n", - "\n", - "The block below is illustrative \u2014 edit the paths/widths for your own reduction\n", - "and run inside CASA with `casa -c scripts/interferometer/casa_reduction.py`.\n", - "It assumes you have already run `split`, `uvcontsub` (optional) and `statwt`\n", - "by hand on the `.ms`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "if __name__ == \"__main__\":\n", - "\n", - " uid = \"my_obs_field_SPT-0418\"\n", - " width = 128\n", - "\n", - " ms_chanavg = f\"{uid}_spw0_chanavg.ms\"\n", - " ms_statwt = f\"{ms_chanavg}.statwt\"\n", - "\n", - " export_uv_wavelengths(ms=ms_chanavg, filename=f\"uv_wavelengths_{uid}_width_{width}\")\n", - " export_visibilities(ms=ms_chanavg, filename=f\"visibilities_{uid}_width_{width}\")\n", - " export_sigma(ms=ms_statwt, filename=f\"sigma_{uid}_width_{width}_statwt\")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Interferometer: CASA Reduction\n", + "==============================\n", + "\n", + "This script is a practical walkthrough of how to reduce an ALMA / JVLA interferometer\n", + "observation with **CASA** and export it into the three FITS files that **PyAutoLens**\n", + "expects for lens modeling:\n", + "\n", + "- `data.fits` \u2014 complex visibilities, shape `(n_vis,)`\n", + "- `noise_map.fits` \u2014 complex per-visibility RMS (sigma), shape `(n_vis,)`\n", + "- `uv_wavelengths.fits` \u2014 `(u, v)` baselines in units of *wavelengths*, shape `(n_vis, 2)`\n", + "\n", + "The script is deliberately a hybrid: the sections that must run **inside the CASA\n", + "environment** (`split`, `statwt`, `uvcontsub`, `tb.open`, etc.) are shown as CASA\n", + "snippets you copy into a CASA session. The plain Python helpers at the bottom use\n", + "`numpy` and `astropy` to convert CASA-exported columns into the shapes PyAutoLens\n", + "requires, and can be run either inside CASA (which has a Python interpreter) or\n", + "in a normal Python environment once the FITS files have been produced.\n", + "\n", + "The script is **a starting point**, not a finished tool. Every real dataset has\n", + "quirks (single-field vs. mosaic, spectral line vs. continuum, different numbers\n", + "of spectral windows, flagged antennas, polarisation products, ...). You will\n", + "almost certainly need to tweak it. If you get stuck, please contact us on the\n", + "PyAutoLens **SLACK** \u2014 interferometry support is actively evolving, and your\n", + "feedback directly shapes these examples.\n", + "\n", + "__Contents__\n", + "\n", + "- **PyAutoLens Requirements:** What the `al.Interferometer` object expects as input.\n", + "- **CASA in Five Steps:** The canonical CASA reduction pipeline for PyAutoLens.\n", + "- **Step 1 \u2014 Split by Field / SPW:** Isolate the lens target and each spectral window.\n", + "- **Step 2 \u2014 Channel Averaging:** Reduce visibility count by averaging channels.\n", + "- **Step 3 \u2014 Continuum Subtraction:** Optional `uvcontsub` for line observations.\n", + "- **Step 4 \u2014 Rescale Sigmas:** `statwt` makes the SIGMA column reflect real scatter.\n", + "- **Step 5 \u2014 Export to FITS:** Python helpers to read the .ms and write .fits.\n", + "- **Combining Spectral Windows:** Concatenate per-SPW arrays into a single dataset.\n", + "- **Polarisations:** Averaging XX/YY (or RR/LL) into a single complex visibility.\n", + "- **Building the Interferometer Object:** Load the FITS files into `al.Interferometer`.\n", + "- **Troubleshooting:** Common pitfalls and what to check.\n", + "- **SLACK:** Where to ask for help.\n", + "\n", + "__PyAutoLens Requirements__\n", + "\n", + "`al.Interferometer.from_fits(...)` loads three FITS files. **On disk**, they are\n", + "stored as plain real-valued arrays with a single `n_vis` axis, after polarisations\n", + "and channels have been collapsed into one long visibility list:\n", + "\n", + "- `data.fits` : real array, shape `(n_vis, 2)` \u2014 col 0 real, col 1 imag\n", + "- `noise_map.fits` : real array, shape `(n_vis, 2)` \u2014 col 0 sigma_real, col 1 sigma_imag\n", + "- `uv_wavelengths.fits` : real array, shape `(n_vis, 2)` \u2014 cols are (u, v) in wavelengths\n", + "\n", + "PyAutoLens converts the first two into complex arrays internally when the\n", + "`Interferometer` object is built.\n", + "\n", + "Critically:\n", + "\n", + "- The `uv` coordinates must be in **wavelengths**, not metres. The CASA `UVW`\n", + " column stores metres, so you must multiply by `frequency / c` per channel.\n", + "- You must flatten *all* polarisation and channel axes into the single `n_vis`\n", + " axis before writing the FITS. The standard order is:\n", + " **(1) average polarisations \u2192 (2) reshape channels into n_vis \u2192 (3) concatenate SPWs**.\n", + " If you skip step (1), the polarisations end up interleaved with the channels\n", + " and the noise-map will be wrong.\n", + "\n", + "__CASA in Five Steps__\n", + "\n", + "A typical ALMA calibrated Measurement Set `my_obs.ms.split.cal` becomes PyAutoLens\n", + "input via this pipeline:\n", + "\n", + " 1. split out the field and (optionally) per-SPW \u2192 field/spw isolated .ms\n", + " 2. channel-average to reduce n_vis \u2192 chanaveraged .ms\n", + " 3. (line only) uvcontsub to remove continuum \u2192 .ms.contsub\n", + " 4. statwt to rescale SIGMA \u2192 .ms.statwt\n", + " 5. export DATA / UVW / CHAN_FREQ / SIGMA to FITS \u2192 data/noise/uv .fits\n", + "\n", + "The rest of this script walks through each of these steps in order. All `split(...)`,\n", + "`statwt(...)`, `uvcontsub(...)`, `tb.open(...)` calls must be executed from within\n", + "a running CASA session (either interactively or via `casa -c my_script.py`).\n", + "\n", + "__Step 1 \u2014 Split by Field / SPW__\n", + "\n", + "ALMA `.ms` datasets typically contain multiple fields (calibrators + science\n", + "target) and multiple spectral windows (SPWs). Start by splitting out just the\n", + "science field. You can also split per-SPW at this stage, which makes later\n", + "channel-averaging simpler because different SPWs can have different numbers\n", + "of channels." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# CASA:\n", + "# split(\n", + "# vis=\"my_obs.ms.split.cal\",\n", + "# outputvis=\"my_obs_field_SPT-0418_spw0.ms\",\n", + "# keepmms=True,\n", + "# field=\"SPT-0418\",\n", + "# spw=\"0\",\n", + "# datacolumn=\"data\", # use \"corrected\" if CORRECTED_DATA exists\n", + "# keepflags=False,\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Repeat the `split` for each SPW you want to include (e.g. spw=\"1\", \"2\", \"3\").\n", + "If your `.ms` only has a single field and you want *all* SPWs in one file, you\n", + "can omit the `spw` argument.\n", + "\n", + "__Step 2 \u2014 Channel Averaging__\n", + "\n", + "ALMA observations often have thousands of channels per SPW, producing huge\n", + "visibility counts. With the JAX-native `TransformerNUFFT` (backed by `nufftax`),\n", + "the light-profile modeling path scales efficiently to ALMA-class datasets with\n", + "many millions of visibilities, so channel averaging is no longer required purely\n", + "for performance. You may still average channels (typically by setting `width`)\n", + "to reduce the dataset size, simplify the analysis, or maximize signal-to-noise\n", + "for continuum-style lens modeling.\n", + "\n", + "Note that `width` must not exceed the number of channels in the SPW \u2014 check with\n", + "the `get_num_chan` helper below if unsure." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# CASA:\n", + "# split(\n", + "# vis=\"my_obs_field_SPT-0418_spw0.ms\",\n", + "# outputvis=\"my_obs_field_SPT-0418_spw0_chanavg.ms\",\n", + "# keepmms=True,\n", + "# width=128, # average 128 channels together\n", + "# datacolumn=\"data\",\n", + "# keepflags=False,\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Step 3 \u2014 Continuum Subtraction (line observations only)__\n", + "\n", + "If you're modeling a **spectral line** (e.g. CO(9-8) from a high-z lensed galaxy),\n", + "you want to remove the underlying continuum first using `uvcontsub`. Specify\n", + "the line-free channel ranges via `fitspw` and the channels to keep via `spw`.\n", + "Skip this step for pure continuum lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# CASA:\n", + "# uvcontsub(\n", + "# vis=\"my_obs_field_SPT-0418_spw0.ms\",\n", + "# outputvis=\"my_obs_field_SPT-0418_spw0.ms.contsub\",\n", + "# fitspw=\"0:10~200;800~1000\", # line-free channel ranges for the fit\n", + "# spw=\"0\", # channels to keep in the output\n", + "# fitorder=1,\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If you ran `uvcontsub`, use the resulting `.ms.contsub` file as input to Step 4\n", + "and Step 5 instead of the original.\n", + "\n", + "__Step 4 \u2014 Rescale Sigmas with statwt__\n", + "\n", + "The `SIGMA` / `WEIGHT` columns that come out of ALMA calibration are often\n", + "set to nominal values rather than reflecting the true scatter in the visibilities.\n", + "`statwt` measures the scatter per-baseline and rescales SIGMA accordingly. This\n", + "is **strongly recommended** before exporting the noise-map \u2014 without it, your\n", + "per-visibility error bars will be wrong and the likelihood will be biased.\n", + "\n", + "`statwt` modifies the `.ms` in place, so copy it first if you want to keep the\n", + "original SIGMAs for comparison." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# CASA (from the shell / inside a CASA session):\n", + "# import os\n", + "# os.system(\"cp -r my_obs_field_SPT-0418_spw0_chanavg.ms \"\n", + "# \"my_obs_field_SPT-0418_spw0_chanavg.ms.statwt\")\n", + "# statwt(\n", + "# vis=\"my_obs_field_SPT-0418_spw0_chanavg.ms.statwt\",\n", + "# datacolumn=\"data\",\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Step 5 \u2014 Export DATA / UVW / CHAN_FREQ / SIGMA to FITS__\n", + "\n", + "With the reduced `.ms` in hand we now read the relevant columns out with the\n", + "CASA `tb` tool and save them as FITS files. The helpers below can be saved\n", + "as a separate `.py` file and executed inside CASA via\n", + "\n", + " casa -c export_to_fits.py\n", + "\n", + "or pasted directly into an interactive CASA session. They use the global `tb`\n", + "object CASA injects, together with numpy and astropy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "import os\n", + "import numpy as np\n", + "\n", + "try:\n", + " from astropy import units, constants\n", + " from astropy.io import fits\n", + "\n", + " astropy_is_imported = True\n", + "except ImportError:\n", + " astropy_is_imported = False\n", + "\n", + "\n", + "def getcol_wrapper(ms, table, colname):\n", + " \"\"\"\n", + " Open a CASA Measurement Set sub-table and return a squeezed column.\n", + "\n", + " Parameters\n", + " ----------\n", + " ms : str\n", + " Path to the .ms directory.\n", + " table : str\n", + " Sub-table name (e.g. \"SPECTRAL_WINDOW\", \"DATA_DESCRIPTION\"). Empty\n", + " string for the main table.\n", + " colname : str\n", + " Column name (e.g. \"DATA\", \"UVW\", \"CHAN_FREQ\", \"SIGMA\", \"NUM_CHAN\").\n", + "\n", + " Returns\n", + " -------\n", + " np.ndarray\n", + " The requested column, with trivial dimensions removed.\n", + " \"\"\"\n", + " if not Path(ms).is_dir():\n", + " raise IOError(f\"{ms} does not exist\")\n", + "\n", + " tb.open(f\"{ms}/{table}\" if table else ms) # noqa: F821 \u2014 `tb` is the CASA tool\n", + " col = np.squeeze(tb.getcol(colname)) # noqa: F821\n", + " tb.close() # noqa: F821\n", + " return col\n", + "\n", + "\n", + "def get_num_chan(ms):\n", + " \"\"\"Number of channels per SPW (shape `(n_spw,)` or scalar).\"\"\"\n", + " return getcol_wrapper(ms=ms, table=\"SPECTRAL_WINDOW\", colname=\"NUM_CHAN\")\n", + "\n", + "\n", + "def get_spw_ids(ms):\n", + " \"\"\"SPW ids present in the main table (shape `(n_spw,)` or scalar).\"\"\"\n", + " return getcol_wrapper(ms=ms, table=\"DATA_DESCRIPTION\", colname=\"SPECTRAL_WINDOW_ID\")\n", + "\n", + "\n", + "def get_frequencies(ms):\n", + " \"\"\"Channel frequencies in Hz (shape `(n_chan,)` or `(n_chan, n_spw)`).\"\"\"\n", + " return getcol_wrapper(ms=ms, table=\"SPECTRAL_WINDOW\", colname=\"CHAN_FREQ\")\n", + "\n", + "\n", + "def _write_array(filename, data):\n", + " \"\"\"Write a numpy array as FITS if astropy is available, else `.npy`.\"\"\"\n", + " if astropy_is_imported:\n", + " fits.writeto(filename=filename + \".fits\", data=data, overwrite=True)\n", + " else:\n", + " with open(filename + \".numpy\", \"wb\") as f:\n", + " np.save(f, data)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visibilities__\n", + "\n", + "The `DATA` column has complex dtype and shape `(n_pol, n_chan, n_vis)` for a\n", + "single-SPW split. We stack the real and imaginary parts along a new trailing\n", + "axis so downstream code can load them as a real-valued FITS image and\n", + "recombine into a complex array via `vis[..., 0] + 1j * vis[..., 1]`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def get_visibilities(ms):\n", + " data = getcol_wrapper(ms=ms, table=\"\", colname=\"DATA\")\n", + " return np.stack(arrays=(data.real, data.imag), axis=-1)\n", + "\n", + "\n", + "def export_visibilities(ms, filename):\n", + " if Path(filename + \".fits\").is_file() or Path(filename + \".numpy\").is_file():\n", + " print(f\"{filename} already exists \u2014 skipping\")\n", + " return\n", + " visibilities = get_visibilities(ms=ms)\n", + " print(\"shape (visibilities):\", visibilities.shape)\n", + " _write_array(filename=filename, data=visibilities)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__UV Wavelengths__\n", + "\n", + "`UVW` in CASA is stored in **metres**, with shape `(3, n_vis)`. PyAutoLens\n", + "needs `(u, v)` in **wavelengths**, so we multiply by `frequency / c` for each\n", + "channel. This produces a `(2, n_chan, n_vis)` array of u, v wavelengths, which\n", + "is later reshaped to `(n_vis_total, 2)` once you flatten over channels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def convert_array_to_wavelengths(array, frequency):\n", + " if astropy_is_imported:\n", + " return (\n", + " ((array * units.m) * (frequency * units.Hz) / constants.c).decompose().value\n", + " )\n", + " return array * frequency / 299792458.0\n", + "\n", + "\n", + "def get_uv_wavelengths(ms):\n", + " uvw = getcol_wrapper(ms=ms, table=\"\", colname=\"UVW\")\n", + " chan_freq = get_frequencies(ms=ms)\n", + "\n", + " if np.shape(chan_freq):\n", + " n_chan = np.shape(chan_freq)[0]\n", + " u_wavelengths = np.zeros((n_chan, uvw.shape[1]))\n", + " v_wavelengths = np.zeros((n_chan, uvw.shape[1]))\n", + " for i in range(n_chan):\n", + " u_wavelengths[i] = convert_array_to_wavelengths(uvw[0], chan_freq[i])\n", + " v_wavelengths[i] = convert_array_to_wavelengths(uvw[1], chan_freq[i])\n", + " else:\n", + " u_wavelengths = convert_array_to_wavelengths(uvw[0], chan_freq)\n", + " v_wavelengths = convert_array_to_wavelengths(uvw[1], chan_freq)\n", + "\n", + " return np.stack(arrays=(u_wavelengths, v_wavelengths), axis=-1)\n", + "\n", + "\n", + "def export_uv_wavelengths(ms, filename):\n", + " if Path(filename + \".fits\").is_file() or Path(filename + \".numpy\").is_file():\n", + " print(f\"{filename} already exists \u2014 skipping\")\n", + " return\n", + " uv_wavelengths = get_uv_wavelengths(ms=ms)\n", + " print(\"shape (uv_wavelengths):\", uv_wavelengths.shape)\n", + " _write_array(filename=filename, data=uv_wavelengths)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sigma (Noise Map)__\n", + "\n", + "The `SIGMA` column in CASA has shape `(n_pol, n_vis)` \u2014 one value per polarisation\n", + "per visibility \u2014 and CASA assigns the *same* sigma to the real and imaginary\n", + "components. We broadcast it across the channel axis and duplicate the last\n", + "axis so the final shape matches the visibilities: `(n_pol, n_chan, n_vis, 2)`.\n", + "\n", + "For this to be meaningful you **must** have run `statwt` first, otherwise the\n", + "sigmas are nominal placeholders rather than real scatter estimates." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def get_sigma(ms):\n", + " sigma = getcol_wrapper(ms=ms, table=\"\", colname=\"SIGMA\")\n", + " chan_freq = np.atleast_1d(get_frequencies(ms=ms))\n", + " n_chan = chan_freq.shape[0]\n", + "\n", + " # Re-introduce the polarisation axis if np.squeeze removed it (n_pol == 1).\n", + " if sigma.ndim == 1:\n", + " sigma = sigma[np.newaxis, :]\n", + "\n", + " sigma = np.tile(sigma[:, np.newaxis, :], (1, n_chan, 1))\n", + " return np.stack(arrays=(sigma, sigma), axis=-1)\n", + "\n", + "\n", + "def export_sigma(ms, filename):\n", + " if Path(filename + \".fits\").is_file() or Path(filename + \".numpy\").is_file():\n", + " print(f\"{filename} already exists \u2014 skipping\")\n", + " return\n", + " sigma = get_sigma(ms=ms)\n", + " print(\"shape (sigma):\", sigma.shape)\n", + " _write_array(filename=filename, data=sigma)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Polarisations (do this FIRST)__\n", + "\n", + "Raw ALMA data has two polarisations (XX, YY) on axis 0 of `DATA` and `SIGMA`.\n", + "For lens modeling the standard approach is to average the two into a single\n", + "complex visibility with combined sigma, since the source emission is not\n", + "polarised for the vast majority of lensed galaxies. **Do this before\n", + "flattening channels or concatenating SPWs**, otherwise the polarisations end\n", + "up interleaved with the channel axis and the noise-map no longer corresponds\n", + "to the visibility it was measured from.\n", + "\n", + " # vis: (n_pol, n_chan, n_vis, 2) \u2014 real/imag on last axis\n", + " vis_avg = 0.5 * (vis[0] + vis[1]) # -> (n_chan, n_vis, 2)\n", + "\n", + " # sig: (n_pol, n_chan, n_vis, 2) \u2014 sigma_real/sigma_imag on last axis\n", + " sig_avg = 0.5 * np.sqrt(sig[0] ** 2 + sig[1] ** 2) # -> (n_chan, n_vis, 2)\n", + "\n", + "If your science *does* care about polarisation, keep each pol as a separate\n", + "dataset and fit them jointly \u2014 contact us on Slack for the current state of\n", + "joint-polarisation modeling.\n", + "\n", + "__Combining Spectral Windows__\n", + "\n", + "Once you have per-SPW pol-averaged arrays, flatten channels into the n_vis\n", + "axis and concatenate across SPWs. All three arrays (visibilities, sigma,\n", + "uv-wavelengths) use the same ordering so they stay aligned.\n", + "\n", + " # Per-SPW (pol-averaged) shapes:\n", + " # vis_avg : (n_chan_spw, n_vis_spw, 2)\n", + " # sig_avg : (n_chan_spw, n_vis_spw, 2)\n", + " # uv : (n_chan_spw, n_vis_spw, 2)\n", + "\n", + " vis_all = np.concatenate(\n", + " [v.reshape(-1, 2) for v in per_spw_vis], axis=0\n", + " )\n", + " sig_all = np.concatenate(\n", + " [s.reshape(-1, 2) for s in per_spw_sig], axis=0\n", + " )\n", + " uv_all = np.concatenate(\n", + " [u.reshape(-1, 2) for u in per_spw_uv], axis=0\n", + " )\n", + "\n", + "After this step all three arrays have shape `(n_vis_total, 2)` \u2014 exactly\n", + "what `al.Interferometer.from_fits` expects.\n", + "\n", + "__Building the Interferometer Object__\n", + "\n", + "Write the three concatenated arrays to FITS (using `astropy.io.fits.writeto`\n", + "or `autoconf.fitsable.output_to_fits`) and load them via the canonical\n", + "workspace pattern used in `start_here.py`:\n", + "\n", + " import autolens as al\n", + "\n", + " # Small, fast real-space mask \u2014 keep shape_native and radius modest while\n", + " # you are iterating on the reduction. Increase later for production fits.\n", + " real_space_mask = al.Mask2D.circular(\n", + " shape_native=(64, 64),\n", + " pixel_scales=0.05,\n", + " radius=1.5,\n", + " )\n", + "\n", + " dataset = al.Interferometer.from_fits(\n", + " data_path=\"data.fits\",\n", + " noise_map_path=\"noise_map.fits\",\n", + " uv_wavelengths_path=\"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT, # NUFFT scales to large n_vis\n", + " )\n", + "\n", + "The 64x64 / 0.05\" mask above keeps the Fourier transform and any dirty-image\n", + "sanity plots fast while you iterate \u2014 a first-pass check of your reduction\n", + "should finish in seconds, not minutes. Bump `shape_native` up (e.g. 256x256)\n", + "and drop `pixel_scales` (e.g. 0.02\") for the real modeling run. Choosing\n", + "these numbers sensibly for your instrument is covered in\n", + "`scripts/interferometer/data_preparation.py`.\n", + "\n", + "__Troubleshooting__\n", + "\n", + "- **`uv_wavelengths` values look way too big/small** \u2014 you forgot the\n", + " metres \u2192 wavelengths conversion, or you used the wrong frequency column.\n", + "- **Dirty image is offset / upside-down** \u2014 the `(u, v)` sign convention or\n", + " the `(y, x)` axis order may disagree with your real-space mask. Plot the\n", + " dirty image with `aplt.subplot_interferometer_dirty_images` as a sanity\n", + " check.\n", + "- **Sigmas produce a flat chi-squared map** \u2014 you probably forgot `statwt`.\n", + "- **n_vis is enormous** \u2014 ensure you are using `TransformerNUFFT` (the default\n", + " recommendation, backed by JAX-native `nufftax`). It scales to many millions\n", + " of visibilities. Channel averaging via a larger `width` in `split` is still\n", + " useful when you want to reduce dataset size for other reasons.\n", + "- **Only one polarisation present** \u2014 if CASA has flagged one pol the\n", + " averaging code above will produce NaNs. Check with `get_visibilities`\n", + " and branch accordingly.\n", + "\n", + "__SLACK__\n", + "\n", + "These scripts are a starting point, not a polished pipeline. If your\n", + "observation has mosaicking, heterogeneous SPWs, unusual correlators, or you\n", + "just can't get the numbers to line up \u2014 ping us on the PyAutoLens Slack.\n", + "Direct user feedback is actively shaping this workflow and we're happy to\n", + "help debug real datasets.\n", + "\n", + "__Running as a Script__\n", + "\n", + "The block below is illustrative \u2014 edit the paths/widths for your own reduction\n", + "and run inside CASA with `casa -c scripts/interferometer/casa_reduction.py`.\n", + "It assumes you have already run `split`, `uvcontsub` (optional) and `statwt`\n", + "by hand on the `.ms`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "if __name__ == \"__main__\":\n", + "\n", + " uid = \"my_obs_field_SPT-0418\"\n", + " width = 128\n", + "\n", + " ms_chanavg = f\"{uid}_spw0_chanavg.ms\"\n", + " ms_statwt = f\"{ms_chanavg}.statwt\"\n", + "\n", + " export_uv_wavelengths(ms=ms_chanavg, filename=f\"uv_wavelengths_{uid}_width_{width}\")\n", + " export_visibilities(ms=ms_chanavg, filename=f\"visibilities_{uid}_width_{width}\")\n", + " export_sigma(ms=ms_statwt, filename=f\"sigma_{uid}_width_{width}_statwt\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/data_preparation.ipynb b/notebooks/interferometer/data_preparation.ipynb index eac433692..7585c4dae 100644 --- a/notebooks/interferometer/data_preparation.ipynb +++ b/notebooks/interferometer/data_preparation.ipynb @@ -1,378 +1,415 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Interferometer: Data Preparation\n", - "================================\n", - "\n", - "When an interferometer dataset is analysed, it must conform to certain standards in order for\n", - "the analysis to be performed correctly. This tutorial describes these standards and links to more detailed scripts\n", - "which will help you prepare your dataset to adhere to them if it does not already.\n", - "\n", - "__Contents__\n", - "\n", - "- **SLACK:** Contact information for help with interferometer data preparation.\n", - "- **Pixel Scale:** Choosing the correct pixel scale for interferometer datasets.\n", - "- **Visibilities:** Loading and inspecting visibility data from FITS files.\n", - "- **Noise-Map:** Loading and inspecting the noise map for the interferometer dataset.\n", - "- **UV Wavelengths:** Loading and inspecting the uv-wavelength baselines.\n", - "- **Real Space Mask:** Setting up the real-space mask for Fourier transform evaluation.\n", - "- **Data Processing Complete:** Summary of required data standards and overview of optional steps.\n", - "- **Positions (Optional):** Marking multiply-imaged source positions to constrain mass models.\n", - "- **Lens Light Centre (Optional):** Marking the lens galaxy light centre to fix or constrain mass model parameters.\n", - "- **Extra Galaxies (Optional):** Marking centres of nearby extra galaxies for inclusion in the model.\n", - "- **Mask Extra Galaxies (Optional):** Creating masks to remove signal from nearby extra galaxies.\n", - "- **Info (Optional):** Storing auxiliary information like redshifts as a JSON file.\n", - "\n", - "__SLACK__\n", - "\n", - "The interferometer data preparation scripts are currently being developed and are not yet complete. If you are\n", - "unsure of how to prepare your dataset, please message us on Slack and we will help you directly!\n", - "\n", - "__Pixel Scale__\n", - "\n", - "When fitting an interferometer dataset, the images of the lens and source galaxies are first evaluated in real-space\n", - "using a grid of pixels, which is then Fourier transformed to the uv-plane.\n", - "\n", - "The \"pixel_scale\" of an interferometer dataset is this pixel-units to arcsecond-units conversion factor. The value\n", - "depends on the instrument used to observe the lens, the wavelength of the light used to observe it and size of\n", - "the baselines used (e.g. longer baselines means higher resolution and therefore a smaller pixel scale).\n", - "\n", - "The pixel scale of some common interferometers is as follows:\n", - "\n", - " - ALMA: 0.02\" - 0.1\" / pixel\n", - " - SMA: 0.05\" - 0.1\" / pixel\n", - " - JVLA: 0.005\" - 0.01\" / pixel\n", - "\n", - "It is absolutely vital you use a sufficently small pixel scale that all structure in the data is resolved after the\n", - "Fourier transform. If the pixel scale is too large, the Fourier transform will smear out the data and the lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / \"simple\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visibilities__\n", - "\n", - "The image is the image of your strong lens, which comes from a telescope like the Hubble Space telescope (HST).\n", - "\n", - "Lets inspect an image which conforms to **PyAutoLens** standards:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "visibilities = al.Visibilities.from_fits(file_path=dataset_path / \"data.fits\", hdu=0)\n", - "\n", - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "These visibilities conforms to **PyAutoLens** standards, because they come from a standard CASA data reduction\n", - "procedure.\n", - "\n", - "More details of this procedure are given in the `examples/casa_reduction.ipynb` notebook.\n", - "\n", - "__Noise-Map__\n", - "\n", - "The noise-map is the real and complex noise in each visiblity of the interferometer dataset. It is used to weight\n", - "the visibilities when a lens model is fitted to the data via a chi-squared statistic.\n", - "\n", - "It is common for all visibilities to have the same noise value, depending on the instrument used to observe the\n", - "the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "visibilities = al.VisibilitiesNoiseMap.from_fits(\n", - " file_path=dataset_path / \"noise_map.fits\", hdu=0\n", - ")\n", - "\n", - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__UV Wavelengths__\n", - "\n", - "The uv-wavelengths define the baselines of the interferometer. They are used to Fourier transform the image to the\n", - "uv-plane, which is where the lens model is evaluated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "uv_wavelengths = al.ndarray_via_fits_from(\n", - " file_path=dataset_path / \"uv_wavelengths.fits\", hdu=0\n", - ")\n", - "\n", - "uv_wavelengths = al.Grid2DIrregular.from_yx_1d(\n", - " y=uv_wavelengths[:, 1] / 10**3.0,\n", - " x=uv_wavelengths[:, 0] / 10**3.0,\n", - ")\n", - "\n", - "aplt.plot_grid(grid=uv_wavelengths, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "These uv wavelengths conform to **PyAutoLens** standards, because they come from a standard CASA data reduction\n", - "procedure.\n", - "\n", - "More details of this procedure are given in the `examples/casa_reduction.ipynb` notebook.\n", - "\n", - "__Real Space Mask__\n", - "\n", - "The `modeling` scripts also define a real-space mask, which defines where the image is evalated in real space\n", - "before it is Fourier transformed.\n", - "\n", - "You must double check that the real-space mask you use:\n", - "\n", - " - Spatially covers the lensed source galaxy, such that the source is not truncated by the mask.\n", - " - Is high enough resolution that the lensed source galaxy is not smeared via the Fourier transform.\n", - "\n", - "__Data Processing Complete__\n", - "\n", - "If your visibilities, noise-map, uv_wavelengths and real space mask conform the standards above, you are ready to analyse\n", - "your dataset!\n", - "\n", - "Below, we provide an overview of optional data preparation steps which prepare other aspects of the analysis.\n", - "\n", - "New users are recommended to skim-read the optional steps below so they are aware of them, but to not perform them\n", - "and instead analyse their dataset now. You can come back to the data preparation scripts below if it becomes necessary.\n", - "\n", - "The following scripts are used to prepare components of an interferometer dataset, however they are used in an\n", - "identical fashion for dataset datasets.\n", - "\n", - "Therefore, they are not located in the `interferometer/data_preparation` package, but instead in the\n", - "`imaging/data_preparation` package, so refer there for a description of their usage.\n", - "\n", - "Note that in order to perform some tasks (e.g. mark on the image where the source is), you will need to use an image\n", - "of the interferometer data even though visibilities are used for the analysis.\n", - "\n", - "__Positions (Optional)__\n", - "\n", - "The script allows you to mark the (y,x) arc-second positions of the multiply imaged lensed source galaxy in\n", - "the image-plane, under the assumption that they originate from the same location in the source-plane.\n", - "\n", - "A non-linear search (e.g. Nautilus) can then use these positions to preferentially choose mass models where these\n", - "positions trace close to one another in the source-plane. This speeding up the initial fitting of lens models and\n", - "removes unwanted solutions from parameter space which have too much or too little mass in the lens galaxy.\n", - "\n", - "If you create positions for your dataset, you must also update your modeling script to use them by loading them\n", - "and passing them to the `Analysis` object via a `PositionsLH` object.\n", - "\n", - "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using positions\n", - "to help the non-linear search find a good lens model.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "Position-based lens model resampling is particularly important for fitting pixelized source models, for the\n", - "reasons disucssed in the following readthedocs\n", - "webapge https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", - "\n", - "- `data_preparation/examples/optional/positions.ipynb`: input the positions manually into a Python script.\n", - "\n", - "- `data_preparation/gui/positions.ipynb` use a Graphical User Interface (GUI) to mark the positions.\n", - "\n", - "- `guides/modeling/customize` for an example.of how to use positions in a `modeling` script.\n", - "\n", - "\n", - "__Lens Light Centre (Optional)__\n", - "\n", - "This script allows you to mark the (y,x) arcsecond locations of the lens galaxy light centre(s) of the strong lens\n", - "you are analysing. These can be used as fixed values for the lens light and mass models in a model-fit.\n", - "\n", - "This reduces the number of free parameters fitted for in a lens model and removes inaccurate solutions where\n", - "the lens mass model centre is unrealistically far from its true centre.\n", - "\n", - "Advanced `chaining` scripts often use these input centres in the early fits to infer an accurate initial lens model,\n", - "amd then make the centres free parameters in later searches to ensure a general and accurate lens model is inferred.\n", - "\n", - "If you create a `light_centre` for your dataset, you must also update your modeling script to use them.\n", - "\n", - "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using a fixed\n", - "lens light and / or mass centre to help the non-linear search find a good lens model.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "The example `data_preparation/examples/optional/lens_light_centre.py` shows how to input the lens galaxy light centre\n", - "manually into a Python script.\n", - "\n", - "The script `data_preparation/gui/lens_light_centre.ipynb` shows how to use a Graphical User Interface (GUI) to mask the\n", - "lens galaxy light centres.\n", - "\n", - "\n", - "__Extra Galaxies (Optional)__\n", - "\n", - "There may be galaxies nearby the lens and source galaxies, whose emission blends with that of the lens and source\n", - "and whose mass may contribute to the ray-tracing and lens model.\n", - "\n", - "We can include these galaxies in the lens model, either as light profiles, mass profiles, or both, using the\n", - "modeling API, where these nearby objects are denoted `extra_galaxies`.\n", - "\n", - "The script `extra_galaxies_centres.py` marks the (y,x) arcsecond locations of these extra galaxies, so that when they\n", - "are included in the lens model the centre of these extra galaxies light and / or mass profiles are fixed to these\n", - "values (or their priors are initialized surrounding these centres).\n", - "\n", - "The example `mask_extra_galaxies.py` (see below) masks the regions of an image where extra galaxies are present.\n", - "This mask is used to remove their signal from the data and increase their noise to make them not impact the fit.\n", - "This means their luminous emission does not need to be included in the model, reducing the number of free parameters\n", - "and speeding up the analysis. It is still a choice whether their mass is included in the model.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/extra_galaxies_centres.py`: input the extra galaxy centres manually into a\n", - " Python script.\n", - "\n", - "- `data_preparation/gui/extra_galaxies_centres.ipynb`: use a Graphical User Interface (GUI) to mark the extra galaxy centres.\n", - "\n", - "- `features/extra_galaxies.py` how to use extra galaxies in a model-fit, including loading the extra galaxy centres.\n", - "\n", - "\n", - "__Mask Extra Galaxies (Optional)__\n", - "\n", - "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated\n", - "with the strong lens we are studying. The emission from these images will impact our model fitting and needs to be\n", - "removed from the analysis.\n", - "\n", - "This script creates a mask of these regions of the image, called the `mask_extra_galaxies`, which can be used to\n", - "prevent them from impacting a fit. This mask may also include emission from objects which are not technically galaxies,\n", - "but blend with the galaxy we are studying in a similar way. Common examples of such objects are foreground stars\n", - "or emission due to the data reduction process.\n", - "\n", - "The mask can be applied in different ways. For example, it could be applied such that the image pixels are discarded\n", - "from the fit entirely, Alternatively the mask could be used to set the image values to (near) zero and increase their\n", - "corresponding noise-map to large values.\n", - "\n", - "The exact method used depends on the nature of the model being fitted. For simple fits like a light profile a mask\n", - "is appropriate, as removing image pixels does not change how the model is fitted. However, for more complex models\n", - "fits, like those using a pixelization, masking regions of the image in a way that removes their image pixels entirely\n", - "from the fit can produce discontinuities in the pixelixation. In this case, scaling the data and noise-map values\n", - "may be a better approach.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/mask_extra_galaxies.py`: create the extra galaxies mask manually via a Python script.\n", - "\n", - "- `data_preparation/gui/extra_galaxies_mask.ipynb` use a Graphical User Interface (GUI) to create the extra galaxies mask.\n", - "\n", - "- `features/extra_galaxies.py` how to use the extra galaxies mask in a model-fit.\n", - "\n", - "__Info (Optional)__\n", - "\n", - "Auxiliary information about a strong lens dataset may used during an analysis or afterwards when interpreting the\n", - " modeling results. For example, the redshifts of the source and lens galaxy.\n", - "\n", - "By storing these as an `info.json` file in the lens's dataset folder, it is straight forward to load the redshifts\n", - "in a modeling script and pass them to a fit, such that **PyAutoLens** can then output results in physical\n", - "units (e.g. kpc instead of arc-seconds).\n", - "\n", - "For analysing large quantities of modeling results, **PyAutoLens** has an sqlite database feature. The info file\n", - "may can also be loaded by the database after a model-fit has completed, such that when one is interpreting\n", - "the results of a model fit additional data on a lens can be used to.\n", - "\n", - "For example, to plot the model-results against other measurements of a lens not made by PyAutoLens. Examples of such\n", - "data might be:\n", - "\n", - "- The velocity dispersion of the lens galaxy.\n", - "- The stellar mass of the lens galaxy.\n", - "- The results of previous strong lens models to the lens performed in previous papers.\n", - "\n", - "**Links / Resources:**\n", - "\n", - "- `data_preparation/examples/optional/info.py`: create the info file manually via a Python script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Interferometer: Data Preparation\n", + "================================\n", + "\n", + "When an interferometer dataset is analysed, it must conform to certain standards in order for\n", + "the analysis to be performed correctly. This tutorial describes these standards and links to more detailed scripts\n", + "which will help you prepare your dataset to adhere to them if it does not already.\n", + "\n", + "__Contents__\n", + "\n", + "- **SLACK:** Contact information for help with interferometer data preparation.\n", + "- **Pixel Scale:** Choosing the correct pixel scale for interferometer datasets.\n", + "- **Visibilities:** Loading and inspecting visibility data from FITS files.\n", + "- **Noise-Map:** Loading and inspecting the noise map for the interferometer dataset.\n", + "- **UV Wavelengths:** Loading and inspecting the uv-wavelength baselines.\n", + "- **Real Space Mask:** Setting up the real-space mask for Fourier transform evaluation.\n", + "- **Data Processing Complete:** Summary of required data standards and overview of optional steps.\n", + "- **Positions (Optional):** Marking multiply-imaged source positions to constrain mass models.\n", + "- **Lens Light Centre (Optional):** Marking the lens galaxy light centre to fix or constrain mass model parameters.\n", + "- **Extra Galaxies (Optional):** Marking centres of nearby extra galaxies for inclusion in the model.\n", + "- **Mask Extra Galaxies (Optional):** Creating masks to remove signal from nearby extra galaxies.\n", + "- **Info (Optional):** Storing auxiliary information like redshifts as a JSON file.\n", + "\n", + "__SLACK__\n", + "\n", + "The interferometer data preparation scripts are currently being developed and are not yet complete. If you are\n", + "unsure of how to prepare your dataset, please message us on Slack and we will help you directly!\n", + "\n", + "__Pixel Scale__\n", + "\n", + "When fitting an interferometer dataset, the images of the lens and source galaxies are first evaluated in real-space\n", + "using a grid of pixels, which is then Fourier transformed to the uv-plane.\n", + "\n", + "The \"pixel_scale\" of an interferometer dataset is this pixel-units to arcsecond-units conversion factor. The value\n", + "depends on the instrument used to observe the lens, the wavelength of the light used to observe it and size of\n", + "the baselines used (e.g. longer baselines means higher resolution and therefore a smaller pixel scale).\n", + "\n", + "The pixel scale of some common interferometers is as follows:\n", + "\n", + " - ALMA: 0.02\" - 0.1\" / pixel\n", + " - SMA: 0.05\" - 0.1\" / pixel\n", + " - JVLA: 0.005\" - 0.01\" / pixel\n", + "\n", + "It is absolutely vital you use a sufficently small pixel scale that all structure in the data is resolved after the\n", + "Fourier transform. If the pixel scale is too large, the Fourier transform will smear out the data and the lens model." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / \"simple\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visibilities__\n", + "\n", + "The image is the image of your strong lens, which comes from a telescope like the Hubble Space telescope (HST).\n", + "\n", + "Lets inspect an image which conforms to **PyAutoLens** standards:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "visibilities = al.Visibilities.from_fits(file_path=dataset_path / \"data.fits\", hdu=0)\n", + "\n", + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "These visibilities conforms to **PyAutoLens** standards, because they come from a standard CASA data reduction\n", + "procedure.\n", + "\n", + "More details of this procedure are given in the `examples/casa_reduction.ipynb` notebook.\n", + "\n", + "__Noise-Map__\n", + "\n", + "The noise-map is the real and complex noise in each visiblity of the interferometer dataset. It is used to weight\n", + "the visibilities when a lens model is fitted to the data via a chi-squared statistic.\n", + "\n", + "It is common for all visibilities to have the same noise value, depending on the instrument used to observe the\n", + "the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "visibilities = al.VisibilitiesNoiseMap.from_fits(\n", + " file_path=dataset_path / \"noise_map.fits\", hdu=0\n", + ")\n", + "\n", + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__UV Wavelengths__\n", + "\n", + "The uv-wavelengths define the baselines of the interferometer. They are used to Fourier transform the image to the\n", + "uv-plane, which is where the lens model is evaluated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "uv_wavelengths = al.ndarray_via_fits_from(\n", + " file_path=dataset_path / \"uv_wavelengths.fits\", hdu=0\n", + ")\n", + "\n", + "uv_wavelengths = al.Grid2DIrregular.from_yx_1d(\n", + " y=uv_wavelengths[:, 1] / 10**3.0,\n", + " x=uv_wavelengths[:, 0] / 10**3.0,\n", + ")\n", + "\n", + "aplt.plot_grid(grid=uv_wavelengths, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "These uv wavelengths conform to **PyAutoLens** standards, because they come from a standard CASA data reduction\n", + "procedure.\n", + "\n", + "More details of this procedure are given in the `examples/casa_reduction.ipynb` notebook.\n", + "\n", + "__Real Space Mask__\n", + "\n", + "The `modeling` scripts also define a real-space mask, which defines where the image is evalated in real space\n", + "before it is Fourier transformed.\n", + "\n", + "You must double check that the real-space mask you use:\n", + "\n", + " - Spatially covers the lensed source galaxy, such that the source is not truncated by the mask.\n", + " - Is high enough resolution that the lensed source galaxy is not smeared via the Fourier transform.\n", + "\n", + "__Data Processing Complete__\n", + "\n", + "If your visibilities, noise-map, uv_wavelengths and real space mask conform the standards above, you are ready to analyse\n", + "your dataset!\n", + "\n", + "Below, we provide an overview of optional data preparation steps which prepare other aspects of the analysis.\n", + "\n", + "New users are recommended to skim-read the optional steps below so they are aware of them, but to not perform them\n", + "and instead analyse their dataset now. You can come back to the data preparation scripts below if it becomes necessary.\n", + "\n", + "The following scripts are used to prepare components of an interferometer dataset, however they are used in an\n", + "identical fashion for dataset datasets.\n", + "\n", + "Therefore, they are not located in the `interferometer/data_preparation` package, but instead in the\n", + "`imaging/data_preparation` package, so refer there for a description of their usage.\n", + "\n", + "Note that in order to perform some tasks (e.g. mark on the image where the source is), you will need to use an image\n", + "of the interferometer data even though visibilities are used for the analysis.\n", + "\n", + "__Positions (Optional)__\n", + "\n", + "The script allows you to mark the (y,x) arc-second positions of the multiply imaged lensed source galaxy in\n", + "the image-plane, under the assumption that they originate from the same location in the source-plane.\n", + "\n", + "A non-linear search (e.g. Nautilus) can then use these positions to preferentially choose mass models where these\n", + "positions trace close to one another in the source-plane. This speeding up the initial fitting of lens models and\n", + "removes unwanted solutions from parameter space which have too much or too little mass in the lens galaxy.\n", + "\n", + "If you create positions for your dataset, you must also update your modeling script to use them by loading them\n", + "and passing them to the `Analysis` object via a `PositionsLH` object.\n", + "\n", + "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using positions\n", + "to help the non-linear search find a good lens model.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "Position-based lens model resampling is particularly important for fitting pixelized source models, for the\n", + "reasons disucssed in the following readthedocs\n", + "webapge https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", + "\n", + "- `data_preparation/examples/optional/positions.ipynb`: input the positions manually into a Python script.\n", + "\n", + "- `data_preparation/gui/positions.ipynb` use a Graphical User Interface (GUI) to mark the positions.\n", + "\n", + "- `guides/modeling/customize` for an example.of how to use positions in a `modeling` script.\n", + "\n", + "\n", + "__Lens Light Centre (Optional)__\n", + "\n", + "This script allows you to mark the (y,x) arcsecond locations of the lens galaxy light centre(s) of the strong lens\n", + "you are analysing. These can be used as fixed values for the lens light and mass models in a model-fit.\n", + "\n", + "This reduces the number of free parameters fitted for in a lens model and removes inaccurate solutions where\n", + "the lens mass model centre is unrealistically far from its true centre.\n", + "\n", + "Advanced `chaining` scripts often use these input centres in the early fits to infer an accurate initial lens model,\n", + "amd then make the centres free parameters in later searches to ensure a general and accurate lens model is inferred.\n", + "\n", + "If you create a `light_centre` for your dataset, you must also update your modeling script to use them.\n", + "\n", + "If your **PyAutoLens** analysis is struggling to converge to a good lens model, you should consider using a fixed\n", + "lens light and / or mass centre to help the non-linear search find a good lens model.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "The example `data_preparation/examples/optional/lens_light_centre.py` shows how to input the lens galaxy light centre\n", + "manually into a Python script.\n", + "\n", + "The script `data_preparation/gui/lens_light_centre.ipynb` shows how to use a Graphical User Interface (GUI) to mask the\n", + "lens galaxy light centres.\n", + "\n", + "\n", + "__Extra Galaxies (Optional)__\n", + "\n", + "There may be galaxies nearby the lens and source galaxies, whose emission blends with that of the lens and source\n", + "and whose mass may contribute to the ray-tracing and lens model.\n", + "\n", + "We can include these galaxies in the lens model, either as light profiles, mass profiles, or both, using the\n", + "modeling API, where these nearby objects are denoted `extra_galaxies`.\n", + "\n", + "The script `extra_galaxies_centres.py` marks the (y,x) arcsecond locations of these extra galaxies, so that when they\n", + "are included in the lens model the centre of these extra galaxies light and / or mass profiles are fixed to these\n", + "values (or their priors are initialized surrounding these centres).\n", + "\n", + "The example `mask_extra_galaxies.py` (see below) masks the regions of an image where extra galaxies are present.\n", + "This mask is used to remove their signal from the data and increase their noise to make them not impact the fit.\n", + "This means their luminous emission does not need to be included in the model, reducing the number of free parameters\n", + "and speeding up the analysis. It is still a choice whether their mass is included in the model.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/extra_galaxies_centres.py`: input the extra galaxy centres manually into a\n", + " Python script.\n", + "\n", + "- `data_preparation/gui/extra_galaxies_centres.ipynb`: use a Graphical User Interface (GUI) to mark the extra galaxy centres.\n", + "\n", + "- `features/extra_galaxies.py` how to use extra galaxies in a model-fit, including loading the extra galaxy centres.\n", + "\n", + "\n", + "__Mask Extra Galaxies (Optional)__\n", + "\n", + "There may be regions of an image that have signal near the lens and source that is from other galaxies not associated\n", + "with the strong lens we are studying. The emission from these images will impact our model fitting and needs to be\n", + "removed from the analysis.\n", + "\n", + "This script creates a mask of these regions of the image, called the `mask_extra_galaxies`, which can be used to\n", + "prevent them from impacting a fit. This mask may also include emission from objects which are not technically galaxies,\n", + "but blend with the galaxy we are studying in a similar way. Common examples of such objects are foreground stars\n", + "or emission due to the data reduction process.\n", + "\n", + "The mask can be applied in different ways. For example, it could be applied such that the image pixels are discarded\n", + "from the fit entirely, Alternatively the mask could be used to set the image values to (near) zero and increase their\n", + "corresponding noise-map to large values.\n", + "\n", + "The exact method used depends on the nature of the model being fitted. For simple fits like a light profile a mask\n", + "is appropriate, as removing image pixels does not change how the model is fitted. However, for more complex models\n", + "fits, like those using a pixelization, masking regions of the image in a way that removes their image pixels entirely\n", + "from the fit can produce discontinuities in the pixelixation. In this case, scaling the data and noise-map values\n", + "may be a better approach.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/mask_extra_galaxies.py`: create the extra galaxies mask manually via a Python script.\n", + "\n", + "- `data_preparation/gui/extra_galaxies_mask.ipynb` use a Graphical User Interface (GUI) to create the extra galaxies mask.\n", + "\n", + "- `features/extra_galaxies.py` how to use the extra galaxies mask in a model-fit.\n", + "\n", + "__Info (Optional)__\n", + "\n", + "Auxiliary information about a strong lens dataset may used during an analysis or afterwards when interpreting the\n", + " modeling results. For example, the redshifts of the source and lens galaxy.\n", + "\n", + "By storing these as an `info.json` file in the lens's dataset folder, it is straight forward to load the redshifts\n", + "in a modeling script and pass them to a fit, such that **PyAutoLens** can then output results in physical\n", + "units (e.g. kpc instead of arc-seconds).\n", + "\n", + "For analysing large quantities of modeling results, **PyAutoLens** has an sqlite database feature. The info file\n", + "may can also be loaded by the database after a model-fit has completed, such that when one is interpreting\n", + "the results of a model fit additional data on a lens can be used to.\n", + "\n", + "For example, to plot the model-results against other measurements of a lens not made by PyAutoLens. Examples of such\n", + "data might be:\n", + "\n", + "- The velocity dispersion of the lens galaxy.\n", + "- The stellar mass of the lens galaxy.\n", + "- The results of previous strong lens models to the lens performed in previous papers.\n", + "\n", + "**Links / Resources:**\n", + "\n", + "- `data_preparation/examples/optional/info.py`: create the info file manually via a Python script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/advanced/shapelets/fit.ipynb b/notebooks/interferometer/features/advanced/shapelets/fit.ipynb index 3e50c0430..280d148d1 100644 --- a/notebooks/interferometer/features/advanced/shapelets/fit.ipynb +++ b/notebooks/interferometer/features/advanced/shapelets/fit.ipynb @@ -1,384 +1,421 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Shapelets Fit (Interferometer)\n", - "==================================================\n", - "\n", - "A shapelet is a basis function appropriate for capturing the exponential / disk-like features of a galaxy.\n", - "The `intensity` of every shapelet is solved for via linear algebra (\"inversion\") rather than being a free\n", - "parameter of the non-linear search.\n", - "\n", - "This script illustrates how to perform a single `FitInterferometer` of a shapelet source model \u2014 that is,\n", - "not the full Nautilus model-fit, but a single likelihood evaluation given known basis parameters. This is\n", - "useful for understanding how the inversion produces the per-shapelet solved-for `intensity` values, and\n", - "how to extract them from the resulting fit.\n", - "\n", - "For an explanation of why shapelet fits to visibility data are now practical thanks to the JAX-native\n", - "NUFFT `nufftax` (https://github.com/GragasLab/nufftax), and why shapelets require the positive-negative\n", - "solver, see the companion `modeling.py` example.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Benefits and drawbacks of a shapelet source for interferometer data.\n", - "- **Positive Negative Solver:** Why shapelets require the positive-negative solver (unlike MGE or linear\n", - " Sersic).\n", - "- **Model:** The lens model whose `intensity` values we solve for via inversion.\n", - "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", - "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", - "- **Basis:** Build the linear polar shapelet basis used as the source bulge.\n", - "- **Fit:** Perform a single `FitInterferometer` and inspect the inversion.\n", - "- **Intensities:** Extract the per-shapelet solved-for `intensity` values.\n", - "- **Visualization:** Build the helper tracer where linear light profiles are replaced with ordinary light\n", - " profiles carrying their solved-for `intensity`, then plot.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data). Interferometer\n", - " convention.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's bulge is a superposition of linear `ShapeletPolar` profiles with shared centre,\n", - " ell_comps, and beta.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, using `TransformerNUFFT`\n", - "backed by `nufftax`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Basis__\n", - "\n", - "We build a `Basis` of linear polar shapelets for the source bulge. All shapelets share centre, ell_comps,\n", - "and a single `beta` size scale; the (n, m) quantum numbers are assigned procedurally.\n", - "\n", - "We use `lp_linear.ShapeletPolar`, which solves for each shapelet's `intensity` analytically via the\n", - "inversion. Linear light profiles are described in detail in the `linear_light_profiles.py` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_n = 10\n", - "total_m = sum(range(2, total_n + 1)) + 1\n", - "\n", - "shapelets_bulge_list = []\n", - "n_count = 1\n", - "m_count = -1\n", - "\n", - "for i in range(total_n + total_m + 1):\n", - " if i == 0:\n", - " n, m = 0, 0\n", - " else:\n", - " n, m = n_count, m_count\n", - " m_count += 2\n", - " if m_count > n_count:\n", - " n_count += 1\n", - " m_count = -n_count\n", - "\n", - " shapelet = al.lp_linear.ShapeletPolar(\n", - " n=n,\n", - " m=m,\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " beta=0.1,\n", - " )\n", - " shapelets_bulge_list.append(shapelet)\n", - "\n", - "source_bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We now illustrate the API for performing a single shapelet fit using standard `Galaxy`, `Tracer` and\n", - "`FitInterferometer` objects. Once we have a `Basis`, we can treat it like any other light profile.\n", - "\n", - "Note `Settings(use_positive_only_solver=False)` is passed to the fit \u2014 shapelets require the\n", - "positive-negative solver to function." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(redshift=1.0, bulge=source_bulge)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitInterferometer(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_positive_only_solver=False),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit's `subplot_fit_interferometer` shows the visibility-plane fit and dirty-image residuals. Because\n", - "the source bulge is a Basis of linear light profiles, the inversion has solved for each shapelet's\n", - "`intensity` to maximize the fit to the observed visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `subplot_fit_dirty_images` provides a real-space view of the data, model and residuals via\n", - "inverse-NUFFT of the visibility-plane quantities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_dirty_images(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The fit contains the solved-for `intensity` value of every shapelet in the basis.\n", - "\n", - "These are computed via `fit.linear_light_profile_intensity_dict`, which maps each linear light profile in\n", - "the model to its inferred `intensity`. Print the first few entries for brevity." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.linear_light_profile_intensity_dict)\n", - "\n", - "for shapelet in shapelets_bulge_list[:5]:\n", - " intensity = fit.linear_light_profile_intensity_dict[shapelet]\n", - " print(f\" n={shapelet.n} m={shapelet.m} intensity = {intensity:+.6e}\")\n", - "\n", - "print(\n", - " f\"\\n number of negative-intensity shapelets: \"\n", - " f\"{sum(1 for s in shapelets_bulge_list if fit.linear_light_profile_intensity_dict[s] < 0)} \"\n", - " f\"/ {len(shapelets_bulge_list)}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the\n", - "solved-for `intensity` values is also accessible from a fit.\n", - "\n", - "The benefit of this helper-tracer is that it can be visualised (linear light profiles cannot be plotted\n", - "by default because they do not have `intensity` values)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization__\n", - "\n", - "The helper-tracer created above replaces every linear `ShapeletPolar` with an ordinary `ShapeletPolar`\n", - "carrying its solved-for `intensity` \u2014 that tracer can be plotted directly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=dataset.grid),\n", - " title=\"Tracer Image (shapelet source)\",\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Shapelets Fit (Interferometer)\n", + "==================================================\n", + "\n", + "A shapelet is a basis function appropriate for capturing the exponential / disk-like features of a galaxy.\n", + "The `intensity` of every shapelet is solved for via linear algebra (\"inversion\") rather than being a free\n", + "parameter of the non-linear search.\n", + "\n", + "This script illustrates how to perform a single `FitInterferometer` of a shapelet source model \u2014 that is,\n", + "not the full Nautilus model-fit, but a single likelihood evaluation given known basis parameters. This is\n", + "useful for understanding how the inversion produces the per-shapelet solved-for `intensity` values, and\n", + "how to extract them from the resulting fit.\n", + "\n", + "For an explanation of why shapelet fits to visibility data are now practical thanks to the JAX-native\n", + "NUFFT `nufftax` (https://github.com/GragasLab/nufftax), and why shapelets require the positive-negative\n", + "solver, see the companion `modeling.py` example.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Benefits and drawbacks of a shapelet source for interferometer data.\n", + "- **Positive Negative Solver:** Why shapelets require the positive-negative solver (unlike MGE or linear\n", + " Sersic).\n", + "- **Model:** The lens model whose `intensity` values we solve for via inversion.\n", + "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", + "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", + "- **Basis:** Build the linear polar shapelet basis used as the source bulge.\n", + "- **Fit:** Perform a single `FitInterferometer` and inspect the inversion.\n", + "- **Intensities:** Extract the per-shapelet solved-for `intensity` values.\n", + "- **Visualization:** Build the helper tracer where linear light profiles are replaced with ordinary light\n", + " profiles carrying their solved-for `intensity`, then plot.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data). Interferometer\n", + " convention.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's bulge is a superposition of linear `ShapeletPolar` profiles with shared centre,\n", + " ell_comps, and beta.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, using `TransformerNUFFT`\n", + "backed by `nufftax`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Basis__\n", + "\n", + "We build a `Basis` of linear polar shapelets for the source bulge. All shapelets share centre, ell_comps,\n", + "and a single `beta` size scale; the (n, m) quantum numbers are assigned procedurally.\n", + "\n", + "We use `lp_linear.ShapeletPolar`, which solves for each shapelet's `intensity` analytically via the\n", + "inversion. Linear light profiles are described in detail in the `linear_light_profiles.py` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_n = 10\n", + "total_m = sum(range(2, total_n + 1)) + 1\n", + "\n", + "shapelets_bulge_list = []\n", + "n_count = 1\n", + "m_count = -1\n", + "\n", + "for i in range(total_n + total_m + 1):\n", + " if i == 0:\n", + " n, m = 0, 0\n", + " else:\n", + " n, m = n_count, m_count\n", + " m_count += 2\n", + " if m_count > n_count:\n", + " n_count += 1\n", + " m_count = -n_count\n", + "\n", + " shapelet = al.lp_linear.ShapeletPolar(\n", + " n=n,\n", + " m=m,\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " beta=0.1,\n", + " )\n", + " shapelets_bulge_list.append(shapelet)\n", + "\n", + "source_bulge = al.lp_basis.Basis(profile_list=shapelets_bulge_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We now illustrate the API for performing a single shapelet fit using standard `Galaxy`, `Tracer` and\n", + "`FitInterferometer` objects. Once we have a `Basis`, we can treat it like any other light profile.\n", + "\n", + "Note `Settings(use_positive_only_solver=False)` is passed to the fit \u2014 shapelets require the\n", + "positive-negative solver to function." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(redshift=1.0, bulge=source_bulge)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitInterferometer(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_positive_only_solver=False),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit's `subplot_fit_interferometer` shows the visibility-plane fit and dirty-image residuals. Because\n", + "the source bulge is a Basis of linear light profiles, the inversion has solved for each shapelet's\n", + "`intensity` to maximize the fit to the observed visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `subplot_fit_dirty_images` provides a real-space view of the data, model and residuals via\n", + "inverse-NUFFT of the visibility-plane quantities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_dirty_images(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The fit contains the solved-for `intensity` value of every shapelet in the basis.\n", + "\n", + "These are computed via `fit.linear_light_profile_intensity_dict`, which maps each linear light profile in\n", + "the model to its inferred `intensity`. Print the first few entries for brevity." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.linear_light_profile_intensity_dict)\n", + "\n", + "for shapelet in shapelets_bulge_list[:5]:\n", + " intensity = fit.linear_light_profile_intensity_dict[shapelet]\n", + " print(f\" n={shapelet.n} m={shapelet.m} intensity = {intensity:+.6e}\")\n", + "\n", + "print(\n", + " f\"\\n number of negative-intensity shapelets: \"\n", + " f\"{sum(1 for s in shapelets_bulge_list if fit.linear_light_profile_intensity_dict[s] < 0)} \"\n", + " f\"/ {len(shapelets_bulge_list)}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the\n", + "solved-for `intensity` values is also accessible from a fit.\n", + "\n", + "The benefit of this helper-tracer is that it can be visualised (linear light profiles cannot be plotted\n", + "by default because they do not have `intensity` values)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization__\n", + "\n", + "The helper-tracer created above replaces every linear `ShapeletPolar` with an ordinary `ShapeletPolar`\n", + "carrying its solved-for `intensity` \u2014 that tracer can be plotted directly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=dataset.grid),\n", + " title=\"Tracer Image (shapelet source)\",\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/advanced/shapelets/modeling.ipynb b/notebooks/interferometer/features/advanced/shapelets/modeling.ipynb index 96d20cb12..7a3d07940 100644 --- a/notebooks/interferometer/features/advanced/shapelets/modeling.ipynb +++ b/notebooks/interferometer/features/advanced/shapelets/modeling.ipynb @@ -1,519 +1,556 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Shapelets (Interferometer)\n", - "==============================================\n", - "\n", - "A shapelet is a basis function appropriate for capturing the exponential / disk-like features of a galaxy.\n", - "It has been employed in many strong lensing studies to model the light of the lensed source galaxy, because\n", - "it can represent features of disky star-forming galaxies that a single Sersic function cannot.\n", - "\n", - "- https://ui.adsabs.harvard.edu/abs/2016MNRAS.457.3066T\n", - "- https://iopscience.iop.org/article/10.1088/0004-637X/813/2/102\n", - "\n", - "Shapelets are described in full in:\n", - "\n", - " https://arxiv.org/abs/astro-ph/0105178\n", - "\n", - "This script performs lens modeling of an `Interferometer` dataset using a polar shapelet basis for the\n", - "source galaxy. The `intensity` of every shapelet is solved for via linear algebra (see the\n", - "`linear_light_profiles` feature for a full description of this).\n", - "\n", - "Shapelet fits to interferometer data were previously impractical because every likelihood evaluation has\n", - "to Fourier-transform each shapelet basis component into the uv-plane, and prior NUFFT backends were not\n", - "JAX-friendly. With `nufftax` (https://github.com/GragasLab/nufftax) \u2014 a JAX-native NUFFT \u2014 the full\n", - "shapelet basis is transformed inside the same jit/vmap pipeline as the rest of the model, amortising the\n", - "per-iteration NUFFT cost on the GPU. Shapelet source fits are now routine even at ALMA-class visibility\n", - "counts.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Benefits and drawbacks of a shapelet source for interferometer data.\n", - "- **NUFFT (nufftax):** Why shapelets-on-visibilities is now practical thanks to nufftax.\n", - "- **Positive Negative Solver:** Why shapelets require a positive-negative solver (unlike MGE or linear\n", - " Sersic).\n", - "- **Model:** Compose the lens model \u2014 `Isothermal` + `ExternalShear` mass and a polar shapelet source\n", - " bulge. Lens light omitted (interferometer convention).\n", - "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", - "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", - "- **Over Sampling:** Interferometer modeling does not use over-sampling.\n", - "- **Search:** Configure the non-linear search (Nautilus).\n", - "- **Analysis:** Create the `AnalysisInterferometer` object with the positive-negative solver enabled.\n", - "- **VRAM:** Memory budget for a multi-component shapelet basis on GPU.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals because they fail to\n", - "capture irregular and asymmetric morphology of source galaxies (e.g. disky star formation, isophotal\n", - "twists). Shapelets capture some of these features and can therefore better represent complex source\n", - "galaxies.\n", - "\n", - "The shapelet model can be composed in a way that has fewer non-linear parameters than an elliptical\n", - "Sersic. In this example, ~10 shapelets which represent the source's `bulge` are composed in a model\n", - "corresponding to just N=3 non-linear parameters (centre + shared beta). A linear Sersic source would have\n", - "N=6.\n", - "\n", - "__Disadvantages__\n", - "\n", - "- There are many types of galaxy structure which shapelets may struggle to represent, such as a bar or\n", - " asymmetric knots of star formation. Shapelets also rely on the galaxy having a distinct centre over\n", - " which the basis can be centred, which is not the case if the galaxy is a multi-component merging system\n", - " or has bright companion galaxies.\n", - "\n", - "- The linear algebra used to solve for the `intensity` of each shapelet has to allow for negative values\n", - " of intensity. Negative surface brightnesses are unphysical, and are often inferred in a shapelet\n", - " decomposition \u2014 for example if the true galaxy has structure that cannot be captured by the shapelet\n", - " basis. Other approaches (MGE, pixelization) can force positive-only intensities on the solution.\n", - "\n", - "- Computationally slower than a single linear `SersicCore` because each shapelet must be NUFFT'd to the\n", - " uv-plane per likelihood. With `nufftax` the per-NUFFT cost is small enough that this is no longer a\n", - " blocker; for many science cases an MGE source is still faster and gives higher quality results.\n", - "\n", - "__NUFFT (nufftax)__\n", - "\n", - "The image-to-visibilities Fourier transform is performed by a Non-Uniform Fast Fourier Transform (NUFFT),\n", - "exposed in **PyAutoLens** as `TransformerNUFFT`. The default backend is `nufftax`, a pure-JAX NUFFT:\n", - "\n", - " https://github.com/GragasLab/nufftax\n", - "\n", - "Because `nufftax` is JAX-native, NUFFT-ing every shapelet basis image happens inside the same compiled\n", - "likelihood that does the inversion, mass model ray-tracing, and chi-squared sum. There is no host\n", - "round-trip between NUFFT calls, so a model with N shapelets costs only N forward-NUFFTs per iteration on\n", - "the GPU \u2014 fast enough that shapelet-on-visibilities is now routinely practical.\n", - "\n", - "If `nufftax` is not installed, install it via `pip install nufftax`.\n", - "\n", - "__Positive Negative Solver__\n", - "\n", - "In other examples which use linear algebra to fit the data \u2014 linear light profiles, the Multi-Gaussian\n", - "Expansion (MGE), and pixelized source reconstructions on CCD imaging \u2014 we use a positive-only solver,\n", - "which forces all solved-for intensities to be positive. This is physical and sensible because the surface\n", - "brightnesses of a galaxy cannot be negative.\n", - "\n", - "Shapelets **cannot** be solved with a positive-only solver. Their ability to decompose the light of a\n", - "galaxy relies on being able to use negative intensities \u2014 shapelets are not physically motivated light\n", - "profiles but a mathematical basis that can represent any light profile, including via cancellations\n", - "between positive and negative basis-function amplitudes.\n", - "\n", - "This means shapelet fits may include negative flux in the reconstructed source galaxy, which is\n", - "unphysical, and is a known disadvantage of using shapelets. The `Settings` object passed to the analysis\n", - "below uses `use_positive_only_solver=False` to allow for negative intensities.\n", - "\n", - "For pixelized source reconstructions on interferometer data this same setting is used for a different\n", - "reason: negative visibility-plane noise can pull individual pixels negative without anything being wrong\n", - "physically.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data). Interferometer\n", - " convention.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's bulge is a superposition of polar `ShapeletPolar` profiles, with all shapelets\n", - " sharing a centre, elliptical components, and a single `beta` size scale.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook.\n", - "\n", - "__Imaging Equivalent__\n", - "\n", - "For the CCD-imaging version of this script, see\n", - "`autolens_workspace/*/imaging/features/advanced/shapelets/modeling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit with\n", - "the lens model.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane\n", - "and compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native NUFFT backed by\n", - "`nufftax`, which is required for fast shapelet modeling and scales efficiently from a few hundred\n", - "visibilities to tens of millions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "If you are familiar with using imaging data, you may have seen that a numerical technique called over\n", - "sampling is used, which evaluates light profiles on a higher resolution grid than the image data to ensure\n", - "the calculation is accurate.\n", - "\n", - "Interferometer data does not observe galaxies in a way where over sampling is necessary, therefore all\n", - "interferometer calculations are performed without over sampling.\n", - "\n", - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's bulge is a superposition of linear `ShapeletPolar` profiles [3 parameters total].\n", - " - All shapelets share a centre, elliptical components, and a single `beta` size scale.\n", - " - The shapelet (n, m) quantum numbers are assigned procedurally.\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=10.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook:\n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source: polar shapelet basis with shared centre/ell_comps/beta.\n", - "total_n = 10\n", - "total_m = sum(range(2, total_n + 1)) + 1\n", - "\n", - "shapelets_bulge_list = af.Collection(\n", - " af.Model(al.lp_linear.ShapeletPolar) for _ in range(total_n + total_m + 1)\n", - ")\n", - "\n", - "n_count = 1\n", - "m_count = -1\n", - "\n", - "for i, shapelet in enumerate(shapelets_bulge_list):\n", - " if i == 0:\n", - " shapelet.n = 0\n", - " shapelet.m = 0\n", - " else:\n", - " shapelet.n = n_count\n", - " shapelet.m = m_count\n", - "\n", - " m_count += 2\n", - "\n", - " if m_count > n_count:\n", - " n_count += 1\n", - " m_count = -n_count\n", - "\n", - " shapelet.centre = shapelets_bulge_list[0].centre\n", - " shapelet.ell_comps = shapelets_bulge_list[0].ell_comps\n", - " shapelet.beta = shapelets_bulge_list[0].beta\n", - "\n", - "source_bulge = af.Model(al.lp_basis.Basis, profile_list=shapelets_bulge_list)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", - "\n", - "# Overall Lens Model:\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "This confirms the source galaxy is made of many `ShapeletPolar` profiles whose centres, elliptical\n", - "components, and beta are all shared." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\") / \"features\" / \"advanced\",\n", - " name=\"shapelets\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisInterferometer` object defining how Nautilus fits the model to the data.\n", - "\n", - "Note `use_positive_only_solver=False` is set on the `Settings` \u2014 shapelets require the positive-negative\n", - "solver, as discussed above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " settings=al.Settings(use_positive_only_solver=False),\n", - " use_jax=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `interferometer/modeling.py` example explains how VRAM is used during GPU-based fitting and how to\n", - "print the estimated VRAM required by a model.\n", - "\n", - "For each linear shapelet, extra VRAM is used to store its NUFFT'd mapping matrix column. For around 30\n", - "shapelets this typically requires a modest amount of VRAM (e.g. 10-50 MB per batched likelihood). Models\n", - "that use hundreds of shapelets, especially in combination with a large batch size, may therefore exceed\n", - "GBs of VRAM and require you to adjust the batch size to fit within your GPU's VRAM.\n", - "\n", - "VRAM on interferometer datasets is driven primarily by the visibility count and the real-space mask size,\n", - "not the number of shapelets in the basis.\n", - "\n", - "__Run Time__\n", - "\n", - "The likelihood evaluation time for a shapelet basis is slower than a single linear `SersicCore` source,\n", - "because the image of every shapelet must be evaluated and NUFFT'd to the uv-plane. With `nufftax`, the\n", - "per-NUFFT cost is small enough that the total slow-down per likelihood is typically 2-5x for a 30-shapelet\n", - "basis compared to a one-component source \u2014 paid back in fewer iterations because the parameter space is\n", - "simpler (only N=3 free non-linear parameters for the source).\n", - "\n", - "Because the shapelet basis has no free `intensity` or `beta`-per-shapelet parameters (only the shared\n", - "beta) and the source intensities are solved by the inversion, Nautilus converges significantly faster\n", - "than for a free-intensity Sersic source.\n", - "\n", - "If shapelets are too slow for your science case, consider the MGE source feature\n", - "(`interferometer/features/multi_gaussian_expansion/modeling.py`), which uses an even simpler basis (no\n", - "quantum-number indexing, just log-spaced sigmas) and is often faster.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the\n", - "output folder for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This script has illustrated how to use shapelets to model the source galaxy of an interferometer-observed\n", - "strong lens. Thanks to nufftax, the per-shapelet NUFFT cost is amortised on the GPU and shapelet fits to\n", - "visibility data are practical at any visibility count.\n", - "\n", - "For most science cases an MGE source (see `features/multi_gaussian_expansion/`) will be faster and give\n", - "higher quality results. Shapelets may perform better for disky / star-forming source morphologies that\n", - "the smoother MGE basis struggles with, but this is not guaranteed \u2014 try both and compare." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Shapelets (Interferometer)\n", + "==============================================\n", + "\n", + "A shapelet is a basis function appropriate for capturing the exponential / disk-like features of a galaxy.\n", + "It has been employed in many strong lensing studies to model the light of the lensed source galaxy, because\n", + "it can represent features of disky star-forming galaxies that a single Sersic function cannot.\n", + "\n", + "- https://ui.adsabs.harvard.edu/abs/2016MNRAS.457.3066T\n", + "- https://iopscience.iop.org/article/10.1088/0004-637X/813/2/102\n", + "\n", + "Shapelets are described in full in:\n", + "\n", + " https://arxiv.org/abs/astro-ph/0105178\n", + "\n", + "This script performs lens modeling of an `Interferometer` dataset using a polar shapelet basis for the\n", + "source galaxy. The `intensity` of every shapelet is solved for via linear algebra (see the\n", + "`linear_light_profiles` feature for a full description of this).\n", + "\n", + "Shapelet fits to interferometer data were previously impractical because every likelihood evaluation has\n", + "to Fourier-transform each shapelet basis component into the uv-plane, and prior NUFFT backends were not\n", + "JAX-friendly. With `nufftax` (https://github.com/GragasLab/nufftax) \u2014 a JAX-native NUFFT \u2014 the full\n", + "shapelet basis is transformed inside the same jit/vmap pipeline as the rest of the model, amortising the\n", + "per-iteration NUFFT cost on the GPU. Shapelet source fits are now routine even at ALMA-class visibility\n", + "counts.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Benefits and drawbacks of a shapelet source for interferometer data.\n", + "- **NUFFT (nufftax):** Why shapelets-on-visibilities is now practical thanks to nufftax.\n", + "- **Positive Negative Solver:** Why shapelets require a positive-negative solver (unlike MGE or linear\n", + " Sersic).\n", + "- **Model:** Compose the lens model \u2014 `Isothermal` + `ExternalShear` mass and a polar shapelet source\n", + " bulge. Lens light omitted (interferometer convention).\n", + "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", + "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", + "- **Over Sampling:** Interferometer modeling does not use over-sampling.\n", + "- **Search:** Configure the non-linear search (Nautilus).\n", + "- **Analysis:** Create the `AnalysisInterferometer` object with the positive-negative solver enabled.\n", + "- **VRAM:** Memory budget for a multi-component shapelet basis on GPU.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals because they fail to\n", + "capture irregular and asymmetric morphology of source galaxies (e.g. disky star formation, isophotal\n", + "twists). Shapelets capture some of these features and can therefore better represent complex source\n", + "galaxies.\n", + "\n", + "The shapelet model can be composed in a way that has fewer non-linear parameters than an elliptical\n", + "Sersic. In this example, ~10 shapelets which represent the source's `bulge` are composed in a model\n", + "corresponding to just N=3 non-linear parameters (centre + shared beta). A linear Sersic source would have\n", + "N=6.\n", + "\n", + "__Disadvantages__\n", + "\n", + "- There are many types of galaxy structure which shapelets may struggle to represent, such as a bar or\n", + " asymmetric knots of star formation. Shapelets also rely on the galaxy having a distinct centre over\n", + " which the basis can be centred, which is not the case if the galaxy is a multi-component merging system\n", + " or has bright companion galaxies.\n", + "\n", + "- The linear algebra used to solve for the `intensity` of each shapelet has to allow for negative values\n", + " of intensity. Negative surface brightnesses are unphysical, and are often inferred in a shapelet\n", + " decomposition \u2014 for example if the true galaxy has structure that cannot be captured by the shapelet\n", + " basis. Other approaches (MGE, pixelization) can force positive-only intensities on the solution.\n", + "\n", + "- Computationally slower than a single linear `SersicCore` because each shapelet must be NUFFT'd to the\n", + " uv-plane per likelihood. With `nufftax` the per-NUFFT cost is small enough that this is no longer a\n", + " blocker; for many science cases an MGE source is still faster and gives higher quality results.\n", + "\n", + "__NUFFT (nufftax)__\n", + "\n", + "The image-to-visibilities Fourier transform is performed by a Non-Uniform Fast Fourier Transform (NUFFT),\n", + "exposed in **PyAutoLens** as `TransformerNUFFT`. The default backend is `nufftax`, a pure-JAX NUFFT:\n", + "\n", + " https://github.com/GragasLab/nufftax\n", + "\n", + "Because `nufftax` is JAX-native, NUFFT-ing every shapelet basis image happens inside the same compiled\n", + "likelihood that does the inversion, mass model ray-tracing, and chi-squared sum. There is no host\n", + "round-trip between NUFFT calls, so a model with N shapelets costs only N forward-NUFFTs per iteration on\n", + "the GPU \u2014 fast enough that shapelet-on-visibilities is now routinely practical.\n", + "\n", + "If `nufftax` is not installed, install it via `pip install nufftax`.\n", + "\n", + "__Positive Negative Solver__\n", + "\n", + "In other examples which use linear algebra to fit the data \u2014 linear light profiles, the Multi-Gaussian\n", + "Expansion (MGE), and pixelized source reconstructions on CCD imaging \u2014 we use a positive-only solver,\n", + "which forces all solved-for intensities to be positive. This is physical and sensible because the surface\n", + "brightnesses of a galaxy cannot be negative.\n", + "\n", + "Shapelets **cannot** be solved with a positive-only solver. Their ability to decompose the light of a\n", + "galaxy relies on being able to use negative intensities \u2014 shapelets are not physically motivated light\n", + "profiles but a mathematical basis that can represent any light profile, including via cancellations\n", + "between positive and negative basis-function amplitudes.\n", + "\n", + "This means shapelet fits may include negative flux in the reconstructed source galaxy, which is\n", + "unphysical, and is a known disadvantage of using shapelets. The `Settings` object passed to the analysis\n", + "below uses `use_positive_only_solver=False` to allow for negative intensities.\n", + "\n", + "For pixelized source reconstructions on interferometer data this same setting is used for a different\n", + "reason: negative visibility-plane noise can pull individual pixels negative without anything being wrong\n", + "physically.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data). Interferometer\n", + " convention.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's bulge is a superposition of polar `ShapeletPolar` profiles, with all shapelets\n", + " sharing a centre, elliptical components, and a single `beta` size scale.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook.\n", + "\n", + "__Imaging Equivalent__\n", + "\n", + "For the CCD-imaging version of this script, see\n", + "`autolens_workspace/*/imaging/features/advanced/shapelets/modeling.py`." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit with\n", + "the lens model.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane\n", + "and compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native NUFFT backed by\n", + "`nufftax`, which is required for fast shapelet modeling and scales efficiently from a few hundred\n", + "visibilities to tens of millions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "If you are familiar with using imaging data, you may have seen that a numerical technique called over\n", + "sampling is used, which evaluates light profiles on a higher resolution grid than the image data to ensure\n", + "the calculation is accurate.\n", + "\n", + "Interferometer data does not observe galaxies in a way where over sampling is necessary, therefore all\n", + "interferometer calculations are performed without over sampling.\n", + "\n", + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's bulge is a superposition of linear `ShapeletPolar` profiles [3 parameters total].\n", + " - All shapelets share a centre, elliptical components, and a single `beta` size scale.\n", + " - The shapelet (n, m) quantum numbers are assigned procedurally.\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=10.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook:\n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source: polar shapelet basis with shared centre/ell_comps/beta.\n", + "total_n = 10\n", + "total_m = sum(range(2, total_n + 1)) + 1\n", + "\n", + "shapelets_bulge_list = af.Collection(\n", + " af.Model(al.lp_linear.ShapeletPolar) for _ in range(total_n + total_m + 1)\n", + ")\n", + "\n", + "n_count = 1\n", + "m_count = -1\n", + "\n", + "for i, shapelet in enumerate(shapelets_bulge_list):\n", + " if i == 0:\n", + " shapelet.n = 0\n", + " shapelet.m = 0\n", + " else:\n", + " shapelet.n = n_count\n", + " shapelet.m = m_count\n", + "\n", + " m_count += 2\n", + "\n", + " if m_count > n_count:\n", + " n_count += 1\n", + " m_count = -n_count\n", + "\n", + " shapelet.centre = shapelets_bulge_list[0].centre\n", + " shapelet.ell_comps = shapelets_bulge_list[0].ell_comps\n", + " shapelet.beta = shapelets_bulge_list[0].beta\n", + "\n", + "source_bulge = af.Model(al.lp_basis.Basis, profile_list=shapelets_bulge_list)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", + "\n", + "# Overall Lens Model:\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "This confirms the source galaxy is made of many `ShapeletPolar` profiles whose centres, elliptical\n", + "components, and beta are all shared." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\") / \"features\" / \"advanced\",\n", + " name=\"shapelets\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisInterferometer` object defining how Nautilus fits the model to the data.\n", + "\n", + "Note `use_positive_only_solver=False` is set on the `Settings` \u2014 shapelets require the positive-negative\n", + "solver, as discussed above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " settings=al.Settings(use_positive_only_solver=False),\n", + " use_jax=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `interferometer/modeling.py` example explains how VRAM is used during GPU-based fitting and how to\n", + "print the estimated VRAM required by a model.\n", + "\n", + "For each linear shapelet, extra VRAM is used to store its NUFFT'd mapping matrix column. For around 30\n", + "shapelets this typically requires a modest amount of VRAM (e.g. 10-50 MB per batched likelihood). Models\n", + "that use hundreds of shapelets, especially in combination with a large batch size, may therefore exceed\n", + "GBs of VRAM and require you to adjust the batch size to fit within your GPU's VRAM.\n", + "\n", + "VRAM on interferometer datasets is driven primarily by the visibility count and the real-space mask size,\n", + "not the number of shapelets in the basis.\n", + "\n", + "__Run Time__\n", + "\n", + "The likelihood evaluation time for a shapelet basis is slower than a single linear `SersicCore` source,\n", + "because the image of every shapelet must be evaluated and NUFFT'd to the uv-plane. With `nufftax`, the\n", + "per-NUFFT cost is small enough that the total slow-down per likelihood is typically 2-5x for a 30-shapelet\n", + "basis compared to a one-component source \u2014 paid back in fewer iterations because the parameter space is\n", + "simpler (only N=3 free non-linear parameters for the source).\n", + "\n", + "Because the shapelet basis has no free `intensity` or `beta`-per-shapelet parameters (only the shared\n", + "beta) and the source intensities are solved by the inversion, Nautilus converges significantly faster\n", + "than for a free-intensity Sersic source.\n", + "\n", + "If shapelets are too slow for your science case, consider the MGE source feature\n", + "(`interferometer/features/multi_gaussian_expansion/modeling.py`), which uses an even simpler basis (no\n", + "quantum-number indexing, just log-spaced sigmas) and is often faster.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the\n", + "output folder for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script has illustrated how to use shapelets to model the source galaxy of an interferometer-observed\n", + "strong lens. Thanks to nufftax, the per-shapelet NUFFT cost is amortised on the GPU and shapelet fits to\n", + "visibility data are practical at any visibility count.\n", + "\n", + "For most science cases an MGE source (see `features/multi_gaussian_expansion/`) will be faster and give\n", + "higher quality results. Shapelets may perform better for disky / star-forming source morphologies that\n", + "the smoother MGE basis struggles with, but this is not guaranteed \u2014 try both and compare." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/datacube/data_preparation.ipynb b/notebooks/interferometer/features/datacube/data_preparation.ipynb index b6dd323ff..247ddeda6 100644 --- a/notebooks/interferometer/features/datacube/data_preparation.ipynb +++ b/notebooks/interferometer/features/datacube/data_preparation.ipynb @@ -1,401 +1,438 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data Preparation: Datacube\n", - "============================\n", - "\n", - "Most users come to datacube modeling with a single 4D FITS file from CASA, with shape\n", - "``(n_pol, n_chan, n_vis, 2)`` \u2014 for example, an ALMA observation of a CO emission line might be\n", - "``(2, 34, 16984, 2)``: two polarisations, 34 spectral channels, ~17k visibilities per channel, real/imag.\n", - "\n", - "This script shows how to bridge from that on-disk shape to the canonical input PyAutoLens expects:\n", - "\n", - " * **In memory**: a Python ``list`` of ``al.Interferometer`` objects, one per channel.\n", - " * **On disk** (canonical 3D layout): three FITS files ``visibilities_cube.fits``,\n", - " ``noise_map_cube.fits`` and ``uv_wavelengths_cube.fits``, all of shape ``(n_chan, n_vis, 2)``.\n", - "\n", - "Two preprocessing steps are usually needed before the 3D layout is reached:\n", - "\n", - " 1. **Polarisation collapse** \u2014 averaging or concatenating the two polarisation entries.\n", - " 2. **Optional `uv_wavelengths` / `noise_map` reduction** \u2014 both quantities change very little\n", - " channel-to-channel, so for the purposes of modeling they're often averaged across channels and\n", - " stored as a single ``(n_vis, 2)`` array shared across the cube.\n", - "\n", - "The script is structured as runnable explanation. The reusable bit is\n", - "``dataset_list_from_3d_fits()`` near the bottom \u2014 copy that into your own script if you have a 3D\n", - "cube already laid out on disk.\n", - "\n", - "__Contents__\n", - "\n", - "- **What ALMA Gives You:** The 4D shape `(n_pol, n_chan, n_vis, 2)` straight from CASA.\n", - "- **Polarisation Handling:** Average vs concatenate; tradeoffs and one-line numpy.\n", - "- **Canonical 3D Shape:** `(n_chan, n_vis, 2)` is what every loader below assumes.\n", - "- **Shared vs Per-Channel uv_wavelengths / noise_map:** When channel-invariant quantities can be a\n", - " single shared `(n_vis, 2)` array, and how the loader handles either case.\n", - "- **3D-FITS Loader:** A self-contained `dataset_list_from_3d_fits()` function that builds the\n", - " in-memory `dataset_list`.\n", - "- **Worked Example:** Loads the reference cube (written by `simulator.py`) and confirms the\n", - " resulting `dataset_list` matches what the per-channel-folder loader would produce." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import numpy as np\n", - "\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__What ALMA Gives You__\n", - "\n", - "A typical ALMA spectral-line `visibilities.fits` from CASA has shape ``(n_pol, n_chan, n_vis, 2)``.\n", - "For Hannah's data this is ``(2, 34, 16984, 2)``: two polarisations, 34 channels of an emission line,\n", - "about 17k visibilities per channel, real/imag.\n", - "\n", - "`noise_map.fits` and `uv_wavelengths.fits` typically share the polarisation and channel dimensions\n", - "(though `uv_wavelengths` may collapse to `(n_pol, n_vis, 2)` if your reduction has decided the\n", - "baselines are channel-invariant \u2014 which is usually a good approximation for narrow emission lines).\n", - "\n", - "The simulator in this folder writes a 4D `(n_pol, n_chan, n_vis, 2)` cube alongside the per-channel\n", - "folders and the 3D `(n_chan, n_vis, 2)` autolens-native cube. We load the simulator's actual 4D\n", - "output below so the polarisation-collapse demonstration runs on the same data the rest of the\n", - "modeling scripts will fit. (Hint: for this synthetic simulator the two polarisations are identical \u2014\n", - "real CASA data has independent noise realisations between pols, which is why averaging real data\n", - "reduces effective noise by `sqrt(2)`. The collapse code below is the same in both cases.)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"datacube\"\n", - "dataset_name = \"sim_simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name\n", - "\n", - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "visibilities_4d = al.ndarray_via_fits_from(\n", - " file_path=dataset_path / \"visibilities_4d_cube.fits\", hdu=0\n", - ")\n", - "noise_map_4d = al.ndarray_via_fits_from(\n", - " file_path=dataset_path / \"noise_map_4d_cube.fits\", hdu=0\n", - ")\n", - "uv_wavelengths_4d = al.ndarray_via_fits_from(\n", - " file_path=dataset_path / \"uv_wavelengths_4d_cube.fits\", hdu=0\n", - ")\n", - "\n", - "print(f\"On-disk shapes (loaded from simulator's *_4d_cube.fits):\")\n", - "print(f\" visibilities_4d: {visibilities_4d.shape} (n_pol, n_chan, n_vis, 2)\")\n", - "print(f\" noise_map_4d: {noise_map_4d.shape} (n_pol, n_chan, n_vis, 2)\")\n", - "print(f\" uv_wavelengths_4d: {uv_wavelengths_4d.shape} (n_pol, n_chan, n_vis, 2)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Polarisation Handling__\n", - "\n", - "Two options for collapsing the polarisation axis. Pick whichever your reduction was designed for \u2014\n", - "the workspace doesn't bake a choice in.\n", - "\n", - "- **Average** preserves ``n_vis``. Best when the two polarisations carry the same source signal and\n", - " you want to suppress polarisation-mode noise:\n", - "\n", - " visibilities_collapsed = visibilities_4d.mean(axis=0) # (n_chan, n_vis, 2)\n", - "\n", - " For the noise map use inverse-variance weighting, which for equal-noise polarisations reduces to\n", - " ``noise_map.mean(axis=0) / sqrt(2)``:\n", - "\n", - " noise_map_collapsed = noise_map_4d.mean(axis=0) / np.sqrt(2) # (n_chan, n_vis, 2)\n", - "\n", - "- **Concatenate** doubles ``n_vis``. Best when you want to keep every visibility independent and let\n", - " the model fit each polarisation separately (or when the polarisations carry slightly different\n", - " signal \u2014 the average would smear that out):\n", - "\n", - " visibilities_collapsed = np.concatenate([visibilities_4d[0], visibilities_4d[1]], axis=1) # (n_chan, 2*n_vis, 2)\n", - " noise_map_collapsed = np.concatenate([noise_map_4d[0], noise_map_4d[1]], axis=1) # (n_chan, 2*n_vis, 2)\n", - " uv_wavelengths_collapsed = np.concatenate([uv_wavelengths_4d[0], uv_wavelengths_4d[1]], axis=1) # (n_chan, 2*n_vis, 2)\n", - "\n", - "We use the averaging path below for clarity. Swap in the concatenate version if that's what your\n", - "reduction expects." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "visibilities_3d = visibilities_4d.mean(axis=0)\n", - "noise_map_3d = noise_map_4d.mean(axis=0) / np.sqrt(2)\n", - "uv_wavelengths_3d = uv_wavelengths_4d.mean(axis=0)\n", - "# uv_wavelengths is channel-invariant in this synthetic simulator, so we can also collapse the\n", - "# channel axis if you prefer the shared 2D form. The loader supports either:\n", - "uv_wavelengths_2d = uv_wavelengths_3d.mean(axis=0)\n", - "\n", - "print(f\"\\nAfter polarisation averaging:\")\n", - "print(f\" visibilities_3d: {visibilities_3d.shape} (n_chan, n_vis, 2)\")\n", - "print(f\" noise_map_3d: {noise_map_3d.shape} (n_chan, n_vis, 2)\")\n", - "print(\n", - " f\" uv_wavelengths_3d: {uv_wavelengths_3d.shape} (n_chan, n_vis, 2) [per-channel]\"\n", - ")\n", - "print(\n", - " f\" uv_wavelengths_2d: {uv_wavelengths_2d.shape} (n_vis, 2) [shared, optional]\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Canonical 3D Shape__\n", - "\n", - "After polarisation collapse, every loader below assumes the canonical post-preprocessing shapes:\n", - "\n", - "- ``visibilities``: ``(n_chan, n_vis, 2)`` \u2014 required.\n", - "- ``noise_map``: ``(n_chan, n_vis, 2)`` \u2014 required, can be channel-invariant if you replicate.\n", - "- ``uv_wavelengths``: ``(n_chan, n_vis, 2)`` for per-channel baselines, **or** ``(n_vis, 2)`` for\n", - " channel-invariant baselines (the loader broadcasts).\n", - "\n", - "If your reduction stores `noise_map` as channel-invariant `(n_vis, 2)`, the loader broadcasts the\n", - "same way." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Shared vs Per-Channel `uv_wavelengths` / `noise_map`__\n", - "\n", - "For most ALMA narrow-line cubes both `uv_wavelengths` and `noise_map` change very little\n", - "channel-to-channel and can safely be stored as a single shared `(n_vis, 2)` array. The deferred\n", - "shared-`L\u1d40 W\u0303 L` optimisation Aris designed (see issue #120) will only fire when these arrays are\n", - "actually shared, so there's some real performance reason to use the shared form when you can.\n", - "\n", - "The loader below supports both: if the input `noise_map` or `uv_wavelengths` array is 2D, it gets\n", - "broadcast to all `n_chan` channels; if it's 3D, each channel gets its own slice.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__3D-FITS Loader__\n", - "\n", - "Self-contained loader function. Copy this into your own script if you have a 3D cube on disk.\n", - "\n", - "The function takes paths to the three FITS files (or numpy arrays directly \u2014 see the `_arrays`\n", - "variant below if you want to skip the disk trip), splits them per-channel, and constructs the list\n", - "of `al.Interferometer` objects.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def dataset_list_from_3d_fits(\n", - " visibilities_path: Path,\n", - " noise_map_path: Path,\n", - " uv_wavelengths_path: Path,\n", - " real_space_mask: al.Mask2D,\n", - " transformer_class=al.TransformerNUFFT,\n", - "):\n", - " \"\"\"Load a 3D-FITS datacube into a list of `al.Interferometer` objects.\n", - "\n", - " Parameters\n", - " ----------\n", - " visibilities_path\n", - " Path to a `.fits` file of shape `(n_chan, n_vis, 2)` storing real/imag visibilities for\n", - " every channel. Polarisations should already be collapsed.\n", - " noise_map_path\n", - " Path to a `.fits` file of shape `(n_chan, n_vis, 2)` (per-channel) or `(n_vis, 2)` (shared\n", - " across channels). The shared form is broadcast to every channel.\n", - " uv_wavelengths_path\n", - " Path to a `.fits` file of shape `(n_chan, n_vis, 2)` (per-channel) or `(n_vis, 2)` (shared).\n", - " real_space_mask\n", - " The 2D real-space mask used by the Fourier transformer.\n", - " transformer_class\n", - " `al.TransformerNUFFT` (default, JAX-native via `nufftax`, scales to many millions of visibilities)\n", - " or `al.TransformerDFT`.\n", - "\n", - " Returns\n", - " -------\n", - " list[al.Interferometer]\n", - " One `Interferometer` per channel, in cube order.\n", - " \"\"\"\n", - " visibilities_3d = al.ndarray_via_fits_from(file_path=visibilities_path, hdu=0)\n", - " noise_map_arr = al.ndarray_via_fits_from(file_path=noise_map_path, hdu=0)\n", - " uv_wavelengths_arr = al.ndarray_via_fits_from(file_path=uv_wavelengths_path, hdu=0)\n", - "\n", - " if visibilities_3d.ndim != 3:\n", - " raise ValueError(\n", - " f\"visibilities array must be 3D (n_chan, n_vis, 2); got shape {visibilities_3d.shape}\"\n", - " )\n", - "\n", - " n_chan = visibilities_3d.shape[0]\n", - "\n", - " # Broadcast 2D shared arrays to (n_chan, n_vis, 2) so the per-channel loop is uniform.\n", - " if noise_map_arr.ndim == 2:\n", - " noise_map_arr = np.broadcast_to(\n", - " noise_map_arr[None], visibilities_3d.shape\n", - " ).copy()\n", - " if uv_wavelengths_arr.ndim == 2:\n", - " uv_wavelengths_arr = np.broadcast_to(\n", - " uv_wavelengths_arr[None], visibilities_3d.shape\n", - " ).copy()\n", - "\n", - " return [\n", - " al.Interferometer(\n", - " data=al.Visibilities(visibilities=visibilities_3d[c]),\n", - " noise_map=al.VisibilitiesNoiseMap(visibilities=noise_map_arr[c]),\n", - " uv_wavelengths=uv_wavelengths_arr[c],\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=transformer_class,\n", - " )\n", - " for c in range(n_chan)\n", - " ]\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Worked Example__\n", - "\n", - "Load the reference cube (already used above for the 4D-collapse demo) and confirm the 3D-FITS\n", - "loader produces a `dataset_list` numerically identical to the per-channel-folder loader. This is\n", - "a sanity check \u2014 once you trust this, you can drop the per-channel-folder pattern entirely if you\n", - "prefer the 3D layout." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=3.5,\n", - ")\n", - "\n", - "dataset_list_3d = dataset_list_from_3d_fits(\n", - " visibilities_path=dataset_path / \"visibilities_cube.fits\",\n", - " noise_map_path=dataset_path / \"noise_map_cube.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths_cube.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerDFT,\n", - ")\n", - "\n", - "# Cross-check against the per-channel-folder loader.\n", - "channel_paths = sorted(\n", - " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", - ")\n", - "dataset_list_folders = [\n", - " al.Interferometer.from_fits(\n", - " data_path=channel_path / \"data.fits\",\n", - " noise_map_path=channel_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerDFT,\n", - " )\n", - " for channel_path in channel_paths\n", - "]\n", - "\n", - "assert len(dataset_list_3d) == len(dataset_list_folders), \"channel count mismatch\"\n", - "\n", - "for c, (d3d, dfolder) in enumerate(zip(dataset_list_3d, dataset_list_folders)):\n", - " np.testing.assert_allclose(\n", - " np.asarray(d3d.data.real),\n", - " np.asarray(dfolder.data.real),\n", - " rtol=1e-12,\n", - " err_msg=f\"channel {c}: data.real mismatch\",\n", - " )\n", - " np.testing.assert_allclose(\n", - " np.asarray(d3d.data.imag),\n", - " np.asarray(dfolder.data.imag),\n", - " rtol=1e-12,\n", - " err_msg=f\"channel {c}: data.imag mismatch\",\n", - " )\n", - " np.testing.assert_allclose(\n", - " np.asarray(d3d.uv_wavelengths),\n", - " np.asarray(dfolder.uv_wavelengths),\n", - " rtol=1e-12,\n", - " err_msg=f\"channel {c}: uv_wavelengths mismatch\",\n", - " )\n", - "\n", - "print(\n", - " f\"\\n3D-FITS loader and per-channel-folder loader agree on {len(dataset_list_3d)} channels.\"\n", - ")\n", - "print(f\" visibilities/channel: {dataset_list_3d[0].data.shape[0]}\")\n", - "print(f\" uv_wavelengths shape: {np.asarray(dataset_list_3d[0].uv_wavelengths).shape}\")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Data Preparation: Datacube\n", + "============================\n", + "\n", + "Most users come to datacube modeling with a single 4D FITS file from CASA, with shape\n", + "``(n_pol, n_chan, n_vis, 2)`` \u2014 for example, an ALMA observation of a CO emission line might be\n", + "``(2, 34, 16984, 2)``: two polarisations, 34 spectral channels, ~17k visibilities per channel, real/imag.\n", + "\n", + "This script shows how to bridge from that on-disk shape to the canonical input PyAutoLens expects:\n", + "\n", + " * **In memory**: a Python ``list`` of ``al.Interferometer`` objects, one per channel.\n", + " * **On disk** (canonical 3D layout): three FITS files ``visibilities_cube.fits``,\n", + " ``noise_map_cube.fits`` and ``uv_wavelengths_cube.fits``, all of shape ``(n_chan, n_vis, 2)``.\n", + "\n", + "Two preprocessing steps are usually needed before the 3D layout is reached:\n", + "\n", + " 1. **Polarisation collapse** \u2014 averaging or concatenating the two polarisation entries.\n", + " 2. **Optional `uv_wavelengths` / `noise_map` reduction** \u2014 both quantities change very little\n", + " channel-to-channel, so for the purposes of modeling they're often averaged across channels and\n", + " stored as a single ``(n_vis, 2)`` array shared across the cube.\n", + "\n", + "The script is structured as runnable explanation. The reusable bit is\n", + "``dataset_list_from_3d_fits()`` near the bottom \u2014 copy that into your own script if you have a 3D\n", + "cube already laid out on disk.\n", + "\n", + "__Contents__\n", + "\n", + "- **What ALMA Gives You:** The 4D shape `(n_pol, n_chan, n_vis, 2)` straight from CASA.\n", + "- **Polarisation Handling:** Average vs concatenate; tradeoffs and one-line numpy.\n", + "- **Canonical 3D Shape:** `(n_chan, n_vis, 2)` is what every loader below assumes.\n", + "- **Shared vs Per-Channel uv_wavelengths / noise_map:** When channel-invariant quantities can be a\n", + " single shared `(n_vis, 2)` array, and how the loader handles either case.\n", + "- **3D-FITS Loader:** A self-contained `dataset_list_from_3d_fits()` function that builds the\n", + " in-memory `dataset_list`.\n", + "- **Worked Example:** Loads the reference cube (written by `simulator.py`) and confirms the\n", + " resulting `dataset_list` matches what the per-channel-folder loader would produce." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__What ALMA Gives You__\n", + "\n", + "A typical ALMA spectral-line `visibilities.fits` from CASA has shape ``(n_pol, n_chan, n_vis, 2)``.\n", + "For Hannah's data this is ``(2, 34, 16984, 2)``: two polarisations, 34 channels of an emission line,\n", + "about 17k visibilities per channel, real/imag.\n", + "\n", + "`noise_map.fits` and `uv_wavelengths.fits` typically share the polarisation and channel dimensions\n", + "(though `uv_wavelengths` may collapse to `(n_pol, n_vis, 2)` if your reduction has decided the\n", + "baselines are channel-invariant \u2014 which is usually a good approximation for narrow emission lines).\n", + "\n", + "The simulator in this folder writes a 4D `(n_pol, n_chan, n_vis, 2)` cube alongside the per-channel\n", + "folders and the 3D `(n_chan, n_vis, 2)` autolens-native cube. We load the simulator's actual 4D\n", + "output below so the polarisation-collapse demonstration runs on the same data the rest of the\n", + "modeling scripts will fit. (Hint: for this synthetic simulator the two polarisations are identical \u2014\n", + "real CASA data has independent noise realisations between pols, which is why averaging real data\n", + "reduces effective noise by `sqrt(2)`. The collapse code below is the same in both cases.)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"datacube\"\n", + "dataset_name = \"sim_simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name\n", + "\n", + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "visibilities_4d = al.ndarray_via_fits_from(\n", + " file_path=dataset_path / \"visibilities_4d_cube.fits\", hdu=0\n", + ")\n", + "noise_map_4d = al.ndarray_via_fits_from(\n", + " file_path=dataset_path / \"noise_map_4d_cube.fits\", hdu=0\n", + ")\n", + "uv_wavelengths_4d = al.ndarray_via_fits_from(\n", + " file_path=dataset_path / \"uv_wavelengths_4d_cube.fits\", hdu=0\n", + ")\n", + "\n", + "print(f\"On-disk shapes (loaded from simulator's *_4d_cube.fits):\")\n", + "print(f\" visibilities_4d: {visibilities_4d.shape} (n_pol, n_chan, n_vis, 2)\")\n", + "print(f\" noise_map_4d: {noise_map_4d.shape} (n_pol, n_chan, n_vis, 2)\")\n", + "print(f\" uv_wavelengths_4d: {uv_wavelengths_4d.shape} (n_pol, n_chan, n_vis, 2)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Polarisation Handling__\n", + "\n", + "Two options for collapsing the polarisation axis. Pick whichever your reduction was designed for \u2014\n", + "the workspace doesn't bake a choice in.\n", + "\n", + "- **Average** preserves ``n_vis``. Best when the two polarisations carry the same source signal and\n", + " you want to suppress polarisation-mode noise:\n", + "\n", + " visibilities_collapsed = visibilities_4d.mean(axis=0) # (n_chan, n_vis, 2)\n", + "\n", + " For the noise map use inverse-variance weighting, which for equal-noise polarisations reduces to\n", + " ``noise_map.mean(axis=0) / sqrt(2)``:\n", + "\n", + " noise_map_collapsed = noise_map_4d.mean(axis=0) / np.sqrt(2) # (n_chan, n_vis, 2)\n", + "\n", + "- **Concatenate** doubles ``n_vis``. Best when you want to keep every visibility independent and let\n", + " the model fit each polarisation separately (or when the polarisations carry slightly different\n", + " signal \u2014 the average would smear that out):\n", + "\n", + " visibilities_collapsed = np.concatenate([visibilities_4d[0], visibilities_4d[1]], axis=1) # (n_chan, 2*n_vis, 2)\n", + " noise_map_collapsed = np.concatenate([noise_map_4d[0], noise_map_4d[1]], axis=1) # (n_chan, 2*n_vis, 2)\n", + " uv_wavelengths_collapsed = np.concatenate([uv_wavelengths_4d[0], uv_wavelengths_4d[1]], axis=1) # (n_chan, 2*n_vis, 2)\n", + "\n", + "We use the averaging path below for clarity. Swap in the concatenate version if that's what your\n", + "reduction expects." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "visibilities_3d = visibilities_4d.mean(axis=0)\n", + "noise_map_3d = noise_map_4d.mean(axis=0) / np.sqrt(2)\n", + "uv_wavelengths_3d = uv_wavelengths_4d.mean(axis=0)\n", + "# uv_wavelengths is channel-invariant in this synthetic simulator, so we can also collapse the\n", + "# channel axis if you prefer the shared 2D form. The loader supports either:\n", + "uv_wavelengths_2d = uv_wavelengths_3d.mean(axis=0)\n", + "\n", + "print(f\"\\nAfter polarisation averaging:\")\n", + "print(f\" visibilities_3d: {visibilities_3d.shape} (n_chan, n_vis, 2)\")\n", + "print(f\" noise_map_3d: {noise_map_3d.shape} (n_chan, n_vis, 2)\")\n", + "print(\n", + " f\" uv_wavelengths_3d: {uv_wavelengths_3d.shape} (n_chan, n_vis, 2) [per-channel]\"\n", + ")\n", + "print(\n", + " f\" uv_wavelengths_2d: {uv_wavelengths_2d.shape} (n_vis, 2) [shared, optional]\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Canonical 3D Shape__\n", + "\n", + "After polarisation collapse, every loader below assumes the canonical post-preprocessing shapes:\n", + "\n", + "- ``visibilities``: ``(n_chan, n_vis, 2)`` \u2014 required.\n", + "- ``noise_map``: ``(n_chan, n_vis, 2)`` \u2014 required, can be channel-invariant if you replicate.\n", + "- ``uv_wavelengths``: ``(n_chan, n_vis, 2)`` for per-channel baselines, **or** ``(n_vis, 2)`` for\n", + " channel-invariant baselines (the loader broadcasts).\n", + "\n", + "If your reduction stores `noise_map` as channel-invariant `(n_vis, 2)`, the loader broadcasts the\n", + "same way." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Shared vs Per-Channel `uv_wavelengths` / `noise_map`__\n", + "\n", + "For most ALMA narrow-line cubes both `uv_wavelengths` and `noise_map` change very little\n", + "channel-to-channel and can safely be stored as a single shared `(n_vis, 2)` array. The deferred\n", + "shared-`L\u1d40 W\u0303 L` optimisation Aris designed (see issue #120) will only fire when these arrays are\n", + "actually shared, so there's some real performance reason to use the shared form when you can.\n", + "\n", + "The loader below supports both: if the input `noise_map` or `uv_wavelengths` array is 2D, it gets\n", + "broadcast to all `n_chan` channels; if it's 3D, each channel gets its own slice.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__3D-FITS Loader__\n", + "\n", + "Self-contained loader function. Copy this into your own script if you have a 3D cube on disk.\n", + "\n", + "The function takes paths to the three FITS files (or numpy arrays directly \u2014 see the `_arrays`\n", + "variant below if you want to skip the disk trip), splits them per-channel, and constructs the list\n", + "of `al.Interferometer` objects.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def dataset_list_from_3d_fits(\n", + " visibilities_path: Path,\n", + " noise_map_path: Path,\n", + " uv_wavelengths_path: Path,\n", + " real_space_mask: al.Mask2D,\n", + " transformer_class=al.TransformerNUFFT,\n", + "):\n", + " \"\"\"Load a 3D-FITS datacube into a list of `al.Interferometer` objects.\n", + "\n", + " Parameters\n", + " ----------\n", + " visibilities_path\n", + " Path to a `.fits` file of shape `(n_chan, n_vis, 2)` storing real/imag visibilities for\n", + " every channel. Polarisations should already be collapsed.\n", + " noise_map_path\n", + " Path to a `.fits` file of shape `(n_chan, n_vis, 2)` (per-channel) or `(n_vis, 2)` (shared\n", + " across channels). The shared form is broadcast to every channel.\n", + " uv_wavelengths_path\n", + " Path to a `.fits` file of shape `(n_chan, n_vis, 2)` (per-channel) or `(n_vis, 2)` (shared).\n", + " real_space_mask\n", + " The 2D real-space mask used by the Fourier transformer.\n", + " transformer_class\n", + " `al.TransformerNUFFT` (default, JAX-native via `nufftax`, scales to many millions of visibilities)\n", + " or `al.TransformerDFT`.\n", + "\n", + " Returns\n", + " -------\n", + " list[al.Interferometer]\n", + " One `Interferometer` per channel, in cube order.\n", + " \"\"\"\n", + " visibilities_3d = al.ndarray_via_fits_from(file_path=visibilities_path, hdu=0)\n", + " noise_map_arr = al.ndarray_via_fits_from(file_path=noise_map_path, hdu=0)\n", + " uv_wavelengths_arr = al.ndarray_via_fits_from(file_path=uv_wavelengths_path, hdu=0)\n", + "\n", + " if visibilities_3d.ndim != 3:\n", + " raise ValueError(\n", + " f\"visibilities array must be 3D (n_chan, n_vis, 2); got shape {visibilities_3d.shape}\"\n", + " )\n", + "\n", + " n_chan = visibilities_3d.shape[0]\n", + "\n", + " # Broadcast 2D shared arrays to (n_chan, n_vis, 2) so the per-channel loop is uniform.\n", + " if noise_map_arr.ndim == 2:\n", + " noise_map_arr = np.broadcast_to(\n", + " noise_map_arr[None], visibilities_3d.shape\n", + " ).copy()\n", + " if uv_wavelengths_arr.ndim == 2:\n", + " uv_wavelengths_arr = np.broadcast_to(\n", + " uv_wavelengths_arr[None], visibilities_3d.shape\n", + " ).copy()\n", + "\n", + " return [\n", + " al.Interferometer(\n", + " data=al.Visibilities(visibilities=visibilities_3d[c]),\n", + " noise_map=al.VisibilitiesNoiseMap(visibilities=noise_map_arr[c]),\n", + " uv_wavelengths=uv_wavelengths_arr[c],\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=transformer_class,\n", + " )\n", + " for c in range(n_chan)\n", + " ]\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Worked Example__\n", + "\n", + "Load the reference cube (already used above for the 4D-collapse demo) and confirm the 3D-FITS\n", + "loader produces a `dataset_list` numerically identical to the per-channel-folder loader. This is\n", + "a sanity check \u2014 once you trust this, you can drop the per-channel-folder pattern entirely if you\n", + "prefer the 3D layout." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=3.5,\n", + ")\n", + "\n", + "dataset_list_3d = dataset_list_from_3d_fits(\n", + " visibilities_path=dataset_path / \"visibilities_cube.fits\",\n", + " noise_map_path=dataset_path / \"noise_map_cube.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths_cube.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerDFT,\n", + ")\n", + "\n", + "# Cross-check against the per-channel-folder loader.\n", + "channel_paths = sorted(\n", + " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", + ")\n", + "dataset_list_folders = [\n", + " al.Interferometer.from_fits(\n", + " data_path=channel_path / \"data.fits\",\n", + " noise_map_path=channel_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerDFT,\n", + " )\n", + " for channel_path in channel_paths\n", + "]\n", + "\n", + "assert len(dataset_list_3d) == len(dataset_list_folders), \"channel count mismatch\"\n", + "\n", + "for c, (d3d, dfolder) in enumerate(zip(dataset_list_3d, dataset_list_folders)):\n", + " np.testing.assert_allclose(\n", + " np.asarray(d3d.data.real),\n", + " np.asarray(dfolder.data.real),\n", + " rtol=1e-12,\n", + " err_msg=f\"channel {c}: data.real mismatch\",\n", + " )\n", + " np.testing.assert_allclose(\n", + " np.asarray(d3d.data.imag),\n", + " np.asarray(dfolder.data.imag),\n", + " rtol=1e-12,\n", + " err_msg=f\"channel {c}: data.imag mismatch\",\n", + " )\n", + " np.testing.assert_allclose(\n", + " np.asarray(d3d.uv_wavelengths),\n", + " np.asarray(dfolder.uv_wavelengths),\n", + " rtol=1e-12,\n", + " err_msg=f\"channel {c}: uv_wavelengths mismatch\",\n", + " )\n", + "\n", + "print(\n", + " f\"\\n3D-FITS loader and per-channel-folder loader agree on {len(dataset_list_3d)} channels.\"\n", + ")\n", + "print(f\" visibilities/channel: {dataset_list_3d[0].data.shape[0]}\")\n", + "print(f\" uv_wavelengths shape: {np.asarray(dataset_list_3d[0].uv_wavelengths).shape}\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/datacube/delaunay.ipynb b/notebooks/interferometer/features/datacube/delaunay.ipynb index 9a8ded331..51ff599ed 100644 --- a/notebooks/interferometer/features/datacube/delaunay.ipynb +++ b/notebooks/interferometer/features/datacube/delaunay.ipynb @@ -1,475 +1,512 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Datacube \u2014 Delaunay Source\n", - "=====================================\n", - "\n", - "This script fits a datacube \u2014 a list of `Interferometer` channels \u2014 with a single shared lens model and a\n", - "per-channel **Delaunay-pixelized** source reconstruction. It is the Delaunay sibling of `modeling.py`, which\n", - "fits the same cube with a `RectangularAdaptDensity` mesh.\n", - "\n", - "A Delaunay mesh adapts the source-plane reconstruction to the lensed source's morphology more flexibly than a\n", - "rectangular mesh: source pixels are placed via a triangulation of (y, x) image-plane points that are ray-traced\n", - "into the source plane for each candidate lens model. This gives more pixels to the highly-magnified parts of\n", - "the source-plane, which is usually where the emission-line signal is concentrated.\n", - "\n", - "For ALMA datacubes Delaunay is often the right default once the simpler rectangular fit has converged on a\n", - "sensible lens model \u2014 the canonical workflow is rectangular for the global fit, Delaunay (or a more adaptive\n", - "image-mesh like `Hilbert`) for the source-science follow-up. This script demonstrates the wiring.\n", - "\n", - "The FactorGraph wiring is identical to `modeling.py`: shared lens, per-channel inversion. Only the source-plane\n", - "mesh (and its regularization) changes.\n", - "\n", - "__Contents__\n", - "\n", - "- **Mask:** Define the 2D real-space mask applied to every channel.\n", - "- **Dataset:** Where the per-channel cube lives on disk and how to point this script at your own.\n", - "- **Dataset Auto-Simulation:** Run `simulator.py` automatically if the cube isn't already on disk.\n", - "- **Dataset Loading:** Loop over the channel folders and load each as an `Interferometer` object.\n", - "- **Sparse Operators:** Pre-compute per-channel sparse-operator matrices used by the Delaunay inversion.\n", - "- **Positions:** Load the cube's multiple-image positions and build a shared `PositionsLH` penalty.\n", - "- **Settings:** Disable the positive-only solver so visibility-space inversions can take negative pixel values.\n", - "- **Image Mesh:** Build the image-plane mesh of (y, x) points that get ray-traced and Delaunay-triangulated.\n", - "- **Edge Zeroing:** Append a ring of edge pixels so the source-plane reconstruction zeroes at the mesh boundary.\n", - "- **Adapt Images:** Pair the image-plane mesh with the source galaxy via `al.AdaptImages`.\n", - "- **Model:** Compose the shared `Isothermal + ExternalShear` lens and Delaunay-pixelized source.\n", - "- **Per-Channel Analyses:** One `AnalysisInterferometer` per channel, with shared `adapt_images` and `PositionsLH`.\n", - "- **FactorGraph:** Wrap each analysis in an `AnalysisFactor` and combine via `af.FactorGraphModel`.\n", - "- **Search:** Configure the `Nautilus` non-linear search.\n", - "- **Model Fit:** Run the fit. Per-channel cost is comparable to the rectangular variant.\n", - "- **Wrap Up:** Pointers to `modeling.py`, `start_here.py`, and the JAX likelihood walkthrough." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import subprocess\n", - "import sys\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "The datacube lives under `dataset/interferometer/datacube//`, with one subfolder per channel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"datacube\"\n", - "dataset_name = \"sim_simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Loading__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "channel_paths = sorted(\n", - " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", - ")\n", - "print(f\"Loading {len(channel_paths)} channels from {dataset_path}\")\n", - "\n", - "dataset_list = [\n", - " al.Interferometer.from_fits(\n", - " data_path=channel_path / \"data.fits\",\n", - " noise_map_path=channel_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - " )\n", - " for channel_path in channel_paths\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " dataset.apply_sparse_operator(use_jax=True, show_progress=False)\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Multiple-image positions + `PositionsLH` are essential for Delaunay fits \u2014 without them, the search routinely\n", - "finds demagnified-source local maxima." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image Mesh__\n", - "\n", - "The Delaunay mesh is built by ray-tracing (y, x) coordinates from the image-plane to the source-plane and\n", - "triangulating the source-plane points. We start with an `Overlay` image-mesh: a regular grid of points spread\n", - "across the image-plane mask. This has a mild adaptive effect \u2014 regions of high lens magnification receive more\n", - "source pixels once they are ray-traced. For more aggressive adaptation, swap `Overlay` for `Hilbert` (which\n", - "weights points by the source's surface brightness).\n", - "\n", - "The number of `pixels` passed to `al.mesh.Delaunay` must equal the number of points in `image_plane_mesh_grid`,\n", - "because JAX uses `pixels` to define static-shape arrays." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=real_space_mask)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Edge Zeroing__\n", - "\n", - "Pixels at the edge of the source-plane mesh are forced to zero brightness during the inversion to prevent\n", - "unphysical solutions where edge pixels reconstruct bright surface brightnesses (often by absorbing residuals).\n", - "For a Delaunay mesh we manually add a ring of edge points to the image-plane mesh and tell `al.mesh.Delaunay`\n", - "how many of those trailing points to zero." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "edge_pixels_total = 30\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=real_space_mask.mask_centre,\n", - " radius=mask_radius + real_space_mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Adapt Images__\n", - "\n", - "The image-plane mesh is passed into modeling via `al.AdaptImages`, keyed on the source galaxy's path in the\n", - "model. The same `adapt_images` object is reused for every channel because the lens model and image-plane mesh\n", - "are shared." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "adapt_images = al.AdaptImages(\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid,\n", - " },\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "Shared `Isothermal + ExternalShear` lens + Delaunay-pixelized source. `ConstantSplit` is the canonical Delaunay\n", - "regularizer (split-prior on inner-vs-edge mesh pixels). The mesh `pixels` is fixed at the number of points in\n", - "`image_plane_mesh_grid`, which is `26*26 + 30` after the edge-ring append." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "mesh = af.Model(\n", - " al.mesh.Delaunay,\n", - " pixels=image_plane_mesh_grid.shape[0],\n", - " zeroed_pixels=edge_pixels_total,\n", - ")\n", - "regularization = af.Model(al.reg.ConstantSplit)\n", - "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Channel Analyses__\n", - "\n", - "Each `AnalysisInterferometer` receives the same `adapt_images` (image-plane mesh is shared) and the same\n", - "`positions_likelihood` (lens model is shared via the FactorGraph)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [\n", - " al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " settings=settings,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " use_jax=True,\n", - " )\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__FactorGraph__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = [\n", - " af.AnalysisFactor(prior_model=model.copy(), analysis=analysis)\n", - " for analysis in analysis_list\n", - "]\n", - "\n", - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - "print(f\" channels in factor graph: {len(analysis_factor_list)}\")\n", - "print(\n", - " f\" global model free parameters: {factor_graph.global_prior_model.total_free_parameters}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\") / \"datacube\",\n", - " name=\"delaunay\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20,\n", - " iterations_per_quick_update=50000,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "Per-channel cost is dominated by the Delaunay inversion + the per-channel NUFFT \u2014 comparable to the rectangular\n", - "variant in `modeling.py`. On CPU expect this to take a few hours for the 4-channel reference cube; on GPU,\n", - "tens of minutes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " Per-channel inversions multiply the per-likelihood cost \u2014 expect this to take longer than a single-channel\n", - " interferometer fit.\n", - " \"\"\"\n", - ")\n", - "\n", - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "For the rectangular variant, see `modeling.py`. For the parametric variant, see `modeling_parametric.py`. For\n", - "the narrative walkthrough, see `start_here.py`. For a step-by-step JAX likelihood walkthrough, see\n", - "`autolens_workspace_developer/datacube/likelihood_function.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Datacube \u2014 Delaunay Source\n", + "=====================================\n", + "\n", + "This script fits a datacube \u2014 a list of `Interferometer` channels \u2014 with a single shared lens model and a\n", + "per-channel **Delaunay-pixelized** source reconstruction. It is the Delaunay sibling of `modeling.py`, which\n", + "fits the same cube with a `RectangularAdaptDensity` mesh.\n", + "\n", + "A Delaunay mesh adapts the source-plane reconstruction to the lensed source's morphology more flexibly than a\n", + "rectangular mesh: source pixels are placed via a triangulation of (y, x) image-plane points that are ray-traced\n", + "into the source plane for each candidate lens model. This gives more pixels to the highly-magnified parts of\n", + "the source-plane, which is usually where the emission-line signal is concentrated.\n", + "\n", + "For ALMA datacubes Delaunay is often the right default once the simpler rectangular fit has converged on a\n", + "sensible lens model \u2014 the canonical workflow is rectangular for the global fit, Delaunay (or a more adaptive\n", + "image-mesh like `Hilbert`) for the source-science follow-up. This script demonstrates the wiring.\n", + "\n", + "The FactorGraph wiring is identical to `modeling.py`: shared lens, per-channel inversion. Only the source-plane\n", + "mesh (and its regularization) changes.\n", + "\n", + "__Contents__\n", + "\n", + "- **Mask:** Define the 2D real-space mask applied to every channel.\n", + "- **Dataset:** Where the per-channel cube lives on disk and how to point this script at your own.\n", + "- **Dataset Auto-Simulation:** Run `simulator.py` automatically if the cube isn't already on disk.\n", + "- **Dataset Loading:** Loop over the channel folders and load each as an `Interferometer` object.\n", + "- **Sparse Operators:** Pre-compute per-channel sparse-operator matrices used by the Delaunay inversion.\n", + "- **Positions:** Load the cube's multiple-image positions and build a shared `PositionsLH` penalty.\n", + "- **Settings:** Disable the positive-only solver so visibility-space inversions can take negative pixel values.\n", + "- **Image Mesh:** Build the image-plane mesh of (y, x) points that get ray-traced and Delaunay-triangulated.\n", + "- **Edge Zeroing:** Append a ring of edge pixels so the source-plane reconstruction zeroes at the mesh boundary.\n", + "- **Adapt Images:** Pair the image-plane mesh with the source galaxy via `al.AdaptImages`.\n", + "- **Model:** Compose the shared `Isothermal + ExternalShear` lens and Delaunay-pixelized source.\n", + "- **Per-Channel Analyses:** One `AnalysisInterferometer` per channel, with shared `adapt_images` and `PositionsLH`.\n", + "- **FactorGraph:** Wrap each analysis in an `AnalysisFactor` and combine via `af.FactorGraphModel`.\n", + "- **Search:** Configure the `Nautilus` non-linear search.\n", + "- **Model Fit:** Run the fit. Per-channel cost is comparable to the rectangular variant.\n", + "- **Wrap Up:** Pointers to `modeling.py`, `start_here.py`, and the JAX likelihood walkthrough." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import subprocess\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "The datacube lives under `dataset/interferometer/datacube//`, with one subfolder per channel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"datacube\"\n", + "dataset_name = \"sim_simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Loading__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "channel_paths = sorted(\n", + " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", + ")\n", + "print(f\"Loading {len(channel_paths)} channels from {dataset_path}\")\n", + "\n", + "dataset_list = [\n", + " al.Interferometer.from_fits(\n", + " data_path=channel_path / \"data.fits\",\n", + " noise_map_path=channel_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + " )\n", + " for channel_path in channel_paths\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " dataset.apply_sparse_operator(use_jax=True, show_progress=False)\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Multiple-image positions + `PositionsLH` are essential for Delaunay fits \u2014 without them, the search routinely\n", + "finds demagnified-source local maxima." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image Mesh__\n", + "\n", + "The Delaunay mesh is built by ray-tracing (y, x) coordinates from the image-plane to the source-plane and\n", + "triangulating the source-plane points. We start with an `Overlay` image-mesh: a regular grid of points spread\n", + "across the image-plane mask. This has a mild adaptive effect \u2014 regions of high lens magnification receive more\n", + "source pixels once they are ray-traced. For more aggressive adaptation, swap `Overlay` for `Hilbert` (which\n", + "weights points by the source's surface brightness).\n", + "\n", + "The number of `pixels` passed to `al.mesh.Delaunay` must equal the number of points in `image_plane_mesh_grid`,\n", + "because JAX uses `pixels` to define static-shape arrays." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(mask=real_space_mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Edge Zeroing__\n", + "\n", + "Pixels at the edge of the source-plane mesh are forced to zero brightness during the inversion to prevent\n", + "unphysical solutions where edge pixels reconstruct bright surface brightnesses (often by absorbing residuals).\n", + "For a Delaunay mesh we manually add a ring of edge points to the image-plane mesh and tell `al.mesh.Delaunay`\n", + "how many of those trailing points to zero." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "edge_pixels_total = 30\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=real_space_mask.mask_centre,\n", + " radius=mask_radius + real_space_mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Adapt Images__\n", + "\n", + "The image-plane mesh is passed into modeling via `al.AdaptImages`, keyed on the source galaxy's path in the\n", + "model. The same `adapt_images` object is reused for every channel because the lens model and image-plane mesh\n", + "are shared." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "adapt_images = al.AdaptImages(\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid,\n", + " },\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "Shared `Isothermal + ExternalShear` lens + Delaunay-pixelized source. `ConstantSplit` is the canonical Delaunay\n", + "regularizer (split-prior on inner-vs-edge mesh pixels). The mesh `pixels` is fixed at the number of points in\n", + "`image_plane_mesh_grid`, which is `26*26 + 30` after the edge-ring append." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "mesh = af.Model(\n", + " al.mesh.Delaunay,\n", + " pixels=image_plane_mesh_grid.shape[0],\n", + " zeroed_pixels=edge_pixels_total,\n", + ")\n", + "regularization = af.Model(al.reg.ConstantSplit)\n", + "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Channel Analyses__\n", + "\n", + "Each `AnalysisInterferometer` receives the same `adapt_images` (image-plane mesh is shared) and the same\n", + "`positions_likelihood` (lens model is shared via the FactorGraph)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [\n", + " al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " settings=settings,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " use_jax=True,\n", + " )\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__FactorGraph__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = [\n", + " af.AnalysisFactor(prior_model=model.copy(), analysis=analysis)\n", + " for analysis in analysis_list\n", + "]\n", + "\n", + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + "print(f\" channels in factor graph: {len(analysis_factor_list)}\")\n", + "print(\n", + " f\" global model free parameters: {factor_graph.global_prior_model.total_free_parameters}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\") / \"datacube\",\n", + " name=\"delaunay\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20,\n", + " iterations_per_quick_update=50000,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "Per-channel cost is dominated by the Delaunay inversion + the per-channel NUFFT \u2014 comparable to the rectangular\n", + "variant in `modeling.py`. On CPU expect this to take a few hours for the 4-channel reference cube; on GPU,\n", + "tens of minutes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " Per-channel inversions multiply the per-likelihood cost \u2014 expect this to take longer than a single-channel\n", + " interferometer fit.\n", + " \"\"\"\n", + ")\n", + "\n", + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "For the rectangular variant, see `modeling.py`. For the parametric variant, see `modeling_parametric.py`. For\n", + "the narrative walkthrough, see `start_here.py`. For a step-by-step JAX likelihood walkthrough, see\n", + "`autolens_workspace_developer/datacube/likelihood_function.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/datacube/likelihood_function.ipynb b/notebooks/interferometer/features/datacube/likelihood_function.ipynb index 5db7cf078..ade644bb9 100644 --- a/notebooks/interferometer/features/datacube/likelihood_function.ipynb +++ b/notebooks/interferometer/features/datacube/likelihood_function.ipynb @@ -1,993 +1,1030 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Datacube__\n", - "\n", - "This script provides a step-by-step guide of the **PyAutoLens** `log_likelihood_function` used to fit a\n", - "**datacube** \u2014 a list of N per-channel `Interferometer` objects sharing a single lens model \u2014 with a per-channel\n", - "pixelized source reconstruction (specifically a `RectangularAdaptDensity` mesh and `Constant` regularization\n", - "scheme).\n", - "\n", - "This script has the same aims as `interferometer/features/pixelization/likelihood_function.py`:\n", - "\n", - " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can understand\n", - " the likelihood function (including references to the previous literature from which it is defined) without\n", - " having to write large quantities of text and equations.\n", - "\n", - " - To make inversions in **PyAutoLens** less of a \"black-box\" to users.\n", - "\n", - "__Comparison To Single-Channel Pixelization Likelihood__\n", - "\n", - "A datacube fit is identical to **N independent single-channel pixelization fits** plus one wrinkle: the lens\n", - "model is shared across all channels, so the same `tracer` is used for every channel's calculation. Inside any\n", - "single channel, every linear-algebra step matches `interferometer/features/pixelization/likelihood_function.py`\n", - "line for line:\n", - "\n", - " - Same ray-tracing (same lens galaxy mass + shear).\n", - " - Same source-plane `Mapper` construction.\n", - " - Same `mapping_matrix` (image-pixel to source-pixel mappings).\n", - " - Same `transformed_mapping_matrix` (NUFFT applied per source pixel), with the *same* memory caveat as the\n", - " single-channel case \u2014 for `n_vis \u2273 10^6` per channel the matrix becomes prohibitive to store and the\n", - " sparse-operator likelihood function is the production path.\n", - " - Same `data_vector`, `curvature_matrix`, `regularization_matrix`, NNLS reconstruction.\n", - " - Same `chi_squared`, regularization term, complexity terms, noise normalisation, per-channel `log_evidence`.\n", - "\n", - "The **cube log-evidence is the sum** of the per-channel log-evidences. That's it.\n", - "\n", - "This script presents the calculation directly:\n", - "\n", - " 1. Build the shared lens galaxy + source pixelization (channel-invariant).\n", - " 2. Walk through **channel 0** in detail (one short section per pixelization-script section, cross-referencing\n", - " that script for the full derivations).\n", - " 3. Loop the per-channel calculation across all `dataset_list` and sum.\n", - " 4. Cross-check against per-channel `al.FitInterferometer.log_evidence`.\n", - "\n", - "If you haven't read `interferometer/features/pixelization/likelihood_function.py` yet, do that first. This\n", - "script defers to it for almost everything that happens inside a single channel.\n", - "\n", - "__Simplifications__\n", - "\n", - "This example uses a `RectangularAdaptDensity` mesh + `Constant` regularization \u2014 the same combination used by\n", - "the rest of the `datacube/` tutorials (`modeling.py`, `start_here.py`). The\n", - "`pixelization/likelihood_function.py` reference uses `RectangularUniform`, which is a thin subclass of\n", - "`RectangularAdaptDensity`; the linear algebra is identical and the construction code is the same. The single\n", - "behaviour difference is that `RectangularAdaptDensity` lets the mesh's pixel density adapt to the source-plane\n", - "magnification map, which gives slightly better resolution in highly-magnified regions but does not change the\n", - "likelihood-function maths at all.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The pixelization likelihood function is the most complex one in **PyAutoLens**. It is strongly advised you read\n", - "through the following first:\n", - "\n", - " - `interferometer/features/pixelization/likelihood_function.py` \u2014 full step-by-step walkthrough of the\n", - " single-channel pixelization likelihood. This script is essentially a per-channel restatement of that one\n", - " plus a sum, so the entire body below uses cross-references to its sections.\n", - " - `interferometer/light_profile/log_likelihood_function.py` \u2014 the simpler light-profile likelihood, which\n", - " introduces visibility-space inner products and the NUFFT without the pixelization linear algebra.\n", - "\n", - "__Contents__\n", - "\n", - "- **Comparison:** datacube = N independent pixelization fits + shared lens; cube log-evidence is the sum.\n", - "- **Simplifications:** `RectangularAdaptDensity` mesh, `Constant` regularization.\n", - "- **Prerequisites:** read `pixelization/likelihood_function.py` first.\n", - "- **Mesh Shape:** identical to the pixelization reference, sized to 14\u00d714 to match `modeling.py`.\n", - "- **Mask:** identical to the pixelization reference, sized to the datacube simulator's 256\u00d7256 / 0.1\u2033 grid.\n", - "- **Dataset:** load a *list* of `Interferometer` objects, one per channel (cube-specific).\n", - "- **Lens Galaxy:** identical to the pixelization reference (channel-invariant).\n", - "- **Source Galaxy Pixelization and Regularization:** identical to the pixelization reference.\n", - "- **One Channel Walkthrough:** run the full pixelization-likelihood calculation on channel 0, cross-referencing\n", - " the single-channel script for shared derivations.\n", - "- **Across All Channels:** loop the per-channel calculation across `dataset_list` and sum \u2014 the headline\n", - " cube-specific section.\n", - "- **Fit:** cross-check the manual sum against per-channel `al.FitInterferometer.log_evidence`.\n", - "- **Lens Modeling:** pointer to `modeling.py`, `start_here.py`, `delaunay.py`, `modeling_parametric.py`.\n", - "- **Log Likelihood Function: Source Code Speed Up:** per-channel fast paths apply, multiplied by N channels;\n", - " the deferred shared-`L\u1d40 W\u0303 L` optimisation recovers a factor-N speed-up when `uv_wavelengths`/`noise_map`\n", - " are nearly channel-invariant.\n", - "- **Wrap Up:** pointers to modeling scripts and the planned JIT-correctness regression tests." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Mesh Shape__`. Sized to 14\u00d714 to match the\n", - "`datacube/modeling.py` mesh." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 14\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Mask__`. Sized to match the datacube simulator's\n", - "256\u00d7256 / 0.1\u2033 real-space grid." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256), pixel_scales=0.1, radius=3.0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "This is the first genuinely cube-specific section. Instead of a single `al.Interferometer.from_fits(...)`\n", - "call, we load a *list* of `Interferometer` objects, one per channel.\n", - "\n", - "The simulator at `simulator.py` writes the cube in two on-disk layouts:\n", - "\n", - " - **Per-channel folders** (used here): one `channel_NNN/` subfolder per channel with its own\n", - " `data.fits`/`noise_map.fits`/`uv_wavelengths.fits`. Convenient when the channels are split-out already.\n", - " - **3D-FITS cube** (`{visibilities,noise_map,uv_wavelengths}_cube.fits`, each `(n_chan, n_vis, 2)`): one\n", - " file containing every channel stacked. `data_preparation.py` shows how to load this form. There is also a\n", - " 4D `(n_pol, n_chan, n_vis, 2)` CASA-like layout that requires a polarisation-collapse pre-processing step.\n", - "\n", - "Whichever loader you use, you end up with the same list `dataset_list` below. Each entry has its own\n", - "`visibilities`, `noise_map`, and `uv_wavelengths`; the lens galaxy is channel-invariant and the per-channel\n", - "sources differ in `intensity` and `centre` (see `simulator.py`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"datacube\"\n", - "dataset_name = \"sim_simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the cube isn't on disk yet, run the simulator. This makes the script runnable on a fresh checkout." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "channel_paths = sorted(\n", - " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", - ")\n", - "\n", - "dataset_list = [\n", - " al.Interferometer.from_fits(\n", - " data_path=channel_path / \"data.fits\",\n", - " noise_map_path=channel_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - " )\n", - " for channel_path in channel_paths\n", - "]\n", - "\n", - "n_channels = len(dataset_list)\n", - "n_vis = int(dataset_list[0].uv_wavelengths.shape[0])\n", - "print(f\"Loaded {n_channels} channels, {n_vis} visibilities per channel.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot the dirty images of the first and last channels to see the per-channel data the cube fit will\n", - "work with. The source intensity and centre vary across the cube, so the dirty images differ even though the\n", - "lens model is the same." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset_list[0])\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset_list[-1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Same as `pixelization/likelihood_function.py:__Over Sampling__`. Interferometer pixelizations do not use\n", - "over-sampling.\n", - "\n", - "__Masked Image Grid__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Masked Image Grid__`. Every channel uses the same\n", - "real-space mask, so `dataset.grids.pixelization` is channel-invariant." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=dataset_list[0].grids.pixelization, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Lens Galaxy__`. The lens galaxy is channel-invariant\n", - "(this is the entire reason datacube modeling can share a single non-linear search across all channels \u2014 the\n", - "mass model doesn't depend on frequency)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Pixelization and Regularization__\n", - "\n", - "Same as `pixelization/likelihood_function.py:__Source Galaxy Pixelization and Regularization__`, with\n", - "`RectangularAdaptDensity` substituted for `RectangularUniform`. The classes share the same construction\n", - "machinery \u2014 `RectangularUniform` is a subclass of `RectangularAdaptDensity` \u2014 so all of the mesh-grid and\n", - "mapper code below is unchanged.\n", - "\n", - "The same source pixelization is used for every channel. Each channel runs its own linear inversion against\n", - "this shared pixelization (which is what gives each channel an independent source-plane reconstruction)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - ")\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Ray Tracing__`. Channel-invariant because the lens is\n", - "channel-invariant; the same `tracer` is used by every per-channel calculation below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "==================================================================================================\n", - " ONE CHANNEL WALKTHROUGH\n", - "==================================================================================================\n", - "\n", - "We now walk through the full pixelization-likelihood calculation for **channel 0**. Every section heading\n", - "below matches a section in `pixelization/likelihood_function.py` and the calculation is the same. Each\n", - "section has a one-line cross-reference, the code (so this script runs end-to-end), and any cube-specific\n", - "notes that don't apply in the single-channel case." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "dataset = dataset_list[0]\n", - "print(f\"\\n=== Channel 0 walkthrough ===\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Border Relocation__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Border Relocation__`. Run per channel because the traced\n", - "grid feeding into it depends on the per-channel `dataset.grids.pixelization` \u2014 but `dataset.grids.pixelization`\n", - "is itself channel-invariant in this cube (every channel uses the same `real_space_mask`), so the relocated\n", - "grid actually only needs to be computed once. We compute it inside the loop below for clarity." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", - "\n", - "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", - " grid=dataset.grids.pixelization\n", - ")[-1]\n", - "\n", - "border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)\n", - "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", - "\n", - "aplt.plot_grid(grid=relocated_grid, title=\"Channel 0 relocated traced grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Pixel Centre Calculation__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Source Pixel Centre Calculation__`. The source-plane mesh\n", - "overlays the relocated traced grid; since the relocated grid is the same for every channel (shared mask,\n", - "shared tracer), the source-plane mesh grid is channel-invariant too." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoarray.inversion.mesh.mesh.rectangular_adapt_density import overlay_grid_from\n", - "\n", - "mesh_grid = overlay_grid_from(\n", - " shape_native=mesh_shape, grid=al.Grid2DIrregular(relocated_grid)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolation / Mapper / Mapping Matrix__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Interpolation__`, `__Mapper__`, and\n", - "`__Mapping Matrix__`. All channel-invariant for the cube \u2014 they depend only on the shared `tracer`,\n", - "mesh shape and mask. We compute them once and reuse across channels in the loop below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "interpolator = pixelization.mesh.interpolator_from(\n", - " source_plane_data_grid=relocated_grid,\n", - " source_plane_mesh_grid=mesh_grid,\n", - ")\n", - "mapper = al.Mapper(interpolator=interpolator)\n", - "\n", - "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", - " pix_indexes_for_sub_slim_index=mapper.pix_indexes_for_sub_slim_index,\n", - " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,\n", - " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,\n", - " pixels=mapper.pixels,\n", - " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", - " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", - " sub_fraction=mapper.over_sampler.sub_fraction,\n", - ")\n", - "\n", - "print(\n", - " f\" mapping_matrix shape: {mapping_matrix.shape} (real-space pixels x source pixels)\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Transformed Mapping Matrix ($f$)__\n", - "\n", - "For a single channel this is identical to `pixelization/likelihood_function.py:__Transformed Mapping Matrix__`.\n", - "\n", - "**This is where per-channel structure starts to appear.** Each channel has its own `uv_wavelengths` and\n", - "therefore its own NUFFT operator (`dataset.transformer`). The shape of the output is the same for every\n", - "channel in our simulator (every channel has the same `n_vis`), but the values differ \u2014 the same source pixel\n", - "maps to different uv-plane visibilities in different channels because the baselines differ.\n", - "\n", - "The same memory caveat applies *per channel*: for `n_vis \u2273 10^6` per channel, this matrix is many GB per\n", - "channel; for an N-channel cube the total memory footprint is N\u00d7 the single-channel one. The sparse-operator\n", - "likelihood function \u2014 see the `apply_sparse_operator` calls in `modeling.py` \u2014 is the production path that\n", - "avoids ever materialising this matrix. The deferred shared-`L\u1d40 W\u0303 L` optimisation goes further: it computes\n", - "the curvature-matrix sandwich once and reuses it across channels when `uv_wavelengths`/`noise_map` are\n", - "nearly channel-invariant (which they typically are for narrow emission lines)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", - " mapping_matrix=mapping_matrix\n", - ")\n", - "\n", - "print(\n", - " f\" transformed_mapping_matrix shape: {transformed_mapping_matrix.shape} \"\n", - " f\"(n_vis x source pixels), dtype: {transformed_mapping_matrix.dtype}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Data Vector (D)__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Data Vector (D)__`. Per channel because each channel has\n", - "its own `dataset.data` (visibilities) and `dataset.noise_map`; the construction formula is unchanged." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_vector = (\n", - " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " visibilities=dataset.data,\n", - " noise_map=dataset.noise_map,\n", - " )\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Matrix (F)__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Curvature Matrix (F)__`. The matrix `F` depends only on\n", - "`transformed_mapping_matrix` and `noise_map` \u2014 both of which are channel-specific because of the per-channel\n", - "NUFFT and noise \u2014 but for the typical narrow-emission-line case where `uv_wavelengths` and `noise_map` change\n", - "very little across the line, `F` is nearly channel-invariant. This is the matrix Aris's deferred\n", - "shared-`L\u1d40 W\u0303 L` optimisation reuses across channels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.real,\n", - " noise_map=dataset.noise_map.real,\n", - ")\n", - "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.imag,\n", - " noise_map=dataset.noise_map.imag,\n", - ")\n", - "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Regularization Matrix (H)__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Regularization Matrix (H)__`. Channel-invariant because\n", - "the regularization scheme + mesh structure are channel-invariant \u2014 `H` only depends on which source pixels\n", - "neighbour which." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", - " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", - " neighbors=mapper.neighbors,\n", - " neighbors_sizes=mapper.neighbors.sizes,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__F + \u03bbH / Galaxy Reconstruction (s)__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__F + Lamdba H__` and `__Galaxy Reconstruction (s)__`. The\n", - "linear system `s = (F + \u03bbH)\u207b\u00b9 D` is solved per channel, producing each channel's source-plane reconstruction.\n", - "\n", - "Per-channel reconstructions are what make the cube fit physically interesting: an emission line that\n", - "brightens-and-fades across the cube produces a sequence of source-plane reconstructions whose total flux\n", - "traces the line profile, while the lens mass model stays fixed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)\n", - "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visibilities Reconstruction__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Visibilities Reconstruction__`. Per channel because the\n", - "model visibilities live in the channel's own uv-plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_visibilities = (\n", - " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " reconstruction=reconstruction,\n", - " )\n", - ")\n", - "mapped_reconstructed_visibilities = al.Visibilities(\n", - " visibilities=mapped_reconstructed_visibilities\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function \u2014 Five Terms__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Likelihood Function__`. The same five-term formula\n", - "applies per channel:\n", - "\n", - " -2 ln \u03b5_c = \u03c7\u00b2_c + s_c\u1d40 H s_c + ln det(F_c + H) - ln det(H) + \u03a3_j ln (2\u03c0 \u03c3\u00b2_{c,j})\n", - "\n", - "where `c` indexes channels. The cube log-evidence is `\u03a3_c log_evidence_c`, computed in the\n", - "\"Across All Channels\" section below.\n", - "\n", - "__Chi Squared__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Chi Squared__`. Per channel \u2014 each channel has its own\n", - "visibilities, noise, and model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_visibilities = mapped_reconstructed_visibilities\n", - "residual_map = dataset.data - model_visibilities\n", - "\n", - "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", - "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", - "chi_squared = np.sum(chi_squared_map_real) + np.sum(chi_squared_map_imag)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Regularization Term__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Regularization Term__`. Per channel via the per-channel\n", - "`reconstruction`; `H` is channel-invariant." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "regularization_term = np.matmul(\n", - " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Complexity Terms__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Complexity Terms__`. The `ln det(F + \u03bbH)` term is\n", - "per-channel (`F` is per-channel); `ln det(H)` is channel-invariant." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", - "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalisation Term__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Noise Normalization Term__`. Per channel.\n", - "\n", - "**Cube-specific note**: if you did polarisation-collapse by *averaging* (see `data_preparation.py`), the\n", - "visibility noise map already incorporates the sqrt(2) noise reduction relative to a single polarisation. The\n", - "noise-normalisation term you compute below matches whatever your data-prep produced \u2014 it doesn't double-count\n", - "the polarisation averaging. If you concatenated polarisations instead, `n_vis` doubled and the term naturally\n", - "scales with it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization_real = np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", - "noise_normalization_imag = np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", - "noise_normalization = noise_normalization_real + noise_normalization_imag" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood (Channel 0)__\n", - "\n", - "Identical to `pixelization/likelihood_function.py:__Calculate The Log Likelihood__`. The result is the\n", - "log-evidence of channel 0 only." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_evidence_channel_0 = float(\n", - " -0.5\n", - " * (\n", - " chi_squared\n", - " + regularization_term\n", - " + log_curvature_reg_matrix_term\n", - " - log_regularization_matrix_term\n", - " + noise_normalization\n", - " )\n", - ")\n", - "print(f\" channel 0 log_evidence: {log_evidence_channel_0:.6f}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "==================================================================================================\n", - " ACROSS ALL CHANNELS\n", - "==================================================================================================\n", - "\n", - "This is the headline cube-specific section. We loop the per-channel calculation across `dataset_list` and\n", - "sum to get the cube log-evidence.\n", - "\n", - "Almost everything inside the loop is channel-invariant \u2014 the `tracer`, `mapper`, `mapping_matrix`, and\n", - "`regularization_matrix`. Only the per-channel `dataset` (visibilities, noise_map, uv_wavelengths) changes, so\n", - "inside the loop we only redo the channel-dependent steps:\n", - "\n", - " - `transformed_mapping_matrix` (depends on `dataset.transformer`).\n", - " - `data_vector` (depends on `dataset.data` and `dataset.noise_map`).\n", - " - `curvature_matrix` (depends on `transformed_mapping_matrix` and `dataset.noise_map`).\n", - " - `reconstruction` (depends on `data_vector` and `curvature_matrix`).\n", - " - Visibilities-space residuals, \u03c7\u00b2, regularization term, complexity terms, noise normalisation.\n", - "\n", - "This is exactly what `af.FactorGraphModel` does internally for the modeling scripts: it feeds the same\n", - "lens-model parameters to every per-channel `AnalysisInterferometer.log_likelihood_function`, and the search\n", - "sees the sum." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def per_channel_log_evidence(dataset):\n", - " \"\"\"Compute the log-evidence of a single channel given the channel-invariant `tracer`, `mapper`,\n", - " `mapping_matrix` and `regularization_matrix` defined above.\n", - "\n", - " All steps mirror `pixelization/likelihood_function.py` line by line; only the dataset (visibilities,\n", - " noise map, uv_wavelengths) varies.\n", - " \"\"\"\n", - "\n", - " transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", - " mapping_matrix=mapping_matrix\n", - " )\n", - "\n", - " data_vector = al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " visibilities=dataset.data,\n", - " noise_map=dataset.noise_map,\n", - " )\n", - "\n", - " real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.real,\n", - " noise_map=dataset.noise_map.real,\n", - " )\n", - " imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.imag,\n", - " noise_map=dataset.noise_map.imag,\n", - " )\n", - " curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", - " curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)\n", - "\n", - " reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)\n", - "\n", - " mapped_reconstructed_visibilities = (\n", - " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " reconstruction=reconstruction,\n", - " )\n", - " )\n", - " model_visibilities = al.Visibilities(visibilities=mapped_reconstructed_visibilities)\n", - "\n", - " residual_map = dataset.data - model_visibilities\n", - " chi_squared = float(\n", - " np.sum((residual_map.real / dataset.noise_map.real) ** 2)\n", - " + np.sum((residual_map.imag / dataset.noise_map.imag) ** 2)\n", - " )\n", - "\n", - " regularization_term = float(\n", - " np.matmul(reconstruction.T, np.matmul(regularization_matrix, reconstruction))\n", - " )\n", - " log_curvature_reg_matrix_term = float(np.linalg.slogdet(curvature_reg_matrix)[1])\n", - " log_regularization_matrix_term = float(np.linalg.slogdet(regularization_matrix)[1])\n", - "\n", - " noise_normalization = float(\n", - " np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", - " + np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", - " )\n", - "\n", - " return -0.5 * (\n", - " chi_squared\n", - " + regularization_term\n", - " + log_curvature_reg_matrix_term\n", - " - log_regularization_matrix_term\n", - " + noise_normalization\n", - " )\n", - "\n", - "\n", - "print(f\"\\n=== Across all channels ===\")\n", - "per_channel_log_evidences = [per_channel_log_evidence(d) for d in dataset_list]\n", - "for c, le in enumerate(per_channel_log_evidences):\n", - " print(f\" channel {c}: log_evidence = {le:.6f}\")\n", - "\n", - "cube_log_evidence = sum(per_channel_log_evidences)\n", - "print(f\" cube log_evidence = sum(per_channel) = {cube_log_evidence:.6f}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "The whole per-channel block above is wrapped inside `al.FitInterferometer` \u2014 exactly as in the single-channel\n", - "case at `pixelization/likelihood_function.py:__Fit__`. We loop the `FitInterferometer` construction across\n", - "`dataset_list` and print the summed `fit.log_evidence` alongside our manual computation.\n", - "\n", - "The two values will agree to ~3 significant figures but typically not exactly. The small residual difference\n", - "(~0.05% relative) comes from source-code internals that this walkthrough deliberately doesn't reproduce \u2014 the\n", - "`__Log Likelihood Function: Source Code Speed Up__` section below describes the production-fast versions of\n", - "`chi_squared` and `curvature_matrix` that bypass the dense `transformed_mapping_matrix`. The pixelization\n", - "reference exhibits the same discrepancy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fits = []\n", - "print(f\"\\n=== Cross-check vs FitInterferometer ===\")\n", - "for c, dataset in enumerate(dataset_list):\n", - " fit = al.FitInterferometer(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_border_relocator=True),\n", - " )\n", - " fits.append(fit)\n", - " print(f\" channel {c}: FitInterferometer.log_evidence = {fit.log_evidence:.6f}\")\n", - "\n", - "fit_total_log_evidence = sum(fit.log_evidence for fit in fits)\n", - "print(f\" summed FitInterferometer.log_evidence = {fit_total_log_evidence:.6f}\")\n", - "print(f\" manual cube_log_evidence = {cube_log_evidence:.6f}\")\n", - "print(\n", - " f\" relative difference = \"\n", - " f\"{abs(cube_log_evidence - fit_total_log_evidence) / abs(fit_total_log_evidence):.2e}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to a datacube, this likelihood function is sampled across many candidate lens-model\n", - "parameters using a non-linear search. For the user-facing modeling story see:\n", - "\n", - " - `modeling.py` \u2014 `RectangularAdaptDensity` pixelization fit with `af.Nautilus`, the canonical entry point.\n", - " - `start_here.py` \u2014 narrative walkthrough wrapping the same fit.\n", - " - `delaunay.py` \u2014 Delaunay-pixelized source variant.\n", - " - `modeling_parametric.py` \u2014 parametric `Sersic` source variant (per-channel intensity).\n", - "\n", - "All four use `af.FactorGraphModel` to wrap a list of `AnalysisInterferometer` objects \u2014 the framework's way\n", - "of expressing the explicit cube sum we just walked through. Internally the FactorGraph routes the same\n", - "lens-model parameters to every per-channel `AnalysisInterferometer.log_likelihood_function` and sums.\n", - "\n", - "__Log Likelihood Function: Source Code Speed Up__\n", - "\n", - "The pixelization-likelihood guide's `__Log Likelihood Function: Source Code Speed Up__` section applies\n", - "unchanged per channel:\n", - "\n", - " - **Fast chi-squared:** the source code never materialises `transformed_mapping_matrix` to compute \u03c7\u00b2.\n", - " - **Sparse-operator curvature matrix:** likewise for `F`.\n", - "\n", - "For an N-channel cube these speed-ups apply per channel, multiplied by N. The deferred shared-`L\u1d40 W\u0303 L`\n", - "optimisation Aris designed goes further: when `uv_wavelengths` and `noise_map` are nearly channel-invariant\n", - "(the typical narrow-emission-line case), the curvature-matrix sandwich `L\u1d40 W\u0303 L` can be computed once and\n", - "reused across all N channels. That recovers a factor-N speed-up on top of the per-channel sparse-operator\n", - "gains, which is what brings ALMA-scale cubes back inside CPU runtime budgets.\n", - "\n", - "__Wrap Up__\n", - "\n", - "This script presented the cube likelihood function as **N independent pixelization likelihoods** + **a shared\n", - "lens model** + **a sum**.\n", - "\n", - "For deeper dives:\n", - "\n", - " - `interferometer/features/pixelization/likelihood_function.py` \u2014 the single-channel pixelization\n", - " walkthrough this script defers to for the per-channel internals.\n", - " - `modeling.py` / `start_here.py` / `delaunay.py` / `modeling_parametric.py` \u2014 user-facing modeling scripts\n", - " that wrap this likelihood in `af.FactorGraphModel` + `af.Nautilus`.\n", - " - `data_preparation.py` \u2014 how to bridge from CASA's 4D `(n_pol, n_chan, n_vis, 2)` output to the per-channel\n", - " `Interferometer` objects this walkthrough loads.\n", - "\n", - "A planned `autolens_workspace_test/scripts/jax_likelihood_functions/datacube/` folder will hold end-to-end\n", - "JAX-JIT correctness tests for the cube likelihood. The JIT-vs-eager `rtol=1e-4` regression that previously\n", - "lived in this file moves there once those test scripts land." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Datacube__\n", + "\n", + "This script provides a step-by-step guide of the **PyAutoLens** `log_likelihood_function` used to fit a\n", + "**datacube** \u2014 a list of N per-channel `Interferometer` objects sharing a single lens model \u2014 with a per-channel\n", + "pixelized source reconstruction (specifically a `RectangularAdaptDensity` mesh and `Constant` regularization\n", + "scheme).\n", + "\n", + "This script has the same aims as `interferometer/features/pixelization/likelihood_function.py`:\n", + "\n", + " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can understand\n", + " the likelihood function (including references to the previous literature from which it is defined) without\n", + " having to write large quantities of text and equations.\n", + "\n", + " - To make inversions in **PyAutoLens** less of a \"black-box\" to users.\n", + "\n", + "__Comparison To Single-Channel Pixelization Likelihood__\n", + "\n", + "A datacube fit is identical to **N independent single-channel pixelization fits** plus one wrinkle: the lens\n", + "model is shared across all channels, so the same `tracer` is used for every channel's calculation. Inside any\n", + "single channel, every linear-algebra step matches `interferometer/features/pixelization/likelihood_function.py`\n", + "line for line:\n", + "\n", + " - Same ray-tracing (same lens galaxy mass + shear).\n", + " - Same source-plane `Mapper` construction.\n", + " - Same `mapping_matrix` (image-pixel to source-pixel mappings).\n", + " - Same `transformed_mapping_matrix` (NUFFT applied per source pixel), with the *same* memory caveat as the\n", + " single-channel case \u2014 for `n_vis \u2273 10^6` per channel the matrix becomes prohibitive to store and the\n", + " sparse-operator likelihood function is the production path.\n", + " - Same `data_vector`, `curvature_matrix`, `regularization_matrix`, NNLS reconstruction.\n", + " - Same `chi_squared`, regularization term, complexity terms, noise normalisation, per-channel `log_evidence`.\n", + "\n", + "The **cube log-evidence is the sum** of the per-channel log-evidences. That's it.\n", + "\n", + "This script presents the calculation directly:\n", + "\n", + " 1. Build the shared lens galaxy + source pixelization (channel-invariant).\n", + " 2. Walk through **channel 0** in detail (one short section per pixelization-script section, cross-referencing\n", + " that script for the full derivations).\n", + " 3. Loop the per-channel calculation across all `dataset_list` and sum.\n", + " 4. Cross-check against per-channel `al.FitInterferometer.log_evidence`.\n", + "\n", + "If you haven't read `interferometer/features/pixelization/likelihood_function.py` yet, do that first. This\n", + "script defers to it for almost everything that happens inside a single channel.\n", + "\n", + "__Simplifications__\n", + "\n", + "This example uses a `RectangularAdaptDensity` mesh + `Constant` regularization \u2014 the same combination used by\n", + "the rest of the `datacube/` tutorials (`modeling.py`, `start_here.py`). The\n", + "`pixelization/likelihood_function.py` reference uses `RectangularUniform`, which is a thin subclass of\n", + "`RectangularAdaptDensity`; the linear algebra is identical and the construction code is the same. The single\n", + "behaviour difference is that `RectangularAdaptDensity` lets the mesh's pixel density adapt to the source-plane\n", + "magnification map, which gives slightly better resolution in highly-magnified regions but does not change the\n", + "likelihood-function maths at all.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The pixelization likelihood function is the most complex one in **PyAutoLens**. It is strongly advised you read\n", + "through the following first:\n", + "\n", + " - `interferometer/features/pixelization/likelihood_function.py` \u2014 full step-by-step walkthrough of the\n", + " single-channel pixelization likelihood. This script is essentially a per-channel restatement of that one\n", + " plus a sum, so the entire body below uses cross-references to its sections.\n", + " - `interferometer/light_profile/log_likelihood_function.py` \u2014 the simpler light-profile likelihood, which\n", + " introduces visibility-space inner products and the NUFFT without the pixelization linear algebra.\n", + "\n", + "__Contents__\n", + "\n", + "- **Comparison:** datacube = N independent pixelization fits + shared lens; cube log-evidence is the sum.\n", + "- **Simplifications:** `RectangularAdaptDensity` mesh, `Constant` regularization.\n", + "- **Prerequisites:** read `pixelization/likelihood_function.py` first.\n", + "- **Mesh Shape:** identical to the pixelization reference, sized to 14\u00d714 to match `modeling.py`.\n", + "- **Mask:** identical to the pixelization reference, sized to the datacube simulator's 256\u00d7256 / 0.1\u2033 grid.\n", + "- **Dataset:** load a *list* of `Interferometer` objects, one per channel (cube-specific).\n", + "- **Lens Galaxy:** identical to the pixelization reference (channel-invariant).\n", + "- **Source Galaxy Pixelization and Regularization:** identical to the pixelization reference.\n", + "- **One Channel Walkthrough:** run the full pixelization-likelihood calculation on channel 0, cross-referencing\n", + " the single-channel script for shared derivations.\n", + "- **Across All Channels:** loop the per-channel calculation across `dataset_list` and sum \u2014 the headline\n", + " cube-specific section.\n", + "- **Fit:** cross-check the manual sum against per-channel `al.FitInterferometer.log_evidence`.\n", + "- **Lens Modeling:** pointer to `modeling.py`, `start_here.py`, `delaunay.py`, `modeling_parametric.py`.\n", + "- **Log Likelihood Function: Source Code Speed Up:** per-channel fast paths apply, multiplied by N channels;\n", + " the deferred shared-`L\u1d40 W\u0303 L` optimisation recovers a factor-N speed-up when `uv_wavelengths`/`noise_map`\n", + " are nearly channel-invariant.\n", + "- **Wrap Up:** pointers to modeling scripts and the planned JIT-correctness regression tests." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Mesh Shape__`. Sized to 14\u00d714 to match the\n", + "`datacube/modeling.py` mesh." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 14\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Mask__`. Sized to match the datacube simulator's\n", + "256\u00d7256 / 0.1\u2033 real-space grid." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256), pixel_scales=0.1, radius=3.0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "This is the first genuinely cube-specific section. Instead of a single `al.Interferometer.from_fits(...)`\n", + "call, we load a *list* of `Interferometer` objects, one per channel.\n", + "\n", + "The simulator at `simulator.py` writes the cube in two on-disk layouts:\n", + "\n", + " - **Per-channel folders** (used here): one `channel_NNN/` subfolder per channel with its own\n", + " `data.fits`/`noise_map.fits`/`uv_wavelengths.fits`. Convenient when the channels are split-out already.\n", + " - **3D-FITS cube** (`{visibilities,noise_map,uv_wavelengths}_cube.fits`, each `(n_chan, n_vis, 2)`): one\n", + " file containing every channel stacked. `data_preparation.py` shows how to load this form. There is also a\n", + " 4D `(n_pol, n_chan, n_vis, 2)` CASA-like layout that requires a polarisation-collapse pre-processing step.\n", + "\n", + "Whichever loader you use, you end up with the same list `dataset_list` below. Each entry has its own\n", + "`visibilities`, `noise_map`, and `uv_wavelengths`; the lens galaxy is channel-invariant and the per-channel\n", + "sources differ in `intensity` and `centre` (see `simulator.py`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"datacube\"\n", + "dataset_name = \"sim_simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the cube isn't on disk yet, run the simulator. This makes the script runnable on a fresh checkout." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "channel_paths = sorted(\n", + " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", + ")\n", + "\n", + "dataset_list = [\n", + " al.Interferometer.from_fits(\n", + " data_path=channel_path / \"data.fits\",\n", + " noise_map_path=channel_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + " )\n", + " for channel_path in channel_paths\n", + "]\n", + "\n", + "n_channels = len(dataset_list)\n", + "n_vis = int(dataset_list[0].uv_wavelengths.shape[0])\n", + "print(f\"Loaded {n_channels} channels, {n_vis} visibilities per channel.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot the dirty images of the first and last channels to see the per-channel data the cube fit will\n", + "work with. The source intensity and centre vary across the cube, so the dirty images differ even though the\n", + "lens model is the same." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset_list[0])\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset_list[-1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Same as `pixelization/likelihood_function.py:__Over Sampling__`. Interferometer pixelizations do not use\n", + "over-sampling.\n", + "\n", + "__Masked Image Grid__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Masked Image Grid__`. Every channel uses the same\n", + "real-space mask, so `dataset.grids.pixelization` is channel-invariant." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=dataset_list[0].grids.pixelization, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Lens Galaxy__`. The lens galaxy is channel-invariant\n", + "(this is the entire reason datacube modeling can share a single non-linear search across all channels \u2014 the\n", + "mass model doesn't depend on frequency)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Pixelization and Regularization__\n", + "\n", + "Same as `pixelization/likelihood_function.py:__Source Galaxy Pixelization and Regularization__`, with\n", + "`RectangularAdaptDensity` substituted for `RectangularUniform`. The classes share the same construction\n", + "machinery \u2014 `RectangularUniform` is a subclass of `RectangularAdaptDensity` \u2014 so all of the mesh-grid and\n", + "mapper code below is unchanged.\n", + "\n", + "The same source pixelization is used for every channel. Each channel runs its own linear inversion against\n", + "this shared pixelization (which is what gives each channel an independent source-plane reconstruction)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.RectangularAdaptDensity(shape=mesh_shape),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + ")\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Ray Tracing__`. Channel-invariant because the lens is\n", + "channel-invariant; the same `tracer` is used by every per-channel calculation below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "==================================================================================================\n", + " ONE CHANNEL WALKTHROUGH\n", + "==================================================================================================\n", + "\n", + "We now walk through the full pixelization-likelihood calculation for **channel 0**. Every section heading\n", + "below matches a section in `pixelization/likelihood_function.py` and the calculation is the same. Each\n", + "section has a one-line cross-reference, the code (so this script runs end-to-end), and any cube-specific\n", + "notes that don't apply in the single-channel case." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "dataset = dataset_list[0]\n", + "print(f\"\\n=== Channel 0 walkthrough ===\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Border Relocation__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Border Relocation__`. Run per channel because the traced\n", + "grid feeding into it depends on the per-channel `dataset.grids.pixelization` \u2014 but `dataset.grids.pixelization`\n", + "is itself channel-invariant in this cube (every channel uses the same `real_space_mask`), so the relocated\n", + "grid actually only needs to be computed once. We compute it inside the loop below for clarity." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", + "\n", + "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", + " grid=dataset.grids.pixelization\n", + ")[-1]\n", + "\n", + "border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)\n", + "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", + "\n", + "aplt.plot_grid(grid=relocated_grid, title=\"Channel 0 relocated traced grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Pixel Centre Calculation__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Source Pixel Centre Calculation__`. The source-plane mesh\n", + "overlays the relocated traced grid; since the relocated grid is the same for every channel (shared mask,\n", + "shared tracer), the source-plane mesh grid is channel-invariant too." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoarray.inversion.mesh.mesh.rectangular_adapt_density import overlay_grid_from\n", + "\n", + "mesh_grid = overlay_grid_from(\n", + " shape_native=mesh_shape, grid=al.Grid2DIrregular(relocated_grid)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolation / Mapper / Mapping Matrix__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Interpolation__`, `__Mapper__`, and\n", + "`__Mapping Matrix__`. All channel-invariant for the cube \u2014 they depend only on the shared `tracer`,\n", + "mesh shape and mask. We compute them once and reuse across channels in the loop below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "interpolator = pixelization.mesh.interpolator_from(\n", + " source_plane_data_grid=relocated_grid,\n", + " source_plane_mesh_grid=mesh_grid,\n", + ")\n", + "mapper = al.Mapper(interpolator=interpolator)\n", + "\n", + "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", + " pix_indexes_for_sub_slim_index=mapper.pix_indexes_for_sub_slim_index,\n", + " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,\n", + " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,\n", + " pixels=mapper.pixels,\n", + " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", + " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", + " sub_fraction=mapper.over_sampler.sub_fraction,\n", + ")\n", + "\n", + "print(\n", + " f\" mapping_matrix shape: {mapping_matrix.shape} (real-space pixels x source pixels)\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Transformed Mapping Matrix ($f$)__\n", + "\n", + "For a single channel this is identical to `pixelization/likelihood_function.py:__Transformed Mapping Matrix__`.\n", + "\n", + "**This is where per-channel structure starts to appear.** Each channel has its own `uv_wavelengths` and\n", + "therefore its own NUFFT operator (`dataset.transformer`). The shape of the output is the same for every\n", + "channel in our simulator (every channel has the same `n_vis`), but the values differ \u2014 the same source pixel\n", + "maps to different uv-plane visibilities in different channels because the baselines differ.\n", + "\n", + "The same memory caveat applies *per channel*: for `n_vis \u2273 10^6` per channel, this matrix is many GB per\n", + "channel; for an N-channel cube the total memory footprint is N\u00d7 the single-channel one. The sparse-operator\n", + "likelihood function \u2014 see the `apply_sparse_operator` calls in `modeling.py` \u2014 is the production path that\n", + "avoids ever materialising this matrix. The deferred shared-`L\u1d40 W\u0303 L` optimisation goes further: it computes\n", + "the curvature-matrix sandwich once and reuses it across channels when `uv_wavelengths`/`noise_map` are\n", + "nearly channel-invariant (which they typically are for narrow emission lines)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", + " mapping_matrix=mapping_matrix\n", + ")\n", + "\n", + "print(\n", + " f\" transformed_mapping_matrix shape: {transformed_mapping_matrix.shape} \"\n", + " f\"(n_vis x source pixels), dtype: {transformed_mapping_matrix.dtype}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Data Vector (D)__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Data Vector (D)__`. Per channel because each channel has\n", + "its own `dataset.data` (visibilities) and `dataset.noise_map`; the construction formula is unchanged." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_vector = (\n", + " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " visibilities=dataset.data,\n", + " noise_map=dataset.noise_map,\n", + " )\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Matrix (F)__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Curvature Matrix (F)__`. The matrix `F` depends only on\n", + "`transformed_mapping_matrix` and `noise_map` \u2014 both of which are channel-specific because of the per-channel\n", + "NUFFT and noise \u2014 but for the typical narrow-emission-line case where `uv_wavelengths` and `noise_map` change\n", + "very little across the line, `F` is nearly channel-invariant. This is the matrix Aris's deferred\n", + "shared-`L\u1d40 W\u0303 L` optimisation reuses across channels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.real,\n", + " noise_map=dataset.noise_map.real,\n", + ")\n", + "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.imag,\n", + " noise_map=dataset.noise_map.imag,\n", + ")\n", + "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Regularization Matrix (H)__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Regularization Matrix (H)__`. Channel-invariant because\n", + "the regularization scheme + mesh structure are channel-invariant \u2014 `H` only depends on which source pixels\n", + "neighbour which." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", + " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", + " neighbors=mapper.neighbors,\n", + " neighbors_sizes=mapper.neighbors.sizes,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__F + \u03bbH / Galaxy Reconstruction (s)__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__F + Lamdba H__` and `__Galaxy Reconstruction (s)__`. The\n", + "linear system `s = (F + \u03bbH)\u207b\u00b9 D` is solved per channel, producing each channel's source-plane reconstruction.\n", + "\n", + "Per-channel reconstructions are what make the cube fit physically interesting: an emission line that\n", + "brightens-and-fades across the cube produces a sequence of source-plane reconstructions whose total flux\n", + "traces the line profile, while the lens mass model stays fixed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)\n", + "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visibilities Reconstruction__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Visibilities Reconstruction__`. Per channel because the\n", + "model visibilities live in the channel's own uv-plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_visibilities = (\n", + " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " reconstruction=reconstruction,\n", + " )\n", + ")\n", + "mapped_reconstructed_visibilities = al.Visibilities(\n", + " visibilities=mapped_reconstructed_visibilities\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function \u2014 Five Terms__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Likelihood Function__`. The same five-term formula\n", + "applies per channel:\n", + "\n", + " -2 ln \u03b5_c = \u03c7\u00b2_c + s_c\u1d40 H s_c + ln det(F_c + H) - ln det(H) + \u03a3_j ln (2\u03c0 \u03c3\u00b2_{c,j})\n", + "\n", + "where `c` indexes channels. The cube log-evidence is `\u03a3_c log_evidence_c`, computed in the\n", + "\"Across All Channels\" section below.\n", + "\n", + "__Chi Squared__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Chi Squared__`. Per channel \u2014 each channel has its own\n", + "visibilities, noise, and model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_visibilities = mapped_reconstructed_visibilities\n", + "residual_map = dataset.data - model_visibilities\n", + "\n", + "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", + "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", + "chi_squared = np.sum(chi_squared_map_real) + np.sum(chi_squared_map_imag)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Regularization Term__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Regularization Term__`. Per channel via the per-channel\n", + "`reconstruction`; `H` is channel-invariant." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "regularization_term = np.matmul(\n", + " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Complexity Terms__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Complexity Terms__`. The `ln det(F + \u03bbH)` term is\n", + "per-channel (`F` is per-channel); `ln det(H)` is channel-invariant." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", + "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalisation Term__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Noise Normalization Term__`. Per channel.\n", + "\n", + "**Cube-specific note**: if you did polarisation-collapse by *averaging* (see `data_preparation.py`), the\n", + "visibility noise map already incorporates the sqrt(2) noise reduction relative to a single polarisation. The\n", + "noise-normalisation term you compute below matches whatever your data-prep produced \u2014 it doesn't double-count\n", + "the polarisation averaging. If you concatenated polarisations instead, `n_vis` doubled and the term naturally\n", + "scales with it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization_real = np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", + "noise_normalization_imag = np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", + "noise_normalization = noise_normalization_real + noise_normalization_imag" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood (Channel 0)__\n", + "\n", + "Identical to `pixelization/likelihood_function.py:__Calculate The Log Likelihood__`. The result is the\n", + "log-evidence of channel 0 only." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_evidence_channel_0 = float(\n", + " -0.5\n", + " * (\n", + " chi_squared\n", + " + regularization_term\n", + " + log_curvature_reg_matrix_term\n", + " - log_regularization_matrix_term\n", + " + noise_normalization\n", + " )\n", + ")\n", + "print(f\" channel 0 log_evidence: {log_evidence_channel_0:.6f}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "==================================================================================================\n", + " ACROSS ALL CHANNELS\n", + "==================================================================================================\n", + "\n", + "This is the headline cube-specific section. We loop the per-channel calculation across `dataset_list` and\n", + "sum to get the cube log-evidence.\n", + "\n", + "Almost everything inside the loop is channel-invariant \u2014 the `tracer`, `mapper`, `mapping_matrix`, and\n", + "`regularization_matrix`. Only the per-channel `dataset` (visibilities, noise_map, uv_wavelengths) changes, so\n", + "inside the loop we only redo the channel-dependent steps:\n", + "\n", + " - `transformed_mapping_matrix` (depends on `dataset.transformer`).\n", + " - `data_vector` (depends on `dataset.data` and `dataset.noise_map`).\n", + " - `curvature_matrix` (depends on `transformed_mapping_matrix` and `dataset.noise_map`).\n", + " - `reconstruction` (depends on `data_vector` and `curvature_matrix`).\n", + " - Visibilities-space residuals, \u03c7\u00b2, regularization term, complexity terms, noise normalisation.\n", + "\n", + "This is exactly what `af.FactorGraphModel` does internally for the modeling scripts: it feeds the same\n", + "lens-model parameters to every per-channel `AnalysisInterferometer.log_likelihood_function`, and the search\n", + "sees the sum." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def per_channel_log_evidence(dataset):\n", + " \"\"\"Compute the log-evidence of a single channel given the channel-invariant `tracer`, `mapper`,\n", + " `mapping_matrix` and `regularization_matrix` defined above.\n", + "\n", + " All steps mirror `pixelization/likelihood_function.py` line by line; only the dataset (visibilities,\n", + " noise map, uv_wavelengths) varies.\n", + " \"\"\"\n", + "\n", + " transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", + " mapping_matrix=mapping_matrix\n", + " )\n", + "\n", + " data_vector = al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " visibilities=dataset.data,\n", + " noise_map=dataset.noise_map,\n", + " )\n", + "\n", + " real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.real,\n", + " noise_map=dataset.noise_map.real,\n", + " )\n", + " imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.imag,\n", + " noise_map=dataset.noise_map.imag,\n", + " )\n", + " curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", + " curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)\n", + "\n", + " reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)\n", + "\n", + " mapped_reconstructed_visibilities = (\n", + " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " reconstruction=reconstruction,\n", + " )\n", + " )\n", + " model_visibilities = al.Visibilities(visibilities=mapped_reconstructed_visibilities)\n", + "\n", + " residual_map = dataset.data - model_visibilities\n", + " chi_squared = float(\n", + " np.sum((residual_map.real / dataset.noise_map.real) ** 2)\n", + " + np.sum((residual_map.imag / dataset.noise_map.imag) ** 2)\n", + " )\n", + "\n", + " regularization_term = float(\n", + " np.matmul(reconstruction.T, np.matmul(regularization_matrix, reconstruction))\n", + " )\n", + " log_curvature_reg_matrix_term = float(np.linalg.slogdet(curvature_reg_matrix)[1])\n", + " log_regularization_matrix_term = float(np.linalg.slogdet(regularization_matrix)[1])\n", + "\n", + " noise_normalization = float(\n", + " np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", + " + np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", + " )\n", + "\n", + " return -0.5 * (\n", + " chi_squared\n", + " + regularization_term\n", + " + log_curvature_reg_matrix_term\n", + " - log_regularization_matrix_term\n", + " + noise_normalization\n", + " )\n", + "\n", + "\n", + "print(f\"\\n=== Across all channels ===\")\n", + "per_channel_log_evidences = [per_channel_log_evidence(d) for d in dataset_list]\n", + "for c, le in enumerate(per_channel_log_evidences):\n", + " print(f\" channel {c}: log_evidence = {le:.6f}\")\n", + "\n", + "cube_log_evidence = sum(per_channel_log_evidences)\n", + "print(f\" cube log_evidence = sum(per_channel) = {cube_log_evidence:.6f}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "The whole per-channel block above is wrapped inside `al.FitInterferometer` \u2014 exactly as in the single-channel\n", + "case at `pixelization/likelihood_function.py:__Fit__`. We loop the `FitInterferometer` construction across\n", + "`dataset_list` and print the summed `fit.log_evidence` alongside our manual computation.\n", + "\n", + "The two values will agree to ~3 significant figures but typically not exactly. The small residual difference\n", + "(~0.05% relative) comes from source-code internals that this walkthrough deliberately doesn't reproduce \u2014 the\n", + "`__Log Likelihood Function: Source Code Speed Up__` section below describes the production-fast versions of\n", + "`chi_squared` and `curvature_matrix` that bypass the dense `transformed_mapping_matrix`. The pixelization\n", + "reference exhibits the same discrepancy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fits = []\n", + "print(f\"\\n=== Cross-check vs FitInterferometer ===\")\n", + "for c, dataset in enumerate(dataset_list):\n", + " fit = al.FitInterferometer(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_border_relocator=True),\n", + " )\n", + " fits.append(fit)\n", + " print(f\" channel {c}: FitInterferometer.log_evidence = {fit.log_evidence:.6f}\")\n", + "\n", + "fit_total_log_evidence = sum(fit.log_evidence for fit in fits)\n", + "print(f\" summed FitInterferometer.log_evidence = {fit_total_log_evidence:.6f}\")\n", + "print(f\" manual cube_log_evidence = {cube_log_evidence:.6f}\")\n", + "print(\n", + " f\" relative difference = \"\n", + " f\"{abs(cube_log_evidence - fit_total_log_evidence) / abs(fit_total_log_evidence):.2e}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to a datacube, this likelihood function is sampled across many candidate lens-model\n", + "parameters using a non-linear search. For the user-facing modeling story see:\n", + "\n", + " - `modeling.py` \u2014 `RectangularAdaptDensity` pixelization fit with `af.Nautilus`, the canonical entry point.\n", + " - `start_here.py` \u2014 narrative walkthrough wrapping the same fit.\n", + " - `delaunay.py` \u2014 Delaunay-pixelized source variant.\n", + " - `modeling_parametric.py` \u2014 parametric `Sersic` source variant (per-channel intensity).\n", + "\n", + "All four use `af.FactorGraphModel` to wrap a list of `AnalysisInterferometer` objects \u2014 the framework's way\n", + "of expressing the explicit cube sum we just walked through. Internally the FactorGraph routes the same\n", + "lens-model parameters to every per-channel `AnalysisInterferometer.log_likelihood_function` and sums.\n", + "\n", + "__Log Likelihood Function: Source Code Speed Up__\n", + "\n", + "The pixelization-likelihood guide's `__Log Likelihood Function: Source Code Speed Up__` section applies\n", + "unchanged per channel:\n", + "\n", + " - **Fast chi-squared:** the source code never materialises `transformed_mapping_matrix` to compute \u03c7\u00b2.\n", + " - **Sparse-operator curvature matrix:** likewise for `F`.\n", + "\n", + "For an N-channel cube these speed-ups apply per channel, multiplied by N. The deferred shared-`L\u1d40 W\u0303 L`\n", + "optimisation Aris designed goes further: when `uv_wavelengths` and `noise_map` are nearly channel-invariant\n", + "(the typical narrow-emission-line case), the curvature-matrix sandwich `L\u1d40 W\u0303 L` can be computed once and\n", + "reused across all N channels. That recovers a factor-N speed-up on top of the per-channel sparse-operator\n", + "gains, which is what brings ALMA-scale cubes back inside CPU runtime budgets.\n", + "\n", + "__Wrap Up__\n", + "\n", + "This script presented the cube likelihood function as **N independent pixelization likelihoods** + **a shared\n", + "lens model** + **a sum**.\n", + "\n", + "For deeper dives:\n", + "\n", + " - `interferometer/features/pixelization/likelihood_function.py` \u2014 the single-channel pixelization\n", + " walkthrough this script defers to for the per-channel internals.\n", + " - `modeling.py` / `start_here.py` / `delaunay.py` / `modeling_parametric.py` \u2014 user-facing modeling scripts\n", + " that wrap this likelihood in `af.FactorGraphModel` + `af.Nautilus`.\n", + " - `data_preparation.py` \u2014 how to bridge from CASA's 4D `(n_pol, n_chan, n_vis, 2)` output to the per-channel\n", + " `Interferometer` objects this walkthrough loads.\n", + "\n", + "A planned `autolens_workspace_test/scripts/jax_likelihood_functions/datacube/` folder will hold end-to-end\n", + "JAX-JIT correctness tests for the cube likelihood. The JIT-vs-eager `rtol=1e-4` regression that previously\n", + "lived in this file moves there once those test scripts land." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/datacube/modeling.ipynb b/notebooks/interferometer/features/datacube/modeling.ipynb index 31956b61e..83b49dbbc 100644 --- a/notebooks/interferometer/features/datacube/modeling.ipynb +++ b/notebooks/interferometer/features/datacube/modeling.ipynb @@ -1,481 +1,518 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Datacube\n", - "===================\n", - "\n", - "This script fits a list of `Interferometer` channels \u2014 a \"datacube\" \u2014 with a single shared lens model and a\n", - "per-channel pixelized source reconstruction. Each channel is an independent `Interferometer` dataset; the\n", - "`af.FactorGraphModel` ties them together by feeding the same lens parameters into every channel's\n", - "`AnalysisInterferometer.log_likelihood_function` and summing the per-channel log-evidences.\n", - "\n", - "A datacube modeled this way captures spatially-resolved spectral-line emission: every channel reconstructs its\n", - "own source-plane pixelization, so an emission line that brightens-and-fades across the cube produces a sequence\n", - "of source-plane reconstructions whose total flux traces the line profile while the lens mass stays fixed.\n", - "\n", - "This script is the focused-modeling sibling of `start_here.py`. Read `start_here.py` first for the narrative\n", - "walkthrough; this file is the one to copy and adapt for your own cube.\n", - "\n", - "__Contents__\n", - "\n", - "- **Mask:** Define the 2D real-space mask applied to every channel.\n", - "- **Dataset:** Where the per-channel cube lives on disk and how to point this script at your own.\n", - "- **Dataset Auto-Simulation:** Run `simulator.py` automatically if the cube isn't already on disk.\n", - "- **Dataset Loading:** Loop over the channel folders and load each as an `Interferometer` object.\n", - "- **Sparse Operators:** Pre-compute per-channel sparse-operator matrices used by the pixelized source inversion.\n", - "- **Settings:** Disable the positive-only solver so visibility-space inversions can take negative pixel values.\n", - "- **Mesh Shape:** Pixelization mesh size \u2014 fixed before modeling because JAX needs static-shape arrays.\n", - "- **Model:** Compose the shared `Isothermal + ExternalShear` lens and pixelized source.\n", - "- **Per-Channel Analyses:** One `AnalysisInterferometer` per channel, with `use_jax=True`.\n", - "- **FactorGraph:** Wrap each analysis in an `AnalysisFactor` and combine via `af.FactorGraphModel`.\n", - "- **Search:** Configure the `Nautilus` non-linear search.\n", - "- **Model Fit:** Run the fit. Per-channel inversion cost dominates runtime \u2014 see notes inline.\n", - "- **Wrap Up:** Summary of the script and pointers to the JAX likelihood walkthrough in\n", - " ``autolens_workspace_developer/datacube/likelihood_function.py``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import subprocess\n", - "import sys\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Every channel uses the same `real_space_mask` \u2014 the lens galaxy and source position don't depend on frequency,\n", - "so masking once is correct. The mask radius is generous enough to contain the lensed source's full extent." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "The datacube lives under `dataset/interferometer/datacube//`, with one subfolder per channel\n", - "(`channel_000/`, `channel_001/`, ...). To point this script at your own cube, drop your channel folders in\n", - "alongside the reference cube and update `dataset_name`. Each channel folder must contain `data.fits`,\n", - "`noise_map.fits` and `uv_wavelengths.fits` in the shape produced by `al.SimulatorInterferometer`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"datacube\"\n", - "dataset_name = \"sim_simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Loading__\n", - "\n", - "Build the cube by loading each channel folder as an `Interferometer` object. The result is a Python list \u2014 no\n", - "new dataset class involved. Channels are discovered by sorted directory listing, so you can add channels by\n", - "simply dropping more `channel_NNN/` folders in." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "channel_paths = sorted(\n", - " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", - ")\n", - "print(f\"Loading {len(channel_paths)} channels from {dataset_path}\")\n", - "\n", - "dataset_list = [\n", - " al.Interferometer.from_fits(\n", - " data_path=channel_path / \"data.fits\",\n", - " noise_map_path=channel_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - " )\n", - " for channel_path in channel_paths\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "Pixelized source modeling uses sparse linear algebra to keep memory and runtime manageable. We pre-compute a\n", - "sparse-operator matrix per channel so each `AnalysisInterferometer.log_likelihood_function` reuses it directly\n", - "during the fit. For SMA-scale data this finishes in seconds per channel; for ALMA-scale cubes it can take\n", - "minutes per channel on CPU, in which case see `pixelization/many_visibilities_preparation.py` for how to\n", - "compute and cache them once." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " dataset.apply_sparse_operator(use_jax=True, show_progress=False)\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Load the cube's multiple-image positions (saved by `simulator.py`) and wrap them in an `al.PositionsLH` penalty.\n", - "For pixelized fits this is essentially required: without the penalty, the search routinely converges on\n", - "demagnified-source local maxima where the source pixels are reconstructed in low-magnification regions of the\n", - "source plane that fit the noise rather than the lensed signal.\n", - "\n", - "The lens model is shared across every channel via the FactorGraph, so a single `PositionsLH` (built once and\n", - "passed to every per-channel analysis) applies the same global constraint everywhere." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "Interferometer pixelizations disable the positive-only inversion solver \u2014 the visibility measurement process\n", - "can produce genuinely negative dirty-image pixel values, so the source-plane reconstruction must be allowed\n", - "to go negative." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The pixelization mesh shape is fixed before modeling because JAX needs static array shapes. We use a\n", - "14 x 14 ``RectangularAdaptDensity`` mesh \u2014 small enough to keep the prototype cheap, large enough to capture the\n", - "emission-line source morphology produced by the simulator. `RectangularAdaptDensity` adapts the source-plane\n", - "pixel density to the lensing magnification map, giving more pixels to the highly-magnified source-plane regions\n", - "where the lensed signal is concentrated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 14\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "The lens galaxy is a shared `Isothermal + ExternalShear`, identical across every channel. The source galaxy is\n", - "a `Pixelization` with a `RectangularAdaptDensity` mesh and `Constant` regularization \u2014 the inversion runs\n", - "independently per channel inside each `AnalysisInterferometer`, giving each channel its own source-plane\n", - "reconstruction without adding any model parameters.\n", - "\n", - "There are no per-channel free parameters: every prior in this base model is identified across factors when the\n", - "`FactorGraph` deduplicates them below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source (pixelization, no free priors):\n", - "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", - "regularization = af.Model(al.reg.Constant)\n", - "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "# Overall lens model:\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Channel Analyses__\n", - "\n", - "One `AnalysisInterferometer` per channel, all with `use_jax=True` so the FactorGraph fit runs on the JAX\n", - "backend. The shared `positions_likelihood` is passed to every analysis \u2014 same penalty, every channel.\n", - "\n", - "__Shared Preloads__\n", - "\n", - "Because every channel shares the same lens model, a large fraction of each channel's likelihood is *identical\n", - "work*. Ray-tracing the lens model, building the source-plane mapper (its mesh and mapping matrix `L`) and the\n", - "curvature matrix `F = L\u1d40W\u0303L` are the dominant inversion-setup costs \u2014 and they are the same for every channel.\n", - "Recomputing them once per channel is pure waste.\n", - "\n", - "Setting `shared_preloads=True` opts each analysis into the `FactorGraphModel` shared-state mechanism: the\n", - "channel-invariant inversion-setup quantities (the mapper and `F`) are computed **once** on the lead channel\n", - "and reused by every other channel, instead of being rebuilt `N` times. For a many-channel cube this collapses\n", - "the dominant inversion-setup cost from `N \u00d7` to `1 \u00d7`, a large speed-up that grows with the number of channels.\n", - "\n", - "This is only correct when those quantities really are channel-invariant \u2014 i.e. when the lens model is shared\n", - "(it is here) **and** the `uv_wavelengths` and `noise_map` are the same for every channel (the curvature matrix\n", - "depends on them through `W\u0303`). The datacube simulated in `simulator.py` is built exactly this way (identical\n", - "`uv_wavelengths` and noise across channels \u2014 the narrow-emission-line regime). If your cube has per-channel\n", - "`uv`/noise that differ significantly, leave `shared_preloads=False` (the default) so each channel computes its\n", - "own inversion \u2014 preloading an invalid quantity would silently corrupt the likelihood. The shared and unshared\n", - "paths are asserted to give identical likelihoods in\n", - "`autolens_workspace_test/scripts/jax_likelihood_functions/datacube/shared_preloads.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [\n", - " al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " settings=settings,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " use_jax=True,\n", - " shared_preloads=True,\n", - " )\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__FactorGraph__\n", - "\n", - "Each analysis is wrapped in an `af.AnalysisFactor` paired with a deep copy of the base model. With no per-factor\n", - "prior overrides, every prior is identified across factors \u2014 so the global model has the same dimensionality as\n", - "the single-channel base model. ``af.FactorGraphModel(..., use_jax=True)`` sums the per-channel log-evidences\n", - "internally, which is exactly the cube log-likelihood you'd write by hand." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = [\n", - " af.AnalysisFactor(prior_model=model.copy(), analysis=analysis)\n", - " for analysis in analysis_list\n", - "]\n", - "\n", - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - "print(f\" channels in factor graph: {len(analysis_factor_list)}\")\n", - "print(\n", - " f\" global model free parameters: {factor_graph.global_prior_model.total_free_parameters}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "`Nautilus` is the standard non-linear search for PyAutoLens. Datacube fits typically need fewer live points\n", - "than imaging fits because the lens dimensionality is unchanged \u2014 only the per-channel inversions multiply.\n", - "Tune `n_live` for your problem." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\") / \"datacube\",\n", - " name=\"modeling\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20,\n", - " iterations_per_quick_update=50000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "Pass the factor graph's `global_prior_model` as the model and the factor graph itself as the analysis \u2014 that's\n", - "the same shape you'd use for any multi-dataset PyAutoFit fit.\n", - "\n", - "**Run time on CPU is dominated by the per-channel inversion.** A 4-channel SMA-scale cube finishes in a few\n", - "hours on CPU; ALMA-scale cubes with 50+ channels need GPU acceleration to complete in reasonable time. The\n", - "``L\u1d40 W\u0303 L`` shared-precompute optimisation (Aris's design \u2014 exploit the fact that ``uv_wavelengths`` and\n", - "``noise_map`` change very little channel-to-channel) is the natural follow-up that brings ALMA-scale cubes\n", - "back inside the budget." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a while\n", - " for a datacube fit, since per-channel inversions multiply the per-likelihood cost.\n", - " \"\"\"\n", - ")\n", - "\n", - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The result returned by `search.fit` is a list of per-factor results \u2014 one entry per channel \u2014 each carrying its\n", - "own `FitInterferometer` against the maximum-likelihood lens model. Use them to inspect the per-channel source\n", - "reconstructions, dirty images, and residuals.\n", - "\n", - "For a step-by-step look at how the per-channel likelihood is summed (and at the JAX JIT pattern that drives\n", - "this fit), see ``autolens_workspace_developer/datacube/likelihood_function.py``." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Datacube\n", + "===================\n", + "\n", + "This script fits a list of `Interferometer` channels \u2014 a \"datacube\" \u2014 with a single shared lens model and a\n", + "per-channel pixelized source reconstruction. Each channel is an independent `Interferometer` dataset; the\n", + "`af.FactorGraphModel` ties them together by feeding the same lens parameters into every channel's\n", + "`AnalysisInterferometer.log_likelihood_function` and summing the per-channel log-evidences.\n", + "\n", + "A datacube modeled this way captures spatially-resolved spectral-line emission: every channel reconstructs its\n", + "own source-plane pixelization, so an emission line that brightens-and-fades across the cube produces a sequence\n", + "of source-plane reconstructions whose total flux traces the line profile while the lens mass stays fixed.\n", + "\n", + "This script is the focused-modeling sibling of `start_here.py`. Read `start_here.py` first for the narrative\n", + "walkthrough; this file is the one to copy and adapt for your own cube.\n", + "\n", + "__Contents__\n", + "\n", + "- **Mask:** Define the 2D real-space mask applied to every channel.\n", + "- **Dataset:** Where the per-channel cube lives on disk and how to point this script at your own.\n", + "- **Dataset Auto-Simulation:** Run `simulator.py` automatically if the cube isn't already on disk.\n", + "- **Dataset Loading:** Loop over the channel folders and load each as an `Interferometer` object.\n", + "- **Sparse Operators:** Pre-compute per-channel sparse-operator matrices used by the pixelized source inversion.\n", + "- **Settings:** Disable the positive-only solver so visibility-space inversions can take negative pixel values.\n", + "- **Mesh Shape:** Pixelization mesh size \u2014 fixed before modeling because JAX needs static-shape arrays.\n", + "- **Model:** Compose the shared `Isothermal + ExternalShear` lens and pixelized source.\n", + "- **Per-Channel Analyses:** One `AnalysisInterferometer` per channel, with `use_jax=True`.\n", + "- **FactorGraph:** Wrap each analysis in an `AnalysisFactor` and combine via `af.FactorGraphModel`.\n", + "- **Search:** Configure the `Nautilus` non-linear search.\n", + "- **Model Fit:** Run the fit. Per-channel inversion cost dominates runtime \u2014 see notes inline.\n", + "- **Wrap Up:** Summary of the script and pointers to the JAX likelihood walkthrough in\n", + " ``autolens_workspace_developer/datacube/likelihood_function.py``." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import subprocess\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Every channel uses the same `real_space_mask` \u2014 the lens galaxy and source position don't depend on frequency,\n", + "so masking once is correct. The mask radius is generous enough to contain the lensed source's full extent." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "The datacube lives under `dataset/interferometer/datacube//`, with one subfolder per channel\n", + "(`channel_000/`, `channel_001/`, ...). To point this script at your own cube, drop your channel folders in\n", + "alongside the reference cube and update `dataset_name`. Each channel folder must contain `data.fits`,\n", + "`noise_map.fits` and `uv_wavelengths.fits` in the shape produced by `al.SimulatorInterferometer`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"datacube\"\n", + "dataset_name = \"sim_simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Loading__\n", + "\n", + "Build the cube by loading each channel folder as an `Interferometer` object. The result is a Python list \u2014 no\n", + "new dataset class involved. Channels are discovered by sorted directory listing, so you can add channels by\n", + "simply dropping more `channel_NNN/` folders in." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "channel_paths = sorted(\n", + " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", + ")\n", + "print(f\"Loading {len(channel_paths)} channels from {dataset_path}\")\n", + "\n", + "dataset_list = [\n", + " al.Interferometer.from_fits(\n", + " data_path=channel_path / \"data.fits\",\n", + " noise_map_path=channel_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + " )\n", + " for channel_path in channel_paths\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "Pixelized source modeling uses sparse linear algebra to keep memory and runtime manageable. We pre-compute a\n", + "sparse-operator matrix per channel so each `AnalysisInterferometer.log_likelihood_function` reuses it directly\n", + "during the fit. For SMA-scale data this finishes in seconds per channel; for ALMA-scale cubes it can take\n", + "minutes per channel on CPU, in which case see `pixelization/many_visibilities_preparation.py` for how to\n", + "compute and cache them once." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " dataset.apply_sparse_operator(use_jax=True, show_progress=False)\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Load the cube's multiple-image positions (saved by `simulator.py`) and wrap them in an `al.PositionsLH` penalty.\n", + "For pixelized fits this is essentially required: without the penalty, the search routinely converges on\n", + "demagnified-source local maxima where the source pixels are reconstructed in low-magnification regions of the\n", + "source plane that fit the noise rather than the lensed signal.\n", + "\n", + "The lens model is shared across every channel via the FactorGraph, so a single `PositionsLH` (built once and\n", + "passed to every per-channel analysis) applies the same global constraint everywhere." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "Interferometer pixelizations disable the positive-only inversion solver \u2014 the visibility measurement process\n", + "can produce genuinely negative dirty-image pixel values, so the source-plane reconstruction must be allowed\n", + "to go negative." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The pixelization mesh shape is fixed before modeling because JAX needs static array shapes. We use a\n", + "14 x 14 ``RectangularAdaptDensity`` mesh \u2014 small enough to keep the prototype cheap, large enough to capture the\n", + "emission-line source morphology produced by the simulator. `RectangularAdaptDensity` adapts the source-plane\n", + "pixel density to the lensing magnification map, giving more pixels to the highly-magnified source-plane regions\n", + "where the lensed signal is concentrated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 14\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "The lens galaxy is a shared `Isothermal + ExternalShear`, identical across every channel. The source galaxy is\n", + "a `Pixelization` with a `RectangularAdaptDensity` mesh and `Constant` regularization \u2014 the inversion runs\n", + "independently per channel inside each `AnalysisInterferometer`, giving each channel its own source-plane\n", + "reconstruction without adding any model parameters.\n", + "\n", + "There are no per-channel free parameters: every prior in this base model is identified across factors when the\n", + "`FactorGraph` deduplicates them below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source (pixelization, no free priors):\n", + "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", + "regularization = af.Model(al.reg.Constant)\n", + "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "# Overall lens model:\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Channel Analyses__\n", + "\n", + "One `AnalysisInterferometer` per channel, all with `use_jax=True` so the FactorGraph fit runs on the JAX\n", + "backend. The shared `positions_likelihood` is passed to every analysis \u2014 same penalty, every channel.\n", + "\n", + "__Shared Preloads__\n", + "\n", + "Because every channel shares the same lens model, a large fraction of each channel's likelihood is *identical\n", + "work*. Ray-tracing the lens model, building the source-plane mapper (its mesh and mapping matrix `L`) and the\n", + "curvature matrix `F = L\u1d40W\u0303L` are the dominant inversion-setup costs \u2014 and they are the same for every channel.\n", + "Recomputing them once per channel is pure waste.\n", + "\n", + "Setting `shared_preloads=True` opts each analysis into the `FactorGraphModel` shared-state mechanism: the\n", + "channel-invariant inversion-setup quantities (the mapper and `F`) are computed **once** on the lead channel\n", + "and reused by every other channel, instead of being rebuilt `N` times. For a many-channel cube this collapses\n", + "the dominant inversion-setup cost from `N \u00d7` to `1 \u00d7`, a large speed-up that grows with the number of channels.\n", + "\n", + "This is only correct when those quantities really are channel-invariant \u2014 i.e. when the lens model is shared\n", + "(it is here) **and** the `uv_wavelengths` and `noise_map` are the same for every channel (the curvature matrix\n", + "depends on them through `W\u0303`). The datacube simulated in `simulator.py` is built exactly this way (identical\n", + "`uv_wavelengths` and noise across channels \u2014 the narrow-emission-line regime). If your cube has per-channel\n", + "`uv`/noise that differ significantly, leave `shared_preloads=False` (the default) so each channel computes its\n", + "own inversion \u2014 preloading an invalid quantity would silently corrupt the likelihood. The shared and unshared\n", + "paths are asserted to give identical likelihoods in\n", + "`autolens_workspace_test/scripts/jax_likelihood_functions/datacube/shared_preloads.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [\n", + " al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " settings=settings,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " use_jax=True,\n", + " shared_preloads=True,\n", + " )\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__FactorGraph__\n", + "\n", + "Each analysis is wrapped in an `af.AnalysisFactor` paired with a deep copy of the base model. With no per-factor\n", + "prior overrides, every prior is identified across factors \u2014 so the global model has the same dimensionality as\n", + "the single-channel base model. ``af.FactorGraphModel(..., use_jax=True)`` sums the per-channel log-evidences\n", + "internally, which is exactly the cube log-likelihood you'd write by hand." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = [\n", + " af.AnalysisFactor(prior_model=model.copy(), analysis=analysis)\n", + " for analysis in analysis_list\n", + "]\n", + "\n", + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + "print(f\" channels in factor graph: {len(analysis_factor_list)}\")\n", + "print(\n", + " f\" global model free parameters: {factor_graph.global_prior_model.total_free_parameters}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "`Nautilus` is the standard non-linear search for PyAutoLens. Datacube fits typically need fewer live points\n", + "than imaging fits because the lens dimensionality is unchanged \u2014 only the per-channel inversions multiply.\n", + "Tune `n_live` for your problem." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\") / \"datacube\",\n", + " name=\"modeling\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20,\n", + " iterations_per_quick_update=50000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "Pass the factor graph's `global_prior_model` as the model and the factor graph itself as the analysis \u2014 that's\n", + "the same shape you'd use for any multi-dataset PyAutoFit fit.\n", + "\n", + "**Run time on CPU is dominated by the per-channel inversion.** A 4-channel SMA-scale cube finishes in a few\n", + "hours on CPU; ALMA-scale cubes with 50+ channels need GPU acceleration to complete in reasonable time. The\n", + "``L\u1d40 W\u0303 L`` shared-precompute optimisation (Aris's design \u2014 exploit the fact that ``uv_wavelengths`` and\n", + "``noise_map`` change very little channel-to-channel) is the natural follow-up that brings ALMA-scale cubes\n", + "back inside the budget." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a while\n", + " for a datacube fit, since per-channel inversions multiply the per-likelihood cost.\n", + " \"\"\"\n", + ")\n", + "\n", + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The result returned by `search.fit` is a list of per-factor results \u2014 one entry per channel \u2014 each carrying its\n", + "own `FitInterferometer` against the maximum-likelihood lens model. Use them to inspect the per-channel source\n", + "reconstructions, dirty images, and residuals.\n", + "\n", + "For a step-by-step look at how the per-channel likelihood is summed (and at the JAX JIT pattern that drives\n", + "this fit), see ``autolens_workspace_developer/datacube/likelihood_function.py``." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/datacube/modeling_parametric.ipynb b/notebooks/interferometer/features/datacube/modeling_parametric.ipynb index ff59ea7e7..770917a79 100644 --- a/notebooks/interferometer/features/datacube/modeling_parametric.ipynb +++ b/notebooks/interferometer/features/datacube/modeling_parametric.ipynb @@ -1,416 +1,453 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Datacube \u2014 Parametric Source\n", - "=======================================\n", - "\n", - "This script fits a datacube \u2014 a list of `Interferometer` channels \u2014 with a single shared lens model and a\n", - "**parametric** source (`al.lp.Sersic`) whose morphology is shared across channels and whose `intensity` varies\n", - "channel-to-channel. It is the parametric-fit sibling of `modeling.py`, which fits the same cube with a\n", - "pixelized source instead.\n", - "\n", - "Use this script when:\n", - "\n", - " - You expect the source emission to be well-described by a single Sersic-like shape across the line \u2014 for\n", - " example, a high-signal-to-noise unresolved emission line in a kinematically simple galaxy.\n", - " - You want a cheaper fit than the pixelization variant. Parametric inversions are 1\u20132 orders of magnitude\n", - " faster than per-channel pixelizations because there is no source-plane linear inversion.\n", - " - You want the per-channel `intensity` posterior directly, with no need to integrate a pixelized\n", - " reconstruction back to a total flux.\n", - "\n", - "Use `modeling.py` (the pixelized variant) instead when the source has complex morphology, internal velocity\n", - "structure that varies across the cube, or signal-to-noise too low for a Sersic fit to converge.\n", - "\n", - "The FactorGraph wiring here is the canonical \"extend the model per dataset\" pattern from\n", - "`autolens_workspace/scripts/multi/modeling.py`: every prior in the base model is shared across channels by\n", - "default, and we explicitly override the source `intensity` prior per `AnalysisFactor` to make it per-channel.\n", - "\n", - "__Contents__\n", - "\n", - "- **Mask:** Define the 2D real-space mask applied to every channel.\n", - "- **Dataset:** Where the per-channel cube lives on disk and how to point this script at your own.\n", - "- **Dataset Auto-Simulation:** Run `simulator.py` automatically if the cube isn't already on disk.\n", - "- **Dataset Loading:** Loop over the channel folders and load each as an `Interferometer` object.\n", - "- **Positions:** Load multiple-image positions and build a shared `PositionsLH` penalty.\n", - "- **Settings:** Default `al.Settings()` \u2014 no positive-only-solver tweak needed (no inversion).\n", - "- **Model:** Compose the shared `Isothermal + ExternalShear` lens and `Sersic` source.\n", - "- **Per-Channel Analyses:** One `AnalysisInterferometer` per channel, with `use_jax=True` and the shared `PositionsLH`.\n", - "- **FactorGraph:** Per-factor `model.copy()` with the source `intensity` prior overridden per channel.\n", - "- **Search:** Configure the `Nautilus` non-linear search.\n", - "- **Model Fit:** Run the fit. Per-channel cost is much cheaper than the pixelization variant.\n", - "- **Wrap Up:** Pointers to `modeling.py`, `start_here.py`, and the JAX likelihood walkthrough." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import subprocess\n", - "import sys\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "The datacube lives under `dataset/interferometer/datacube//`, with one subfolder per channel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"datacube\"\n", - "dataset_name = \"sim_simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Loading__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "channel_paths = sorted(\n", - " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", - ")\n", - "print(f\"Loading {len(channel_paths)} channels from {dataset_path}\")\n", - "\n", - "dataset_list = [\n", - " al.Interferometer.from_fits(\n", - " data_path=channel_path / \"data.fits\",\n", - " noise_map_path=channel_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - " )\n", - " for channel_path in channel_paths\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Load the cube's multiple-image positions and wrap them in an `al.PositionsLH` penalty. Pixelized fits really\n", - "need this; parametric Sersic fits less so, but the penalty still helps the search avoid local maxima where the\n", - "mass model places multiple images far apart in the source plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "Parametric sources don't run a source-plane inversion, so we don't need to disable a positive-only solver\n", - "here. Default settings are fine." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "The cube model has two ingredients:\n", - "\n", - " - A shared `Isothermal + ExternalShear` lens (7 free parameters).\n", - " - A parametric `al.lp.Sersic` source. Its morphology \u2014 `centre`, `ell_comps`, `effective_radius`,\n", - " `sersic_index` \u2014 is shared across channels (4 free parameters). Its `intensity` will be overridden\n", - " per-factor below to give each channel its own free `intensity` parameter, capturing the emission-line\n", - " spectrum.\n", - "\n", - "Total free parameters:\n", - " 7 (lens) + 4 (source morphology, shared) + N_channels * 1 (per-channel intensity)\n", - " = 11 + N_channels\n", - "\n", - "For the 4-channel reference cube that's 15 free parameters total \u2014 tractable for `Nautilus`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source (parametric Sersic \u2014 base prior on intensity gets overridden per factor below):\n", - "bulge = af.Model(al.lp.Sersic)\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall lens model:\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Channel Analyses__\n", - "\n", - "The shared `positions_likelihood` is passed to every per-channel analysis so the multiple-image penalty applies\n", - "globally to the lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [\n", - " al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " settings=settings,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " use_jax=True,\n", - " )\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__FactorGraph__\n", - "\n", - "Each analysis is wrapped in an `af.AnalysisFactor` paired with a deep copy of the base model. We then override\n", - "the source `intensity` prior per factor \u2014 overwriting the prior with a fresh `LogUniformPrior` makes that\n", - "parameter per-channel rather than identified across factors. Every other source parameter (centre, ell_comps,\n", - "effective_radius, sersic_index) is left untouched, so the FactorGraph identifies them across channels.\n", - "\n", - "This is the canonical \"extend the model per dataset\" pattern from `autolens_workspace/scripts/multi/modeling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = []\n", - "\n", - "for analysis in analysis_list:\n", - " model_analysis = model.copy()\n", - "\n", - " # Per-channel intensity prior \u2014 overrides the shared default with a fresh prior object,\n", - " # which the FactorGraph treats as a distinct (per-factor) parameter.\n", - " model_analysis.galaxies.source.bulge.intensity = af.LogUniformPrior(\n", - " lower_limit=1e-3, upper_limit=10.0\n", - " )\n", - "\n", - " analysis_factor_list.append(\n", - " af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", - " )\n", - "\n", - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - "print(f\" channels in factor graph: {len(analysis_factor_list)}\")\n", - "print(\n", - " f\" global model free parameters: {factor_graph.global_prior_model.total_free_parameters}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "Parametric datacube fits have a higher non-linear dimensionality than the pixelization variant (per-channel\n", - "intensity adds one parameter per channel) but a much cheaper per-likelihood cost. Tuning depends on your cube;\n", - "`n_live=150` is a reasonable starting point." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\") / \"datacube\",\n", - " name=\"modeling_parametric\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=20,\n", - " iterations_per_quick_update=50000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "Run the fit. Per-channel cost is much cheaper than `modeling.py` because there is no source-plane inversion,\n", - "just a forward Sersic evaluation per channel followed by the Fourier transform." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " Parametric datacube fits are typically faster than pixelized ones \u2014 the per-likelihood cost is dominated\n", - " by the per-channel NUFFT, with no per-channel inversion on top.\n", - " \"\"\"\n", - ")\n", - "\n", - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The result returned by `search.fit` is a list of per-factor results \u2014 one entry per channel \u2014 each carrying\n", - "its own `FitInterferometer` against the maximum-likelihood lens + source-morphology model and the per-channel\n", - "intensity. The recovered per-channel intensities should trace the input emission-line spectrum stored in\n", - "`dataset/interferometer/datacube//cube_summary.json`.\n", - "\n", - "For the pixelized variant of this fit (free-form per-channel source reconstruction), see `modeling.py`. For\n", - "the narrative walkthrough, see `start_here.py`. For a step-by-step JAX likelihood walkthrough of how the\n", - "per-channel log-evidences are summed inside the FactorGraph, see\n", - "`autolens_workspace_developer/datacube/likelihood_function.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Datacube \u2014 Parametric Source\n", + "=======================================\n", + "\n", + "This script fits a datacube \u2014 a list of `Interferometer` channels \u2014 with a single shared lens model and a\n", + "**parametric** source (`al.lp.Sersic`) whose morphology is shared across channels and whose `intensity` varies\n", + "channel-to-channel. It is the parametric-fit sibling of `modeling.py`, which fits the same cube with a\n", + "pixelized source instead.\n", + "\n", + "Use this script when:\n", + "\n", + " - You expect the source emission to be well-described by a single Sersic-like shape across the line \u2014 for\n", + " example, a high-signal-to-noise unresolved emission line in a kinematically simple galaxy.\n", + " - You want a cheaper fit than the pixelization variant. Parametric inversions are 1\u20132 orders of magnitude\n", + " faster than per-channel pixelizations because there is no source-plane linear inversion.\n", + " - You want the per-channel `intensity` posterior directly, with no need to integrate a pixelized\n", + " reconstruction back to a total flux.\n", + "\n", + "Use `modeling.py` (the pixelized variant) instead when the source has complex morphology, internal velocity\n", + "structure that varies across the cube, or signal-to-noise too low for a Sersic fit to converge.\n", + "\n", + "The FactorGraph wiring here is the canonical \"extend the model per dataset\" pattern from\n", + "`autolens_workspace/scripts/multi/modeling.py`: every prior in the base model is shared across channels by\n", + "default, and we explicitly override the source `intensity` prior per `AnalysisFactor` to make it per-channel.\n", + "\n", + "__Contents__\n", + "\n", + "- **Mask:** Define the 2D real-space mask applied to every channel.\n", + "- **Dataset:** Where the per-channel cube lives on disk and how to point this script at your own.\n", + "- **Dataset Auto-Simulation:** Run `simulator.py` automatically if the cube isn't already on disk.\n", + "- **Dataset Loading:** Loop over the channel folders and load each as an `Interferometer` object.\n", + "- **Positions:** Load multiple-image positions and build a shared `PositionsLH` penalty.\n", + "- **Settings:** Default `al.Settings()` \u2014 no positive-only-solver tweak needed (no inversion).\n", + "- **Model:** Compose the shared `Isothermal + ExternalShear` lens and `Sersic` source.\n", + "- **Per-Channel Analyses:** One `AnalysisInterferometer` per channel, with `use_jax=True` and the shared `PositionsLH`.\n", + "- **FactorGraph:** Per-factor `model.copy()` with the source `intensity` prior overridden per channel.\n", + "- **Search:** Configure the `Nautilus` non-linear search.\n", + "- **Model Fit:** Run the fit. Per-channel cost is much cheaper than the pixelization variant.\n", + "- **Wrap Up:** Pointers to `modeling.py`, `start_here.py`, and the JAX likelihood walkthrough." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import subprocess\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "The datacube lives under `dataset/interferometer/datacube//`, with one subfolder per channel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"datacube\"\n", + "dataset_name = \"sim_simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Loading__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "channel_paths = sorted(\n", + " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", + ")\n", + "print(f\"Loading {len(channel_paths)} channels from {dataset_path}\")\n", + "\n", + "dataset_list = [\n", + " al.Interferometer.from_fits(\n", + " data_path=channel_path / \"data.fits\",\n", + " noise_map_path=channel_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + " )\n", + " for channel_path in channel_paths\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Load the cube's multiple-image positions and wrap them in an `al.PositionsLH` penalty. Pixelized fits really\n", + "need this; parametric Sersic fits less so, but the penalty still helps the search avoid local maxima where the\n", + "mass model places multiple images far apart in the source plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "Parametric sources don't run a source-plane inversion, so we don't need to disable a positive-only solver\n", + "here. Default settings are fine." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "The cube model has two ingredients:\n", + "\n", + " - A shared `Isothermal + ExternalShear` lens (7 free parameters).\n", + " - A parametric `al.lp.Sersic` source. Its morphology \u2014 `centre`, `ell_comps`, `effective_radius`,\n", + " `sersic_index` \u2014 is shared across channels (4 free parameters). Its `intensity` will be overridden\n", + " per-factor below to give each channel its own free `intensity` parameter, capturing the emission-line\n", + " spectrum.\n", + "\n", + "Total free parameters:\n", + " 7 (lens) + 4 (source morphology, shared) + N_channels * 1 (per-channel intensity)\n", + " = 11 + N_channels\n", + "\n", + "For the 4-channel reference cube that's 15 free parameters total \u2014 tractable for `Nautilus`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source (parametric Sersic \u2014 base prior on intensity gets overridden per factor below):\n", + "bulge = af.Model(al.lp.Sersic)\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall lens model:\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Channel Analyses__\n", + "\n", + "The shared `positions_likelihood` is passed to every per-channel analysis so the multiple-image penalty applies\n", + "globally to the lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [\n", + " al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " settings=settings,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " use_jax=True,\n", + " )\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__FactorGraph__\n", + "\n", + "Each analysis is wrapped in an `af.AnalysisFactor` paired with a deep copy of the base model. We then override\n", + "the source `intensity` prior per factor \u2014 overwriting the prior with a fresh `LogUniformPrior` makes that\n", + "parameter per-channel rather than identified across factors. Every other source parameter (centre, ell_comps,\n", + "effective_radius, sersic_index) is left untouched, so the FactorGraph identifies them across channels.\n", + "\n", + "This is the canonical \"extend the model per dataset\" pattern from `autolens_workspace/scripts/multi/modeling.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = []\n", + "\n", + "for analysis in analysis_list:\n", + " model_analysis = model.copy()\n", + "\n", + " # Per-channel intensity prior \u2014 overrides the shared default with a fresh prior object,\n", + " # which the FactorGraph treats as a distinct (per-factor) parameter.\n", + " model_analysis.galaxies.source.bulge.intensity = af.LogUniformPrior(\n", + " lower_limit=1e-3, upper_limit=10.0\n", + " )\n", + "\n", + " analysis_factor_list.append(\n", + " af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", + " )\n", + "\n", + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + "print(f\" channels in factor graph: {len(analysis_factor_list)}\")\n", + "print(\n", + " f\" global model free parameters: {factor_graph.global_prior_model.total_free_parameters}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "Parametric datacube fits have a higher non-linear dimensionality than the pixelization variant (per-channel\n", + "intensity adds one parameter per channel) but a much cheaper per-likelihood cost. Tuning depends on your cube;\n", + "`n_live=150` is a reasonable starting point." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\") / \"datacube\",\n", + " name=\"modeling_parametric\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=20,\n", + " iterations_per_quick_update=50000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "Run the fit. Per-channel cost is much cheaper than `modeling.py` because there is no source-plane inversion,\n", + "just a forward Sersic evaluation per channel followed by the Fourier transform." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " Parametric datacube fits are typically faster than pixelized ones \u2014 the per-likelihood cost is dominated\n", + " by the per-channel NUFFT, with no per-channel inversion on top.\n", + " \"\"\"\n", + ")\n", + "\n", + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The result returned by `search.fit` is a list of per-factor results \u2014 one entry per channel \u2014 each carrying\n", + "its own `FitInterferometer` against the maximum-likelihood lens + source-morphology model and the per-channel\n", + "intensity. The recovered per-channel intensities should trace the input emission-line spectrum stored in\n", + "`dataset/interferometer/datacube//cube_summary.json`.\n", + "\n", + "For the pixelized variant of this fit (free-form per-channel source reconstruction), see `modeling.py`. For\n", + "the narrative walkthrough, see `start_here.py`. For a step-by-step JAX likelihood walkthrough of how the\n", + "per-channel log-evidences are summed inside the FactorGraph, see\n", + "`autolens_workspace_developer/datacube/likelihood_function.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/datacube/simulator.ipynb b/notebooks/interferometer/features/datacube/simulator.ipynb index a9cec9f8d..db31ceae0 100644 --- a/notebooks/interferometer/features/datacube/simulator.ipynb +++ b/notebooks/interferometer/features/datacube/simulator.ipynb @@ -1,596 +1,633 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Datacube\n", - "====================\n", - "\n", - "This script simulates an ALMA-style spectral-line `Interferometer` datacube of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`, identical across all channels.\n", - " - The source galaxy's light is an `Sersic`, whose `intensity` follows a Gaussian emission-line profile in the\n", - " channel index (peaking near the cube centre and falling off at the edges).\n", - "\n", - "A \"datacube\" here is a Python list of `Interferometer` objects, one per spectral channel. Each channel writes its\n", - "own ``data.fits`` / ``noise_map.fits`` / ``uv_wavelengths.fits`` into a separate folder, alongside the true tracer\n", - "for that channel. The modeling examples in this folder load the cube by listing those folders.\n", - "\n", - "For Phase 1 of datacube modeling we keep things deliberately simple: every channel uses the same `uv_wavelengths`\n", - "and the same noise level. Two source properties vary across channels \u2014 `intensity` (a Gaussian emission-line\n", - "profile) and `centre` (a small linear shift along the y axis, simulating the kinematic gradient that real\n", - "emission lines almost always exhibit).\n", - "\n", - "__Contents__\n", - "\n", - "- **Cube Configuration:** Number of channels and the emission-line shape used to drive the per-channel intensity.\n", - "- **Dataset Paths:** Where the per-channel datasets, summary JSON and overview plots are written.\n", - "- **uv_wavelengths:** Reuse the SMA `uv_wavelengths.fits` shipped with the workspace as a stand-in for ALMA coverage.\n", - "- **Real-Space Grid:** The 2D image-plane grid each channel is evaluated on before the Fourier transform.\n", - "- **Lens Galaxy:** Shared `Isothermal + ExternalShear` lens, identical for every channel.\n", - "- **Per-Channel Source:** A Gaussian emission line drives the per-channel `Sersic.intensity`; the `centre` shifts linearly along y across channels to mimic a kinematic gradient.\n", - "- **Per-Channel Simulate:** Loop over channels: build the tracer, simulate, write FITS + tracer.json to disk.\n", - "- **3D-FITS Cube:** Stack the per-channel arrays into single `(n_chan, n_vis, 2)` FITS files \u2014 the autolens-canonical post-polarisation-collapse shape.\n", - "- **4D CASA-like Cube:** Wrap the 3D cubes in a polarisation axis to produce `(n_pol, n_chan, n_vis, 2)` files matching CASA's native output. Polarisations are identical here for simplicity.\n", - "- **Multiple Images:** Compute and save the lensed multiple-image positions used by the modeling scripts' `PositionsLH` penalty.\n", - "- **Cube Summary:** Dump the emission-line parameters and per-channel intensities to `cube_summary.json`.\n", - "- **Cube Overview Plot:** Row-per-channel sanity figure (lensed image, uv-plane Re/Im, |vis| vs baseline length).\n", - "- **Spectrum Plot:** Source intensity as a function of channel index." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import json\n", - "from pathlib import Path\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "\n", - "import autolens as al" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cube Configuration__\n", - "\n", - "The cube shape is set by ``N_CHANNELS`` and the emission-line shape by the Gaussian profile parameters below. Keep\n", - "``N_CHANNELS`` small for the prototype: per-channel inversion cost grows linearly, and 4 channels are enough to\n", - "exercise the FactorGraph wiring end-to-end." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "N_CHANNELS = 4\n", - "PEAK_CHANNEL = 1.5\n", - "SIGMA_CHANNEL = 1.2\n", - "PEAK_INTENSITY = 0.6\n", - "\n", - "# Total source-centre shift along the y axis across the whole cube, in arcsec.\n", - "# At the simulator's 0.1\"/pixel grid, 0.12\" is about 1.2 pixels end-to-end \u2014\n", - "# visible in the per-channel lensed image without the source jumping outside the Einstein ring.\n", - "CENTRE_SHIFT_TOTAL = 0.12\n", - "SOURCE_CENTRE_BASE = (0.1, 0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "Each channel lives in its own subfolder so the modeling scripts can iterate over them with a simple loop. The\n", - "``cube_summary.json`` records the emission-line parameters used to simulate the cube \u2014 modeling can compare the\n", - "recovered per-channel source amplitudes against those true values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"interferometer\"\n", - "dataset_label = \"datacube\"\n", - "dataset_name = \"sim_simple\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name\n", - "dataset_path.mkdir(parents=True, exist_ok=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__uv_wavelengths__\n", - "\n", - "Load the SMA `uv_wavelengths.fits` shipped with the workspace at ``dataset/interferometer/uv_wavelengths/sma.fits``.\n", - "SMA's coverage is small (~190 visibilities) but it makes the prototype fast to iterate on without paying the ALMA\n", - "runtime cost. To swap in a real ALMA cube, point this path at your own per-channel `uv_wavelengths.fits` files\n", - "inside the simulation loop below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "uv_wavelengths_path = Path(\"dataset\", dataset_type, \"uv_wavelengths\")\n", - "uv_wavelengths = al.ndarray_via_fits_from(\n", - " file_path=uv_wavelengths_path / \"sma.fits\", hdu=0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Real-Space Grid__\n", - "\n", - "For interferometer data this image is evaluated in real space and then Fourier-transformed. Interferometer\n", - "calculations don't need over-sampling \u2014 the Fourier transform is the dominant numerical cost." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(256, 256), pixel_scales=0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "The lens mass + external shear is shared across channels \u2014 the lens doesn't change with frequency." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Channel Source__\n", - "\n", - "Two things vary channel-to-channel:\n", - "\n", - " - ``intensity`` follows a Gaussian profile in channel index \u2014 the simplest model of an emission line peaking at\n", - " one velocity and falling off symmetrically.\n", - " - ``centre`` shifts linearly along the y axis from ``-CENTRE_SHIFT_TOTAL / 2`` to ``+CENTRE_SHIFT_TOTAL / 2``\n", - " relative to ``SOURCE_CENTRE_BASE``. This mimics a kinematic gradient (e.g. a rotating disk's centroid drifting\n", - " across the emission line) and gives each channel a visually distinct lensed image. A real cube might also shift\n", - " in x or rotate; one-axis is enough to demonstrate the principle.\n", - "\n", - "The other shape parameters (ellipticity, effective radius, Sersic index) are held fixed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def channel_intensity(channel: int) -> float:\n", - " return PEAK_INTENSITY * float(\n", - " np.exp(-0.5 * ((channel - PEAK_CHANNEL) / SIGMA_CHANNEL) ** 2)\n", - " )\n", - "\n", - "\n", - "def channel_centre(channel: int) -> tuple:\n", - " \"\"\"Source centre for a given channel.\n", - "\n", - " Linearly shifts along the y axis from channel 0 to channel N-1, centred on\n", - " ``SOURCE_CENTRE_BASE``. For ``N_CHANNELS=4`` and ``CENTRE_SHIFT_TOTAL=0.12\"`` the per-channel\n", - " centres are ``(0.04, 0.1)``, ``(0.08, 0.1)``, ``(0.12, 0.1)``, ``(0.16, 0.1)``.\n", - " \"\"\"\n", - " if N_CHANNELS == 1:\n", - " return SOURCE_CENTRE_BASE\n", - " fractional = (channel - (N_CHANNELS - 1) / 2.0) / (N_CHANNELS - 1)\n", - " offset_y = CENTRE_SHIFT_TOTAL * fractional\n", - " return (SOURCE_CENTRE_BASE[0] + offset_y, SOURCE_CENTRE_BASE[1])\n", - "\n", - "\n", - "def source_galaxy_for(channel: int) -> al.Galaxy:\n", - " return al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=channel_centre(channel),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=channel_intensity(channel),\n", - " effective_radius=1.0,\n", - " sersic_index=2.5,\n", - " ),\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Channel Simulate__\n", - "\n", - "For each channel we build a tracer with the shared lens + the per-channel source, run\n", - "``SimulatorInterferometer.via_tracer_from``, and write the resulting visibilities, noise map, baselines and true\n", - "tracer into ``channel_NNN/``. The `noise_seed` is incremented per channel so each channel sees an independent noise\n", - "realisation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "channel_intensities = []\n", - "channel_centres = []\n", - "datasets = []\n", - "tracers = []\n", - "\n", - "for channel in range(N_CHANNELS):\n", - " intensity = channel_intensity(channel)\n", - " centre = channel_centre(channel)\n", - " channel_intensities.append(intensity)\n", - " channel_centres.append(list(centre))\n", - "\n", - " channel_path = dataset_path / f\"channel_{channel:03d}\"\n", - " channel_path.mkdir(parents=True, exist_ok=True)\n", - "\n", - " source_galaxy = source_galaxy_for(channel)\n", - " tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - " simulator = al.SimulatorInterferometer(\n", - " uv_wavelengths=uv_wavelengths,\n", - " exposure_time=300.0,\n", - " noise_sigma=1000.0,\n", - " transformer_class=al.TransformerDFT,\n", - " noise_seed=1 + channel,\n", - " )\n", - "\n", - " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - " datasets.append(dataset)\n", - " tracers.append(tracer)\n", - "\n", - " al.output_to_fits(\n", - " values=np.stack([dataset.data.real, dataset.data.imag], axis=-1),\n", - " file_path=channel_path / \"data.fits\",\n", - " overwrite=True,\n", - " )\n", - " al.output_to_fits(\n", - " values=np.stack([dataset.noise_map.real, dataset.noise_map.imag], axis=-1),\n", - " file_path=channel_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - " )\n", - " al.output_to_fits(\n", - " values=dataset.uv_wavelengths,\n", - " file_path=channel_path / \"uv_wavelengths.fits\",\n", - " overwrite=True,\n", - " )\n", - " al.output_to_json(\n", - " obj=tracer,\n", - " file_path=channel_path / \"tracer.json\",\n", - " )\n", - "\n", - " print(\n", - " f\" channel {channel:03d}: intensity={intensity:.4f}, \"\n", - " f\"centre={centre}, \"\n", - " f\"|vis|_max={np.max(np.abs(dataset.data)):.3e}\"\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__3D-FITS Cube__\n", - "\n", - "ALMA visibilities exit CASA as a single 4D FITS of shape `(n_pol, n_chan, n_vis, 2)`. After the user collapses\n", - "the polarisation axis (averaging or concatenating \u2014 see `data_preparation.py`), the canonical input shape is\n", - "`(n_chan, n_vis, 2)`. We write three additional files at the cube root in that shape so users with CASA-native\n", - "data can load the cube via `data_preparation.dataset_list_from_3d_fits` without first splitting into per-channel\n", - "folders. The per-channel `channel_NNN/` folders above are kept as well \u2014 both layouts coexist for now." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "visibilities_cube = np.stack(\n", - " [np.stack([d.data.real, d.data.imag], axis=-1) for d in datasets], axis=0\n", - ")\n", - "noise_map_cube = np.stack(\n", - " [np.stack([d.noise_map.real, d.noise_map.imag], axis=-1) for d in datasets], axis=0\n", - ")\n", - "uv_wavelengths_cube = np.stack([np.asarray(d.uv_wavelengths) for d in datasets], axis=0)\n", - "\n", - "al.output_to_fits(\n", - " values=visibilities_cube,\n", - " file_path=dataset_path / \"visibilities_cube.fits\",\n", - " overwrite=True,\n", - ")\n", - "al.output_to_fits(\n", - " values=noise_map_cube,\n", - " file_path=dataset_path / \"noise_map_cube.fits\",\n", - " overwrite=True,\n", - ")\n", - "al.output_to_fits(\n", - " values=uv_wavelengths_cube,\n", - " file_path=dataset_path / \"uv_wavelengths_cube.fits\",\n", - " overwrite=True,\n", - ")\n", - "\n", - "print(f\" 3D-FITS cubes: shape {visibilities_cube.shape} (n_chan, n_vis, 2)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__4D CASA-like Cube__\n", - "\n", - "ALMA visibilities arrive from CASA as a single 4D FITS of shape `(n_pol, n_chan, n_vis, 2)` \u2014 for example\n", - "`(2, 34, 16984, 2)` for a typical narrow-line observation. Users typically run a polarisation-collapse step\n", - "(averaging or concatenating the two pols \u2014 see `data_preparation.py`) before the data is fed to autolens.\n", - "\n", - "We write a 4D version of each cube here so users can practice the polarisation-collapse step against the\n", - "simulator's actual output rather than synthetic random arrays. For this synthetic simulator both polarisations\n", - "carry **identical** data (same visibilities, same noise_map, same baselines) \u2014 that's a pedagogical\n", - "simplification: real CASA data has independent noise realisations between polarisations, which is why\n", - "averaging real data reduces the effective noise by `sqrt(2)`. The baselines being identical across pols is\n", - "astronomically correct (both pols share the same antenna pairs)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "visibilities_4d_cube = np.stack([visibilities_cube, visibilities_cube], axis=0)\n", - "noise_map_4d_cube = np.stack([noise_map_cube, noise_map_cube], axis=0)\n", - "uv_wavelengths_4d_cube = np.stack([uv_wavelengths_cube, uv_wavelengths_cube], axis=0)\n", - "\n", - "al.output_to_fits(\n", - " values=visibilities_4d_cube,\n", - " file_path=dataset_path / \"visibilities_4d_cube.fits\",\n", - " overwrite=True,\n", - ")\n", - "al.output_to_fits(\n", - " values=noise_map_4d_cube,\n", - " file_path=dataset_path / \"noise_map_4d_cube.fits\",\n", - " overwrite=True,\n", - ")\n", - "al.output_to_fits(\n", - " values=uv_wavelengths_4d_cube,\n", - " file_path=dataset_path / \"uv_wavelengths_4d_cube.fits\",\n", - " overwrite=True,\n", - ")\n", - "\n", - "print(\n", - " f\" 4D CASA-like cubes: shape {visibilities_4d_cube.shape} (n_pol, n_chan, n_vis, 2)\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multiple Images__\n", - "\n", - "Pixelized source modeling can drift toward unphysical demagnified-source local maxima \u2014 the source pixels are\n", - "reconstructed in low-magnification regions of the source plane that fit the noise rather than the lensed signal.\n", - "The `PositionsLH` penalty defends against that by reading a small set of multiple-image positions from disk and\n", - "adding a likelihood penalty for any candidate lens model whose source-plane back-projection of those positions\n", - "spreads them apart.\n", - "\n", - "For simulated data we can compute the multiple-image positions automatically with `al.PointSolver`. The lens\n", - "model is channel-invariant; the source centre shifts slightly across channels, so we compute positions for the\n", - "*mean* source centre (`SOURCE_CENTRE_BASE`) \u2014 the resulting positions are within the `PositionsLH` threshold\n", - "for every individual channel. A single `positions.json` covers the entire cube." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "positions = solver.solve(\n", - " tracer=tracers[0],\n", - " source_plane_coordinate=SOURCE_CENTRE_BASE,\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=positions,\n", - " file_path=dataset_path / \"positions.json\",\n", - ")\n", - "\n", - "print(f\" multiple-image positions: {len(positions)}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cube Summary__\n", - "\n", - "A small JSON sidecar lets downstream scripts (and the user) recover the emission-line parameters used during\n", - "simulation without re-running the script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "summary = {\n", - " \"n_channels\": N_CHANNELS,\n", - " \"peak_channel\": PEAK_CHANNEL,\n", - " \"sigma_channel\": SIGMA_CHANNEL,\n", - " \"peak_intensity\": PEAK_INTENSITY,\n", - " \"centre_shift_total\": CENTRE_SHIFT_TOTAL,\n", - " \"source_centre_base\": list(SOURCE_CENTRE_BASE),\n", - " \"channel_intensities\": channel_intensities,\n", - " \"channel_centres\": channel_centres,\n", - "}\n", - "with open(dataset_path / \"cube_summary.json\", \"w\") as f:\n", - " json.dump(summary, f, indent=2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cube Overview Plot__\n", - "\n", - "Row-per-channel sanity figure: lensed real-space image, Re(visibilities) and Im(visibilities) in the uv-plane,\n", - "and |vis| as a function of baseline length. Useful for eyeballing whether the emission-line modulation has\n", - "propagated through the simulator end-to-end." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fig, axes = plt.subplots(N_CHANNELS, 4, figsize=(16, 3.4 * N_CHANNELS), squeeze=False)\n", - "\n", - "for c in range(N_CHANNELS):\n", - " dataset = datasets[c]\n", - " tracer = tracers[c]\n", - " uv = np.asarray(dataset.uv_wavelengths)\n", - " re = np.asarray(dataset.data.real)\n", - " im = np.asarray(dataset.data.imag)\n", - " amp = np.hypot(re, im)\n", - " baseline = np.hypot(uv[:, 0], uv[:, 1])\n", - "\n", - " image = tracer.image_2d_from(grid=grid)\n", - " axes[c, 0].imshow(np.asarray(image.native), origin=\"lower\", cmap=\"hot\")\n", - " axes[c, 0].set_title(f\"channel {c}: lensed image (I={channel_intensities[c]:.3f})\")\n", - " axes[c, 0].set_axis_off()\n", - "\n", - " sc = axes[c, 1].scatter(uv[:, 0], uv[:, 1], c=re, s=8, cmap=\"RdBu_r\")\n", - " axes[c, 1].set_title(\"Re(visibilities)\")\n", - " axes[c, 1].set_aspect(\"equal\")\n", - " axes[c, 1].set_xlabel(\"u\")\n", - " axes[c, 1].set_ylabel(\"v\")\n", - " plt.colorbar(sc, ax=axes[c, 1], fraction=0.046)\n", - "\n", - " sc = axes[c, 2].scatter(uv[:, 0], uv[:, 1], c=im, s=8, cmap=\"RdBu_r\")\n", - " axes[c, 2].set_title(\"Im(visibilities)\")\n", - " axes[c, 2].set_aspect(\"equal\")\n", - " axes[c, 2].set_xlabel(\"u\")\n", - " axes[c, 2].set_ylabel(\"v\")\n", - " plt.colorbar(sc, ax=axes[c, 2], fraction=0.046)\n", - "\n", - " axes[c, 3].scatter(baseline, amp, s=8, alpha=0.6)\n", - " axes[c, 3].set_title(\"|vis| vs baseline length\")\n", - " axes[c, 3].set_xlabel(r\"$\\sqrt{u^2 + v^2}$\")\n", - " axes[c, 3].set_ylabel(\"|visibilities|\")\n", - "\n", - "fig.tight_layout()\n", - "fig.savefig(dataset_path / \"cube_overview.png\", dpi=120, bbox_inches=\"tight\")\n", - "plt.close(fig)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Spectrum Plot__\n", - "\n", - "Source intensity vs channel index \u2014 the input emission-line profile this cube was simulated from. Modeling\n", - "scripts can compare the recovered per-channel inversion magnitude against this curve to sanity-check the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fig, ax = plt.subplots(figsize=(6, 4))\n", - "ax.plot(range(N_CHANNELS), channel_intensities, \"o-\")\n", - "ax.set_xlabel(\"channel index\")\n", - "ax.set_ylabel(\"source intensity\")\n", - "ax.set_title(f\"emission-line spectrum (peak={PEAK_CHANNEL}, sigma={SIGMA_CHANNEL})\")\n", - "ax.grid(True, alpha=0.3)\n", - "fig.tight_layout()\n", - "fig.savefig(dataset_path / \"spectrum.png\", dpi=120, bbox_inches=\"tight\")\n", - "plt.close(fig)\n", - "\n", - "print(f\" cube simulated: {dataset_path}\")\n", - "print(f\" channels: {N_CHANNELS}\")\n", - "print(f\" visibilities/chan: {uv_wavelengths.shape[0]}\")\n", - "print(f\" real-space grid: 256 x 256\")\n", - "print(f\" peak intensity: {PEAK_INTENSITY}\")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Datacube\n", + "====================\n", + "\n", + "This script simulates an ALMA-style spectral-line `Interferometer` datacube of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`, identical across all channels.\n", + " - The source galaxy's light is an `Sersic`, whose `intensity` follows a Gaussian emission-line profile in the\n", + " channel index (peaking near the cube centre and falling off at the edges).\n", + "\n", + "A \"datacube\" here is a Python list of `Interferometer` objects, one per spectral channel. Each channel writes its\n", + "own ``data.fits`` / ``noise_map.fits`` / ``uv_wavelengths.fits`` into a separate folder, alongside the true tracer\n", + "for that channel. The modeling examples in this folder load the cube by listing those folders.\n", + "\n", + "For Phase 1 of datacube modeling we keep things deliberately simple: every channel uses the same `uv_wavelengths`\n", + "and the same noise level. Two source properties vary across channels \u2014 `intensity` (a Gaussian emission-line\n", + "profile) and `centre` (a small linear shift along the y axis, simulating the kinematic gradient that real\n", + "emission lines almost always exhibit).\n", + "\n", + "__Contents__\n", + "\n", + "- **Cube Configuration:** Number of channels and the emission-line shape used to drive the per-channel intensity.\n", + "- **Dataset Paths:** Where the per-channel datasets, summary JSON and overview plots are written.\n", + "- **uv_wavelengths:** Reuse the SMA `uv_wavelengths.fits` shipped with the workspace as a stand-in for ALMA coverage.\n", + "- **Real-Space Grid:** The 2D image-plane grid each channel is evaluated on before the Fourier transform.\n", + "- **Lens Galaxy:** Shared `Isothermal + ExternalShear` lens, identical for every channel.\n", + "- **Per-Channel Source:** A Gaussian emission line drives the per-channel `Sersic.intensity`; the `centre` shifts linearly along y across channels to mimic a kinematic gradient.\n", + "- **Per-Channel Simulate:** Loop over channels: build the tracer, simulate, write FITS + tracer.json to disk.\n", + "- **3D-FITS Cube:** Stack the per-channel arrays into single `(n_chan, n_vis, 2)` FITS files \u2014 the autolens-canonical post-polarisation-collapse shape.\n", + "- **4D CASA-like Cube:** Wrap the 3D cubes in a polarisation axis to produce `(n_pol, n_chan, n_vis, 2)` files matching CASA's native output. Polarisations are identical here for simplicity.\n", + "- **Multiple Images:** Compute and save the lensed multiple-image positions used by the modeling scripts' `PositionsLH` penalty.\n", + "- **Cube Summary:** Dump the emission-line parameters and per-channel intensities to `cube_summary.json`.\n", + "- **Cube Overview Plot:** Row-per-channel sanity figure (lensed image, uv-plane Re/Im, |vis| vs baseline length).\n", + "- **Spectrum Plot:** Source intensity as a function of channel index." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import json\n", + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cube Configuration__\n", + "\n", + "The cube shape is set by ``N_CHANNELS`` and the emission-line shape by the Gaussian profile parameters below. Keep\n", + "``N_CHANNELS`` small for the prototype: per-channel inversion cost grows linearly, and 4 channels are enough to\n", + "exercise the FactorGraph wiring end-to-end." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "N_CHANNELS = 4\n", + "PEAK_CHANNEL = 1.5\n", + "SIGMA_CHANNEL = 1.2\n", + "PEAK_INTENSITY = 0.6\n", + "\n", + "# Total source-centre shift along the y axis across the whole cube, in arcsec.\n", + "# At the simulator's 0.1\"/pixel grid, 0.12\" is about 1.2 pixels end-to-end \u2014\n", + "# visible in the per-channel lensed image without the source jumping outside the Einstein ring.\n", + "CENTRE_SHIFT_TOTAL = 0.12\n", + "SOURCE_CENTRE_BASE = (0.1, 0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "Each channel lives in its own subfolder so the modeling scripts can iterate over them with a simple loop. The\n", + "``cube_summary.json`` records the emission-line parameters used to simulate the cube \u2014 modeling can compare the\n", + "recovered per-channel source amplitudes against those true values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"interferometer\"\n", + "dataset_label = \"datacube\"\n", + "dataset_name = \"sim_simple\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name\n", + "dataset_path.mkdir(parents=True, exist_ok=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__uv_wavelengths__\n", + "\n", + "Load the SMA `uv_wavelengths.fits` shipped with the workspace at ``dataset/interferometer/uv_wavelengths/sma.fits``.\n", + "SMA's coverage is small (~190 visibilities) but it makes the prototype fast to iterate on without paying the ALMA\n", + "runtime cost. To swap in a real ALMA cube, point this path at your own per-channel `uv_wavelengths.fits` files\n", + "inside the simulation loop below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "uv_wavelengths_path = Path(\"dataset\", dataset_type, \"uv_wavelengths\")\n", + "uv_wavelengths = al.ndarray_via_fits_from(\n", + " file_path=uv_wavelengths_path / \"sma.fits\", hdu=0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Real-Space Grid__\n", + "\n", + "For interferometer data this image is evaluated in real space and then Fourier-transformed. Interferometer\n", + "calculations don't need over-sampling \u2014 the Fourier transform is the dominant numerical cost." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(256, 256), pixel_scales=0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "The lens mass + external shear is shared across channels \u2014 the lens doesn't change with frequency." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Channel Source__\n", + "\n", + "Two things vary channel-to-channel:\n", + "\n", + " - ``intensity`` follows a Gaussian profile in channel index \u2014 the simplest model of an emission line peaking at\n", + " one velocity and falling off symmetrically.\n", + " - ``centre`` shifts linearly along the y axis from ``-CENTRE_SHIFT_TOTAL / 2`` to ``+CENTRE_SHIFT_TOTAL / 2``\n", + " relative to ``SOURCE_CENTRE_BASE``. This mimics a kinematic gradient (e.g. a rotating disk's centroid drifting\n", + " across the emission line) and gives each channel a visually distinct lensed image. A real cube might also shift\n", + " in x or rotate; one-axis is enough to demonstrate the principle.\n", + "\n", + "The other shape parameters (ellipticity, effective radius, Sersic index) are held fixed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def channel_intensity(channel: int) -> float:\n", + " return PEAK_INTENSITY * float(\n", + " np.exp(-0.5 * ((channel - PEAK_CHANNEL) / SIGMA_CHANNEL) ** 2)\n", + " )\n", + "\n", + "\n", + "def channel_centre(channel: int) -> tuple:\n", + " \"\"\"Source centre for a given channel.\n", + "\n", + " Linearly shifts along the y axis from channel 0 to channel N-1, centred on\n", + " ``SOURCE_CENTRE_BASE``. For ``N_CHANNELS=4`` and ``CENTRE_SHIFT_TOTAL=0.12\"`` the per-channel\n", + " centres are ``(0.04, 0.1)``, ``(0.08, 0.1)``, ``(0.12, 0.1)``, ``(0.16, 0.1)``.\n", + " \"\"\"\n", + " if N_CHANNELS == 1:\n", + " return SOURCE_CENTRE_BASE\n", + " fractional = (channel - (N_CHANNELS - 1) / 2.0) / (N_CHANNELS - 1)\n", + " offset_y = CENTRE_SHIFT_TOTAL * fractional\n", + " return (SOURCE_CENTRE_BASE[0] + offset_y, SOURCE_CENTRE_BASE[1])\n", + "\n", + "\n", + "def source_galaxy_for(channel: int) -> al.Galaxy:\n", + " return al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=channel_centre(channel),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=channel_intensity(channel),\n", + " effective_radius=1.0,\n", + " sersic_index=2.5,\n", + " ),\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Channel Simulate__\n", + "\n", + "For each channel we build a tracer with the shared lens + the per-channel source, run\n", + "``SimulatorInterferometer.via_tracer_from``, and write the resulting visibilities, noise map, baselines and true\n", + "tracer into ``channel_NNN/``. The `noise_seed` is incremented per channel so each channel sees an independent noise\n", + "realisation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "channel_intensities = []\n", + "channel_centres = []\n", + "datasets = []\n", + "tracers = []\n", + "\n", + "for channel in range(N_CHANNELS):\n", + " intensity = channel_intensity(channel)\n", + " centre = channel_centre(channel)\n", + " channel_intensities.append(intensity)\n", + " channel_centres.append(list(centre))\n", + "\n", + " channel_path = dataset_path / f\"channel_{channel:03d}\"\n", + " channel_path.mkdir(parents=True, exist_ok=True)\n", + "\n", + " source_galaxy = source_galaxy_for(channel)\n", + " tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + " simulator = al.SimulatorInterferometer(\n", + " uv_wavelengths=uv_wavelengths,\n", + " exposure_time=300.0,\n", + " noise_sigma=1000.0,\n", + " transformer_class=al.TransformerDFT,\n", + " noise_seed=1 + channel,\n", + " )\n", + "\n", + " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + " datasets.append(dataset)\n", + " tracers.append(tracer)\n", + "\n", + " al.output_to_fits(\n", + " values=np.stack([dataset.data.real, dataset.data.imag], axis=-1),\n", + " file_path=channel_path / \"data.fits\",\n", + " overwrite=True,\n", + " )\n", + " al.output_to_fits(\n", + " values=np.stack([dataset.noise_map.real, dataset.noise_map.imag], axis=-1),\n", + " file_path=channel_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + " )\n", + " al.output_to_fits(\n", + " values=dataset.uv_wavelengths,\n", + " file_path=channel_path / \"uv_wavelengths.fits\",\n", + " overwrite=True,\n", + " )\n", + " al.output_to_json(\n", + " obj=tracer,\n", + " file_path=channel_path / \"tracer.json\",\n", + " )\n", + "\n", + " print(\n", + " f\" channel {channel:03d}: intensity={intensity:.4f}, \"\n", + " f\"centre={centre}, \"\n", + " f\"|vis|_max={np.max(np.abs(dataset.data)):.3e}\"\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__3D-FITS Cube__\n", + "\n", + "ALMA visibilities exit CASA as a single 4D FITS of shape `(n_pol, n_chan, n_vis, 2)`. After the user collapses\n", + "the polarisation axis (averaging or concatenating \u2014 see `data_preparation.py`), the canonical input shape is\n", + "`(n_chan, n_vis, 2)`. We write three additional files at the cube root in that shape so users with CASA-native\n", + "data can load the cube via `data_preparation.dataset_list_from_3d_fits` without first splitting into per-channel\n", + "folders. The per-channel `channel_NNN/` folders above are kept as well \u2014 both layouts coexist for now." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "visibilities_cube = np.stack(\n", + " [np.stack([d.data.real, d.data.imag], axis=-1) for d in datasets], axis=0\n", + ")\n", + "noise_map_cube = np.stack(\n", + " [np.stack([d.noise_map.real, d.noise_map.imag], axis=-1) for d in datasets], axis=0\n", + ")\n", + "uv_wavelengths_cube = np.stack([np.asarray(d.uv_wavelengths) for d in datasets], axis=0)\n", + "\n", + "al.output_to_fits(\n", + " values=visibilities_cube,\n", + " file_path=dataset_path / \"visibilities_cube.fits\",\n", + " overwrite=True,\n", + ")\n", + "al.output_to_fits(\n", + " values=noise_map_cube,\n", + " file_path=dataset_path / \"noise_map_cube.fits\",\n", + " overwrite=True,\n", + ")\n", + "al.output_to_fits(\n", + " values=uv_wavelengths_cube,\n", + " file_path=dataset_path / \"uv_wavelengths_cube.fits\",\n", + " overwrite=True,\n", + ")\n", + "\n", + "print(f\" 3D-FITS cubes: shape {visibilities_cube.shape} (n_chan, n_vis, 2)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__4D CASA-like Cube__\n", + "\n", + "ALMA visibilities arrive from CASA as a single 4D FITS of shape `(n_pol, n_chan, n_vis, 2)` \u2014 for example\n", + "`(2, 34, 16984, 2)` for a typical narrow-line observation. Users typically run a polarisation-collapse step\n", + "(averaging or concatenating the two pols \u2014 see `data_preparation.py`) before the data is fed to autolens.\n", + "\n", + "We write a 4D version of each cube here so users can practice the polarisation-collapse step against the\n", + "simulator's actual output rather than synthetic random arrays. For this synthetic simulator both polarisations\n", + "carry **identical** data (same visibilities, same noise_map, same baselines) \u2014 that's a pedagogical\n", + "simplification: real CASA data has independent noise realisations between polarisations, which is why\n", + "averaging real data reduces the effective noise by `sqrt(2)`. The baselines being identical across pols is\n", + "astronomically correct (both pols share the same antenna pairs)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "visibilities_4d_cube = np.stack([visibilities_cube, visibilities_cube], axis=0)\n", + "noise_map_4d_cube = np.stack([noise_map_cube, noise_map_cube], axis=0)\n", + "uv_wavelengths_4d_cube = np.stack([uv_wavelengths_cube, uv_wavelengths_cube], axis=0)\n", + "\n", + "al.output_to_fits(\n", + " values=visibilities_4d_cube,\n", + " file_path=dataset_path / \"visibilities_4d_cube.fits\",\n", + " overwrite=True,\n", + ")\n", + "al.output_to_fits(\n", + " values=noise_map_4d_cube,\n", + " file_path=dataset_path / \"noise_map_4d_cube.fits\",\n", + " overwrite=True,\n", + ")\n", + "al.output_to_fits(\n", + " values=uv_wavelengths_4d_cube,\n", + " file_path=dataset_path / \"uv_wavelengths_4d_cube.fits\",\n", + " overwrite=True,\n", + ")\n", + "\n", + "print(\n", + " f\" 4D CASA-like cubes: shape {visibilities_4d_cube.shape} (n_pol, n_chan, n_vis, 2)\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multiple Images__\n", + "\n", + "Pixelized source modeling can drift toward unphysical demagnified-source local maxima \u2014 the source pixels are\n", + "reconstructed in low-magnification regions of the source plane that fit the noise rather than the lensed signal.\n", + "The `PositionsLH` penalty defends against that by reading a small set of multiple-image positions from disk and\n", + "adding a likelihood penalty for any candidate lens model whose source-plane back-projection of those positions\n", + "spreads them apart.\n", + "\n", + "For simulated data we can compute the multiple-image positions automatically with `al.PointSolver`. The lens\n", + "model is channel-invariant; the source centre shifts slightly across channels, so we compute positions for the\n", + "*mean* source centre (`SOURCE_CENTRE_BASE`) \u2014 the resulting positions are within the `PositionsLH` threshold\n", + "for every individual channel. A single `positions.json` covers the entire cube." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "positions = solver.solve(\n", + " tracer=tracers[0],\n", + " source_plane_coordinate=SOURCE_CENTRE_BASE,\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=positions,\n", + " file_path=dataset_path / \"positions.json\",\n", + ")\n", + "\n", + "print(f\" multiple-image positions: {len(positions)}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cube Summary__\n", + "\n", + "A small JSON sidecar lets downstream scripts (and the user) recover the emission-line parameters used during\n", + "simulation without re-running the script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "summary = {\n", + " \"n_channels\": N_CHANNELS,\n", + " \"peak_channel\": PEAK_CHANNEL,\n", + " \"sigma_channel\": SIGMA_CHANNEL,\n", + " \"peak_intensity\": PEAK_INTENSITY,\n", + " \"centre_shift_total\": CENTRE_SHIFT_TOTAL,\n", + " \"source_centre_base\": list(SOURCE_CENTRE_BASE),\n", + " \"channel_intensities\": channel_intensities,\n", + " \"channel_centres\": channel_centres,\n", + "}\n", + "with open(dataset_path / \"cube_summary.json\", \"w\") as f:\n", + " json.dump(summary, f, indent=2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cube Overview Plot__\n", + "\n", + "Row-per-channel sanity figure: lensed real-space image, Re(visibilities) and Im(visibilities) in the uv-plane,\n", + "and |vis| as a function of baseline length. Useful for eyeballing whether the emission-line modulation has\n", + "propagated through the simulator end-to-end." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fig, axes = plt.subplots(N_CHANNELS, 4, figsize=(16, 3.4 * N_CHANNELS), squeeze=False)\n", + "\n", + "for c in range(N_CHANNELS):\n", + " dataset = datasets[c]\n", + " tracer = tracers[c]\n", + " uv = np.asarray(dataset.uv_wavelengths)\n", + " re = np.asarray(dataset.data.real)\n", + " im = np.asarray(dataset.data.imag)\n", + " amp = np.hypot(re, im)\n", + " baseline = np.hypot(uv[:, 0], uv[:, 1])\n", + "\n", + " image = tracer.image_2d_from(grid=grid)\n", + " axes[c, 0].imshow(np.asarray(image.native), origin=\"lower\", cmap=\"hot\")\n", + " axes[c, 0].set_title(f\"channel {c}: lensed image (I={channel_intensities[c]:.3f})\")\n", + " axes[c, 0].set_axis_off()\n", + "\n", + " sc = axes[c, 1].scatter(uv[:, 0], uv[:, 1], c=re, s=8, cmap=\"RdBu_r\")\n", + " axes[c, 1].set_title(\"Re(visibilities)\")\n", + " axes[c, 1].set_aspect(\"equal\")\n", + " axes[c, 1].set_xlabel(\"u\")\n", + " axes[c, 1].set_ylabel(\"v\")\n", + " plt.colorbar(sc, ax=axes[c, 1], fraction=0.046)\n", + "\n", + " sc = axes[c, 2].scatter(uv[:, 0], uv[:, 1], c=im, s=8, cmap=\"RdBu_r\")\n", + " axes[c, 2].set_title(\"Im(visibilities)\")\n", + " axes[c, 2].set_aspect(\"equal\")\n", + " axes[c, 2].set_xlabel(\"u\")\n", + " axes[c, 2].set_ylabel(\"v\")\n", + " plt.colorbar(sc, ax=axes[c, 2], fraction=0.046)\n", + "\n", + " axes[c, 3].scatter(baseline, amp, s=8, alpha=0.6)\n", + " axes[c, 3].set_title(\"|vis| vs baseline length\")\n", + " axes[c, 3].set_xlabel(r\"$\\sqrt{u^2 + v^2}$\")\n", + " axes[c, 3].set_ylabel(\"|visibilities|\")\n", + "\n", + "fig.tight_layout()\n", + "fig.savefig(dataset_path / \"cube_overview.png\", dpi=120, bbox_inches=\"tight\")\n", + "plt.close(fig)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Spectrum Plot__\n", + "\n", + "Source intensity vs channel index \u2014 the input emission-line profile this cube was simulated from. Modeling\n", + "scripts can compare the recovered per-channel inversion magnitude against this curve to sanity-check the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fig, ax = plt.subplots(figsize=(6, 4))\n", + "ax.plot(range(N_CHANNELS), channel_intensities, \"o-\")\n", + "ax.set_xlabel(\"channel index\")\n", + "ax.set_ylabel(\"source intensity\")\n", + "ax.set_title(f\"emission-line spectrum (peak={PEAK_CHANNEL}, sigma={SIGMA_CHANNEL})\")\n", + "ax.grid(True, alpha=0.3)\n", + "fig.tight_layout()\n", + "fig.savefig(dataset_path / \"spectrum.png\", dpi=120, bbox_inches=\"tight\")\n", + "plt.close(fig)\n", + "\n", + "print(f\" cube simulated: {dataset_path}\")\n", + "print(f\" channels: {N_CHANNELS}\")\n", + "print(f\" visibilities/chan: {uv_wavelengths.shape[0]}\")\n", + "print(f\" real-space grid: 256 x 256\")\n", + "print(f\" peak intensity: {PEAK_INTENSITY}\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/datacube/start_here.ipynb b/notebooks/interferometer/features/datacube/start_here.ipynb index aaebe17ee..315329124 100644 --- a/notebooks/interferometer/features/datacube/start_here.ipynb +++ b/notebooks/interferometer/features/datacube/start_here.ipynb @@ -1,516 +1,553 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Start Here: Datacube\n", - "====================\n", - "\n", - "A datacube is a stack of interferometer observations of the same lens taken across many spectral channels \u2014 for\n", - "example, an ALMA observation of a CO emission line in a high-redshift lensed galaxy, where each channel records\n", - "visibilities at a slightly different frequency. The lens galaxy is the same in every channel, but the source's\n", - "emission-line morphology varies across the cube: bright in channels close to the line peak, faint in channels at\n", - "the line wings, and shifting in shape if the source has internal velocity structure.\n", - "\n", - "This script shows how to model such a cube in PyAutoLens by treating it as exactly that \u2014 a Python list of\n", - "`Interferometer` objects, one per channel \u2014 and tying them together with `af.FactorGraphModel`. The lens model is\n", - "shared across every channel, and each channel reconstructs its own pixelized source. The result is a per-channel\n", - "sequence of source-plane reconstructions whose combined likelihood drives a single, global lens model fit.\n", - "\n", - "This is Phase 1 of datacube modeling. It deliberately runs each channel's NUFFT and source inversion\n", - "independently, because reusing existing single-channel code is the fastest path to a working end-to-end\n", - "demonstration. The faster shared-`L\u1d40 W\u0303 L` design \u2014 which exploits the fact that `uv_wavelengths` and\n", - "`noise_map` change very little across an emission line \u2014 is a follow-up that lands once this prototype is proven.\n", - "\n", - "If you've already read `interferometer/start_here.py`, the bulk of this script will look familiar. The new\n", - "ingredients are the loop that loads each channel, the `AnalysisFactor` per channel, and the `FactorGraphModel`\n", - "that sums them.\n", - "\n", - "__Contents__\n", - "\n", - "- **JAX:** GPU/CPU acceleration via JAX \u2014 the same backend that single-channel interferometer fits use.\n", - "- **Imports:** Standard PyAutoLens imports + `autofit` for the FactorGraph wiring.\n", - "- **Mask:** A single 2D real-space mask shared across all channels.\n", - "- **Dataset:** Where the per-channel cube lives on disk and how to point this script at your own.\n", - "- **Dataset Auto-Simulation:** Run `simulator.py` automatically if the cube isn't already on disk.\n", - "- **Dataset Loading:** Loop over channel folders to build a `dataset_list` of `Interferometer` objects.\n", - "- **Sparse Operators:** Per-channel sparse-operator pre-compute used by the pixelized source inversion.\n", - "- **Positions:** Load the cube's multiple-image positions and build a shared `PositionsLH` penalty.\n", - "- **Settings:** Disable the positive-only solver (visibility inversions can take negative pixel values).\n", - "- **Mesh Shape:** The pixelization mesh shape \u2014 fixed before modeling because JAX needs static shapes.\n", - "- **Model:** Shared lens galaxy + pixelized source. The same model is reused unchanged across every channel.\n", - "- **Per-Channel Analyses:** One `AnalysisInterferometer` per channel, all sharing the same `PositionsLH`.\n", - "- **FactorGraph:** Wrap each analysis in an `AnalysisFactor`; combine via `af.FactorGraphModel`.\n", - "- **Search:** Configure the `Nautilus` non-linear search.\n", - "- **Model Fit:** Fit the cube \u2014 the FactorGraph routes shared lens parameters into every channel's likelihood.\n", - "- **Result:** What the returned `result_list` contains and how to inspect per-channel reconstructions.\n", - "- **Wrap Up:** Pointers to `modeling.py`, `simulator.py`, and the JAX likelihood walkthrough.\n", - "\n", - "__JAX__\n", - "\n", - "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU support, your\n", - "fits will run much faster (tens of minutes instead of several hours for a 4-channel cube). On CPU, JAX still\n", - "provides a meaningful speed-up via multithreading, but datacube fits are inherently more expensive than\n", - "single-channel fits because the per-channel inversion cost multiplies by the number of channels.\n", - "\n", - "If you don't have a GPU locally, consider Google Colab, which provides free GPUs." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import subprocess\n", - "import sys\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "A single 2D circular mask shared across every channel. The lens galaxy and the source emission live in the same\n", - "sky region in every frequency channel, so masking once and reusing the mask is correct." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "The reference cube ships in `dataset/interferometer/datacube/sim_simple/`, with one subfolder per channel\n", - "(`channel_000/`, `channel_001/`, ...) each containing `data.fits`, `noise_map.fits`, `uv_wavelengths.fits` and\n", - "the true `tracer.json`.\n", - "\n", - "To point this script at your own cube, drop your channel folders in alongside the reference cube and update\n", - "`dataset_name`. Each channel folder must contain `data.fits`, `noise_map.fits` and `uv_wavelengths.fits` in the\n", - "shape produced by `al.SimulatorInterferometer` (visibilities and noise stored as ``(n_vis, 2)`` real/imag pairs;\n", - "baselines as ``(n_vis, 2)`` u/v pairs)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"datacube\"\n", - "dataset_name = \"sim_simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Loading__\n", - "\n", - "Build the cube by loading each channel folder as an `Interferometer` object. The result is a Python list \u2014 no\n", - "new dataset class involved. Every downstream component (analyses, factors, fits, plotters) operates on this\n", - "list directly, which is the whole reason we picked the list-of-Interferometer design over a bespoke\n", - "`Datacube3D` class." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "channel_paths = sorted(\n", - " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", - ")\n", - "print(f\"Found {len(channel_paths)} channels in {dataset_path}\")\n", - "\n", - "dataset_list = [\n", - " al.Interferometer.from_fits(\n", - " data_path=channel_path / \"data.fits\",\n", - " noise_map_path=channel_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - " )\n", - " for channel_path in channel_paths\n", - "]\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset_list[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "Pixelized source modeling uses sparse linear algebra to keep memory and runtime manageable. We pre-compute a\n", - "sparse-operator matrix per channel \u2014 `apply_sparse_operator()` does this work in seconds for SMA-scale data and\n", - "in minutes per channel on CPU for ALMA-scale data. For very large cubes you'll want to compute these once and\n", - "cache them; see `pixelization/many_visibilities_preparation.py` for the pattern." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " dataset.apply_sparse_operator(use_jax=True, show_progress=False)\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Pixelized source modeling has a known failure mode: without a position-likelihood penalty, the search routinely\n", - "converges on demagnified-source local maxima where the source pixels are reconstructed in low-magnification\n", - "regions of the source plane that fit the noise rather than the lensed signal. The `PositionsLH` penalty defends\n", - "against that by reading a small set of multiple-image positions from disk and adding a likelihood penalty for\n", - "any candidate lens model whose source-plane back-projection of those positions spreads them apart.\n", - "\n", - "For the cube we load `positions.json` (written by `simulator.py`) and build one `PositionsLH` that gets passed to\n", - "every per-channel analysis below. The lens model is shared across channels via the FactorGraph, so applying the\n", - "same penalty in every analysis enforces a single global constraint.\n", - "\n", - "The threshold of 0.3\" is generous; for a real fit you'd tighten it (typically < 0.05\") once the lens model has\n", - "settled into the right region of parameter space." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "Interferometer pixelizations disable the positive-only inversion solver. The visibility measurement process\n", - "can produce genuinely negative dirty-image pixel values, so the source-plane reconstruction must be allowed to\n", - "go negative \u2014 forcing positivity here would create unphysical bias." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The pixelization mesh shape is fixed before modeling because JAX needs static-shape arrays for its source-plane\n", - "linear algebra. We use a 14 x 14 `RectangularAdaptDensity` mesh \u2014 small enough to make the prototype iteration\n", - "cheap, large enough to capture the emission-line source morphology produced by the simulator.\n", - "`RectangularAdaptDensity` adapts the source-plane pixel density to the lensing magnification map, giving more\n", - "pixels to the highly-magnified regions of the source plane where the lensed signal is concentrated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 14\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "The cube model has two ingredients:\n", - "\n", - " - A shared `Isothermal + ExternalShear` lens. There are 7 free parameters (mass centre, ellipticity components,\n", - " einstein radius, two shear components). The lens does not change with frequency, so a single set of priors is\n", - " used for every channel.\n", - " - A pixelized source: a `RectangularAdaptDensity` mesh with `Constant` regularization (1 free parameter \u2014 the\n", - " regularization coefficient). The pixelization itself has no per-pixel priors; the source-plane fluxes are a\n", - " linear inversion output computed by each channel's `AnalysisInterferometer` at fit time. That is what makes\n", - " each channel an independent linear solve while sharing all of the non-linear parameters.\n", - "\n", - "The total dimensionality of the non-linear parameter space is therefore 8." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source (pixelization, no per-pixel priors):\n", - "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", - "regularization = af.Model(al.reg.Constant)\n", - "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "# Overall lens model:\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Per-Channel Analyses__\n", - "\n", - "One `AnalysisInterferometer` per channel \u2014 there are no per-channel parameters here, only per-channel data.\n", - "Each analysis runs its own NUFFT, builds its own visibility-space inversion, and returns its own log-evidence\n", - "when called with a candidate lens model. The shared `positions_likelihood` is passed to every analysis to apply\n", - "the same global multiple-image penalty across the cube." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [\n", - " al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " settings=settings,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " use_jax=True,\n", - " )\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__FactorGraph__\n", - "\n", - "`af.AnalysisFactor` pairs each analysis with a deep copy of the base model, and `af.FactorGraphModel` combines\n", - "the factors into a single global model whose log-likelihood is the sum of the per-factor log-likelihoods. With\n", - "no per-factor prior overrides, every prior is *identified* across factors \u2014 the factor-graph machinery\n", - "deduplicates them \u2014 so the global model has the same dimensionality as the single-channel base model.\n", - "\n", - "If you wanted per-channel free parameters (for example, a per-channel `intensity` for a parametric source),\n", - "you'd override that prior on each `model.copy()` before wrapping it in an `AnalysisFactor`. See\n", - "`autolens_workspace/scripts/multi/modeling.py` for how that works in the multi-band case." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = [\n", - " af.AnalysisFactor(prior_model=model.copy(), analysis=analysis)\n", - " for analysis in analysis_list\n", - "]\n", - "\n", - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - "print(f\" channels in factor graph: {len(analysis_factor_list)}\")\n", - "print(\n", - " f\" global model free parameters: {factor_graph.global_prior_model.total_free_parameters}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "`Nautilus` is the standard non-linear search for PyAutoLens. The lens dimensionality is unchanged from a\n", - "single-channel fit, so `n_live=100` is a reasonable starting point \u2014 but the per-likelihood cost is N times\n", - "larger because each likelihood call runs N inversions, so wall-clock time scales linearly with the number of\n", - "channels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\") / \"datacube\",\n", - " name=\"start_here\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20,\n", - " iterations_per_quick_update=50000,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "We pass the factor graph's `global_prior_model` as the model and the factor graph itself as the analysis \u2014 this\n", - "is the same shape used for any multi-dataset PyAutoFit fit. Internally the search proposes lens parameters from\n", - "the global prior, hands them to the FactorGraph's `log_likelihood_function`, and the FactorGraph routes them\n", - "into every channel's `AnalysisInterferometer.log_likelihood_function` and sums the per-channel log-evidences.\n", - "\n", - "**Run time on CPU is dominated by the per-channel inversion.** A 4-channel SMA-scale cube finishes in a few\n", - "hours on CPU; ALMA-scale cubes with 50+ channels need GPU acceleration to complete in reasonable time." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " Per-channel inversions multiply the per-likelihood cost \u2014 expect this to take longer than a single-channel\n", - " interferometer fit. On CPU plan for hours; on GPU, tens of minutes.\n", - " \"\"\"\n", - ")\n", - "\n", - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "`result_list` contains one entry per factor \u2014 one per channel \u2014 each carrying its own `FitInterferometer`\n", - "against the maximum-likelihood lens model. Use them to inspect:\n", - "\n", - " - Per-channel source reconstructions (each channel's pixelized source is independent).\n", - " - Per-channel dirty images and residuals.\n", - " - The shared maximum-likelihood lens parameters (identical across factors by construction).\n", - "\n", - "Compared against `dataset/interferometer/datacube//cube_summary.json`, the per-channel reconstructed total\n", - "flux should trace the input emission-line spectrum.\n", - "\n", - "__Wrap Up__\n", - "\n", - "This script walks through the full datacube modeling pipeline. For deeper dives:\n", - "\n", - " - `modeling.py` \u2014 the focused FactorGraph + Nautilus example, ready to copy and adapt for your own cube.\n", - " - `simulator.py` \u2014 how the reference cube is generated.\n", - " - `autolens_workspace_developer/datacube/likelihood_function.py` \u2014 a step-by-step JAX walkthrough of how the\n", - " per-channel log-evidences are built and summed inside the FactorGraph, with an explicit eager-vs-JIT\n", - " correctness check.\n", - "\n", - "Phase 2 work \u2014 the shared-`L\u1d40 W\u0303 L` optimisation that exploits channel-invariant `uv_wavelengths` and\n", - "`noise_map` to bring ALMA-scale cubes inside CPU runtime budgets \u2014 is a separate follow-up issue." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Start Here: Datacube\n", + "====================\n", + "\n", + "A datacube is a stack of interferometer observations of the same lens taken across many spectral channels \u2014 for\n", + "example, an ALMA observation of a CO emission line in a high-redshift lensed galaxy, where each channel records\n", + "visibilities at a slightly different frequency. The lens galaxy is the same in every channel, but the source's\n", + "emission-line morphology varies across the cube: bright in channels close to the line peak, faint in channels at\n", + "the line wings, and shifting in shape if the source has internal velocity structure.\n", + "\n", + "This script shows how to model such a cube in PyAutoLens by treating it as exactly that \u2014 a Python list of\n", + "`Interferometer` objects, one per channel \u2014 and tying them together with `af.FactorGraphModel`. The lens model is\n", + "shared across every channel, and each channel reconstructs its own pixelized source. The result is a per-channel\n", + "sequence of source-plane reconstructions whose combined likelihood drives a single, global lens model fit.\n", + "\n", + "This is Phase 1 of datacube modeling. It deliberately runs each channel's NUFFT and source inversion\n", + "independently, because reusing existing single-channel code is the fastest path to a working end-to-end\n", + "demonstration. The faster shared-`L\u1d40 W\u0303 L` design \u2014 which exploits the fact that `uv_wavelengths` and\n", + "`noise_map` change very little across an emission line \u2014 is a follow-up that lands once this prototype is proven.\n", + "\n", + "If you've already read `interferometer/start_here.py`, the bulk of this script will look familiar. The new\n", + "ingredients are the loop that loads each channel, the `AnalysisFactor` per channel, and the `FactorGraphModel`\n", + "that sums them.\n", + "\n", + "__Contents__\n", + "\n", + "- **JAX:** GPU/CPU acceleration via JAX \u2014 the same backend that single-channel interferometer fits use.\n", + "- **Imports:** Standard PyAutoLens imports + `autofit` for the FactorGraph wiring.\n", + "- **Mask:** A single 2D real-space mask shared across all channels.\n", + "- **Dataset:** Where the per-channel cube lives on disk and how to point this script at your own.\n", + "- **Dataset Auto-Simulation:** Run `simulator.py` automatically if the cube isn't already on disk.\n", + "- **Dataset Loading:** Loop over channel folders to build a `dataset_list` of `Interferometer` objects.\n", + "- **Sparse Operators:** Per-channel sparse-operator pre-compute used by the pixelized source inversion.\n", + "- **Positions:** Load the cube's multiple-image positions and build a shared `PositionsLH` penalty.\n", + "- **Settings:** Disable the positive-only solver (visibility inversions can take negative pixel values).\n", + "- **Mesh Shape:** The pixelization mesh shape \u2014 fixed before modeling because JAX needs static shapes.\n", + "- **Model:** Shared lens galaxy + pixelized source. The same model is reused unchanged across every channel.\n", + "- **Per-Channel Analyses:** One `AnalysisInterferometer` per channel, all sharing the same `PositionsLH`.\n", + "- **FactorGraph:** Wrap each analysis in an `AnalysisFactor`; combine via `af.FactorGraphModel`.\n", + "- **Search:** Configure the `Nautilus` non-linear search.\n", + "- **Model Fit:** Fit the cube \u2014 the FactorGraph routes shared lens parameters into every channel's likelihood.\n", + "- **Result:** What the returned `result_list` contains and how to inspect per-channel reconstructions.\n", + "- **Wrap Up:** Pointers to `modeling.py`, `simulator.py`, and the JAX likelihood walkthrough.\n", + "\n", + "__JAX__\n", + "\n", + "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU support, your\n", + "fits will run much faster (tens of minutes instead of several hours for a 4-channel cube). On CPU, JAX still\n", + "provides a meaningful speed-up via multithreading, but datacube fits are inherently more expensive than\n", + "single-channel fits because the per-channel inversion cost multiplies by the number of channels.\n", + "\n", + "If you don't have a GPU locally, consider Google Colab, which provides free GPUs." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import subprocess\n", + "import sys\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "A single 2D circular mask shared across every channel. The lens galaxy and the source emission live in the same\n", + "sky region in every frequency channel, so masking once and reusing the mask is correct." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "The reference cube ships in `dataset/interferometer/datacube/sim_simple/`, with one subfolder per channel\n", + "(`channel_000/`, `channel_001/`, ...) each containing `data.fits`, `noise_map.fits`, `uv_wavelengths.fits` and\n", + "the true `tracer.json`.\n", + "\n", + "To point this script at your own cube, drop your channel folders in alongside the reference cube and update\n", + "`dataset_name`. Each channel folder must contain `data.fits`, `noise_map.fits` and `uv_wavelengths.fits` in the\n", + "shape produced by `al.SimulatorInterferometer` (visibilities and noise stored as ``(n_vis, 2)`` real/imag pairs;\n", + "baselines as ``(n_vis, 2)`` u/v pairs)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"datacube\"\n", + "dataset_name = \"sim_simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/features/datacube/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Loading__\n", + "\n", + "Build the cube by loading each channel folder as an `Interferometer` object. The result is a Python list \u2014 no\n", + "new dataset class involved. Every downstream component (analyses, factors, fits, plotters) operates on this\n", + "list directly, which is the whole reason we picked the list-of-Interferometer design over a bespoke\n", + "`Datacube3D` class." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "channel_paths = sorted(\n", + " p for p in dataset_path.iterdir() if p.is_dir() and p.name.startswith(\"channel_\")\n", + ")\n", + "print(f\"Found {len(channel_paths)} channels in {dataset_path}\")\n", + "\n", + "dataset_list = [\n", + " al.Interferometer.from_fits(\n", + " data_path=channel_path / \"data.fits\",\n", + " noise_map_path=channel_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=channel_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + " )\n", + " for channel_path in channel_paths\n", + "]\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset_list[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "Pixelized source modeling uses sparse linear algebra to keep memory and runtime manageable. We pre-compute a\n", + "sparse-operator matrix per channel \u2014 `apply_sparse_operator()` does this work in seconds for SMA-scale data and\n", + "in minutes per channel on CPU for ALMA-scale data. For very large cubes you'll want to compute these once and\n", + "cache them; see `pixelization/many_visibilities_preparation.py` for the pattern." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " dataset.apply_sparse_operator(use_jax=True, show_progress=False)\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Pixelized source modeling has a known failure mode: without a position-likelihood penalty, the search routinely\n", + "converges on demagnified-source local maxima where the source pixels are reconstructed in low-magnification\n", + "regions of the source plane that fit the noise rather than the lensed signal. The `PositionsLH` penalty defends\n", + "against that by reading a small set of multiple-image positions from disk and adding a likelihood penalty for\n", + "any candidate lens model whose source-plane back-projection of those positions spreads them apart.\n", + "\n", + "For the cube we load `positions.json` (written by `simulator.py`) and build one `PositionsLH` that gets passed to\n", + "every per-channel analysis below. The lens model is shared across channels via the FactorGraph, so applying the\n", + "same penalty in every analysis enforces a single global constraint.\n", + "\n", + "The threshold of 0.3\" is generous; for a real fit you'd tighten it (typically < 0.05\") once the lens model has\n", + "settled into the right region of parameter space." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(al.from_json(file_path=dataset_path / \"positions.json\"))\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "Interferometer pixelizations disable the positive-only inversion solver. The visibility measurement process\n", + "can produce genuinely negative dirty-image pixel values, so the source-plane reconstruction must be allowed to\n", + "go negative \u2014 forcing positivity here would create unphysical bias." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The pixelization mesh shape is fixed before modeling because JAX needs static-shape arrays for its source-plane\n", + "linear algebra. We use a 14 x 14 `RectangularAdaptDensity` mesh \u2014 small enough to make the prototype iteration\n", + "cheap, large enough to capture the emission-line source morphology produced by the simulator.\n", + "`RectangularAdaptDensity` adapts the source-plane pixel density to the lensing magnification map, giving more\n", + "pixels to the highly-magnified regions of the source plane where the lensed signal is concentrated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 14\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "The cube model has two ingredients:\n", + "\n", + " - A shared `Isothermal + ExternalShear` lens. There are 7 free parameters (mass centre, ellipticity components,\n", + " einstein radius, two shear components). The lens does not change with frequency, so a single set of priors is\n", + " used for every channel.\n", + " - A pixelized source: a `RectangularAdaptDensity` mesh with `Constant` regularization (1 free parameter \u2014 the\n", + " regularization coefficient). The pixelization itself has no per-pixel priors; the source-plane fluxes are a\n", + " linear inversion output computed by each channel's `AnalysisInterferometer` at fit time. That is what makes\n", + " each channel an independent linear solve while sharing all of the non-linear parameters.\n", + "\n", + "The total dimensionality of the non-linear parameter space is therefore 8." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source (pixelization, no per-pixel priors):\n", + "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", + "regularization = af.Model(al.reg.Constant)\n", + "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "# Overall lens model:\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Per-Channel Analyses__\n", + "\n", + "One `AnalysisInterferometer` per channel \u2014 there are no per-channel parameters here, only per-channel data.\n", + "Each analysis runs its own NUFFT, builds its own visibility-space inversion, and returns its own log-evidence\n", + "when called with a candidate lens model. The shared `positions_likelihood` is passed to every analysis to apply\n", + "the same global multiple-image penalty across the cube." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [\n", + " al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " settings=settings,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " use_jax=True,\n", + " )\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__FactorGraph__\n", + "\n", + "`af.AnalysisFactor` pairs each analysis with a deep copy of the base model, and `af.FactorGraphModel` combines\n", + "the factors into a single global model whose log-likelihood is the sum of the per-factor log-likelihoods. With\n", + "no per-factor prior overrides, every prior is *identified* across factors \u2014 the factor-graph machinery\n", + "deduplicates them \u2014 so the global model has the same dimensionality as the single-channel base model.\n", + "\n", + "If you wanted per-channel free parameters (for example, a per-channel `intensity` for a parametric source),\n", + "you'd override that prior on each `model.copy()` before wrapping it in an `AnalysisFactor`. See\n", + "`autolens_workspace/scripts/multi/modeling.py` for how that works in the multi-band case." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = [\n", + " af.AnalysisFactor(prior_model=model.copy(), analysis=analysis)\n", + " for analysis in analysis_list\n", + "]\n", + "\n", + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + "print(f\" channels in factor graph: {len(analysis_factor_list)}\")\n", + "print(\n", + " f\" global model free parameters: {factor_graph.global_prior_model.total_free_parameters}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "`Nautilus` is the standard non-linear search for PyAutoLens. The lens dimensionality is unchanged from a\n", + "single-channel fit, so `n_live=100` is a reasonable starting point \u2014 but the per-likelihood cost is N times\n", + "larger because each likelihood call runs N inversions, so wall-clock time scales linearly with the number of\n", + "channels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\") / \"datacube\",\n", + " name=\"start_here\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20,\n", + " iterations_per_quick_update=50000,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "We pass the factor graph's `global_prior_model` as the model and the factor graph itself as the analysis \u2014 this\n", + "is the same shape used for any multi-dataset PyAutoFit fit. Internally the search proposes lens parameters from\n", + "the global prior, hands them to the FactorGraph's `log_likelihood_function`, and the FactorGraph routes them\n", + "into every channel's `AnalysisInterferometer.log_likelihood_function` and sums the per-channel log-evidences.\n", + "\n", + "**Run time on CPU is dominated by the per-channel inversion.** A 4-channel SMA-scale cube finishes in a few\n", + "hours on CPU; ALMA-scale cubes with 50+ channels need GPU acceleration to complete in reasonable time." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " Per-channel inversions multiply the per-likelihood cost \u2014 expect this to take longer than a single-channel\n", + " interferometer fit. On CPU plan for hours; on GPU, tens of minutes.\n", + " \"\"\"\n", + ")\n", + "\n", + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "`result_list` contains one entry per factor \u2014 one per channel \u2014 each carrying its own `FitInterferometer`\n", + "against the maximum-likelihood lens model. Use them to inspect:\n", + "\n", + " - Per-channel source reconstructions (each channel's pixelized source is independent).\n", + " - Per-channel dirty images and residuals.\n", + " - The shared maximum-likelihood lens parameters (identical across factors by construction).\n", + "\n", + "Compared against `dataset/interferometer/datacube//cube_summary.json`, the per-channel reconstructed total\n", + "flux should trace the input emission-line spectrum.\n", + "\n", + "__Wrap Up__\n", + "\n", + "This script walks through the full datacube modeling pipeline. For deeper dives:\n", + "\n", + " - `modeling.py` \u2014 the focused FactorGraph + Nautilus example, ready to copy and adapt for your own cube.\n", + " - `simulator.py` \u2014 how the reference cube is generated.\n", + " - `autolens_workspace_developer/datacube/likelihood_function.py` \u2014 a step-by-step JAX walkthrough of how the\n", + " per-channel log-evidences are built and summed inside the FactorGraph, with an explicit eager-vs-JIT\n", + " correctness check.\n", + "\n", + "Phase 2 work \u2014 the shared-`L\u1d40 W\u0303 L` optimisation that exploits channel-invariant `uv_wavelengths` and\n", + "`noise_map` to bring ALMA-scale cubes inside CPU runtime budgets \u2014 is a separate follow-up issue." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/extra_galaxies/modeling.ipynb b/notebooks/interferometer/features/extra_galaxies/modeling.ipynb index e380c79bd..2a3abf743 100644 --- a/notebooks/interferometer/features/extra_galaxies/modeling.ipynb +++ b/notebooks/interferometer/features/extra_galaxies/modeling.ipynb @@ -1,473 +1,510 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Extra Galaxies\n", - "=================================\n", - "\n", - "\n", - "There may be extra galaxies nearby the lens and source galaxies, whose mass may contribute to the ray-tracing and\n", - "lens model.\n", - "\n", - "They may also emit luminous emission, but for interferometer datasets they are rarely detected at the long wavelengths\n", - "probed. The extra galaxies examples for interferometer throughout the workspace therefore do not include their light\n", - "(unlike the CCD image examples).\n", - "\n", - "This example shows how to perform lens modeling which accounts for the mass of these extra galaxies. The centres of\n", - "each galaxy (e.g. their brightest pixels observed from imaging data) are used as the centre of the mass profiles of\n", - "these galaxies, in order to reduce model complexity.\n", - "\n", - "__Contents__\n", - "\n", - "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Extra Galaxies Dataset:** We are now going to model the dataset with extra galaxies included in the model, where these.\n", - "- **Extra Galaxies Centres:** To set up a lens model including each extra galaxy with a mass profile, we input manually the.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Extra Galaxies Model:** We now use the modeling API to create the model for the extra galaxies.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Scaling Relations:** The modeling API has full support for composing the extra galaxies such that their mass follows.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Data Preparation__\n", - "\n", - "To perform modeling which accounts for extra galaxies, a list of the centre of each extra galaxy are used to set up\n", - "the model-fit. For the example dataset used here, these tasks have already been performed and the\n", - "metadata (`mask_extra_galaxies.fits` and `extra_galaxies_centres.json` are already included in results folder.\n", - "\n", - "The tutorial `autolens_workspace/*/imaging/data_preparation/optional/extra_galaxies_centres.py`\n", - "describes how to create these centres and output them to a `.json` file. You will need to use imaging data\n", - "to do not, as interferometer data rarely detects the light of these extra galaxies. If this data is not available,\n", - "you probably dont have any evidence of there being multiple galaxies in the system!\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit \n", - "with the lens model.\n", - "\n", - "We load the `extra_galaxies` dataset, which includes two extra galaxies either side of the main lens galaxy.\n", - "\n", - "Load and plot the strong lens dataset `extra_galaxies` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/features/extra_galaxies/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We do not perform a model-fit using this dataset, as using a mask like this requires that we use a pixelization\n", - "to fit the lensed source, which you may not be familiar with yet.\n", - "\n", - "In the `features/pixelization` example we perform a fit using this noise scaling scheme and a pixelization,\n", - "so check this out if you are interested in how to do this.\n", - "\n", - "__Extra Galaxies Dataset__\n", - "\n", - "We are now going to model the dataset with extra galaxies included in the model, where these galaxies include\n", - "the mass profiles of the extra galaxies.\n", - "\n", - "__Extra Galaxies Centres__\n", - "\n", - "To set up a lens model including each extra galaxy with a mass profile, we input manually the centres of\n", - "the extra galaxies.\n", - "\n", - "In principle, a lens model including the extra galaxies could be composed without these centres. For example, if \n", - "there were two extra galaxies in the data, we could simply add two additional mass profiles into the model. \n", - "The modeling API does support this, but we will not use it in this example.\n", - "\n", - "This is because models where the extra galaxies have free centres are often too complex to fit. It is likely the fit \n", - "will infer an inaccurate lens model and local maxima, because the parameter space is too complex.\n", - "\n", - "For example, a common problem is that one of the extra galaxy light profiles intended to model a nearby galaxy instead \n", - "fits one of the lensed source's multiple images. Alternatively, an extra galaxy's mass profile may recenter itself and \n", - "act as part of the main lens galaxy's mass distribution.\n", - "\n", - "Therefore, when modeling extra galaxies we input the centre of each, in order to fix their mass profile \n", - "centres or set up priors centre around these values.\n", - "\n", - "The `data_preparation` tutorial `autolens_workspace/*/imaging/data_preparation/examples/optional/extra_galaxies_centres.py` \n", - "describes how to create these centres. Using this script they have been output to the `.json` file we load below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_centres = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"extra_galaxies_centres.json\"))\n", - ")\n", - "\n", - "print(extra_galaxies_centres)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__ \n", - "\n", - "Perform the normal steps to set up the main model of the lens galaxy and source.\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=5,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Model__ \n", - "\n", - "We now use the modeling API to create the model for the extra galaxies.\n", - "\n", - "Currently, the extra galaxies API require that the centres of the mass profiles are fixed to the input centres\n", - "(but the other parameters of the mass profiles remain free). \n", - "\n", - "Therefore, in this example fits a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", - "\n", - " - The source galaxy's light is a Multi Gaussian Expansion [4 parameters].\n", - "\n", - " - Each extra galaxy's total mass distribution is a `IsothermalSph` profile with fixed \n", - " centre [2 extra galaxies x 1 parameters = 2 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=11.\n", - "\n", - "Extra galaxy mass profiles often to go unphysically high `einstein_radius` values, degrading the fit. The \n", - "`einstein_radius` parameter is set a `UniformPrior` with an upper limit of 0.1\" to prevent this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Extra Galaxies:\n", - "\n", - "extra_galaxies_list = []\n", - "\n", - "for extra_galaxy_centre in extra_galaxies_centres:\n", - "\n", - " # Extra Galaxy Mass\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - "\n", - " mass.centre = extra_galaxy_centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", - "\n", - " # Extra Galaxy\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source=source), extra_galaxies=extra_galaxies\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute confirms the model includes extra galaxies that we defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Analysis__ \n", - "\n", - "The code below performs the normal steps to set up a model-fit.\n", - "\n", - "Given the extra model parameters due to the extra galaxies, we increase the number of live points to 200." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\") / \"features\",\n", - " name=\"extra_galaxies_model\",\n", - " unique_tag=dataset_name,\n", - " n_live=200,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=20000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", - "required by a model.\n", - "\n", - "Adding extra galaxies with just mass profiles has a negligible effect on VRAM, because mass profiles are fast to\n", - "compute and do not require large images to be stored in VRAM.\n", - "\n", - "__Run Time__\n", - "\n", - "Adding extra galaxies to the model increases the likelihood evaluation times, because their mass profiles need their \n", - "deflection angles computed. These calculations are pretty fast, so only a small increase in time is expected.\n", - "\n", - "The bigger hit on run time is due to the extra free parameters, 1 `einstein_radius` per extra galaxy for its mass. \n", - "This increases the dimensionality of non-linear parameter space. This means Nautilus takes longer to converge on \n", - "the highest likelihood regions of parameter space.\n", - "\n", - "The Source, Light and Mass (SLaM) pipelines support extra galaxies but in a way that reduces the number of free\n", - "parameters they add to the model. This is described in the `slam` examples. The `group` package, which models systems\n", - "with 10+ extra galaxies, introduces even more clever parameterizations which add 0 free parameters per extra galaxy,\n", - "so if your model has many extra galaxies you should check out the `group` package.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We can now begin the model-fit by passing the model and analysis object to the search, which performs a non-linear\n", - "search to find which models fit the data with the highest likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "By plotting the maximum log likelihood `FitInterferometer` object we can confirm the extra galaxies contribute to the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "These examples show how the results API can be extended to investigate extra galaxies in the results.\n", - "\n", - "__Scaling Relations__\n", - "\n", - "The modeling API has full support for composing the extra galaxies such that their mass follows scaling\n", - "relations. For example, you could assume that the mass of the extra galaxies is related to their luminosity via a\n", - "constant mass-to-light ratio.\n", - "\n", - "This is documented in the `autolens_workspace/*/imaging/features/scaling_relation` example.\n", - "\n", - "__Wrap Up__\n", - "\n", - "The extra galaxies API makes it straight forward for us to model galaxy-scale strong lenses with additional components\n", - "for mass of nearby objects.\n", - "\n", - "The `autolens_workspace` includes a `group` package, for modeling group scale strong lenses which have multiple lens \n", - "galaxies. When you should use the extra galaxies API as shown here, and when you should use the group package? \n", - "\n", - "The distinction is as follows:\n", - "\n", - " - A galaxy scale lens is a system which can be modeled to a high level of accuracy using a single light and mass \n", - " distribution for the main lens galaxy. Including additional galaxies in the model via the extra galaxies API makes small \n", - " improvements on the lens model, but a good fit is possible without them. \n", - "\n", - " - A group scale lens is a system which cannot be modeled to a high level of accuracy using a single light and mass \n", - " distribution. Defining a 'main' lens galaxy is unclear and two or more main lens galaxies are required to fit an \n", - " accurate model. \n", - "\n", - "The `group` package also uses the extra galaxies API for model composition, but does so to compose and fit more \n", - "complex lens models." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Extra Galaxies\n", + "=================================\n", + "\n", + "\n", + "There may be extra galaxies nearby the lens and source galaxies, whose mass may contribute to the ray-tracing and\n", + "lens model.\n", + "\n", + "They may also emit luminous emission, but for interferometer datasets they are rarely detected at the long wavelengths\n", + "probed. The extra galaxies examples for interferometer throughout the workspace therefore do not include their light\n", + "(unlike the CCD image examples).\n", + "\n", + "This example shows how to perform lens modeling which accounts for the mass of these extra galaxies. The centres of\n", + "each galaxy (e.g. their brightest pixels observed from imaging data) are used as the centre of the mass profiles of\n", + "these galaxies, in order to reduce model complexity.\n", + "\n", + "__Contents__\n", + "\n", + "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Extra Galaxies Dataset:** We are now going to model the dataset with extra galaxies included in the model, where these.\n", + "- **Extra Galaxies Centres:** To set up a lens model including each extra galaxy with a mass profile, we input manually the.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Extra Galaxies Model:** We now use the modeling API to create the model for the extra galaxies.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Scaling Relations:** The modeling API has full support for composing the extra galaxies such that their mass follows.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Data Preparation__\n", + "\n", + "To perform modeling which accounts for extra galaxies, a list of the centre of each extra galaxy are used to set up\n", + "the model-fit. For the example dataset used here, these tasks have already been performed and the\n", + "metadata (`mask_extra_galaxies.fits` and `extra_galaxies_centres.json` are already included in results folder.\n", + "\n", + "The tutorial `autolens_workspace/*/imaging/data_preparation/optional/extra_galaxies_centres.py`\n", + "describes how to create these centres and output them to a `.json` file. You will need to use imaging data\n", + "to do not, as interferometer data rarely detects the light of these extra galaxies. If this data is not available,\n", + "you probably dont have any evidence of there being multiple galaxies in the system!\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit \n", + "with the lens model.\n", + "\n", + "We load the `extra_galaxies` dataset, which includes two extra galaxies either side of the main lens galaxy.\n", + "\n", + "Load and plot the strong lens dataset `extra_galaxies` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/features/extra_galaxies/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We do not perform a model-fit using this dataset, as using a mask like this requires that we use a pixelization\n", + "to fit the lensed source, which you may not be familiar with yet.\n", + "\n", + "In the `features/pixelization` example we perform a fit using this noise scaling scheme and a pixelization,\n", + "so check this out if you are interested in how to do this.\n", + "\n", + "__Extra Galaxies Dataset__\n", + "\n", + "We are now going to model the dataset with extra galaxies included in the model, where these galaxies include\n", + "the mass profiles of the extra galaxies.\n", + "\n", + "__Extra Galaxies Centres__\n", + "\n", + "To set up a lens model including each extra galaxy with a mass profile, we input manually the centres of\n", + "the extra galaxies.\n", + "\n", + "In principle, a lens model including the extra galaxies could be composed without these centres. For example, if \n", + "there were two extra galaxies in the data, we could simply add two additional mass profiles into the model. \n", + "The modeling API does support this, but we will not use it in this example.\n", + "\n", + "This is because models where the extra galaxies have free centres are often too complex to fit. It is likely the fit \n", + "will infer an inaccurate lens model and local maxima, because the parameter space is too complex.\n", + "\n", + "For example, a common problem is that one of the extra galaxy light profiles intended to model a nearby galaxy instead \n", + "fits one of the lensed source's multiple images. Alternatively, an extra galaxy's mass profile may recenter itself and \n", + "act as part of the main lens galaxy's mass distribution.\n", + "\n", + "Therefore, when modeling extra galaxies we input the centre of each, in order to fix their mass profile \n", + "centres or set up priors centre around these values.\n", + "\n", + "The `data_preparation` tutorial `autolens_workspace/*/imaging/data_preparation/examples/optional/extra_galaxies_centres.py` \n", + "describes how to create these centres. Using this script they have been output to the `.json` file we load below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_centres = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"extra_galaxies_centres.json\"))\n", + ")\n", + "\n", + "print(extra_galaxies_centres)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__ \n", + "\n", + "Perform the normal steps to set up the main model of the lens galaxy and source.\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=5,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Model__ \n", + "\n", + "We now use the modeling API to create the model for the extra galaxies.\n", + "\n", + "Currently, the extra galaxies API require that the centres of the mass profiles are fixed to the input centres\n", + "(but the other parameters of the mass profiles remain free). \n", + "\n", + "Therefore, in this example fits a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + "\n", + " - The source galaxy's light is a Multi Gaussian Expansion [4 parameters].\n", + "\n", + " - Each extra galaxy's total mass distribution is a `IsothermalSph` profile with fixed \n", + " centre [2 extra galaxies x 1 parameters = 2 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=11.\n", + "\n", + "Extra galaxy mass profiles often to go unphysically high `einstein_radius` values, degrading the fit. The \n", + "`einstein_radius` parameter is set a `UniformPrior` with an upper limit of 0.1\" to prevent this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Extra Galaxies:\n", + "\n", + "extra_galaxies_list = []\n", + "\n", + "for extra_galaxy_centre in extra_galaxies_centres:\n", + "\n", + " # Extra Galaxy Mass\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + "\n", + " mass.centre = extra_galaxy_centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.5)\n", + "\n", + " # Extra Galaxy\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source=source), extra_galaxies=extra_galaxies\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute confirms the model includes extra galaxies that we defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Analysis__ \n", + "\n", + "The code below performs the normal steps to set up a model-fit.\n", + "\n", + "Given the extra model parameters due to the extra galaxies, we increase the number of live points to 200." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\") / \"features\",\n", + " name=\"extra_galaxies_model\",\n", + " unique_tag=dataset_name,\n", + " n_live=200,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=20000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", + "required by a model.\n", + "\n", + "Adding extra galaxies with just mass profiles has a negligible effect on VRAM, because mass profiles are fast to\n", + "compute and do not require large images to be stored in VRAM.\n", + "\n", + "__Run Time__\n", + "\n", + "Adding extra galaxies to the model increases the likelihood evaluation times, because their mass profiles need their \n", + "deflection angles computed. These calculations are pretty fast, so only a small increase in time is expected.\n", + "\n", + "The bigger hit on run time is due to the extra free parameters, 1 `einstein_radius` per extra galaxy for its mass. \n", + "This increases the dimensionality of non-linear parameter space. This means Nautilus takes longer to converge on \n", + "the highest likelihood regions of parameter space.\n", + "\n", + "The Source, Light and Mass (SLaM) pipelines support extra galaxies but in a way that reduces the number of free\n", + "parameters they add to the model. This is described in the `slam` examples. The `group` package, which models systems\n", + "with 10+ extra galaxies, introduces even more clever parameterizations which add 0 free parameters per extra galaxy,\n", + "so if your model has many extra galaxies you should check out the `group` package.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We can now begin the model-fit by passing the model and analysis object to the search, which performs a non-linear\n", + "search to find which models fit the data with the highest likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "By plotting the maximum log likelihood `FitInterferometer` object we can confirm the extra galaxies contribute to the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "These examples show how the results API can be extended to investigate extra galaxies in the results.\n", + "\n", + "__Scaling Relations__\n", + "\n", + "The modeling API has full support for composing the extra galaxies such that their mass follows scaling\n", + "relations. For example, you could assume that the mass of the extra galaxies is related to their luminosity via a\n", + "constant mass-to-light ratio.\n", + "\n", + "This is documented in the `autolens_workspace/*/imaging/features/scaling_relation` example.\n", + "\n", + "__Wrap Up__\n", + "\n", + "The extra galaxies API makes it straight forward for us to model galaxy-scale strong lenses with additional components\n", + "for mass of nearby objects.\n", + "\n", + "The `autolens_workspace` includes a `group` package, for modeling group scale strong lenses which have multiple lens \n", + "galaxies. When you should use the extra galaxies API as shown here, and when you should use the group package? \n", + "\n", + "The distinction is as follows:\n", + "\n", + " - A galaxy scale lens is a system which can be modeled to a high level of accuracy using a single light and mass \n", + " distribution for the main lens galaxy. Including additional galaxies in the model via the extra galaxies API makes small \n", + " improvements on the lens model, but a good fit is possible without them. \n", + "\n", + " - A group scale lens is a system which cannot be modeled to a high level of accuracy using a single light and mass \n", + " distribution. Defining a 'main' lens galaxy is unclear and two or more main lens galaxies are required to fit an \n", + " accurate model. \n", + "\n", + "The `group` package also uses the extra galaxies API for model composition, but does so to compose and fit more \n", + "complex lens models." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/extra_galaxies/simulator.ipynb b/notebooks/interferometer/features/extra_galaxies/simulator.ipynb index 35813e2ba..ee157b346 100644 --- a/notebooks/interferometer/features/extra_galaxies/simulator.ipynb +++ b/notebooks/interferometer/features/extra_galaxies/simulator.ipynb @@ -1,455 +1,492 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Extra Galaxies\n", - "=========================\n", - "\n", - "There may be extra galaxies nearby the lens and source galaxies, whose mass may contribute to the ray-tracing and\n", - "lens model.\n", - "\n", - "They may also emit luminous emission, but for interferometer datasets they are rarely detected at the long wavelengths\n", - "probed. The extra galaxies examples for interferometer throughout the workspace therefore do not include their light\n", - "(unlike the CCD image examples).\n", - "\n", - "The emission of these galaxies may overlap the lensed source emission, and their mass may contribute to the lensing\n", - "of the source.\n", - "\n", - "We therefore will include these galaxies as mass profiles in the lens model, accounting for their lensing effects\n", - "via ray-tracing.\n", - "\n", - "This uses the modeling API, which is illustrated in the script `autolens_workspace/*/features/extra_galaxies/modeling`.\n", - "\n", - "This script simulates an interferometer dataset which includes extra galaxies near the lens and source\n", - "galaxies. This is used to illustrate the extra galaxies API in the script above.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Other Scripts:** To illustrate how compose and fit a lens model which includes the extra galaxies as light and mass.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Simulate:** Simulate the image using a (y,x) grid.\n", - "- **Galaxies:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", - "- **Extra Galaxies:** Includes two extra galaxies, which must be modeled to ensure the lens model is accurate.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "- **Multiple Images:** Output the multiple image positions of the source galaxy which can help with lens modeling.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Interferometer` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The source galaxy's light is an `Sersic`.\n", - " - There are two extra galaxies whose mass perturbs the lensed source's emission.\n", - "\n", - "__Other Scripts__\n", - "\n", - "To illustrate how compose and fit a lens model which includes the extra galaxies as light and mass profiles.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"interferometer\"\n", - "dataset_name = \"extra_galaxies\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image using a (y,x) grid.\n", - "\n", - "This simulated galaxy has additional galaxies and light profiles which are offset from the main galaxy centre \n", - "of (0.0\", 0.0\"). The adaptive over sampling grid has all centres input to account for this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(256, 256), pixel_scales=0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To perform the Fourier transform we need the wavelengths of the baselines." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "uv_wavelengths_path = Path(\"dataset\", dataset_type, \"uv_wavelengths\")\n", - "uv_wavelengths = al.ndarray_via_fits_from(\n", - " file_path=Path(uv_wavelengths_path, \"sma.fits\"), hdu=0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the interferometer data, which defines the exposure time, noise levels and transformer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorInterferometer(\n", - " uv_wavelengths=uv_wavelengths,\n", - " exposure_time=300.0,\n", - " noise_sigma=1000.0,\n", - " transformer_class=al.TransformerDFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxies__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.1, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=1.0,\n", - " sersic_index=2.5,\n", - " ),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "Includes two extra galaxies, which must be modeled to ensure the lens model is accurate.\n", - "\n", - "Note that their redshift is the same as the main galaxy, which is not necessarily the case in real observations. \n", - "\n", - "If they are at a different redshift, the tools for masking or modeling the luminous emission of the extra galaxies \n", - "are equipped to handle this.\n", - "\n", - "For mass modeling, their redshifts being different to the main galaxy will lead to multi-plane ray-tracing being\n", - "performed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_0_centre = (1.0, 3.5)\n", - "\n", - "extra_galaxy_0 = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=extra_galaxy_0_centre, einstein_radius=0.1),\n", - ")\n", - "\n", - "extra_galaxy_1_centre = (-2.0, -3.5)\n", - "\n", - "extra_galaxy_1 = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.IsothermalSph(centre=extra_galaxy_1_centre, einstein_radius=0.2),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated `Interferometer` dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an interferometer dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Interferometer` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_interferometer(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_interferometer_dirty_images(\n", - " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Centres__\n", - "\n", - "Output the centres of the extra galaxies to a .json file, so that they can be used to set up the model\n", - "in the modeling scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=al.Grid2DIrregular(values=[extra_galaxy_0_centre, extra_galaxy_1_centre]),\n", - " file_path=dataset_path / \"extra_galaxies_centres.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multiple Images__\n", - "\n", - "Output the multiple image positions of the source galaxy which can help with lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " file_path=dataset_path / \"positions.json\",\n", - " obj=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/interferometer/extra_galaxies`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Extra Galaxies\n", + "=========================\n", + "\n", + "There may be extra galaxies nearby the lens and source galaxies, whose mass may contribute to the ray-tracing and\n", + "lens model.\n", + "\n", + "They may also emit luminous emission, but for interferometer datasets they are rarely detected at the long wavelengths\n", + "probed. The extra galaxies examples for interferometer throughout the workspace therefore do not include their light\n", + "(unlike the CCD image examples).\n", + "\n", + "The emission of these galaxies may overlap the lensed source emission, and their mass may contribute to the lensing\n", + "of the source.\n", + "\n", + "We therefore will include these galaxies as mass profiles in the lens model, accounting for their lensing effects\n", + "via ray-tracing.\n", + "\n", + "This uses the modeling API, which is illustrated in the script `autolens_workspace/*/features/extra_galaxies/modeling`.\n", + "\n", + "This script simulates an interferometer dataset which includes extra galaxies near the lens and source\n", + "galaxies. This is used to illustrate the extra galaxies API in the script above.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Other Scripts:** To illustrate how compose and fit a lens model which includes the extra galaxies as light and mass.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Simulate:** Simulate the image using a (y,x) grid.\n", + "- **Galaxies:** Setup the lens galaxy's light, mass and source galaxy light for this simulated lens.\n", + "- **Extra Galaxies:** Includes two extra galaxies, which must be modeled to ensure the lens model is accurate.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "- **Multiple Images:** Output the multiple image positions of the source galaxy which can help with lens modeling.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Interferometer` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The source galaxy's light is an `Sersic`.\n", + " - There are two extra galaxies whose mass perturbs the lensed source's emission.\n", + "\n", + "__Other Scripts__\n", + "\n", + "To illustrate how compose and fit a lens model which includes the extra galaxies as light and mass profiles.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"interferometer\"\n", + "dataset_name = \"extra_galaxies\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image using a (y,x) grid.\n", + "\n", + "This simulated galaxy has additional galaxies and light profiles which are offset from the main galaxy centre \n", + "of (0.0\", 0.0\"). The adaptive over sampling grid has all centres input to account for this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(256, 256), pixel_scales=0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To perform the Fourier transform we need the wavelengths of the baselines." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "uv_wavelengths_path = Path(\"dataset\", dataset_type, \"uv_wavelengths\")\n", + "uv_wavelengths = al.ndarray_via_fits_from(\n", + " file_path=Path(uv_wavelengths_path, \"sma.fits\"), hdu=0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the interferometer data, which defines the exposure time, noise levels and transformer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorInterferometer(\n", + " uv_wavelengths=uv_wavelengths,\n", + " exposure_time=300.0,\n", + " noise_sigma=1000.0,\n", + " transformer_class=al.TransformerDFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxies__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.1, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=1.0,\n", + " sersic_index=2.5,\n", + " ),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "Includes two extra galaxies, which must be modeled to ensure the lens model is accurate.\n", + "\n", + "Note that their redshift is the same as the main galaxy, which is not necessarily the case in real observations. \n", + "\n", + "If they are at a different redshift, the tools for masking or modeling the luminous emission of the extra galaxies \n", + "are equipped to handle this.\n", + "\n", + "For mass modeling, their redshifts being different to the main galaxy will lead to multi-plane ray-tracing being\n", + "performed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_0_centre = (1.0, 3.5)\n", + "\n", + "extra_galaxy_0 = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=extra_galaxy_0_centre, einstein_radius=0.1),\n", + ")\n", + "\n", + "extra_galaxy_1_centre = (-2.0, -3.5)\n", + "\n", + "extra_galaxy_1 = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.IsothermalSph(centre=extra_galaxy_1_centre, einstein_radius=0.2),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated `Interferometer` dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[lens_galaxy, extra_galaxy_0, extra_galaxy_1, source_galaxy]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an interferometer dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Interferometer` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_interferometer(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_interferometer_dirty_images(\n", + " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Centres__\n", + "\n", + "Output the centres of the extra galaxies to a .json file, so that they can be used to set up the model\n", + "in the modeling scripts." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=al.Grid2DIrregular(values=[extra_galaxy_0_centre, extra_galaxy_1_centre]),\n", + " file_path=dataset_path / \"extra_galaxies_centres.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multiple Images__\n", + "\n", + "Output the multiple image positions of the source galaxy which can help with lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " file_path=dataset_path / \"positions.json\",\n", + " obj=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/interferometer/extra_galaxies`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/extra_galaxies/slam.ipynb b/notebooks/interferometer/features/extra_galaxies/slam.ipynb index dd23f3004..14f5eaee7 100644 --- a/notebooks/interferometer/features/extra_galaxies/slam.ipynb +++ b/notebooks/interferometer/features/extra_galaxies/slam.ipynb @@ -1,795 +1,832 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Extra Galaxies: SLaM\n", - "=====================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for fitting a\n", - "lens model where extra galaxies surrounding the lens are included in the lens model.\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", - "guide before working through this example.\n", - "\n", - "This example only provides documentation specific to the extra galaxies, describing how the pipeline\n", - "differs from the standard SLaM pipelines described in the SLaM start here guide.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **Group SLaM:** This SLaM pipeline is designed for the regime where one is modeling galaxy scale lenses with nearby.\n", - "- **This Script:** Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this.\n", - "- **SOURCE LP PIPELINE:** Initializes the mass model + source-light using MGE light profiles via `TransformerNUFFT`.\n", - "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP result.\n", - "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included as interferometer data.\n", - "- **Extra Galaxies Centres:** This is the same API as described in the `features/extra_galaxies.ipynb` example, where the centres.\n", - "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with `TransformerNUFFT` + sparse operator (source_pix onwards).\n", - "- **Sparse Operators:** The `pixelization/modeling` example describes how the sparse operator formalism speeds up.\n", - "- **Settings:** Disable the default position only linear algebra solver so the source reconstruction can have.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **Extra Galaxies:** Build the extra galaxies model: each extra galaxy has an `IsothermalSph` mass profile with its.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "- **Output:** The `start_here.ipynb` example describes how results can be output to hard-disk after the SLaM.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "- **Extra Galaxies** (`features/extra_galaxies.ipynb`):\n", - " How we include extra galaxies in the lens model, by using the centres of the galaxies\n", - " which have been determined beforehand.\n", - "\n", - "You can still run the script without fully understanding the guide, but reviewing it later will\n", - "make the structure and choices of the SLaM workflow clearer.\n", - "\n", - "__Group SLaM__\n", - "\n", - "This SLaM pipeline is designed for the regime where one is modeling galaxy scale lenses with nearby surrounding\n", - "extra galaxies.\n", - "\n", - "However, these systems can often become close to the group scale lensing regime, for which PyAutoLens has a dedicated\n", - "package for modeling (`autolens_workspace/*/group`) and its own dedicated SLaM pipelines.\n", - "\n", - "The main difference between this SLaM pipeline and the group SLaM pipelines is that in the latter, the masses of\n", - "the extra galaxies are modeled using scaling relations tied to their light profiles. The group SLaM pipeline has\n", - "additional searches in the SOURCE LP PIPELINE to measure the luminosities of the extra galaxies for this purpose.\n", - "\n", - "Which SLaM pipeline you should use depends on your particular strong lens, but as a rule of thumb if you are\n", - "including a lot of extra galaxies (e.g. more than 5) and your model complexity is increasing significantly, you should\n", - "consider using the group SLaM pipelines.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this SLaM modeling\n", - "script fits `Imaging` dataset of a strong lens system where in the final model:\n", - "\n", - " - The lens galaxy's light is a bulge with Multiple Gaussian Expansion (MGE) light profile.\n", - " - The lens galaxy's total mass distribution is an `PowerLaw` plus an `ExternalShear`.\n", - " - The source galaxy's light is a `Pixelization`.\n", - " - Two extra galaxies are included in the model, each with their mass as a `IsothermalSph` profile.\n", - "\n", - "This modeling script uses the SLaM pipelines:\n", - "\n", - " `source_pix`\n", - " `light_lp`\n", - " `mass_total`\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Initializes a robust mass + source-light model using a Multi Gaussian Expansion (MGE), fit with light profiles.\n", - "\n", - "This stage uses `dataset_nufft` (built with `TransformerNUFFT`, backed by JAX-native `nufftax`), which makes\n", - "light-profile fitting fast even on ALMA-class datasets with millions of visibilities. The result provides the\n", - "adapt image and position likelihood threaded into the pixelized pipelines that follow.\n", - "\n", - "Each extra galaxy carries an `IsothermalSph` mass profile centred on its known position. Extra galaxies do not\n", - "have light components in the interferometer pipeline, because interferometer data does not contain lens light\n", - "emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " extra_galaxies,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "The first search of the SOURCE PIX PIPELINE fits a pixelization whose purpose is to generate a high-quality\n", - "adapt image used in search 2. It uses the adapt image computed from the SOURCE LP result, with the position\n", - "likelihood derived automatically via `source_lp_result.positions_likelihood_from(...)`.\n", - "\n", - "This stage uses `dataset_sparse` (built with `TransformerNUFFT` + `apply_sparse_operator`). Pixelizations\n", - "exploit sparsity in the linear inversion rather than the NUFFT path.\n", - "\n", - "The `extra_galaxies` mass priors are carried forward from the SOURCE LP result as free model parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " extra_galaxies = source_lp_result.model.extra_galaxies\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=source_lp_result.model.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " extra_galaxies=extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", - "pixelization and regularization.\n", - "\n", - "The extra galaxies are passed from `source_pix_result_1` as fixed instances.\n", - "\n", - "Note that the LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not\n", - "contain lens light emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " settings=settings,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " # interferometry does not support lens light\n", - " bulge=None,\n", - " disk=None,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", - "contain lens light emission.\n", - "\n", - "The extra galaxies are passed from `source_pix_result_1` as free model parameters, so their masses are\n", - "updated during this search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " source_pix_result_2: af.Result,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_1.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=source_pix_result_1.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " extra_galaxies=source_pix_result_1.model.extra_galaxies,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load the `Interferometer` data, define the visibility and real-space masks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"extra_galaxies\"\n", - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", - ")\n", - "\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/features/extra_galaxies/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "# dataset_name = \"alma\"\n", - "\n", - "# if dataset_name == \"alma\":\n", - "#\n", - "# real_space_mask = al.Mask2D.circular(\n", - "# shape_native=(800, 800),\n", - "# pixel_scales=0.01,\n", - "# radius=mask_radius,\n", - "# )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Two Datasets__\n", - "\n", - "The SLaM pipeline runs in two phases that prefer different transformers:\n", - "\n", - "- `dataset_nufft` uses `TransformerNUFFT` (backed by JAX-native `nufftax`) for the `source_lp` stage.\n", - "- `dataset_sparse` uses `TransformerNUFFT` + `apply_sparse_operator(...)` for `source_pix_1`,\n", - " `source_pix_2` and `mass_total`.\n", - "\n", - "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload) differ." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_nufft = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "dataset_sparse = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Centres__\n", - "\n", - "This is the same API as described in the `features/extra_galaxies.ipynb` example, where the centres of the extra\n", - "galaxies are loaded from a `.json` file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_centres = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"extra_galaxies_centres.json\"))\n", - ")\n", - "\n", - "print(extra_galaxies_centres)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", - "pixelized source modeling, especially for many visibilities.\n", - "\n", - "We use a try / except to load the pre-computed curvature preload, which is necessary to use\n", - "the sparse operator formalism. If this file does not exist (e.g. you have not made it manually via\n", - "the `many_visibilities_preparation` example) it is made here.\n", - "\n", - "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", - "`source_lp` does not need it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " nufft_precision_operator = np.load(\n", - " file=dataset_path / \"nufft_precision_operator.npy\",\n", - " )\n", - "except FileNotFoundError:\n", - " nufft_precision_operator = None\n", - "\n", - "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", - " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "Disable the default position only linear algebra solver so the source reconstruction can have\n", - "negative pixel values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"interferometer\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies__\n", - "\n", - "Build the extra galaxies model: each extra galaxy has an `IsothermalSph` mass profile with its centre fixed\n", - "to the known extra galaxy centre and a free `einstein_radius`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxies_list = []\n", - "\n", - "for extra_galaxy_centre in extra_galaxies_centres:\n", - "\n", - " mass = af.Model(al.mp.IsothermalSph)\n", - " mass.centre = extra_galaxy_centre\n", - " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.1)\n", - "\n", - " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - " extra_galaxy.mass.centre = extra_galaxy_centre\n", - "\n", - " extra_galaxies_list.append(extra_galaxy)\n", - "\n", - "extra_galaxies = af.Collection(extra_galaxies_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step.\n", - "\n", - "Note the transformer split: `source_lp` is passed `dataset_nufft` (TransformerNUFFT), while every later stage\n", - "is passed `dataset_sparse` (TransformerNUFFT + sparse operator)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset_nufft,\n", - " mask_radius=mask_radius,\n", - " extra_galaxies=extra_galaxies,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " settings=settings,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "The `start_here.ipynb` example describes how results can be output to hard-disk after the SLaM pipelines have been run.\n", - "Checkout that script for a complete description of the output of this script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Extra Galaxies: SLaM\n", + "=====================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for fitting a\n", + "lens model where extra galaxies surrounding the lens are included in the lens model.\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", + "guide before working through this example.\n", + "\n", + "This example only provides documentation specific to the extra galaxies, describing how the pipeline\n", + "differs from the standard SLaM pipelines described in the SLaM start here guide.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **Group SLaM:** This SLaM pipeline is designed for the regime where one is modeling galaxy scale lenses with nearby.\n", + "- **This Script:** Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this.\n", + "- **SOURCE LP PIPELINE:** Initializes the mass model + source-light using MGE light profiles via `TransformerNUFFT`.\n", + "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP result.\n", + "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included as interferometer data.\n", + "- **Extra Galaxies Centres:** This is the same API as described in the `features/extra_galaxies.ipynb` example, where the centres.\n", + "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with `TransformerNUFFT` + sparse operator (source_pix onwards).\n", + "- **Sparse Operators:** The `pixelization/modeling` example describes how the sparse operator formalism speeds up.\n", + "- **Settings:** Disable the default position only linear algebra solver so the source reconstruction can have.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **Extra Galaxies:** Build the extra galaxies model: each extra galaxy has an `IsothermalSph` mass profile with its.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "- **Output:** The `start_here.ipynb` example describes how results can be output to hard-disk after the SLaM.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "- **Extra Galaxies** (`features/extra_galaxies.ipynb`):\n", + " How we include extra galaxies in the lens model, by using the centres of the galaxies\n", + " which have been determined beforehand.\n", + "\n", + "You can still run the script without fully understanding the guide, but reviewing it later will\n", + "make the structure and choices of the SLaM workflow clearer.\n", + "\n", + "__Group SLaM__\n", + "\n", + "This SLaM pipeline is designed for the regime where one is modeling galaxy scale lenses with nearby surrounding\n", + "extra galaxies.\n", + "\n", + "However, these systems can often become close to the group scale lensing regime, for which PyAutoLens has a dedicated\n", + "package for modeling (`autolens_workspace/*/group`) and its own dedicated SLaM pipelines.\n", + "\n", + "The main difference between this SLaM pipeline and the group SLaM pipelines is that in the latter, the masses of\n", + "the extra galaxies are modeled using scaling relations tied to their light profiles. The group SLaM pipeline has\n", + "additional searches in the SOURCE LP PIPELINE to measure the luminosities of the extra galaxies for this purpose.\n", + "\n", + "Which SLaM pipeline you should use depends on your particular strong lens, but as a rule of thumb if you are\n", + "including a lot of extra galaxies (e.g. more than 5) and your model complexity is increasing significantly, you should\n", + "consider using the group SLaM pipelines.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this SLaM modeling\n", + "script fits `Imaging` dataset of a strong lens system where in the final model:\n", + "\n", + " - The lens galaxy's light is a bulge with Multiple Gaussian Expansion (MGE) light profile.\n", + " - The lens galaxy's total mass distribution is an `PowerLaw` plus an `ExternalShear`.\n", + " - The source galaxy's light is a `Pixelization`.\n", + " - Two extra galaxies are included in the model, each with their mass as a `IsothermalSph` profile.\n", + "\n", + "This modeling script uses the SLaM pipelines:\n", + "\n", + " `source_pix`\n", + " `light_lp`\n", + " `mass_total`\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Initializes a robust mass + source-light model using a Multi Gaussian Expansion (MGE), fit with light profiles.\n", + "\n", + "This stage uses `dataset_nufft` (built with `TransformerNUFFT`, backed by JAX-native `nufftax`), which makes\n", + "light-profile fitting fast even on ALMA-class datasets with millions of visibilities. The result provides the\n", + "adapt image and position likelihood threaded into the pixelized pipelines that follow.\n", + "\n", + "Each extra galaxy carries an `IsothermalSph` mass profile centred on its known position. Extra galaxies do not\n", + "have light components in the interferometer pipeline, because interferometer data does not contain lens light\n", + "emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " extra_galaxies,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "The first search of the SOURCE PIX PIPELINE fits a pixelization whose purpose is to generate a high-quality\n", + "adapt image used in search 2. It uses the adapt image computed from the SOURCE LP result, with the position\n", + "likelihood derived automatically via `source_lp_result.positions_likelihood_from(...)`.\n", + "\n", + "This stage uses `dataset_sparse` (built with `TransformerNUFFT` + `apply_sparse_operator`). Pixelizations\n", + "exploit sparsity in the linear inversion rather than the NUFFT path.\n", + "\n", + "The `extra_galaxies` mass priors are carried forward from the SOURCE LP result as free model parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " extra_galaxies = source_lp_result.model.extra_galaxies\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=source_lp_result.model.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " extra_galaxies=extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", + "pixelization and regularization.\n", + "\n", + "The extra galaxies are passed from `source_pix_result_1` as fixed instances.\n", + "\n", + "Note that the LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not\n", + "contain lens light emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " settings=settings,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " # interferometry does not support lens light\n", + " bulge=None,\n", + " disk=None,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " extra_galaxies=source_pix_result_1.instance.extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", + "contain lens light emission.\n", + "\n", + "The extra galaxies are passed from `source_pix_result_1` as free model parameters, so their masses are\n", + "updated during this search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " source_pix_result_2: af.Result,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_1.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=source_pix_result_1.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " extra_galaxies=source_pix_result_1.model.extra_galaxies,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load the `Interferometer` data, define the visibility and real-space masks." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"extra_galaxies\"\n", + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", + ")\n", + "\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/features/extra_galaxies/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "# dataset_name = \"alma\"\n", + "\n", + "# if dataset_name == \"alma\":\n", + "#\n", + "# real_space_mask = al.Mask2D.circular(\n", + "# shape_native=(800, 800),\n", + "# pixel_scales=0.01,\n", + "# radius=mask_radius,\n", + "# )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Two Datasets__\n", + "\n", + "The SLaM pipeline runs in two phases that prefer different transformers:\n", + "\n", + "- `dataset_nufft` uses `TransformerNUFFT` (backed by JAX-native `nufftax`) for the `source_lp` stage.\n", + "- `dataset_sparse` uses `TransformerNUFFT` + `apply_sparse_operator(...)` for `source_pix_1`,\n", + " `source_pix_2` and `mass_total`.\n", + "\n", + "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload) differ." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_nufft = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "dataset_sparse = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Centres__\n", + "\n", + "This is the same API as described in the `features/extra_galaxies.ipynb` example, where the centres of the extra\n", + "galaxies are loaded from a `.json` file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_centres = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"extra_galaxies_centres.json\"))\n", + ")\n", + "\n", + "print(extra_galaxies_centres)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", + "pixelized source modeling, especially for many visibilities.\n", + "\n", + "We use a try / except to load the pre-computed curvature preload, which is necessary to use\n", + "the sparse operator formalism. If this file does not exist (e.g. you have not made it manually via\n", + "the `many_visibilities_preparation` example) it is made here.\n", + "\n", + "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", + "`source_lp` does not need it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " nufft_precision_operator = np.load(\n", + " file=dataset_path / \"nufft_precision_operator.npy\",\n", + " )\n", + "except FileNotFoundError:\n", + " nufft_precision_operator = None\n", + "\n", + "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", + " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "Disable the default position only linear algebra solver so the source reconstruction can have\n", + "negative pixel values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"interferometer\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies__\n", + "\n", + "Build the extra galaxies model: each extra galaxy has an `IsothermalSph` mass profile with its centre fixed\n", + "to the known extra galaxy centre and a free `einstein_radius`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxies_list = []\n", + "\n", + "for extra_galaxy_centre in extra_galaxies_centres:\n", + "\n", + " mass = af.Model(al.mp.IsothermalSph)\n", + " mass.centre = extra_galaxy_centre\n", + " mass.einstein_radius = af.UniformPrior(lower_limit=0.0, upper_limit=0.1)\n", + "\n", + " extra_galaxy = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + " extra_galaxy.mass.centre = extra_galaxy_centre\n", + "\n", + " extra_galaxies_list.append(extra_galaxy)\n", + "\n", + "extra_galaxies = af.Collection(extra_galaxies_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step.\n", + "\n", + "Note the transformer split: `source_lp` is passed `dataset_nufft` (TransformerNUFFT), while every later stage\n", + "is passed `dataset_sparse` (TransformerNUFFT + sparse operator)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset_nufft,\n", + " mask_radius=mask_radius,\n", + " extra_galaxies=extra_galaxies,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " settings=settings,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "The `start_here.ipynb` example describes how results can be output to hard-disk after the SLaM pipelines have been run.\n", + "Checkout that script for a complete description of the output of this script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/linear_light_profiles/fit.ipynb b/notebooks/interferometer/features/linear_light_profiles/fit.ipynb index 2f9249412..ac6e829f9 100644 --- a/notebooks/interferometer/features/linear_light_profiles/fit.ipynb +++ b/notebooks/interferometer/features/linear_light_profiles/fit.ipynb @@ -1,376 +1,413 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Linear Light Profiles Fit (Interferometer)\n", - "==============================================================\n", - "\n", - "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", - "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", - "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood) given\n", - "the light profile's other parameters.\n", - "\n", - "This script illustrates how to perform a single `FitInterferometer` of a linear light profile model \u2014 that is,\n", - "not the full Nautilus model-fit, but a single likelihood evaluation given known light/mass profile parameters.\n", - "This is useful for understanding how the inversion produces the solved-for `intensity`, and how to extract that\n", - "value from the resulting fit.\n", - "\n", - "For an explanation of why linear light profile fits are now practical against visibility data thanks to the\n", - "JAX-native NUFFT `nufftax` (https://github.com/GragasLab/nufftax), see the companion `modeling.py` example.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Benefits and drawbacks of linear light profiles for interferometer data.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Model:** The lens model whose `intensity` we solve for via inversion.\n", - "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", - "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", - "- **Fit:** Perform a single `FitInterferometer` using the model and inspect the inversion.\n", - "- **Intensities:** Extract the solved-for `intensity` via `fit.linear_light_profile_intensity_dict`.\n", - "- **Visualization:** Build the helper tracer where linear light profiles are replaced with ordinary light\n", - " profiles carrying their solved-for `intensity`, then plot.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "The source galaxy's `intensity` parameter is therefore not a free parameter in the model-fit, reducing the\n", - "dimensionality of non-linear parameter space by one. The lens light is already omitted for interferometer data,\n", - "so the saving is smaller than the imaging case (where lens and source both contribute) \u2014 but the inversion still\n", - "removes the degeneracies between `intensity` and the source's shape parameters (e.g. `effective_radius`,\n", - "`sersic_index`), which are difficult degeneracies for the non-linear search to map out accurately.\n", - "\n", - "The inversion has a relatively small computational cost on top of the NUFFT, so we reduce the model complexity\n", - "without much slow-down.\n", - "\n", - "__Disadvantages__\n", - "\n", - "Although the computation time of the inversion is small, it is not non-negligible. It is approximately 3-4x\n", - "slower per inversion-only term than using a standard light profile with a fixed `intensity`. The NUFFT typically\n", - "dominates the total per-likelihood cost on interferometer data, so the overall slow-down is usually smaller.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and\n", - "negative values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a\n", - "galaxy's light, which is clearly unphysical.\n", - "\n", - "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is\n", - "as fast as positive-negative solvers. This ensures that all light profile intensities are positive and\n", - "therefore physical.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data). This is the standard\n", - " convention for interferometer modeling.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a linear `SersicCore`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, using `TransformerNUFFT`\n", - "backed by `nufftax`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We now illustrate how to perform a fit to the dataset using a linear light profile, with the lens mass and the\n", - "source's shape parameters fixed to known values.\n", - "\n", - "The API follows closely the standard use of a `FitInterferometer` object, but simply uses a linear light\n", - "profile (via the `lp_linear` module) instead of a standard light profile.\n", - "\n", - "Note that the linear light profile below does not have an `intensity` parameter input \u2014 we let the inversion\n", - "solve for it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit's `subplot_fit_interferometer` shows the visibility-plane fit and dirty-image residuals. Because the\n", - "source bulge is a linear light profile, the inversion has solved for its `intensity` to maximize the fit to\n", - "the observed visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `subplot_fit_dirty_images` provides a real-space view of the data, model and residuals via inverse-NUFFT\n", - "of the visibility-plane quantities. This is generally more interpretable to the human eye than the uv-plane\n", - "plots above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_dirty_images(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The fit contains the solved-for `intensity` value.\n", - "\n", - "This is computed using a fit's `linear_light_profile_intensity_dict`, which maps each linear light profile in\n", - "the model parameterization above to its `intensity`.\n", - "\n", - "The code below shows how to use this dictionary, as an alternative to using the `max_log_likelihood` quantities\n", - "covered in `modeling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_bulge = tracer.galaxies[-1].bulge\n", - "\n", - "print(fit.linear_light_profile_intensity_dict)\n", - "\n", - "print(\n", - " f\"\\n Intensity of source bulge (lp_linear.SersicCore) = \"\n", - " f\"{fit.linear_light_profile_intensity_dict[source_bulge]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the\n", - "solved-for `intensity` values is also accessible from a fit.\n", - "\n", - "For example, the linear `SersicCore` of the source `bulge` component above has a solved-for `intensity` of\n", - "~0.3.\n", - "\n", - "The `tracer` created below instead has an ordinary `SersicCore` light profile with `intensity` ~0.3. The\n", - "benefit of this tracer is that it can be visualised (linear light profiles cannot be plotted by default\n", - "because they do not have `intensity` values)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", - "\n", - "print(tracer.galaxies[-1].bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization__\n", - "\n", - "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they\n", - "do not have an `intensity` value.\n", - "\n", - "Therefore, the helper-tracer created above (with all linear light profiles replaced by ordinary light profiles\n", - "carrying their solved-for `intensity`) must be used for visualization:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Linear Light Profiles Fit (Interferometer)\n", + "==============================================================\n", + "\n", + "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", + "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", + "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood) given\n", + "the light profile's other parameters.\n", + "\n", + "This script illustrates how to perform a single `FitInterferometer` of a linear light profile model \u2014 that is,\n", + "not the full Nautilus model-fit, but a single likelihood evaluation given known light/mass profile parameters.\n", + "This is useful for understanding how the inversion produces the solved-for `intensity`, and how to extract that\n", + "value from the resulting fit.\n", + "\n", + "For an explanation of why linear light profile fits are now practical against visibility data thanks to the\n", + "JAX-native NUFFT `nufftax` (https://github.com/GragasLab/nufftax), see the companion `modeling.py` example.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Benefits and drawbacks of linear light profiles for interferometer data.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Model:** The lens model whose `intensity` we solve for via inversion.\n", + "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", + "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", + "- **Fit:** Perform a single `FitInterferometer` using the model and inspect the inversion.\n", + "- **Intensities:** Extract the solved-for `intensity` via `fit.linear_light_profile_intensity_dict`.\n", + "- **Visualization:** Build the helper tracer where linear light profiles are replaced with ordinary light\n", + " profiles carrying their solved-for `intensity`, then plot.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "The source galaxy's `intensity` parameter is therefore not a free parameter in the model-fit, reducing the\n", + "dimensionality of non-linear parameter space by one. The lens light is already omitted for interferometer data,\n", + "so the saving is smaller than the imaging case (where lens and source both contribute) \u2014 but the inversion still\n", + "removes the degeneracies between `intensity` and the source's shape parameters (e.g. `effective_radius`,\n", + "`sersic_index`), which are difficult degeneracies for the non-linear search to map out accurately.\n", + "\n", + "The inversion has a relatively small computational cost on top of the NUFFT, so we reduce the model complexity\n", + "without much slow-down.\n", + "\n", + "__Disadvantages__\n", + "\n", + "Although the computation time of the inversion is small, it is not non-negligible. It is approximately 3-4x\n", + "slower per inversion-only term than using a standard light profile with a fixed `intensity`. The NUFFT typically\n", + "dominates the total per-likelihood cost on interferometer data, so the overall slow-down is usually smaller.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and\n", + "negative values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a\n", + "galaxy's light, which is clearly unphysical.\n", + "\n", + "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is\n", + "as fast as positive-negative solvers. This ensures that all light profile intensities are positive and\n", + "therefore physical.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data). This is the standard\n", + " convention for interferometer modeling.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a linear `SersicCore`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, using `TransformerNUFFT`\n", + "backed by `nufftax`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We now illustrate how to perform a fit to the dataset using a linear light profile, with the lens mass and the\n", + "source's shape parameters fixed to known values.\n", + "\n", + "The API follows closely the standard use of a `FitInterferometer` object, but simply uses a linear light\n", + "profile (via the `lp_linear` module) instead of a standard light profile.\n", + "\n", + "Note that the linear light profile below does not have an `intensity` parameter input \u2014 we let the inversion\n", + "solve for it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit's `subplot_fit_interferometer` shows the visibility-plane fit and dirty-image residuals. Because the\n", + "source bulge is a linear light profile, the inversion has solved for its `intensity` to maximize the fit to\n", + "the observed visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `subplot_fit_dirty_images` provides a real-space view of the data, model and residuals via inverse-NUFFT\n", + "of the visibility-plane quantities. This is generally more interpretable to the human eye than the uv-plane\n", + "plots above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_dirty_images(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The fit contains the solved-for `intensity` value.\n", + "\n", + "This is computed using a fit's `linear_light_profile_intensity_dict`, which maps each linear light profile in\n", + "the model parameterization above to its `intensity`.\n", + "\n", + "The code below shows how to use this dictionary, as an alternative to using the `max_log_likelihood` quantities\n", + "covered in `modeling.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_bulge = tracer.galaxies[-1].bulge\n", + "\n", + "print(fit.linear_light_profile_intensity_dict)\n", + "\n", + "print(\n", + " f\"\\n Intensity of source bulge (lp_linear.SersicCore) = \"\n", + " f\"{fit.linear_light_profile_intensity_dict[source_bulge]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the\n", + "solved-for `intensity` values is also accessible from a fit.\n", + "\n", + "For example, the linear `SersicCore` of the source `bulge` component above has a solved-for `intensity` of\n", + "~0.3.\n", + "\n", + "The `tracer` created below instead has an ordinary `SersicCore` light profile with `intensity` ~0.3. The\n", + "benefit of this tracer is that it can be visualised (linear light profiles cannot be plotted by default\n", + "because they do not have `intensity` values)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles\n", + "\n", + "print(tracer.galaxies[-1].bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization__\n", + "\n", + "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they\n", + "do not have an `intensity` value.\n", + "\n", + "Therefore, the helper-tracer created above (with all linear light profiles replaced by ordinary light profiles\n", + "carrying their solved-for `intensity`) must be used for visualization:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/linear_light_profiles/likelihood_function.ipynb b/notebooks/interferometer/features/linear_light_profiles/likelihood_function.ipynb index 6e8793115..7b32cc795 100644 --- a/notebooks/interferometer/features/linear_light_profiles/likelihood_function.ipynb +++ b/notebooks/interferometer/features/linear_light_profiles/likelihood_function.ipynb @@ -1,839 +1,876 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Linear Light Profile (Interferometer)__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit\n", - "`Interferometer` data with a linear light profile (e.g. a linear `SersicCore` source).\n", - "\n", - "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", - "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and\n", - "it always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood)\n", - "given the light profile's other parameters.\n", - "\n", - "For interferometer data this is now practical thanks to the JAX-native NUFFT `nufftax`\n", - "(https://github.com/GragasLab/nufftax) \u2014 the image-to-uv Fourier transform of each linear basis component\n", - "happens inside the same jit/vmap pipeline as the rest of the model, so per-iteration NUFFT cost is amortised\n", - "on the GPU.\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can\n", - " understand the likelihood function (including references to the previous literature from which it is\n", - " defined) without having to write large quantities of text and equations.\n", - "\n", - " - To make linear inversions in **PyAutoLens** less of a \"black-box\" to users.\n", - "\n", - "Accompanying this script is the imaging-version `imaging/features/linear_light_profiles/likelihood_function.py`,\n", - "which walks through the same calculation for CCD imaging data (PSF convolution rather than NUFFT). Most of the\n", - "linear algebra is identical \u2014 only the operation that produces the columns of the `operated_mapping_matrix`\n", - "differs.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function of a linear light profile builds on the standard parametric likelihood function and on\n", - "the interferometer-specific NUFFT step, so you must read the following notebooks before this script:\n", - "\n", - " - `interferometer/log_likelihood_function.ipynb` \u2014 the standard interferometer parametric likelihood\n", - " function (NUFFT of a real-space image, visibility-plane $\\\\chi^2$).\n", - " - `imaging/features/linear_light_profiles/likelihood_function.ipynb` \u2014 the linear-inversion linear algebra\n", - " (data vector, curvature matrix, positive-only solver) for CCD imaging.\n", - "\n", - "This script repeats just enough setup that you can follow it without rereading those two \u2014 but if anything is\n", - "unclear, those are the places to look first.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script.\n", - "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", - "- **Dataset:** Load and plot the strong lens `Interferometer` dataset using `TransformerNUFFT` (nufftax).\n", - "- **Lens Galaxy:** A mass-only lens galaxy (no light \u2014 interferometer convention).\n", - "- **Source Galaxy Linear Light Profile:** A linear `SersicCore` source with no `intensity` parameter.\n", - "- **Internal Intensity:** Why the linear light profile carries an internal `intensity=1.0`.\n", - "- **Ray Tracing:** Image-plane to source-plane grid.\n", - "- **Source Image (Internal):** Source-plane image of the linear profile with `intensity=1.0`.\n", - "- **Mapping Matrix:** Real-space mapping matrix \u2014 one column per linear profile, equal to the lensed image\n", - " of each linear basis component evaluated with `intensity=1.0`.\n", - "- **Transformed Mapping Matrix ($f$):** NUFFT each column to give the visibility-space mapping matrix used\n", - " in the inversion.\n", - "- **Data Vector (D):** Compute $D$ from the transformed mapping matrix, visibilities, and noise map.\n", - "- **Curvature Matrix (F):** Compute $F$ separately for real and imaginary components, then sum.\n", - "- **Reconstruction (Positive-Negative):** Solve $s = F^{-1} D$ via NumPy.\n", - "- **Reconstruction (Positive Only):** Solve with the fast non-negative least squares (`fnnls`) algorithm.\n", - "- **Visibilities Reconstruction:** Map $s$ back to visibility space.\n", - "- **Likelihood Function:** Visibility-plane $\\\\chi^2$ and noise normalization.\n", - "- **Chi Squared:** Sum chi-squared contributions over real and imaginary components.\n", - "- **Noise Normalization Term:** The fixed noise normalization term.\n", - "- **Calculate The Log Likelihood:** Combine into the final log likelihood.\n", - "- **Fit:** Cross-check via `FitInterferometer`.\n", - "- **Lens Modeling:** How this likelihood is sampled in the full Nautilus fit.\n", - "- **Wrap Up:** Summary and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit with the\n", - "model. We use `TransformerNUFFT` (backed by `nufftax`), the JAX-native NUFFT that makes linear inversions in\n", - "the visibility plane practical at any visibility count." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "For interferometer data the lens galaxy's optical/IR emission is typically below detection, so we model only\n", - "its mass:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Linear Light Profile__\n", - "\n", - "The source galaxy is fitted using a **linear** `SersicCore`. Compared to the standard `SersicCore` of the\n", - "parametric interferometer likelihood guide, we drop the `intensity` argument \u2014 it is solved for via the\n", - "inversion below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_bulge = al.lp_linear.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, bulge=source_bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Internal Intensity__\n", - "\n", - "Internally in the source code, linear light profiles still have an `intensity` parameter \u2014 its value is fixed\n", - "to 1.0. The inversion later rescales the column of the mapping matrix by the solved-for intensity. Without an\n", - "internal value of 1.0 the image evaluated below would be ambiguous in normalisation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Source Bulge Internal Intensity:\")\n", - "print(source_bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\\\theta$ from the image-plane to its\n", - "(y,x) source-plane coordinate $\\\\beta$ using the summed deflection angles $\\\\alpha$ of the mass profiles:\n", - "\n", - " $\\\\beta = \\\\theta - \\\\alpha(\\\\theta)$\n", - "\n", - "For interferometer data the only grid we need to ray-trace is the real-space grid associated with the\n", - "`real_space_mask` \u2014 there is no blurring grid (that is an imaging-only concept used to account for flux\n", - "outside the mask convolving into it via the PSF)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane); we only need the source-plane grid (index -1).\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=dataset.grids.lp)[-1]\n", - "\n", - "aplt.plot_grid(grid=traced_grid, title=\"Traced Source-Plane Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Image (Internal)__\n", - "\n", - "Evaluate the source bulge on the ray-traced source-plane grid. Because the linear profile carries\n", - "`intensity=1.0` internally, the absolute normalization of this image is arbitrary \u2014 it represents the\n", - "*shape* of the source's contribution to the model. The inversion will scale it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_2d_source_bulge_internal = source_galaxy.bulge.image_2d_from(grid=traced_grid)\n", - "\n", - "aplt.plot_array(\n", - " array=image_2d_source_bulge_internal, title=\"Source Bulge Image (intensity=1.0)\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapping Matrix__\n", - "\n", - "For interferometer data the mapping matrix is the same conceptually as for imaging data: each column\n", - "holds the real-space image of one linear basis component, evaluated with `intensity=1.0`. The lensing of\n", - "that image is already baked in (because we evaluated it on the ray-traced source-plane grid).\n", - "\n", - "We have only one linear light profile (the source bulge), so the mapping matrix has dimensions\n", - "`(total_real_space_pixels, 1)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lp_linear_func_source = al.LightProfileLinearObjFuncList(\n", - " grid=traced_grid,\n", - " blurring_grid=None,\n", - " psf=None,\n", - " light_profile_list=[source_galaxy.bulge],\n", - " regularization=None,\n", - ")\n", - "\n", - "mapping_matrix = lp_linear_func_source.mapping_matrix\n", - "\n", - "print(\"Mapping matrix shape (real-space pixels, linear basis count):\")\n", - "print(mapping_matrix.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `mapping_matrix` shows the source bulge image in 1D (column 0). For the single-source case\n", - "the matrix is a single column, so the plot looks like a tall stripe." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Transformed Mapping Matrix ($f$)__\n", - "\n", - "To fit visibilities, every column of the real-space `mapping_matrix` must be NUFFT'd to the uv-plane. The\n", - "result is the `transformed_mapping_matrix` \u2014 a *complex-valued* matrix with dimensions\n", - "`(total_visibilities, total_linear_basis_components)`.\n", - "\n", - "In our case it is a single column: the visibilities the source bulge contributes per unit `intensity`. The\n", - "inversion's job is to find the scalar that, when multiplied by this column, best fits the observed\n", - "visibilities.\n", - "\n", - "The NUFFT here uses `TransformerNUFFT` (nufftax) \u2014 this is exactly the operation that used to be slow and\n", - "which nufftax has made fast. The whole pipeline (deflections \u2192 ray-trace \u2192 source-plane image \u2192\n", - "`transform_mapping_matrix` \u2192 linear inversion) is JIT-compilable under JAX." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", - " mapping_matrix=mapping_matrix\n", - ")\n", - "\n", - "print(\"Transformed mapping matrix shape (visibilities, linear basis count):\")\n", - "print(transformed_mapping_matrix.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the real and imaginary components of the (single column of the) transformed mapping matrix. Each row is\n", - "one observed visibility; the value is the contribution of the source per unit intensity at that uv point." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " transformed_mapping_matrix.real,\n", - " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.title(\"Re(f) \u2014 real component of transformed mapping matrix\")\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "plt.imshow(\n", - " transformed_mapping_matrix.imag,\n", - " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.title(\"Im(f) \u2014 imaginary component of transformed mapping matrix\")\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) (hereafter WD03) introduce the linear inversion\n", - "formalism used to compute the intensity values of the linear light profiles. WD03 indexes the transformed\n", - "mapping matrix as $f_{ij}$ where $i$ maps over all $I$ linear basis components and $j$ maps over all $J$\n", - "visibilities.\n", - "\n", - "The indexing of the `transformed_mapping_matrix` array is reversed compared to the WD03 convention (rows are\n", - "visibilities, columns are basis components).\n", - "\n", - "__Data Vector (D)__\n", - "\n", - "To solve for the linear light profile intensities we pose the problem as a linear inversion.\n", - "\n", - "This requires us to convert the `transformed_mapping_matrix` and our `data` and `noise_map` into matrices of\n", - "the right dimensions. The `data_vector`, $D$, has dimensions `(total_linear_basis_components,)`.\n", - "\n", - "In WD03 the data vector is given by:\n", - "\n", - " $\\\\vec{D}_{i} = \\\\sum_{\\\\rm j=1}^{J} f_{ij}\\\\, d_{j} / \\\\sigma_{j}^2 \\\\, ,$\n", - "\n", - "where $d_j$ are the observed visibility values, $\\\\sigma_j^2$ are the visibility variances, and the sum runs\n", - "over real and imaginary components. The interferometer helper handles the real/imaginary split internally.\n", - "\n", - "For a single-basis model $D$ is a length-1 vector \u2014 just the noise-weighted overlap of the source's\n", - "contribution with the observed visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_vector = (\n", - " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " visibilities=dataset.data,\n", - " noise_map=dataset.noise_map,\n", - " )\n", - ")\n", - "\n", - "print(\"Data Vector D:\")\n", - "print(data_vector)\n", - "print(data_vector.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Matrix (F)__\n", - "\n", - "The `curvature_matrix` $F$ has dimensions\n", - "`(total_linear_basis_components, total_linear_basis_components)`.\n", - "\n", - "In WD03 the curvature matrix is given by:\n", - "\n", - " ${F}_{ik} = \\\\sum_{\\\\rm j=1}^{J} f_{ij}\\\\, f_{kj} / \\\\sigma_{j}^2 \\\\, .$\n", - "\n", - "Because visibilities (and therefore $f$) are complex-valued, the curvature is computed separately for the\n", - "real and imaginary parts and summed. For a single-basis model $F$ is a 1x1 matrix." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.real,\n", - " noise_map=dataset.noise_map.real,\n", - ")\n", - "\n", - "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.imag,\n", - " noise_map=dataset.noise_map.imag,\n", - ")\n", - "\n", - "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", - "\n", - "print(\"Curvature Matrix F:\")\n", - "print(curvature_matrix)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction (Positive-Negative)__\n", - "\n", - "The following chi-squared is minimized when we perform the inversion and solve for the source intensity:\n", - "\n", - "$\\\\chi^2 = \\\\sum_{\\\\rm j=1}^{J} \\\\bigg[ \\\\frac{(\\\\sum_{\\\\rm i=1}^{I} s_{i} f_{ij}) - d_{j}}{\\\\sigma_{j}} \\\\bigg]^2$\n", - "\n", - "Where $s_i$ is the solved-for `intensity` of the $i$-th linear basis component. The solution is given by\n", - "(equation 5 WD03):\n", - "\n", - " $s = F^{-1} D$\n", - "\n", - "Computed with NumPy:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = np.linalg.solve(curvature_matrix, data_vector)\n", - "\n", - "print(\"Reconstruction s (positive-negative solver):\")\n", - "print(reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For linear light profiles fit to a clean dataset like ours, the solved-for `intensity` is positive and the\n", - "positive-negative solution is physical. For more complex models (e.g. many components or noisy data) the\n", - "positive-negative solver can return negative `intensity` values \u2014 unphysical for a light profile.\n", - "\n", - "__Reconstruction (Positive Only)__\n", - "\n", - "The linear algebra can be solved with the constraint that all `intensity` values are positive. The naive\n", - "approach is `scipy.optimize.nnls`, which is iterative and works directly on the transformed mapping matrix \u2014\n", - "it does not use `data_vector` or `curvature_matrix`. This is slow, especially with many linear components.\n", - "\n", - "The source code therefore uses a \"fast nnls\" algorithm \u2014 an adaptation of:\n", - " https://github.com/jvendrow/fnnls\n", - "\n", - "`fnnls` *does* use `data_vector` $D$ and `curvature_matrix` $F$, which is why it is much faster. The function\n", - "`reconstruction_positive_only_from` wraps `fnnls`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", - " data_vector=data_vector,\n", - " curvature_reg_matrix=curvature_matrix, # ignore the _reg_ tag in this guide\n", - ")\n", - "\n", - "print(\"Reconstruction s (positive-only solver):\")\n", - "print(reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visibilities Reconstruction__\n", - "\n", - "Using the reconstructed `intensity` value(s) we can map the reconstruction back to the visibility plane via\n", - "the `transformed_mapping_matrix`, producing the model visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_visibilities = (\n", - " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " reconstruction=reconstruction,\n", - " )\n", - ")\n", - "\n", - "mapped_reconstructed_visibilities = al.Visibilities(\n", - " visibilities=mapped_reconstructed_visibilities\n", - ")\n", - "\n", - "aplt.plot_grid(\n", - " grid=mapped_reconstructed_visibilities.in_grid, title=\"Model Visibilities\"\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our linear-light-profile reconstruction.\n", - "\n", - "The likelihood function for parametric galaxy modeling, even with linear light profiles, consists of two\n", - "terms in the visibility plane:\n", - "\n", - " $-2 \\\\mathrm{ln} \\\\, \\\\epsilon = \\\\chi^2 + \\\\sum_{\\\\rm j=1}^{J} { \\\\mathrm{ln}} \\\\left [2 \\\\pi (\\\\sigma_j)^2 \\\\right] \\\\, .$\n", - "\n", - "We now explain each term.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\\\chi^2$ statistic computed in the visibility plane:\n", - "\n", - " - `model_data` = `mapped_reconstructed_visibilities`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = `normalized_residual_map` ** 2.0\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "Visibilities are complex-valued, so we split into real and imaginary components, compute $\\\\chi^2$ for each,\n", - "and sum." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_visibilities = mapped_reconstructed_visibilities\n", - "\n", - "residual_map = dataset.data - model_visibilities\n", - "\n", - "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", - "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", - "\n", - "chi_squared_real = np.sum(chi_squared_map_real)\n", - "chi_squared_imag = np.sum(chi_squared_map_imag)\n", - "chi_squared = chi_squared_real + chi_squared_imag\n", - "\n", - "print(f\"chi_squared (real + imag) = {chi_squared}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "The likelihood function assumes the visibility data consists of independent Gaussian noise on every\n", - "visibility (real and imaginary parts treated independently).\n", - "\n", - "The `noise_normalization` term is the sum of the log of every noise-map value squared. Because the\n", - "`noise_map` is fixed, this term does not change during modeling and has no impact on the inferred model \u2014 it\n", - "is included so that the absolute value of `log_likelihood` has the correct calibration for model comparison." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization_real = float(\n", - " np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", - ")\n", - "noise_normalization_imag = float(\n", - " np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", - ")\n", - "noise_normalization = noise_normalization_real + noise_normalization_imag" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "Combine the two terms to compute the `log_likelihood`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(f\"log_likelihood (figure of merit) = {figure_of_merit}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "The exact same likelihood evaluation is performed inside the `FitInterferometer` object. We construct one,\n", - "print its `figure_of_merit`, and confirm it matches the value we computed by hand above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", - "print(f\"FitInterferometer.figure_of_merit = {fit.figure_of_merit}\")\n", - "\n", - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit contains an `Inversion` object, which handles all the linear algebra we have covered in this script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.inversion)\n", - "print(fit.inversion.data_vector)\n", - "print(fit.inversion.curvature_matrix)\n", - "print(fit.inversion.reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Inversion` can also be computed from the tracer and dataset directly, by passing them to the\n", - "`TracerToInversion` object. This object additionally handles cases where the tracer contains a mix of linear\n", - "light profiles and pixelization-based components." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(\n", - " tracer=tracer,\n", - " dataset=dataset,\n", - ")\n", - "\n", - "inversion = tracer_to_inversion.inversion" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `nautilus`\n", - "(https://github.com/johannesulf/nautilus); multiple MCMC and optimization algorithms are also supported.\n", - "\n", - "For linear light profiles, the reduced number of free parameters (the `intensity` is solved for via the\n", - "inversion instead of being a non-linear search dimension) means that the sampler converges in fewer\n", - "iterations and is less likely to be confused by intensity-shape degeneracies.\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the linear light profile interferometer likelihood\n", - "function. The pipeline:\n", - "\n", - " ray-trace \u2192 source-plane image (linear profile, intensity=1) \u2192 `mapping_matrix`\n", - " \u2192 NUFFT \u2192 `transformed_mapping_matrix` \u2192 $D$ and $F$ \u2192 solve $s = F^{-1} D$\n", - " \u2192 `mapped_reconstructed_visibilities` \u2192 visibility-plane $\\\\chi^2$ \u2192 log likelihood\n", - "\n", - "is the same as for CCD imaging linear light profiles, with the PSF convolution step replaced by the NUFFT\n", - "step. The NUFFT is the operation that nufftax made fast on the GPU, which is why this entire workflow is\n", - "now practical on interferometer data at any visibility count.\n", - "\n", - "There are a number of other inputs which slightly change the behaviour of this likelihood function, which\n", - "are described in additional notebooks found in the `guides` package:\n", - "\n", - " - `over_sampling`: Not applicable to interferometer data (over-sampling is an imaging-only technique).\n", - " - `pixelization`: For sources whose morphology cannot be captured by analytic light profiles, see the\n", - " interferometer pixelization examples in `interferometer/features/pixelization/`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Linear Light Profile (Interferometer)__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit\n", + "`Interferometer` data with a linear light profile (e.g. a linear `SersicCore` source).\n", + "\n", + "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", + "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and\n", + "it always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood)\n", + "given the light profile's other parameters.\n", + "\n", + "For interferometer data this is now practical thanks to the JAX-native NUFFT `nufftax`\n", + "(https://github.com/GragasLab/nufftax) \u2014 the image-to-uv Fourier transform of each linear basis component\n", + "happens inside the same jit/vmap pipeline as the rest of the model, so per-iteration NUFFT cost is amortised\n", + "on the GPU.\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can\n", + " understand the likelihood function (including references to the previous literature from which it is\n", + " defined) without having to write large quantities of text and equations.\n", + "\n", + " - To make linear inversions in **PyAutoLens** less of a \"black-box\" to users.\n", + "\n", + "Accompanying this script is the imaging-version `imaging/features/linear_light_profiles/likelihood_function.py`,\n", + "which walks through the same calculation for CCD imaging data (PSF convolution rather than NUFFT). Most of the\n", + "linear algebra is identical \u2014 only the operation that produces the columns of the `operated_mapping_matrix`\n", + "differs.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function of a linear light profile builds on the standard parametric likelihood function and on\n", + "the interferometer-specific NUFFT step, so you must read the following notebooks before this script:\n", + "\n", + " - `interferometer/log_likelihood_function.ipynb` \u2014 the standard interferometer parametric likelihood\n", + " function (NUFFT of a real-space image, visibility-plane $\\\\chi^2$).\n", + " - `imaging/features/linear_light_profiles/likelihood_function.ipynb` \u2014 the linear-inversion linear algebra\n", + " (data vector, curvature matrix, positive-only solver) for CCD imaging.\n", + "\n", + "This script repeats just enough setup that you can follow it without rereading those two \u2014 but if anything is\n", + "unclear, those are the places to look first.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script.\n", + "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", + "- **Dataset:** Load and plot the strong lens `Interferometer` dataset using `TransformerNUFFT` (nufftax).\n", + "- **Lens Galaxy:** A mass-only lens galaxy (no light \u2014 interferometer convention).\n", + "- **Source Galaxy Linear Light Profile:** A linear `SersicCore` source with no `intensity` parameter.\n", + "- **Internal Intensity:** Why the linear light profile carries an internal `intensity=1.0`.\n", + "- **Ray Tracing:** Image-plane to source-plane grid.\n", + "- **Source Image (Internal):** Source-plane image of the linear profile with `intensity=1.0`.\n", + "- **Mapping Matrix:** Real-space mapping matrix \u2014 one column per linear profile, equal to the lensed image\n", + " of each linear basis component evaluated with `intensity=1.0`.\n", + "- **Transformed Mapping Matrix ($f$):** NUFFT each column to give the visibility-space mapping matrix used\n", + " in the inversion.\n", + "- **Data Vector (D):** Compute $D$ from the transformed mapping matrix, visibilities, and noise map.\n", + "- **Curvature Matrix (F):** Compute $F$ separately for real and imaginary components, then sum.\n", + "- **Reconstruction (Positive-Negative):** Solve $s = F^{-1} D$ via NumPy.\n", + "- **Reconstruction (Positive Only):** Solve with the fast non-negative least squares (`fnnls`) algorithm.\n", + "- **Visibilities Reconstruction:** Map $s$ back to visibility space.\n", + "- **Likelihood Function:** Visibility-plane $\\\\chi^2$ and noise normalization.\n", + "- **Chi Squared:** Sum chi-squared contributions over real and imaginary components.\n", + "- **Noise Normalization Term:** The fixed noise normalization term.\n", + "- **Calculate The Log Likelihood:** Combine into the final log likelihood.\n", + "- **Fit:** Cross-check via `FitInterferometer`.\n", + "- **Lens Modeling:** How this likelihood is sampled in the full Nautilus fit.\n", + "- **Wrap Up:** Summary and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit with the\n", + "model. We use `TransformerNUFFT` (backed by `nufftax`), the JAX-native NUFFT that makes linear inversions in\n", + "the visibility plane practical at any visibility count." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "For interferometer data the lens galaxy's optical/IR emission is typically below detection, so we model only\n", + "its mass:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Linear Light Profile__\n", + "\n", + "The source galaxy is fitted using a **linear** `SersicCore`. Compared to the standard `SersicCore` of the\n", + "parametric interferometer likelihood guide, we drop the `intensity` argument \u2014 it is solved for via the\n", + "inversion below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_bulge = al.lp_linear.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, bulge=source_bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Internal Intensity__\n", + "\n", + "Internally in the source code, linear light profiles still have an `intensity` parameter \u2014 its value is fixed\n", + "to 1.0. The inversion later rescales the column of the mapping matrix by the solved-for intensity. Without an\n", + "internal value of 1.0 the image evaluated below would be ambiguous in normalisation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Source Bulge Internal Intensity:\")\n", + "print(source_bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\\\theta$ from the image-plane to its\n", + "(y,x) source-plane coordinate $\\\\beta$ using the summed deflection angles $\\\\alpha$ of the mass profiles:\n", + "\n", + " $\\\\beta = \\\\theta - \\\\alpha(\\\\theta)$\n", + "\n", + "For interferometer data the only grid we need to ray-trace is the real-space grid associated with the\n", + "`real_space_mask` \u2014 there is no blurring grid (that is an imaging-only concept used to account for flux\n", + "outside the mask convolving into it via the PSF)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane); we only need the source-plane grid (index -1).\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=dataset.grids.lp)[-1]\n", + "\n", + "aplt.plot_grid(grid=traced_grid, title=\"Traced Source-Plane Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Image (Internal)__\n", + "\n", + "Evaluate the source bulge on the ray-traced source-plane grid. Because the linear profile carries\n", + "`intensity=1.0` internally, the absolute normalization of this image is arbitrary \u2014 it represents the\n", + "*shape* of the source's contribution to the model. The inversion will scale it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_2d_source_bulge_internal = source_galaxy.bulge.image_2d_from(grid=traced_grid)\n", + "\n", + "aplt.plot_array(\n", + " array=image_2d_source_bulge_internal, title=\"Source Bulge Image (intensity=1.0)\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapping Matrix__\n", + "\n", + "For interferometer data the mapping matrix is the same conceptually as for imaging data: each column\n", + "holds the real-space image of one linear basis component, evaluated with `intensity=1.0`. The lensing of\n", + "that image is already baked in (because we evaluated it on the ray-traced source-plane grid).\n", + "\n", + "We have only one linear light profile (the source bulge), so the mapping matrix has dimensions\n", + "`(total_real_space_pixels, 1)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lp_linear_func_source = al.LightProfileLinearObjFuncList(\n", + " grid=traced_grid,\n", + " blurring_grid=None,\n", + " psf=None,\n", + " light_profile_list=[source_galaxy.bulge],\n", + " regularization=None,\n", + ")\n", + "\n", + "mapping_matrix = lp_linear_func_source.mapping_matrix\n", + "\n", + "print(\"Mapping matrix shape (real-space pixels, linear basis count):\")\n", + "print(mapping_matrix.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `mapping_matrix` shows the source bulge image in 1D (column 0). For the single-source case\n", + "the matrix is a single column, so the plot looks like a tall stripe." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Transformed Mapping Matrix ($f$)__\n", + "\n", + "To fit visibilities, every column of the real-space `mapping_matrix` must be NUFFT'd to the uv-plane. The\n", + "result is the `transformed_mapping_matrix` \u2014 a *complex-valued* matrix with dimensions\n", + "`(total_visibilities, total_linear_basis_components)`.\n", + "\n", + "In our case it is a single column: the visibilities the source bulge contributes per unit `intensity`. The\n", + "inversion's job is to find the scalar that, when multiplied by this column, best fits the observed\n", + "visibilities.\n", + "\n", + "The NUFFT here uses `TransformerNUFFT` (nufftax) \u2014 this is exactly the operation that used to be slow and\n", + "which nufftax has made fast. The whole pipeline (deflections \u2192 ray-trace \u2192 source-plane image \u2192\n", + "`transform_mapping_matrix` \u2192 linear inversion) is JIT-compilable under JAX." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", + " mapping_matrix=mapping_matrix\n", + ")\n", + "\n", + "print(\"Transformed mapping matrix shape (visibilities, linear basis count):\")\n", + "print(transformed_mapping_matrix.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the real and imaginary components of the (single column of the) transformed mapping matrix. Each row is\n", + "one observed visibility; the value is the contribution of the source per unit intensity at that uv point." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " transformed_mapping_matrix.real,\n", + " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.title(\"Re(f) \u2014 real component of transformed mapping matrix\")\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "plt.imshow(\n", + " transformed_mapping_matrix.imag,\n", + " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.title(\"Im(f) \u2014 imaginary component of transformed mapping matrix\")\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) (hereafter WD03) introduce the linear inversion\n", + "formalism used to compute the intensity values of the linear light profiles. WD03 indexes the transformed\n", + "mapping matrix as $f_{ij}$ where $i$ maps over all $I$ linear basis components and $j$ maps over all $J$\n", + "visibilities.\n", + "\n", + "The indexing of the `transformed_mapping_matrix` array is reversed compared to the WD03 convention (rows are\n", + "visibilities, columns are basis components).\n", + "\n", + "__Data Vector (D)__\n", + "\n", + "To solve for the linear light profile intensities we pose the problem as a linear inversion.\n", + "\n", + "This requires us to convert the `transformed_mapping_matrix` and our `data` and `noise_map` into matrices of\n", + "the right dimensions. The `data_vector`, $D$, has dimensions `(total_linear_basis_components,)`.\n", + "\n", + "In WD03 the data vector is given by:\n", + "\n", + " $\\\\vec{D}_{i} = \\\\sum_{\\\\rm j=1}^{J} f_{ij}\\\\, d_{j} / \\\\sigma_{j}^2 \\\\, ,$\n", + "\n", + "where $d_j$ are the observed visibility values, $\\\\sigma_j^2$ are the visibility variances, and the sum runs\n", + "over real and imaginary components. The interferometer helper handles the real/imaginary split internally.\n", + "\n", + "For a single-basis model $D$ is a length-1 vector \u2014 just the noise-weighted overlap of the source's\n", + "contribution with the observed visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_vector = (\n", + " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " visibilities=dataset.data,\n", + " noise_map=dataset.noise_map,\n", + " )\n", + ")\n", + "\n", + "print(\"Data Vector D:\")\n", + "print(data_vector)\n", + "print(data_vector.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Matrix (F)__\n", + "\n", + "The `curvature_matrix` $F$ has dimensions\n", + "`(total_linear_basis_components, total_linear_basis_components)`.\n", + "\n", + "In WD03 the curvature matrix is given by:\n", + "\n", + " ${F}_{ik} = \\\\sum_{\\\\rm j=1}^{J} f_{ij}\\\\, f_{kj} / \\\\sigma_{j}^2 \\\\, .$\n", + "\n", + "Because visibilities (and therefore $f$) are complex-valued, the curvature is computed separately for the\n", + "real and imaginary parts and summed. For a single-basis model $F$ is a 1x1 matrix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.real,\n", + " noise_map=dataset.noise_map.real,\n", + ")\n", + "\n", + "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.imag,\n", + " noise_map=dataset.noise_map.imag,\n", + ")\n", + "\n", + "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", + "\n", + "print(\"Curvature Matrix F:\")\n", + "print(curvature_matrix)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction (Positive-Negative)__\n", + "\n", + "The following chi-squared is minimized when we perform the inversion and solve for the source intensity:\n", + "\n", + "$\\\\chi^2 = \\\\sum_{\\\\rm j=1}^{J} \\\\bigg[ \\\\frac{(\\\\sum_{\\\\rm i=1}^{I} s_{i} f_{ij}) - d_{j}}{\\\\sigma_{j}} \\\\bigg]^2$\n", + "\n", + "Where $s_i$ is the solved-for `intensity` of the $i$-th linear basis component. The solution is given by\n", + "(equation 5 WD03):\n", + "\n", + " $s = F^{-1} D$\n", + "\n", + "Computed with NumPy:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = np.linalg.solve(curvature_matrix, data_vector)\n", + "\n", + "print(\"Reconstruction s (positive-negative solver):\")\n", + "print(reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For linear light profiles fit to a clean dataset like ours, the solved-for `intensity` is positive and the\n", + "positive-negative solution is physical. For more complex models (e.g. many components or noisy data) the\n", + "positive-negative solver can return negative `intensity` values \u2014 unphysical for a light profile.\n", + "\n", + "__Reconstruction (Positive Only)__\n", + "\n", + "The linear algebra can be solved with the constraint that all `intensity` values are positive. The naive\n", + "approach is `scipy.optimize.nnls`, which is iterative and works directly on the transformed mapping matrix \u2014\n", + "it does not use `data_vector` or `curvature_matrix`. This is slow, especially with many linear components.\n", + "\n", + "The source code therefore uses a \"fast nnls\" algorithm \u2014 an adaptation of:\n", + " https://github.com/jvendrow/fnnls\n", + "\n", + "`fnnls` *does* use `data_vector` $D$ and `curvature_matrix` $F$, which is why it is much faster. The function\n", + "`reconstruction_positive_only_from` wraps `fnnls`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", + " data_vector=data_vector,\n", + " curvature_reg_matrix=curvature_matrix, # ignore the _reg_ tag in this guide\n", + ")\n", + "\n", + "print(\"Reconstruction s (positive-only solver):\")\n", + "print(reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visibilities Reconstruction__\n", + "\n", + "Using the reconstructed `intensity` value(s) we can map the reconstruction back to the visibility plane via\n", + "the `transformed_mapping_matrix`, producing the model visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_visibilities = (\n", + " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " reconstruction=reconstruction,\n", + " )\n", + ")\n", + "\n", + "mapped_reconstructed_visibilities = al.Visibilities(\n", + " visibilities=mapped_reconstructed_visibilities\n", + ")\n", + "\n", + "aplt.plot_grid(\n", + " grid=mapped_reconstructed_visibilities.in_grid, title=\"Model Visibilities\"\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our linear-light-profile reconstruction.\n", + "\n", + "The likelihood function for parametric galaxy modeling, even with linear light profiles, consists of two\n", + "terms in the visibility plane:\n", + "\n", + " $-2 \\\\mathrm{ln} \\\\, \\\\epsilon = \\\\chi^2 + \\\\sum_{\\\\rm j=1}^{J} { \\\\mathrm{ln}} \\\\left [2 \\\\pi (\\\\sigma_j)^2 \\\\right] \\\\, .$\n", + "\n", + "We now explain each term.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\\\chi^2$ statistic computed in the visibility plane:\n", + "\n", + " - `model_data` = `mapped_reconstructed_visibilities`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = `normalized_residual_map` ** 2.0\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "Visibilities are complex-valued, so we split into real and imaginary components, compute $\\\\chi^2$ for each,\n", + "and sum." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_visibilities = mapped_reconstructed_visibilities\n", + "\n", + "residual_map = dataset.data - model_visibilities\n", + "\n", + "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", + "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", + "\n", + "chi_squared_real = np.sum(chi_squared_map_real)\n", + "chi_squared_imag = np.sum(chi_squared_map_imag)\n", + "chi_squared = chi_squared_real + chi_squared_imag\n", + "\n", + "print(f\"chi_squared (real + imag) = {chi_squared}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "The likelihood function assumes the visibility data consists of independent Gaussian noise on every\n", + "visibility (real and imaginary parts treated independently).\n", + "\n", + "The `noise_normalization` term is the sum of the log of every noise-map value squared. Because the\n", + "`noise_map` is fixed, this term does not change during modeling and has no impact on the inferred model \u2014 it\n", + "is included so that the absolute value of `log_likelihood` has the correct calibration for model comparison." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization_real = float(\n", + " np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", + ")\n", + "noise_normalization_imag = float(\n", + " np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", + ")\n", + "noise_normalization = noise_normalization_real + noise_normalization_imag" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "Combine the two terms to compute the `log_likelihood`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(f\"log_likelihood (figure of merit) = {figure_of_merit}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "The exact same likelihood evaluation is performed inside the `FitInterferometer` object. We construct one,\n", + "print its `figure_of_merit`, and confirm it matches the value we computed by hand above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", + "print(f\"FitInterferometer.figure_of_merit = {fit.figure_of_merit}\")\n", + "\n", + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit contains an `Inversion` object, which handles all the linear algebra we have covered in this script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.inversion)\n", + "print(fit.inversion.data_vector)\n", + "print(fit.inversion.curvature_matrix)\n", + "print(fit.inversion.reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Inversion` can also be computed from the tracer and dataset directly, by passing them to the\n", + "`TracerToInversion` object. This object additionally handles cases where the tracer contains a mix of linear\n", + "light profiles and pixelization-based components." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(\n", + " tracer=tracer,\n", + " dataset=dataset,\n", + ")\n", + "\n", + "inversion = tracer_to_inversion.inversion" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `nautilus`\n", + "(https://github.com/johannesulf/nautilus); multiple MCMC and optimization algorithms are also supported.\n", + "\n", + "For linear light profiles, the reduced number of free parameters (the `intensity` is solved for via the\n", + "inversion instead of being a non-linear search dimension) means that the sampler converges in fewer\n", + "iterations and is less likely to be confused by intensity-shape degeneracies.\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the linear light profile interferometer likelihood\n", + "function. The pipeline:\n", + "\n", + " ray-trace \u2192 source-plane image (linear profile, intensity=1) \u2192 `mapping_matrix`\n", + " \u2192 NUFFT \u2192 `transformed_mapping_matrix` \u2192 $D$ and $F$ \u2192 solve $s = F^{-1} D$\n", + " \u2192 `mapped_reconstructed_visibilities` \u2192 visibility-plane $\\\\chi^2$ \u2192 log likelihood\n", + "\n", + "is the same as for CCD imaging linear light profiles, with the PSF convolution step replaced by the NUFFT\n", + "step. The NUFFT is the operation that nufftax made fast on the GPU, which is why this entire workflow is\n", + "now practical on interferometer data at any visibility count.\n", + "\n", + "There are a number of other inputs which slightly change the behaviour of this likelihood function, which\n", + "are described in additional notebooks found in the `guides` package:\n", + "\n", + " - `over_sampling`: Not applicable to interferometer data (over-sampling is an imaging-only technique).\n", + " - `pixelization`: For sources whose morphology cannot be captured by analytic light profiles, see the\n", + " interferometer pixelization examples in `interferometer/features/pixelization/`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/linear_light_profiles/modeling.ipynb b/notebooks/interferometer/features/linear_light_profiles/modeling.ipynb index f49118680..8e1169d34 100644 --- a/notebooks/interferometer/features/linear_light_profiles/modeling.ipynb +++ b/notebooks/interferometer/features/linear_light_profiles/modeling.ipynb @@ -1,656 +1,693 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Linear Light Profiles (Interferometer)\n", - "==========================================================\n", - "\n", - "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", - "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", - "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood) given\n", - "the light profile's other parameters.\n", - "\n", - "Linear light profiles have been a standard tool for fitting CCD imaging data for a long time. For interferometer\n", - "data they used to be impractical, because every likelihood evaluation has to Fourier-transform each basis component\n", - "into the uv-plane, and prior NUFFT backends were not JAX-friendly. With `nufftax`\n", - "(https://github.com/GragasLab/nufftax) \u2014 a JAX-native Non-Uniform Fast Fourier Transform \u2014 the image-to-uv\n", - "transform now runs inside the same jit/vmap pipeline as the rest of the model, so the per-iteration overhead of\n", - "NUFFT-ing each basis component is amortised on the GPU. Linear light profile fits are therefore practical for\n", - "interferometer data at any visibility count, including ALMA-class datasets with tens of millions of visibilities.\n", - "\n", - "Based on the advantages below, we recommend you use linear light profiles whenever fitting light profiles to\n", - "interferometer data.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Benefits and drawbacks of linear light profiles, and how they apply to\n", - " interferometer data specifically.\n", - "- **NUFFT (nufftax):** Why linear light profile fits to visibilities are now practical thanks to nufftax.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Model:** Compose the lens model fitted to the data \u2014 `Isothermal` + `ExternalShear` lens mass and a\n", - " linear `SersicCore` source. The lens light is omitted (interferometer convention).\n", - "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", - "- **Dataset:** Load the strong lens `Interferometer` dataset, using `TransformerNUFFT` (backed by `nufftax`).\n", - "- **Over Sampling:** Interferometer modeling does not use over-sampling (covered briefly here for users\n", - " familiar with imaging).\n", - "- **Search:** Configure the non-linear search (Nautilus).\n", - "- **Analysis:** Create the `AnalysisInterferometer` object.\n", - "- **VRAM:** Linear light profiles add negligible VRAM compared to standard light profiles.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Intensities:** How to extract solved-for `intensity` values from the result.\n", - "- **Visualization:** Visualising fits with linear light profiles requires the\n", - " `model_obj_linear_light_profiles_to_light_profiles` helper.\n", - "- **Max Likelihood Inversion:** Access the `Inversion` object from the result.\n", - "- **Linear Objects (Internal Source Code):** The internal `linear_obj_list` representation used by the\n", - " inversion.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "The source galaxy's `intensity` parameter is therefore not a free parameter in the model-fit, reducing the\n", - "dimensionality of non-linear parameter space by one. The lens light is already omitted for interferometer data,\n", - "so the saving is smaller than the imaging case (where lens and source both contribute) \u2014 but the inversion still\n", - "removes the degeneracies between `intensity` and the source's shape parameters (e.g. `effective_radius`,\n", - "`sersic_index`), which are difficult degeneracies for the non-linear search to map out accurately. This produces\n", - "more reliable lens model results and the fit converges in fewer iterations.\n", - "\n", - "The inversion has a relatively small computational cost on top of the NUFFT, so we reduce the model complexity\n", - "without much slow-down.\n", - "\n", - "__Disadvantages__\n", - "\n", - "Although the computation time of the inversion is small, it is not non-negligible. It is approximately 3-4x\n", - "slower per likelihood than using a standard light profile with a fixed `intensity`.\n", - "\n", - "The gains in run times from the simpler parameter space therefore broadly balance the slower per-likelihood\n", - "evaluation. The headline benefit is reliability, not raw speed.\n", - "\n", - "__NUFFT (nufftax)__\n", - "\n", - "The image-to-visibilities Fourier transform is performed by a Non-Uniform Fast Fourier Transform (NUFFT),\n", - "exposed in **PyAutoLens** as `TransformerNUFFT`. The default backend is `nufftax`, a pure-JAX NUFFT that\n", - "jit-compiles and vmap-batches like the rest of the library:\n", - "\n", - " https://github.com/GragasLab/nufftax\n", - "\n", - "Because `nufftax` is JAX-native, NUFFT-ing each linear basis image happens inside the same compiled likelihood\n", - "that does the inversion, mass model ray-tracing, and chi-squared sum. There is no host round-trip between\n", - "NUFFT calls, so a model with N linear light profiles costs only N forward-NUFFTs per iteration on the GPU \u2014\n", - "fast enough that linear inversions in the visibility plane are now routinely practical.\n", - "\n", - "If `nufftax` is not installed, install it via `pip install nufftax`. A legacy pynufft-backed transformer\n", - "(`TransformerNUFFTPyNUFFT`) is available as a non-JAX fallback but is not recommended for linear light profiles.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and\n", - "negative values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical.\n", - "\n", - "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is\n", - "as fast as positive-negative solvers. This ensures that all light profile intensities are positive and therefore\n", - "physical.\n", - "\n", - "For pixelized source reconstructions on interferometer data this solver is often disabled because negative\n", - "visibility-plane noise can pull individual pixels negative without anything being wrong physically. For linear\n", - "*light profiles*, the intensity is a single physical normalisation of an extended profile, so we keep the\n", - "positive-only solver enabled.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data). This is the standard\n", - " convention for interferometer modeling, as the lens galaxy's optical/IR emission is typically below the\n", - " detection threshold of mm/sub-mm interferometers.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a linear `SersicCore`.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook.\n", - "\n", - "__Imaging Equivalent__\n", - "\n", - "For the CCD-imaging version of this script, which also fits a linear `Sersic` for the lens light, see\n", - "`autolens_workspace/*/imaging/features/linear_light_profiles/modeling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit with\n", - "the lens model.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and\n", - "compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native Non-Uniform Fast Fourier\n", - "Transform backed by `nufftax`, which is the required choice for fast linear light profile modeling and\n", - "scales efficiently from a few hundred visibilities to tens of millions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling\n", - "is used, which evaluates light profiles on a higher resolution grid than the image data to ensure the\n", - "calculation is accurate.\n", - "\n", - "Interferometer data does not observe galaxies in a way where over sampling is necessary, therefore all\n", - "interferometer calculations are performed without over sampling.\n", - "\n", - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is a linear `SersicCore` [5 parameters \u2014 `intensity` is solved for analytically].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=12.\n", - "\n", - "Note how the source galaxy uses a linear light profile, meaning that its `intensity` parameter is no longer a\n", - "free parameter in the fit. There is no lens-light component (interferometer convention).\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook:\n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", - "refer to `start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the source galaxy's light profile does not include an `intensity` parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", - "full description).\n", - "\n", - "In the `interferometer/modeling.py` example 75 live points (`n_live=75`) were used to sample parameter space.\n", - "For this linear light profile fit we keep `n_live=75` \u2014 the saving from one fewer free parameter is modest, and\n", - "the run-time benefit on interferometer data comes mostly from the reliability of the linear inversion rather\n", - "than a reduction in live points." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\") / \"features\",\n", - " name=\"linear_light_profiles\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisInterferometer` object defining how Nautilus fits the model to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `interferometer/modeling.py` example explains how VRAM is used during GPU-based fitting and how to print\n", - "the estimated VRAM required by a model.\n", - "\n", - "For each linear light profile in the model a small additional amount of VRAM is used to store its NUFFT'd\n", - "mapping matrix column. For 1-10 linear light profiles this is a tiny amount of VRAM (e.g. < 10MB per batched\n", - "likelihood). Even for large batch sizes you almost certainly will not use enough VRAM to require monitoring.\n", - "\n", - "VRAM on interferometer datasets is driven primarily by the visibility count and the real-space mask size, not\n", - "the number of linear light profiles in the model.\n", - "\n", - "__Run Time__\n", - "\n", - "For standard light profiles fitting interferometer data, the log likelihood evaluation time is dominated by the\n", - "NUFFT step.\n", - "\n", - "For linear light profiles, the per-evaluation cost is the NUFFT plus a small additional cost from the linear\n", - "inversion. The inversion adds approximately 3-4x the cost of the inversion-only term compared to the\n", - "fixed-intensity case, but because the NUFFT typically dominates the total cost, the overall slow-down per\n", - "likelihood is usually closer to 1.1-1.5x for a model with a single linear source profile.\n", - "\n", - "Because one free parameter has been removed from the model (the source `intensity`) and the parameter-space\n", - "degeneracy between `intensity` and shape parameters is broken, the total number of likelihood evaluations needed\n", - "for convergence is usually reduced. Fits using standard light profiles and linear light profiles therefore take\n", - "roughly the same wall-clock time to run. The simpler parameter space of linear light profiles means the\n", - "model-fit is more reliable, less susceptible to converging to a local maximum, and scales better if more linear\n", - "light profiles are added (e.g. an MGE source).\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output\n", - "folder for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", - "refer to `start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that `intensity` parameters are not inferred by the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "The source galaxy appears similar to that in the data, confirming that the `intensity` value inferred by the\n", - "inversion process is accurate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The intensity of a linear light profile is not part of the model parameterization, and is therefore not\n", - "displayed in the `model.results` file.\n", - "\n", - "To extract the `intensity` value of a specific component in the model, we use the `max_log_likelihood_tracer`,\n", - "which has already performed the inversion and therefore the galaxy light profiles have their solved-for\n", - "`intensity` values associated with them." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "# The source is the only galaxy with a light profile in the interferometer model \u2014 index -1 grabs it\n", - "# regardless of tracer ordering.\n", - "print(tracer.galaxies[-1].bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Tracer` contained in the `max_log_likelihood_fit` also has the solved for `intensity` value:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = result.max_log_likelihood_fit\n", - "\n", - "tracer = fit.tracer\n", - "\n", - "print(tracer.galaxies[-1].bulge.intensity)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization__\n", - "\n", - "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they do\n", - "not have an `intensity` value.\n", - "\n", - "Therefore, a helper produces an equivalent tracer in which every linear light profile has been replaced with an\n", - "ordinary light profile carrying its solved-for `intensity`. That helper-tracer can then be visualised:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = result.max_log_likelihood_tracer\n", - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "__Result (Advanced)__\n", - "\n", - "The code below shows additional results that can be computed from a `Result` object following a fit with a\n", - "linear light profile.\n", - "\n", - "__Max Likelihood Inversion__\n", - "\n", - "As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit`, which contains the\n", - "`Inversion` object we need." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = result.max_log_likelihood_fit.inversion" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This `Inversion` is what handled the linear algebra that produced the source `intensity` value above.\n", - "\n", - "__Linear Objects (Internal Source Code)__\n", - "\n", - "An `Inversion` contains all of the linear objects used to reconstruct the data in its `linear_obj_list`.\n", - "\n", - "This list may include the following objects:\n", - "\n", - " - `LightProfileLinearObjFuncList`: Holds a list of linear light profiles and the functionality used to\n", - " reconstruct data in an inversion. It may contain a single light profile (e.g. `lp_linear.SersicCore`) or\n", - " many light profiles combined in a `Basis` (e.g. `lp_basis.Basis`).\n", - "\n", - " - `Mapper`: The linear object used by a `Pixelization` to reconstruct data via an `Inversion`. The `Mapper`\n", - " is specific to the `Pixelization`'s `Mesh` (e.g. a `RectangularMapper` is used for a `RectangularAdaptDensity`\n", - " mesh).\n", - "\n", - "In this example, the model has one linear `SersicCore` for the source galaxy's bulge and no lens-light\n", - "component. The inversion therefore has a single `LightProfileLinearObjFuncList` entry, which holds the source\n", - "bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.linear_obj_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To extract results from an inversion many quantities come in lists or require us to specify the linear object\n", - "we wish to use. Knowing what linear objects are in the `linear_obj_list`, and what indexes they correspond to,\n", - "is therefore important.\n", - "\n", - "The single entry in this example is the source bulge." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"LightProfileLinearObjFuncList (Source SersicCore) = {inversion.linear_obj_list[0]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `LightProfileLinearObjFuncList` contains a `light_profile_list`. In this example the list has a single\n", - "light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"Linear Light Profile list (Source SersicCore) = {inversion.linear_obj_list[0].light_profile_list}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Linear Light Profiles (Interferometer)\n", + "==========================================================\n", + "\n", + "A \"linear light profile\" is a variant of a standard light profile where the `intensity` parameter is solved for\n", + "via linear algebra every time the model is fitted to the data. This uses a process called an \"inversion\" and it\n", + "always computes the `intensity` values that give the best fit to the data (e.g. maximize the likelihood) given\n", + "the light profile's other parameters.\n", + "\n", + "Linear light profiles have been a standard tool for fitting CCD imaging data for a long time. For interferometer\n", + "data they used to be impractical, because every likelihood evaluation has to Fourier-transform each basis component\n", + "into the uv-plane, and prior NUFFT backends were not JAX-friendly. With `nufftax`\n", + "(https://github.com/GragasLab/nufftax) \u2014 a JAX-native Non-Uniform Fast Fourier Transform \u2014 the image-to-uv\n", + "transform now runs inside the same jit/vmap pipeline as the rest of the model, so the per-iteration overhead of\n", + "NUFFT-ing each basis component is amortised on the GPU. Linear light profile fits are therefore practical for\n", + "interferometer data at any visibility count, including ALMA-class datasets with tens of millions of visibilities.\n", + "\n", + "Based on the advantages below, we recommend you use linear light profiles whenever fitting light profiles to\n", + "interferometer data.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Benefits and drawbacks of linear light profiles, and how they apply to\n", + " interferometer data specifically.\n", + "- **NUFFT (nufftax):** Why linear light profile fits to visibilities are now practical thanks to nufftax.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Model:** Compose the lens model fitted to the data \u2014 `Isothermal` + `ExternalShear` lens mass and a\n", + " linear `SersicCore` source. The lens light is omitted (interferometer convention).\n", + "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", + "- **Dataset:** Load the strong lens `Interferometer` dataset, using `TransformerNUFFT` (backed by `nufftax`).\n", + "- **Over Sampling:** Interferometer modeling does not use over-sampling (covered briefly here for users\n", + " familiar with imaging).\n", + "- **Search:** Configure the non-linear search (Nautilus).\n", + "- **Analysis:** Create the `AnalysisInterferometer` object.\n", + "- **VRAM:** Linear light profiles add negligible VRAM compared to standard light profiles.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Intensities:** How to extract solved-for `intensity` values from the result.\n", + "- **Visualization:** Visualising fits with linear light profiles requires the\n", + " `model_obj_linear_light_profiles_to_light_profiles` helper.\n", + "- **Max Likelihood Inversion:** Access the `Inversion` object from the result.\n", + "- **Linear Objects (Internal Source Code):** The internal `linear_obj_list` representation used by the\n", + " inversion.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "The source galaxy's `intensity` parameter is therefore not a free parameter in the model-fit, reducing the\n", + "dimensionality of non-linear parameter space by one. The lens light is already omitted for interferometer data,\n", + "so the saving is smaller than the imaging case (where lens and source both contribute) \u2014 but the inversion still\n", + "removes the degeneracies between `intensity` and the source's shape parameters (e.g. `effective_radius`,\n", + "`sersic_index`), which are difficult degeneracies for the non-linear search to map out accurately. This produces\n", + "more reliable lens model results and the fit converges in fewer iterations.\n", + "\n", + "The inversion has a relatively small computational cost on top of the NUFFT, so we reduce the model complexity\n", + "without much slow-down.\n", + "\n", + "__Disadvantages__\n", + "\n", + "Although the computation time of the inversion is small, it is not non-negligible. It is approximately 3-4x\n", + "slower per likelihood than using a standard light profile with a fixed `intensity`.\n", + "\n", + "The gains in run times from the simpler parameter space therefore broadly balance the slower per-likelihood\n", + "evaluation. The headline benefit is reliability, not raw speed.\n", + "\n", + "__NUFFT (nufftax)__\n", + "\n", + "The image-to-visibilities Fourier transform is performed by a Non-Uniform Fast Fourier Transform (NUFFT),\n", + "exposed in **PyAutoLens** as `TransformerNUFFT`. The default backend is `nufftax`, a pure-JAX NUFFT that\n", + "jit-compiles and vmap-batches like the rest of the library:\n", + "\n", + " https://github.com/GragasLab/nufftax\n", + "\n", + "Because `nufftax` is JAX-native, NUFFT-ing each linear basis image happens inside the same compiled likelihood\n", + "that does the inversion, mass model ray-tracing, and chi-squared sum. There is no host round-trip between\n", + "NUFFT calls, so a model with N linear light profiles costs only N forward-NUFFTs per iteration on the GPU \u2014\n", + "fast enough that linear inversions in the visibility plane are now routinely practical.\n", + "\n", + "If `nufftax` is not installed, install it via `pip install nufftax`. A legacy pynufft-backed transformer\n", + "(`TransformerNUFFTPyNUFFT`) is available as a non-JAX fallback but is not recommended for linear light profiles.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and\n", + "negative values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical.\n", + "\n", + "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it is\n", + "as fast as positive-negative solvers. This ensures that all light profile intensities are positive and therefore\n", + "physical.\n", + "\n", + "For pixelized source reconstructions on interferometer data this solver is often disabled because negative\n", + "visibility-plane noise can pull individual pixels negative without anything being wrong physically. For linear\n", + "*light profiles*, the intensity is a single physical normalisation of an extended profile, so we keep the\n", + "positive-only solver enabled.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data). This is the standard\n", + " convention for interferometer modeling, as the lens galaxy's optical/IR emission is typically below the\n", + " detection threshold of mm/sub-mm interferometers.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a linear `SersicCore`.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook.\n", + "\n", + "__Imaging Equivalent__\n", + "\n", + "For the CCD-imaging version of this script, which also fits a linear `Sersic` for the lens light, see\n", + "`autolens_workspace/*/imaging/features/linear_light_profiles/modeling.py`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit with\n", + "the lens model.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and\n", + "compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native Non-Uniform Fast Fourier\n", + "Transform backed by `nufftax`, which is the required choice for fast linear light profile modeling and\n", + "scales efficiently from a few hundred visibilities to tens of millions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling\n", + "is used, which evaluates light profiles on a higher resolution grid than the image data to ensure the\n", + "calculation is accurate.\n", + "\n", + "Interferometer data does not observe galaxies in a way where over sampling is necessary, therefore all\n", + "interferometer calculations are performed without over sampling.\n", + "\n", + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is a linear `SersicCore` [5 parameters \u2014 `intensity` is solved for analytically].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=12.\n", + "\n", + "Note how the source galaxy uses a linear light profile, meaning that its `intensity` parameter is no longer a\n", + "free parameter in the fit. There is no lens-light component (interferometer convention).\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook:\n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", + "refer to `start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the source galaxy's light profile does not include an `intensity` parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", + "full description).\n", + "\n", + "In the `interferometer/modeling.py` example 75 live points (`n_live=75`) were used to sample parameter space.\n", + "For this linear light profile fit we keep `n_live=75` \u2014 the saving from one fewer free parameter is modest, and\n", + "the run-time benefit on interferometer data comes mostly from the reliability of the linear inversion rather\n", + "than a reduction in live points." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\") / \"features\",\n", + " name=\"linear_light_profiles\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisInterferometer` object defining how Nautilus fits the model to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `interferometer/modeling.py` example explains how VRAM is used during GPU-based fitting and how to print\n", + "the estimated VRAM required by a model.\n", + "\n", + "For each linear light profile in the model a small additional amount of VRAM is used to store its NUFFT'd\n", + "mapping matrix column. For 1-10 linear light profiles this is a tiny amount of VRAM (e.g. < 10MB per batched\n", + "likelihood). Even for large batch sizes you almost certainly will not use enough VRAM to require monitoring.\n", + "\n", + "VRAM on interferometer datasets is driven primarily by the visibility count and the real-space mask size, not\n", + "the number of linear light profiles in the model.\n", + "\n", + "__Run Time__\n", + "\n", + "For standard light profiles fitting interferometer data, the log likelihood evaluation time is dominated by the\n", + "NUFFT step.\n", + "\n", + "For linear light profiles, the per-evaluation cost is the NUFFT plus a small additional cost from the linear\n", + "inversion. The inversion adds approximately 3-4x the cost of the inversion-only term compared to the\n", + "fixed-intensity case, but because the NUFFT typically dominates the total cost, the overall slow-down per\n", + "likelihood is usually closer to 1.1-1.5x for a model with a single linear source profile.\n", + "\n", + "Because one free parameter has been removed from the model (the source `intensity`) and the parameter-space\n", + "degeneracy between `intensity` and shape parameters is broken, the total number of likelihood evaluations needed\n", + "for convergence is usually reduced. Fits using standard light profiles and linear light profiles therefore take\n", + "roughly the same wall-clock time to run. The simpler parameter space of linear light profiles means the\n", + "model-fit is more reliable, less susceptible to converging to a local maximum, and scales better if more linear\n", + "light profiles are added (e.g. an MGE source).\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output\n", + "folder for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", + "refer to `start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that `intensity` parameters are not inferred by the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "The source galaxy appears similar to that in the data, confirming that the `intensity` value inferred by the\n", + "inversion process is accurate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The intensity of a linear light profile is not part of the model parameterization, and is therefore not\n", + "displayed in the `model.results` file.\n", + "\n", + "To extract the `intensity` value of a specific component in the model, we use the `max_log_likelihood_tracer`,\n", + "which has already performed the inversion and therefore the galaxy light profiles have their solved-for\n", + "`intensity` values associated with them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "# The source is the only galaxy with a light profile in the interferometer model \u2014 index -1 grabs it\n", + "# regardless of tracer ordering.\n", + "print(tracer.galaxies[-1].bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Tracer` contained in the `max_log_likelihood_fit` also has the solved for `intensity` value:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = result.max_log_likelihood_fit\n", + "\n", + "tracer = fit.tracer\n", + "\n", + "print(tracer.galaxies[-1].bulge.intensity)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization__\n", + "\n", + "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they do\n", + "not have an `intensity` value.\n", + "\n", + "Therefore, a helper produces an equivalent tracer in which every linear light profile has been replaced with an\n", + "ordinary light profile carrying its solved-for `intensity`. That helper-tracer can then be visualised:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = result.max_log_likelihood_tracer\n", + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "__Result (Advanced)__\n", + "\n", + "The code below shows additional results that can be computed from a `Result` object following a fit with a\n", + "linear light profile.\n", + "\n", + "__Max Likelihood Inversion__\n", + "\n", + "As seen elsewhere in the workspace, the result contains a `max_log_likelihood_fit`, which contains the\n", + "`Inversion` object we need." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = result.max_log_likelihood_fit.inversion" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `Inversion` is what handled the linear algebra that produced the source `intensity` value above.\n", + "\n", + "__Linear Objects (Internal Source Code)__\n", + "\n", + "An `Inversion` contains all of the linear objects used to reconstruct the data in its `linear_obj_list`.\n", + "\n", + "This list may include the following objects:\n", + "\n", + " - `LightProfileLinearObjFuncList`: Holds a list of linear light profiles and the functionality used to\n", + " reconstruct data in an inversion. It may contain a single light profile (e.g. `lp_linear.SersicCore`) or\n", + " many light profiles combined in a `Basis` (e.g. `lp_basis.Basis`).\n", + "\n", + " - `Mapper`: The linear object used by a `Pixelization` to reconstruct data via an `Inversion`. The `Mapper`\n", + " is specific to the `Pixelization`'s `Mesh` (e.g. a `RectangularMapper` is used for a `RectangularAdaptDensity`\n", + " mesh).\n", + "\n", + "In this example, the model has one linear `SersicCore` for the source galaxy's bulge and no lens-light\n", + "component. The inversion therefore has a single `LightProfileLinearObjFuncList` entry, which holds the source\n", + "bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.linear_obj_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To extract results from an inversion many quantities come in lists or require us to specify the linear object\n", + "we wish to use. Knowing what linear objects are in the `linear_obj_list`, and what indexes they correspond to,\n", + "is therefore important.\n", + "\n", + "The single entry in this example is the source bulge." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"LightProfileLinearObjFuncList (Source SersicCore) = {inversion.linear_obj_list[0]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `LightProfileLinearObjFuncList` contains a `light_profile_list`. In this example the list has a single\n", + "light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"Linear Light Profile list (Source SersicCore) = {inversion.linear_obj_list[0].light_profile_list}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/linear_light_profiles/slam.ipynb b/notebooks/interferometer/features/linear_light_profiles/slam.ipynb index d05a99428..fee75f7c7 100644 --- a/notebooks/interferometer/features/linear_light_profiles/slam.ipynb +++ b/notebooks/interferometer/features/linear_light_profiles/slam.ipynb @@ -1,738 +1,775 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Linear Light Profiles: SLaM (Interferometer)\n", - "=============================================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for `Interferometer`\n", - "data, using a **linear light profile** in the SOURCE LP stage instead of a Multi-Gaussian Expansion (MGE).\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that guide before\n", - "working through this example.\n", - "\n", - "The interferometer pixelization SLaM pipeline at `interferometer/features/pixelization/slam.py` is the closest\n", - "analogue \u2014 this script keeps the same pipeline structure (SOURCE LP \u2192 SOURCE PIX 1 \u2192 SOURCE PIX 2 \u2192 MASS\n", - "TOTAL) and only changes the SOURCE LP source bulge from an MGE to a linear `SersicCore`. The LIGHT LP\n", - "pipeline from `slam_start_here.py` is omitted because interferometer data does not contain lens light\n", - "emission.\n", - "\n", - "The differences from the interferometer pixelization SLaM are:\n", - "\n", - " - The SOURCE LP PIPELINE uses a single linear `SersicCore` profile (`al.lp_linear.SersicCore`) for the\n", - " source galaxy's bulge instead of a multi-Gaussian expansion.\n", - "\n", - "Linear light profiles solve for the `intensity` analytically via linear algebra, removing it from the\n", - "non-linear parameter space. This reduces the dimensionality of the SOURCE LP search and eliminates\n", - "intensity-shape degeneracies on the source bulge. The pixelized source-reconstruction stages\n", - "(SOURCE PIX 1 and 2) and the MASS TOTAL stage are unchanged.\n", - "\n", - "The SOURCE LP stage uses `TransformerNUFFT` (backed by JAX-native `nufftax`,\n", - "https://github.com/GragasLab/nufftax). Linear light profile fits to visibilities are now practical at any\n", - "visibility count because the per-iteration NUFFT of each linear basis component is fast on the GPU.\n", - "\n", - "The SOURCE PIX and MASS TOTAL stages switch to `TransformerNUFFT` combined with the pre-computed sparse\n", - "operator, because pixelized source reconstructions exploit sparsity rather than the NUFFT path.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with `slam_start_here`.\n", - "- **SOURCE LP PIPELINE:** Initializes the mass and source-light model with a linear `SersicCore` profile,\n", - " fitted via `TransformerNUFFT`.\n", - "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP result.\n", - "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included\n", - " (interferometer data).\n", - "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with\n", - " `TransformerNUFFT` + sparse operator (source_pix onwards).\n", - "- **Sparse Operators:** Pre-compute the sparse operator for the pixelized stages.\n", - "- **Settings:** Disable the positive-only solver so the pixelized source reconstruction can have negative\n", - " pixel values.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database\n", - " use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in `features/pixelization/modeling`, the mesh shape is fixed before modeling.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines and how they\n", - " integrate into strong-lens modeling.\n", - "\n", - "- **Linear Light Profiles** (`interferometer/features/linear_light_profiles/modeling`)\n", - " How linear light profiles work for interferometer data, why nufftax makes them practical, and what they\n", - " add over fixed-intensity profiles.\n", - "\n", - "- **Interferometer Pixelization SLaM** (`interferometer/features/pixelization/slam`)\n", - " The canonical interferometer SLaM pipeline. This script is essentially that pipeline with the SOURCE LP\n", - " source bulge swapped from an MGE to a linear `SersicCore`.\n", - "\n", - "You can still run the script without fully understanding these guides, but reviewing them later will make\n", - "the structure and choices of the SLaM workflow clearer.\n", - "\n", - "__High Resolution Dataset__\n", - "\n", - "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large\n", - "files which are too big to include in the main `autolens_workspace` repository:\n", - "\n", - " https://github.com/PyAutoLabs/autolens_workspace_large_files\n", - "\n", - "After downloading the file, place it in the directory:\n", - "\n", - " `autolens_workspace/dataset/interferometer/alma`\n", - "\n", - "You can then perform modeling using this high-resolution dataset by uncommenting the relevant line of code\n", - "below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light using a\n", - "**linear** `SersicCore` profile. The lens galaxy mass and external shear are fitted at the same time.\n", - "\n", - "The linear `SersicCore` has its `intensity` solved for analytically via the linear inversion rather than as\n", - "a free parameter of the non-linear search. This makes the SOURCE LP search faster and more reliable: the\n", - "intensity-shape degeneracies (between `intensity` and `effective_radius` / `sersic_index`) are eliminated.\n", - "\n", - "This stage uses the `dataset_nufft` Interferometer (built with `TransformerNUFFT`, backed by `nufftax`),\n", - "which makes the per-iteration NUFFT of the linear basis component fast even for ALMA-class datasets with\n", - "millions of visibilities.\n", - "\n", - "The mass and source models from this search initialize the SOURCE PIX PIPELINE searches that follow, and\n", - "the result also provides the adapt image and position likelihood used by those later stages.\n", - "\n", - "Note that no lens light is fitted: interferometer data does not contain lens light emission, so\n", - "`lens.bulge` and `lens.disk` are kept at `None`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", - "\n", - " source_bulge = af.Model(al.lp_linear.SersicCore)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " # interferometer data does not contain lens light emission\n", - " bulge=None,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source galaxy.\n", - "\n", - "The first search fits a pixelization whose purpose is to generate a high-quality adapt image used in\n", - "search 2. It uses the adapt image computed from the SOURCE LP result and the position likelihood is also\n", - "derived automatically from the SOURCE LP result \u2014 no manual positions input is required.\n", - "\n", - "This stage uses the `dataset_sparse` Interferometer (built with `TransformerNUFFT` +\n", - "`apply_sparse_operator`). The NUFFT keeps the one-time dirty-image setup tractable at ALMA-scale\n", - "visibility counts, and the precomputed sparse operator makes per-likelihood curvature assembly use the\n", - "FFT-based W\u0303 precision matrix instead of the dense `transformed_mapping_matrix`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", - "pixelization and regularization.\n", - "\n", - "Note that the LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not\n", - "contain lens light emission.\n", - "\n", - "Like SOURCE PIX PIPELINE 1, this stage uses the `dataset_sparse` Interferometer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " settings=settings,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", - "contain lens light emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " source_pix_result_2: af.Result,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_1.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=source_pix_result_1.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load the `Interferometer` data and define the real-space mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", - ")\n", - "\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "# dataset_name = \"alma\"\n", - "#\n", - "# if dataset_name == \"alma\":\n", - "#\n", - "# real_space_mask = al.Mask2D.circular(\n", - "# shape_native=(800, 800),\n", - "# pixel_scales=0.01,\n", - "# radius=mask_radius,\n", - "# )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Two Datasets__\n", - "\n", - "The SLaM pipeline runs in two phases that prefer different transformers:\n", - "\n", - "- `dataset_nufft` uses `TransformerNUFFT` (backed by JAX-native `nufftax`) for the `source_lp` stage. With\n", - " the linear `SersicCore` source bulge this is the fast path at any visibility count.\n", - "- `dataset_sparse` uses `TransformerNUFFT` combined with `apply_sparse_operator(...)` for\n", - " `source_pix_1`, `source_pix_2`, and `mass_total`. Pixelized source reconstructions exploit sparsity in\n", - " the linear inversion rather than the NUFFT, so this combination is the right choice for the pixelized\n", - " stages.\n", - "\n", - "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload) differ." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_nufft = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "dataset_sparse = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", - "pixelized source modeling, especially for many visibilities.\n", - "\n", - "We use a try / except to load the pre-computed curvature preload, which is necessary to use the sparse\n", - "operator formalism. If this file does not exist (e.g. you have not made it manually via the\n", - "`many_visibilities_preparation` example) it is made here.\n", - "\n", - "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", - "`source_lp` does not need it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " nufft_precision_operator = np.load(\n", - " file=dataset_path / \"nufft_precision_operator.npy\",\n", - " )\n", - "except FileNotFoundError:\n", - " nufft_precision_operator = None\n", - "\n", - "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", - " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "Disable the default positive-only linear algebra solver so the pixelized source reconstruction can have\n", - "negative pixel values. (The linear `SersicCore` in SOURCE LP is still constrained to positive `intensity`\n", - "internally because that is a physical normalization.)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"interferometer\") / \"slam_linear_light_profiles\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for a\n", - "description of each pipeline step.\n", - "\n", - "Note the transformer split: `source_lp` is passed `dataset_nufft` (TransformerNUFFT), while every later\n", - "stage is passed `dataset_sparse` (TransformerNUFFT + sparse operator)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset_nufft,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " settings=settings,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Linear Light Profiles: SLaM (Interferometer)\n", + "=============================================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for `Interferometer`\n", + "data, using a **linear light profile** in the SOURCE LP stage instead of a Multi-Gaussian Expansion (MGE).\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that guide before\n", + "working through this example.\n", + "\n", + "The interferometer pixelization SLaM pipeline at `interferometer/features/pixelization/slam.py` is the closest\n", + "analogue \u2014 this script keeps the same pipeline structure (SOURCE LP \u2192 SOURCE PIX 1 \u2192 SOURCE PIX 2 \u2192 MASS\n", + "TOTAL) and only changes the SOURCE LP source bulge from an MGE to a linear `SersicCore`. The LIGHT LP\n", + "pipeline from `slam_start_here.py` is omitted because interferometer data does not contain lens light\n", + "emission.\n", + "\n", + "The differences from the interferometer pixelization SLaM are:\n", + "\n", + " - The SOURCE LP PIPELINE uses a single linear `SersicCore` profile (`al.lp_linear.SersicCore`) for the\n", + " source galaxy's bulge instead of a multi-Gaussian expansion.\n", + "\n", + "Linear light profiles solve for the `intensity` analytically via linear algebra, removing it from the\n", + "non-linear parameter space. This reduces the dimensionality of the SOURCE LP search and eliminates\n", + "intensity-shape degeneracies on the source bulge. The pixelized source-reconstruction stages\n", + "(SOURCE PIX 1 and 2) and the MASS TOTAL stage are unchanged.\n", + "\n", + "The SOURCE LP stage uses `TransformerNUFFT` (backed by JAX-native `nufftax`,\n", + "https://github.com/GragasLab/nufftax). Linear light profile fits to visibilities are now practical at any\n", + "visibility count because the per-iteration NUFFT of each linear basis component is fast on the GPU.\n", + "\n", + "The SOURCE PIX and MASS TOTAL stages switch to `TransformerNUFFT` combined with the pre-computed sparse\n", + "operator, because pixelized source reconstructions exploit sparsity rather than the NUFFT path.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with `slam_start_here`.\n", + "- **SOURCE LP PIPELINE:** Initializes the mass and source-light model with a linear `SersicCore` profile,\n", + " fitted via `TransformerNUFFT`.\n", + "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP result.\n", + "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included\n", + " (interferometer data).\n", + "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with\n", + " `TransformerNUFFT` + sparse operator (source_pix onwards).\n", + "- **Sparse Operators:** Pre-compute the sparse operator for the pixelized stages.\n", + "- **Settings:** Disable the positive-only solver so the pixelized source reconstruction can have negative\n", + " pixel values.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database\n", + " use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in `features/pixelization/modeling`, the mesh shape is fixed before modeling.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines and how they\n", + " integrate into strong-lens modeling.\n", + "\n", + "- **Linear Light Profiles** (`interferometer/features/linear_light_profiles/modeling`)\n", + " How linear light profiles work for interferometer data, why nufftax makes them practical, and what they\n", + " add over fixed-intensity profiles.\n", + "\n", + "- **Interferometer Pixelization SLaM** (`interferometer/features/pixelization/slam`)\n", + " The canonical interferometer SLaM pipeline. This script is essentially that pipeline with the SOURCE LP\n", + " source bulge swapped from an MGE to a linear `SersicCore`.\n", + "\n", + "You can still run the script without fully understanding these guides, but reviewing them later will make\n", + "the structure and choices of the SLaM workflow clearer.\n", + "\n", + "__High Resolution Dataset__\n", + "\n", + "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large\n", + "files which are too big to include in the main `autolens_workspace` repository:\n", + "\n", + " https://github.com/PyAutoLabs/autolens_workspace_large_files\n", + "\n", + "After downloading the file, place it in the directory:\n", + "\n", + " `autolens_workspace/dataset/interferometer/alma`\n", + "\n", + "You can then perform modeling using this high-resolution dataset by uncommenting the relevant line of code\n", + "below." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light using a\n", + "**linear** `SersicCore` profile. The lens galaxy mass and external shear are fitted at the same time.\n", + "\n", + "The linear `SersicCore` has its `intensity` solved for analytically via the linear inversion rather than as\n", + "a free parameter of the non-linear search. This makes the SOURCE LP search faster and more reliable: the\n", + "intensity-shape degeneracies (between `intensity` and `effective_radius` / `sersic_index`) are eliminated.\n", + "\n", + "This stage uses the `dataset_nufft` Interferometer (built with `TransformerNUFFT`, backed by `nufftax`),\n", + "which makes the per-iteration NUFFT of the linear basis component fast even for ALMA-class datasets with\n", + "millions of visibilities.\n", + "\n", + "The mass and source models from this search initialize the SOURCE PIX PIPELINE searches that follow, and\n", + "the result also provides the adapt image and position likelihood used by those later stages.\n", + "\n", + "Note that no lens light is fitted: interferometer data does not contain lens light emission, so\n", + "`lens.bulge` and `lens.disk` are kept at `None`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", + "\n", + " source_bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " # interferometer data does not contain lens light emission\n", + " bulge=None,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source galaxy.\n", + "\n", + "The first search fits a pixelization whose purpose is to generate a high-quality adapt image used in\n", + "search 2. It uses the adapt image computed from the SOURCE LP result and the position likelihood is also\n", + "derived automatically from the SOURCE LP result \u2014 no manual positions input is required.\n", + "\n", + "This stage uses the `dataset_sparse` Interferometer (built with `TransformerNUFFT` +\n", + "`apply_sparse_operator`). The NUFFT keeps the one-time dirty-image setup tractable at ALMA-scale\n", + "visibility counts, and the precomputed sparse operator makes per-likelihood curvature assembly use the\n", + "FFT-based W\u0303 precision matrix instead of the dense `transformed_mapping_matrix`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", + "pixelization and regularization.\n", + "\n", + "Note that the LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not\n", + "contain lens light emission.\n", + "\n", + "Like SOURCE PIX PIPELINE 1, this stage uses the `dataset_sparse` Interferometer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " settings=settings,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", + "contain lens light emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " source_pix_result_2: af.Result,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_1.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=source_pix_result_1.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load the `Interferometer` data and define the real-space mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", + ")\n", + "\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "# dataset_name = \"alma\"\n", + "#\n", + "# if dataset_name == \"alma\":\n", + "#\n", + "# real_space_mask = al.Mask2D.circular(\n", + "# shape_native=(800, 800),\n", + "# pixel_scales=0.01,\n", + "# radius=mask_radius,\n", + "# )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Two Datasets__\n", + "\n", + "The SLaM pipeline runs in two phases that prefer different transformers:\n", + "\n", + "- `dataset_nufft` uses `TransformerNUFFT` (backed by JAX-native `nufftax`) for the `source_lp` stage. With\n", + " the linear `SersicCore` source bulge this is the fast path at any visibility count.\n", + "- `dataset_sparse` uses `TransformerNUFFT` combined with `apply_sparse_operator(...)` for\n", + " `source_pix_1`, `source_pix_2`, and `mass_total`. Pixelized source reconstructions exploit sparsity in\n", + " the linear inversion rather than the NUFFT, so this combination is the right choice for the pixelized\n", + " stages.\n", + "\n", + "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload) differ." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_nufft = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "dataset_sparse = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", + "pixelized source modeling, especially for many visibilities.\n", + "\n", + "We use a try / except to load the pre-computed curvature preload, which is necessary to use the sparse\n", + "operator formalism. If this file does not exist (e.g. you have not made it manually via the\n", + "`many_visibilities_preparation` example) it is made here.\n", + "\n", + "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", + "`source_lp` does not need it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " nufft_precision_operator = np.load(\n", + " file=dataset_path / \"nufft_precision_operator.npy\",\n", + " )\n", + "except FileNotFoundError:\n", + " nufft_precision_operator = None\n", + "\n", + "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", + " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "Disable the default positive-only linear algebra solver so the pixelized source reconstruction can have\n", + "negative pixel values. (The linear `SersicCore` in SOURCE LP is still constrained to positive `intensity`\n", + "internally because that is a physical normalization.)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"interferometer\") / \"slam_linear_light_profiles\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for a\n", + "description of each pipeline step.\n", + "\n", + "Note the transformer split: `source_lp` is passed `dataset_nufft` (TransformerNUFFT), while every later\n", + "stage is passed `dataset_sparse` (TransformerNUFFT + sparse operator)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset_nufft,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " settings=settings,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/multi_gaussian_expansion/fit.ipynb b/notebooks/interferometer/features/multi_gaussian_expansion/fit.ipynb index 4b4592712..9acf9642b 100644 --- a/notebooks/interferometer/features/multi_gaussian_expansion/fit.ipynb +++ b/notebooks/interferometer/features/multi_gaussian_expansion/fit.ipynb @@ -1,386 +1,423 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Multi Gaussian Expansion Fit (Interferometer)\n", - "=================================================================\n", - "\n", - "A multi-Gaussian expansion (MGE) decomposes a galaxy's light into ~15-100 Gaussians, where the `intensity`\n", - "of every Gaussian is solved for via linear algebra using a process called an \"inversion\".\n", - "\n", - "This script illustrates how to perform a single `FitInterferometer` of an MGE source model \u2014 that is, not\n", - "the full Nautilus model-fit, but a single likelihood evaluation given known basis parameters. This is\n", - "useful for understanding how the inversion produces the per-Gaussian solved-for `intensity` values, and\n", - "how to extract them from the resulting fit.\n", - "\n", - "**Role swap vs the imaging MGE example:** the imaging script fits the *lens* galaxy's complex morphology\n", - "with the MGE. For interferometer data the lens light is omitted (no detection in mm/sub-mm), so the MGE is\n", - "instead applied to the *source* galaxy.\n", - "\n", - "For an explanation of why MGE fits to visibility data are now practical thanks to the JAX-native NUFFT\n", - "`nufftax` (https://github.com/GragasLab/nufftax), see the companion `modeling.py` example.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Benefits and drawbacks of an MGE source for interferometer data.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Model:** The lens model whose `intensity` values we solve for via inversion \u2014 `Isothermal +\n", - " ExternalShear` mass with a 30-Gaussian MGE source bulge.\n", - "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", - "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", - "- **Basis:** Build the linear Gaussian basis used as the source bulge.\n", - "- **Fit:** Perform a single `FitInterferometer` and inspect the inversion.\n", - "- **Intensities:** Extract the per-Gaussian solved-for `intensity` values via\n", - " `fit.linear_light_profile_intensity_dict`.\n", - "- **Visualization:** Build the helper tracer where linear light profiles are replaced with ordinary light\n", - " profiles carrying their solved-for `intensity`, then plot.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "**PyAutoLens** uses a positive only linear algebra solver for the MGE inversion which has been extensively\n", - "optimized to ensure it is as fast as positive-negative solvers. This ensures that all Gaussian intensities\n", - "are positive and therefore physical \u2014 without it, an unconstrained MGE inversion can produce a\n", - "positive-negative \"ringing\" pattern across the basis.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data). Interferometer\n", - " convention.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's bulge is a multi-Gaussian expansion of 5 linear `Gaussian` profiles, all sharing\n", - " the same centre and `ell_comps`, with `sigma` values spanning 0.01\" to the mask radius in log-spaced\n", - " increments.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, using `TransformerNUFFT`\n", - "backed by `nufftax`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Basis__\n", - "\n", - "We build a `Basis` of 5 linear `Gaussian` profiles, all sharing the same centre and `ell_comps`, with\n", - "`sigma` values spanning 0.01\" to the mask radius in log-spaced increments.\n", - "\n", - "We use linear light profile Gaussians (`lp_linear.Gaussian`), which solve for each Gaussian's `intensity`\n", - "analytically via the inversion. This is essential for MGE \u2014 a wide range of positive `intensity` values\n", - "are needed to decompose the source's morphology, and we cannot guess them by eye.\n", - "\n", - "Linear light profiles are described in detail in the `linear_light_profiles.py` example; you should\n", - "familiarize yourself with that example before using the multi-Gaussian expansion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 5\n", - "\n", - "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius.\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "bulge_gaussian_list = []\n", - "for i in range(total_gaussians):\n", - " gaussian = al.lp_linear.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " bulge_gaussian_list.append(gaussian)\n", - "\n", - "source_bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "We now illustrate the API for performing a single MGE fit using standard `Galaxy`, `Tracer` and\n", - "`FitInterferometer` objects. Once we have a `Basis`, we can treat it like any other light profile." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=source_bulge,\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit's `subplot_fit_interferometer` shows the visibility-plane fit and dirty-image residuals. Because\n", - "the source bulge is a Basis of linear light profiles, the inversion has solved for each Gaussian's\n", - "`intensity` to maximize the fit to the observed visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `subplot_fit_dirty_images` provides a real-space view of the data, model and residuals via inverse-NUFFT\n", - "of the visibility-plane quantities. This is generally more interpretable to the human eye than the uv-plane\n", - "plots above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_dirty_images(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensities__\n", - "\n", - "The fit contains the solved-for `intensity` value of every Gaussian in the MGE basis.\n", - "\n", - "These are computed via `fit.linear_light_profile_intensity_dict`, which maps each linear light profile in\n", - "the model to its inferred `intensity`.\n", - "\n", - "The code below prints the first five Gaussian intensities for brevity \u2014 for an MGE of N Gaussians the full\n", - "dict has N entries." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.linear_light_profile_intensity_dict)\n", - "\n", - "for gaussian in bulge_gaussian_list[:5]:\n", - " print(\n", - " f\" sigma = {gaussian.sigma:.4f} \"\n", - " f\"intensity = {fit.linear_light_profile_intensity_dict[gaussian]:.6e}\"\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the\n", - "solved-for `intensity` values is also accessible from a fit.\n", - "\n", - "The benefit of this helper-tracer is that it can be visualised (linear light profiles cannot be plotted by\n", - "default because they do not have `intensity` values)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = fit.model_obj_linear_light_profiles_to_light_profiles" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization__\n", - "\n", - "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they\n", - "do not have an `intensity` value.\n", - "\n", - "The helper-tracer created above replaces every linear `Gaussian` with an ordinary `Gaussian` carrying its\n", - "solved-for `intensity` \u2014 that tracer can be plotted directly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image (MGE source)\"\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Multi Gaussian Expansion Fit (Interferometer)\n", + "=================================================================\n", + "\n", + "A multi-Gaussian expansion (MGE) decomposes a galaxy's light into ~15-100 Gaussians, where the `intensity`\n", + "of every Gaussian is solved for via linear algebra using a process called an \"inversion\".\n", + "\n", + "This script illustrates how to perform a single `FitInterferometer` of an MGE source model \u2014 that is, not\n", + "the full Nautilus model-fit, but a single likelihood evaluation given known basis parameters. This is\n", + "useful for understanding how the inversion produces the per-Gaussian solved-for `intensity` values, and\n", + "how to extract them from the resulting fit.\n", + "\n", + "**Role swap vs the imaging MGE example:** the imaging script fits the *lens* galaxy's complex morphology\n", + "with the MGE. For interferometer data the lens light is omitted (no detection in mm/sub-mm), so the MGE is\n", + "instead applied to the *source* galaxy.\n", + "\n", + "For an explanation of why MGE fits to visibility data are now practical thanks to the JAX-native NUFFT\n", + "`nufftax` (https://github.com/GragasLab/nufftax), see the companion `modeling.py` example.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Benefits and drawbacks of an MGE source for interferometer data.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Model:** The lens model whose `intensity` values we solve for via inversion \u2014 `Isothermal +\n", + " ExternalShear` mass with a 30-Gaussian MGE source bulge.\n", + "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", + "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", + "- **Basis:** Build the linear Gaussian basis used as the source bulge.\n", + "- **Fit:** Perform a single `FitInterferometer` and inspect the inversion.\n", + "- **Intensities:** Extract the per-Gaussian solved-for `intensity` values via\n", + " `fit.linear_light_profile_intensity_dict`.\n", + "- **Visualization:** Build the helper tracer where linear light profiles are replaced with ordinary light\n", + " profiles carrying their solved-for `intensity`, then plot.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "**PyAutoLens** uses a positive only linear algebra solver for the MGE inversion which has been extensively\n", + "optimized to ensure it is as fast as positive-negative solvers. This ensures that all Gaussian intensities\n", + "are positive and therefore physical \u2014 without it, an unconstrained MGE inversion can produce a\n", + "positive-negative \"ringing\" pattern across the basis.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data). Interferometer\n", + " convention.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's bulge is a multi-Gaussian expansion of 5 linear `Gaussian` profiles, all sharing\n", + " the same centre and `ell_comps`, with `sigma` values spanning 0.01\" to the mask radius in log-spaced\n", + " increments.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, using `TransformerNUFFT`\n", + "backed by `nufftax`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Basis__\n", + "\n", + "We build a `Basis` of 5 linear `Gaussian` profiles, all sharing the same centre and `ell_comps`, with\n", + "`sigma` values spanning 0.01\" to the mask radius in log-spaced increments.\n", + "\n", + "We use linear light profile Gaussians (`lp_linear.Gaussian`), which solve for each Gaussian's `intensity`\n", + "analytically via the inversion. This is essential for MGE \u2014 a wide range of positive `intensity` values\n", + "are needed to decompose the source's morphology, and we cannot guess them by eye.\n", + "\n", + "Linear light profiles are described in detail in the `linear_light_profiles.py` example; you should\n", + "familiarize yourself with that example before using the multi-Gaussian expansion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 5\n", + "\n", + "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius.\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "bulge_gaussian_list = []\n", + "for i in range(total_gaussians):\n", + " gaussian = al.lp_linear.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " bulge_gaussian_list.append(gaussian)\n", + "\n", + "source_bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "We now illustrate the API for performing a single MGE fit using standard `Galaxy`, `Tracer` and\n", + "`FitInterferometer` objects. Once we have a `Basis`, we can treat it like any other light profile." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=source_bulge,\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit's `subplot_fit_interferometer` shows the visibility-plane fit and dirty-image residuals. Because\n", + "the source bulge is a Basis of linear light profiles, the inversion has solved for each Gaussian's\n", + "`intensity` to maximize the fit to the observed visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `subplot_fit_dirty_images` provides a real-space view of the data, model and residuals via inverse-NUFFT\n", + "of the visibility-plane quantities. This is generally more interpretable to the human eye than the uv-plane\n", + "plots above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_dirty_images(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensities__\n", + "\n", + "The fit contains the solved-for `intensity` value of every Gaussian in the MGE basis.\n", + "\n", + "These are computed via `fit.linear_light_profile_intensity_dict`, which maps each linear light profile in\n", + "the model to its inferred `intensity`.\n", + "\n", + "The code below prints the first five Gaussian intensities for brevity \u2014 for an MGE of N Gaussians the full\n", + "dict has N entries." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.linear_light_profile_intensity_dict)\n", + "\n", + "for gaussian in bulge_gaussian_list[:5]:\n", + " print(\n", + " f\" sigma = {gaussian.sigma:.4f} \"\n", + " f\"intensity = {fit.linear_light_profile_intensity_dict[gaussian]:.6e}\"\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Tracer` where all linear light profile objects are replaced with ordinary light profiles using the\n", + "solved-for `intensity` values is also accessible from a fit.\n", + "\n", + "The benefit of this helper-tracer is that it can be visualised (linear light profiles cannot be plotted by\n", + "default because they do not have `intensity` values)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = fit.model_obj_linear_light_profiles_to_light_profiles" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization__\n", + "\n", + "Linear light profiles and objects containing them (e.g. galaxies, a tracer) cannot be plotted because they\n", + "do not have an `intensity` value.\n", + "\n", + "The helper-tracer created above replaces every linear `Gaussian` with an ordinary `Gaussian` carrying its\n", + "solved-for `intensity` \u2014 that tracer can be plotted directly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image (MGE source)\"\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/multi_gaussian_expansion/likelihood_function.ipynb b/notebooks/interferometer/features/multi_gaussian_expansion/likelihood_function.ipynb index 4f2854154..71c7a7f3c 100644 --- a/notebooks/interferometer/features/multi_gaussian_expansion/likelihood_function.ipynb +++ b/notebooks/interferometer/features/multi_gaussian_expansion/likelihood_function.ipynb @@ -1,766 +1,803 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Multi Gaussian Expansion (Interferometer)__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit\n", - "`Interferometer` data with a multi-Gaussian expansion (MGE) source: a `Basis` of many linear `Gaussian`\n", - "profiles whose `intensity` values are solved for analytically via a linear inversion.\n", - "\n", - "The visibility-plane linear inversion is **mathematically identical** to the single-component linear light\n", - "profile case (see `interferometer/features/linear_light_profiles/likelihood_function.py`). The only\n", - "difference is the number of columns in the mapping matrix \u2014 one column per Gaussian in the basis instead\n", - "of one column total. The data vector `D` and curvature matrix `F` therefore have dimensions equal to the\n", - "basis size, and the solved `s = F^{-1} D` gives one `intensity` per Gaussian.\n", - "\n", - "**Role swap vs the imaging MGE example:** the imaging script's basis is on the lens galaxy. Here it's on\n", - "the source galaxy, because interferometer data does not contain lens light. The source bulge image is\n", - "evaluated on the *ray-traced* source-plane grid; the rest of the inversion machinery is unchanged.\n", - "\n", - "For interferometer data this is now practical thanks to the JAX-native NUFFT `nufftax`\n", - "(https://github.com/GragasLab/nufftax) \u2014 every Gaussian's NUFFT happens inside the same jit/vmap pipeline\n", - "as the rest of the model, so the per-iteration NUFFT cost scales linearly with the basis size and is\n", - "amortised on the GPU.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The MGE likelihood function builds on:\n", - "\n", - " - `interferometer/log_likelihood_function.ipynb` \u2014 the standard interferometer parametric likelihood\n", - " function (NUFFT of a real-space image, visibility-plane $\\\\chi^2$).\n", - " - `interferometer/features/linear_light_profiles/likelihood_function.ipynb` \u2014 the single-component\n", - " visibility-plane linear inversion (data vector, curvature matrix, positive-only solver). The MGE is\n", - " a direct generalisation to N-component bases.\n", - "\n", - "This script repeats just enough setup that you can follow it without rereading those two \u2014 but if anything\n", - "is unclear, those are the places to look first.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Reading order before this script.\n", - "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", - "- **Dataset:** Load and plot the strong lens `Interferometer` dataset using `TransformerNUFFT` (nufftax).\n", - "- **Lens Galaxy:** A mass-only lens galaxy (no light \u2014 interferometer convention).\n", - "- **Source Galaxy MGE Basis:** Build the linear `Gaussian` basis used as the source bulge.\n", - "- **Ray Tracing:** Image-plane to source-plane grid.\n", - "- **Mapping Matrix:** Real-space mapping matrix \u2014 one column per Gaussian in the basis.\n", - "- **Transformed Mapping Matrix ($f$):** NUFFT each column to give the visibility-space mapping matrix.\n", - "- **Data Vector (D):** Compute $D$ from the transformed mapping matrix, visibilities, and noise map.\n", - "- **Curvature Matrix (F):** Compute $F$ separately for real and imaginary components, then sum.\n", - "- **Reconstruction (Positive-Negative):** Solve $s = F^{-1} D$ via NumPy.\n", - "- **Reconstruction (Positive Only):** Solve with the fast non-negative least squares (`fnnls`) algorithm.\n", - "- **Visibilities Reconstruction:** Map $s$ back to visibility space.\n", - "- **Likelihood Function:** Visibility-plane $\\\\chi^2$ and noise normalization.\n", - "- **Chi Squared:** Sum chi-squared contributions over real and imaginary components.\n", - "- **Noise Normalization Term:** The fixed noise normalization term.\n", - "- **Calculate The Log Likelihood:** Combine into the final log likelihood.\n", - "- **Fit:** Cross-check via `FitInterferometer`.\n", - "- **Lens Modeling:** How this likelihood is sampled in the full Nautilus fit.\n", - "- **Wrap Up:** Summary and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files using `TransformerNUFFT`\n", - "backed by `nufftax`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "For interferometer data the lens galaxy's optical/IR emission is typically below detection, so we model\n", - "only its mass:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy MGE Basis__\n", - "\n", - "We build a `Basis` of 5 linear `Gaussian` profiles for the source bulge, all sharing the same centre and\n", - "`ell_comps`, with `sigma` values spanning 0.01\" to the mask radius in log-spaced increments.\n", - "\n", - "Each Gaussian is a linear light profile \u2014 its `intensity` is solved for analytically via the inversion\n", - "below. Internally each linear profile carries `intensity=1.0`, which the inversion later rescales." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 5\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "bulge_gaussian_list = []\n", - "for i in range(total_gaussians):\n", - " gaussian = al.lp_linear.Gaussian(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " sigma=10 ** log10_sigma_list[i],\n", - " )\n", - " bulge_gaussian_list.append(gaussian)\n", - "\n", - "source_bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)\n", - "source_galaxy = al.Galaxy(redshift=1.0, bulge=source_bulge)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\\\theta$ from the image-plane to\n", - "its (y,x) source-plane coordinate $\\\\beta$ using the summed deflection angles $\\\\alpha$ of the mass\n", - "profiles:\n", - "\n", - " $\\\\beta = \\\\theta - \\\\alpha(\\\\theta)$\n", - "\n", - "For interferometer data the only grid we need to ray-trace is the real-space grid associated with the\n", - "`real_space_mask`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane); we only need the source-plane grid (index -1).\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=dataset.grids.lp)[-1]\n", - "\n", - "aplt.plot_grid(grid=traced_grid, title=\"Traced Source-Plane Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapping Matrix__\n", - "\n", - "For interferometer MGE modeling, the mapping matrix has one column per Gaussian in the basis. Each column\n", - "holds the real-space image of one Gaussian, evaluated on the ray-traced source-plane grid with\n", - "`intensity=1.0` internally. The lensing is already baked in because we evaluated on the traced grid.\n", - "\n", - "The dimensions are `(total_real_space_pixels, total_gaussians)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lp_linear_func_source = al.LightProfileLinearObjFuncList(\n", - " grid=traced_grid,\n", - " blurring_grid=None,\n", - " psf=None,\n", - " light_profile_list=bulge_gaussian_list,\n", - " regularization=None,\n", - ")\n", - "\n", - "mapping_matrix = lp_linear_func_source.mapping_matrix\n", - "\n", - "print(\"Mapping matrix shape (real-space pixels, total Gaussians):\")\n", - "print(mapping_matrix.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `mapping_matrix` shows each Gaussian's lensed source-plane image as one column." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.colorbar()\n", - "plt.title(\"Real-space mapping matrix \u2014 one column per Gaussian\")\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Transformed Mapping Matrix ($f$)__\n", - "\n", - "Every column of the real-space `mapping_matrix` must be NUFFT'd to the uv-plane. The result is the\n", - "`transformed_mapping_matrix` \u2014 a *complex-valued* matrix with dimensions\n", - "`(total_visibilities, total_gaussians)`.\n", - "\n", - "For an N-Gaussian MGE this matrix has N columns. The inversion's job is to find the N scalars that, when\n", - "multiplied by their respective columns and summed, best fit the observed visibilities.\n", - "\n", - "This NUFFT-each-column operation is exactly what nufftax made fast. The whole pipeline (ray-trace \u2192\n", - "source-plane Gaussian images \u2192 `transform_mapping_matrix` \u2192 linear inversion) is JIT-compilable under JAX." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", - " mapping_matrix=mapping_matrix\n", - ")\n", - "\n", - "print(\"Transformed mapping matrix shape (visibilities, total Gaussians):\")\n", - "print(transformed_mapping_matrix.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the real and imaginary components of the transformed mapping matrix. Each row is one observed\n", - "visibility; the values across that row are the per-unit-intensity contributions of each Gaussian at that\n", - "uv point." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " transformed_mapping_matrix.real,\n", - " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.title(\"Re(f) \u2014 real component of transformed mapping matrix\")\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "plt.imshow(\n", - " transformed_mapping_matrix.imag,\n", - " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.title(\"Im(f) \u2014 imaginary component of transformed mapping matrix\")\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) (hereafter WD03) introduce the linear inversion\n", - "formalism used to compute the intensity values of the linear light profiles. WD03 indexes the transformed\n", - "mapping matrix as $f_{ij}$ where $i$ maps over all $I$ linear basis components and $j$ maps over all $J$\n", - "visibilities. For our MGE, $I = $ `total_gaussians`.\n", - "\n", - "The indexing of the `transformed_mapping_matrix` array is reversed compared to the WD03 convention (rows\n", - "are visibilities, columns are basis components).\n", - "\n", - "__Data Vector (D)__\n", - "\n", - "To solve for the per-Gaussian intensities we pose the problem as a linear inversion.\n", - "\n", - "The `data_vector`, $D$, has dimensions `(total_gaussians,)`. In WD03 the data vector is given by:\n", - "\n", - " $\\\\vec{D}_{i} = \\\\sum_{\\\\rm j=1}^{J} f_{ij}\\\\, d_{j} / \\\\sigma_{j}^2 \\\\, ,$\n", - "\n", - "where $d_j$ are the observed visibility values, $\\\\sigma_j^2$ are the visibility variances, and the sum\n", - "runs over real and imaginary components. The interferometer helper handles the real/imaginary split\n", - "internally." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_vector = (\n", - " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " visibilities=dataset.data,\n", - " noise_map=dataset.noise_map,\n", - " )\n", - ")\n", - "\n", - "print(\"Data Vector D shape:\")\n", - "print(data_vector.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Matrix (F)__\n", - "\n", - "The `curvature_matrix` $F$ has dimensions `(total_gaussians, total_gaussians)`.\n", - "\n", - "In WD03 the curvature matrix is given by:\n", - "\n", - " ${F}_{ik} = \\\\sum_{\\\\rm j=1}^{J} f_{ij}\\\\, f_{kj} / \\\\sigma_{j}^2 \\\\, .$\n", - "\n", - "Because visibilities (and therefore $f$) are complex-valued, the curvature is computed separately for the\n", - "real and imaginary parts and summed.\n", - "\n", - "For an N-Gaussian MGE $F$ is an NxN matrix; the off-diagonal entries $F_{ik}$ quantify how much Gaussians\n", - "$i$ and $k$ overlap in the visibility plane. Adjacent Gaussians (similar sigma) have large overlaps; very\n", - "different sigmas overlap less." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.real,\n", - " noise_map=dataset.noise_map.real,\n", - ")\n", - "\n", - "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.imag,\n", - " noise_map=dataset.noise_map.imag,\n", - ")\n", - "\n", - "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", - "\n", - "print(\"Curvature Matrix F shape:\")\n", - "print(curvature_matrix.shape)\n", - "\n", - "plt.imshow(curvature_matrix)\n", - "plt.colorbar()\n", - "plt.title(\"Curvature Matrix F\")\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction (Positive-Negative)__\n", - "\n", - "The chi-squared minimised by the inversion is:\n", - "\n", - "$\\\\chi^2 = \\\\sum_{\\\\rm j=1}^{J} \\\\bigg[ \\\\frac{(\\\\sum_{\\\\rm i=1}^{I} s_{i} f_{ij}) - d_{j}}{\\\\sigma_{j}} \\\\bigg]^2$\n", - "\n", - "Where $s_i$ is the solved-for `intensity` of the $i$-th Gaussian. The solution is given by (equation 5\n", - "WD03):\n", - "\n", - " $s = F^{-1} D$\n", - "\n", - "For an MGE this often returns negative `intensity` values for some Gaussians \u2014 the \"ringing\" pattern\n", - "where alternating Gaussians take large positive/negative values. This is unphysical." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction_positive_negative = np.linalg.solve(curvature_matrix, data_vector)\n", - "\n", - "print(\"Reconstruction s (positive-negative solver, first 5):\")\n", - "print(reconstruction_positive_negative[:5])\n", - "print(\n", - " f\" number of negative entries: {(reconstruction_positive_negative < 0).sum()} / {total_gaussians}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction (Positive Only)__\n", - "\n", - "The linear algebra can be solved with the constraint that all `intensity` values are positive. The naive\n", - "approach is `scipy.optimize.nnls`, which is iterative and works directly on the transformed mapping matrix\n", - "\u2014 slow for many Gaussians.\n", - "\n", - "The source code therefore uses a \"fast nnls\" algorithm \u2014 an adaptation of:\n", - " https://github.com/jvendrow/fnnls\n", - "\n", - "`fnnls` uses `data_vector` $D$ and `curvature_matrix` $F$ (not the full mapping matrix), which makes it\n", - "much faster than scipy's nnls. This is essential for MGE because the basis size can be large." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", - " data_vector=data_vector,\n", - " curvature_reg_matrix=curvature_matrix, # ignore the _reg_ tag in this guide\n", - ")\n", - "\n", - "print(\"Reconstruction s (positive-only solver, first 5):\")\n", - "print(reconstruction[:5])\n", - "print(f\" number of zero entries: {(reconstruction == 0).sum()} / {total_gaussians}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visibilities Reconstruction__\n", - "\n", - "Using the reconstructed `intensity` values we can map the reconstruction back to the visibility plane via\n", - "the `transformed_mapping_matrix`, producing the model visibilities.\n", - "\n", - "This sums the per-Gaussian visibility-plane contributions weighted by their solved-for intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_visibilities = (\n", - " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " reconstruction=reconstruction,\n", - " )\n", - ")\n", - "\n", - "mapped_reconstructed_visibilities = al.Visibilities(\n", - " visibilities=mapped_reconstructed_visibilities\n", - ")\n", - "\n", - "aplt.plot_grid(\n", - " grid=mapped_reconstructed_visibilities.in_grid,\n", - " title=\"Model Visibilities (MGE source)\",\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our MGE source reconstruction.\n", - "\n", - "The likelihood function for parametric galaxy modeling, even with linear light profile bases, consists of\n", - "two terms in the visibility plane:\n", - "\n", - " $-2 \\\\mathrm{ln} \\\\, \\\\epsilon = \\\\chi^2 + \\\\sum_{\\\\rm j=1}^{J} { \\\\mathrm{ln}} \\\\left [2 \\\\pi (\\\\sigma_j)^2 \\\\right] \\\\, .$\n", - "\n", - "(Note: for a *pixelization* there are additional regularisation terms. The MGE inversion does not use\n", - "regularisation \u2014 its smoothness comes from the basis-function form, not from a regularisation matrix.)\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\\\chi^2$ statistic computed in the visibility plane:\n", - "\n", - " - `model_data` = `mapped_reconstructed_visibilities`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = `normalized_residual_map` ** 2.0\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "Visibilities are complex-valued, so we split into real and imaginary components, compute $\\\\chi^2$ for\n", - "each, and sum." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_visibilities = mapped_reconstructed_visibilities\n", - "\n", - "residual_map = dataset.data - model_visibilities\n", - "\n", - "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", - "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", - "\n", - "chi_squared_real = np.sum(chi_squared_map_real)\n", - "chi_squared_imag = np.sum(chi_squared_map_imag)\n", - "chi_squared = chi_squared_real + chi_squared_imag\n", - "\n", - "print(f\"chi_squared (real + imag) = {chi_squared}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "The likelihood function assumes the visibility data consists of independent Gaussian noise on every\n", - "visibility (real and imaginary parts treated independently)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization_real = float(\n", - " np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", - ")\n", - "noise_normalization_imag = float(\n", - " np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", - ")\n", - "noise_normalization = noise_normalization_real + noise_normalization_imag" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "Combine the two terms to compute the `log_likelihood`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(f\"log_likelihood (figure of merit) = {figure_of_merit}\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "The exact same likelihood evaluation is performed inside the `FitInterferometer` object. We construct one,\n", - "print its `figure_of_merit`, and confirm it matches the value we computed by hand above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", - "print(f\"FitInterferometer.figure_of_merit = {fit.figure_of_merit}\")\n", - "\n", - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit contains an `Inversion` object, which handles all the linear algebra we have covered in this\n", - "script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.inversion)\n", - "print(f\"data_vector shape: {fit.inversion.data_vector.shape}\")\n", - "print(f\"curvature_matrix shape: {fit.inversion.curvature_matrix.shape}\")\n", - "print(f\"reconstruction shape: {fit.inversion.reconstruction.shape}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `nautilus`\n", - "(https://github.com/johannesulf/nautilus); multiple MCMC and optimization algorithms are also supported.\n", - "\n", - "For an MGE source, the reduced number of free parameters (intensities are solved analytically, sigmas are\n", - "fixed, centres and ell_comps are shared) means that the sampler converges in fewer iterations than a\n", - "free-intensity Sersic source \u2014 even though each likelihood evaluation is slower per-iteration because of\n", - "the larger basis.\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the multi-Gaussian expansion interferometer likelihood\n", - "function. The pipeline:\n", - "\n", - " ray-trace \u2192 source-plane Gaussian images (intensity=1 each) \u2192 `mapping_matrix` (N columns)\n", - " \u2192 NUFFT \u2192 `transformed_mapping_matrix` \u2192 $D$ (length N) and $F$ (NxN) \u2192 solve $s = F^{-1} D$\n", - " \u2192 `mapped_reconstructed_visibilities` \u2192 visibility-plane $\\\\chi^2$ \u2192 log likelihood\n", - "\n", - "is the same as for the single-component linear light profile case\n", - "(`features/linear_light_profiles/likelihood_function.py`), just generalised to N basis components. The\n", - "NUFFT is the operation that nufftax made fast on the GPU, which is why this entire workflow is now\n", - "practical even for large MGE bases on visibility data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Multi Gaussian Expansion (Interferometer)__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit\n", + "`Interferometer` data with a multi-Gaussian expansion (MGE) source: a `Basis` of many linear `Gaussian`\n", + "profiles whose `intensity` values are solved for analytically via a linear inversion.\n", + "\n", + "The visibility-plane linear inversion is **mathematically identical** to the single-component linear light\n", + "profile case (see `interferometer/features/linear_light_profiles/likelihood_function.py`). The only\n", + "difference is the number of columns in the mapping matrix \u2014 one column per Gaussian in the basis instead\n", + "of one column total. The data vector `D` and curvature matrix `F` therefore have dimensions equal to the\n", + "basis size, and the solved `s = F^{-1} D` gives one `intensity` per Gaussian.\n", + "\n", + "**Role swap vs the imaging MGE example:** the imaging script's basis is on the lens galaxy. Here it's on\n", + "the source galaxy, because interferometer data does not contain lens light. The source bulge image is\n", + "evaluated on the *ray-traced* source-plane grid; the rest of the inversion machinery is unchanged.\n", + "\n", + "For interferometer data this is now practical thanks to the JAX-native NUFFT `nufftax`\n", + "(https://github.com/GragasLab/nufftax) \u2014 every Gaussian's NUFFT happens inside the same jit/vmap pipeline\n", + "as the rest of the model, so the per-iteration NUFFT cost scales linearly with the basis size and is\n", + "amortised on the GPU.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The MGE likelihood function builds on:\n", + "\n", + " - `interferometer/log_likelihood_function.ipynb` \u2014 the standard interferometer parametric likelihood\n", + " function (NUFFT of a real-space image, visibility-plane $\\\\chi^2$).\n", + " - `interferometer/features/linear_light_profiles/likelihood_function.ipynb` \u2014 the single-component\n", + " visibility-plane linear inversion (data vector, curvature matrix, positive-only solver). The MGE is\n", + " a direct generalisation to N-component bases.\n", + "\n", + "This script repeats just enough setup that you can follow it without rereading those two \u2014 but if anything\n", + "is unclear, those are the places to look first.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Reading order before this script.\n", + "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", + "- **Dataset:** Load and plot the strong lens `Interferometer` dataset using `TransformerNUFFT` (nufftax).\n", + "- **Lens Galaxy:** A mass-only lens galaxy (no light \u2014 interferometer convention).\n", + "- **Source Galaxy MGE Basis:** Build the linear `Gaussian` basis used as the source bulge.\n", + "- **Ray Tracing:** Image-plane to source-plane grid.\n", + "- **Mapping Matrix:** Real-space mapping matrix \u2014 one column per Gaussian in the basis.\n", + "- **Transformed Mapping Matrix ($f$):** NUFFT each column to give the visibility-space mapping matrix.\n", + "- **Data Vector (D):** Compute $D$ from the transformed mapping matrix, visibilities, and noise map.\n", + "- **Curvature Matrix (F):** Compute $F$ separately for real and imaginary components, then sum.\n", + "- **Reconstruction (Positive-Negative):** Solve $s = F^{-1} D$ via NumPy.\n", + "- **Reconstruction (Positive Only):** Solve with the fast non-negative least squares (`fnnls`) algorithm.\n", + "- **Visibilities Reconstruction:** Map $s$ back to visibility space.\n", + "- **Likelihood Function:** Visibility-plane $\\\\chi^2$ and noise normalization.\n", + "- **Chi Squared:** Sum chi-squared contributions over real and imaginary components.\n", + "- **Noise Normalization Term:** The fixed noise normalization term.\n", + "- **Calculate The Log Likelihood:** Combine into the final log likelihood.\n", + "- **Fit:** Cross-check via `FitInterferometer`.\n", + "- **Lens Modeling:** How this likelihood is sampled in the full Nautilus fit.\n", + "- **Wrap Up:** Summary and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files using `TransformerNUFFT`\n", + "backed by `nufftax`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "For interferometer data the lens galaxy's optical/IR emission is typically below detection, so we model\n", + "only its mass:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy MGE Basis__\n", + "\n", + "We build a `Basis` of 5 linear `Gaussian` profiles for the source bulge, all sharing the same centre and\n", + "`ell_comps`, with `sigma` values spanning 0.01\" to the mask radius in log-spaced increments.\n", + "\n", + "Each Gaussian is a linear light profile \u2014 its `intensity` is solved for analytically via the inversion\n", + "below. Internally each linear profile carries `intensity=1.0`, which the inversion later rescales." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 5\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "bulge_gaussian_list = []\n", + "for i in range(total_gaussians):\n", + " gaussian = al.lp_linear.Gaussian(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " sigma=10 ** log10_sigma_list[i],\n", + " )\n", + " bulge_gaussian_list.append(gaussian)\n", + "\n", + "source_bulge = al.lp_basis.Basis(profile_list=bulge_gaussian_list)\n", + "source_galaxy = al.Galaxy(redshift=1.0, bulge=source_bulge)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\\\theta$ from the image-plane to\n", + "its (y,x) source-plane coordinate $\\\\beta$ using the summed deflection angles $\\\\alpha$ of the mass\n", + "profiles:\n", + "\n", + " $\\\\beta = \\\\theta - \\\\alpha(\\\\theta)$\n", + "\n", + "For interferometer data the only grid we need to ray-trace is the real-space grid associated with the\n", + "`real_space_mask`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane); we only need the source-plane grid (index -1).\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=dataset.grids.lp)[-1]\n", + "\n", + "aplt.plot_grid(grid=traced_grid, title=\"Traced Source-Plane Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapping Matrix__\n", + "\n", + "For interferometer MGE modeling, the mapping matrix has one column per Gaussian in the basis. Each column\n", + "holds the real-space image of one Gaussian, evaluated on the ray-traced source-plane grid with\n", + "`intensity=1.0` internally. The lensing is already baked in because we evaluated on the traced grid.\n", + "\n", + "The dimensions are `(total_real_space_pixels, total_gaussians)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lp_linear_func_source = al.LightProfileLinearObjFuncList(\n", + " grid=traced_grid,\n", + " blurring_grid=None,\n", + " psf=None,\n", + " light_profile_list=bulge_gaussian_list,\n", + " regularization=None,\n", + ")\n", + "\n", + "mapping_matrix = lp_linear_func_source.mapping_matrix\n", + "\n", + "print(\"Mapping matrix shape (real-space pixels, total Gaussians):\")\n", + "print(mapping_matrix.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `mapping_matrix` shows each Gaussian's lensed source-plane image as one column." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.colorbar()\n", + "plt.title(\"Real-space mapping matrix \u2014 one column per Gaussian\")\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Transformed Mapping Matrix ($f$)__\n", + "\n", + "Every column of the real-space `mapping_matrix` must be NUFFT'd to the uv-plane. The result is the\n", + "`transformed_mapping_matrix` \u2014 a *complex-valued* matrix with dimensions\n", + "`(total_visibilities, total_gaussians)`.\n", + "\n", + "For an N-Gaussian MGE this matrix has N columns. The inversion's job is to find the N scalars that, when\n", + "multiplied by their respective columns and summed, best fit the observed visibilities.\n", + "\n", + "This NUFFT-each-column operation is exactly what nufftax made fast. The whole pipeline (ray-trace \u2192\n", + "source-plane Gaussian images \u2192 `transform_mapping_matrix` \u2192 linear inversion) is JIT-compilable under JAX." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", + " mapping_matrix=mapping_matrix\n", + ")\n", + "\n", + "print(\"Transformed mapping matrix shape (visibilities, total Gaussians):\")\n", + "print(transformed_mapping_matrix.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the real and imaginary components of the transformed mapping matrix. Each row is one observed\n", + "visibility; the values across that row are the per-unit-intensity contributions of each Gaussian at that\n", + "uv point." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " transformed_mapping_matrix.real,\n", + " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.title(\"Re(f) \u2014 real component of transformed mapping matrix\")\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "plt.imshow(\n", + " transformed_mapping_matrix.imag,\n", + " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.title(\"Im(f) \u2014 imaginary component of transformed mapping matrix\")\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) (hereafter WD03) introduce the linear inversion\n", + "formalism used to compute the intensity values of the linear light profiles. WD03 indexes the transformed\n", + "mapping matrix as $f_{ij}$ where $i$ maps over all $I$ linear basis components and $j$ maps over all $J$\n", + "visibilities. For our MGE, $I = $ `total_gaussians`.\n", + "\n", + "The indexing of the `transformed_mapping_matrix` array is reversed compared to the WD03 convention (rows\n", + "are visibilities, columns are basis components).\n", + "\n", + "__Data Vector (D)__\n", + "\n", + "To solve for the per-Gaussian intensities we pose the problem as a linear inversion.\n", + "\n", + "The `data_vector`, $D$, has dimensions `(total_gaussians,)`. In WD03 the data vector is given by:\n", + "\n", + " $\\\\vec{D}_{i} = \\\\sum_{\\\\rm j=1}^{J} f_{ij}\\\\, d_{j} / \\\\sigma_{j}^2 \\\\, ,$\n", + "\n", + "where $d_j$ are the observed visibility values, $\\\\sigma_j^2$ are the visibility variances, and the sum\n", + "runs over real and imaginary components. The interferometer helper handles the real/imaginary split\n", + "internally." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_vector = (\n", + " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " visibilities=dataset.data,\n", + " noise_map=dataset.noise_map,\n", + " )\n", + ")\n", + "\n", + "print(\"Data Vector D shape:\")\n", + "print(data_vector.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Matrix (F)__\n", + "\n", + "The `curvature_matrix` $F$ has dimensions `(total_gaussians, total_gaussians)`.\n", + "\n", + "In WD03 the curvature matrix is given by:\n", + "\n", + " ${F}_{ik} = \\\\sum_{\\\\rm j=1}^{J} f_{ij}\\\\, f_{kj} / \\\\sigma_{j}^2 \\\\, .$\n", + "\n", + "Because visibilities (and therefore $f$) are complex-valued, the curvature is computed separately for the\n", + "real and imaginary parts and summed.\n", + "\n", + "For an N-Gaussian MGE $F$ is an NxN matrix; the off-diagonal entries $F_{ik}$ quantify how much Gaussians\n", + "$i$ and $k$ overlap in the visibility plane. Adjacent Gaussians (similar sigma) have large overlaps; very\n", + "different sigmas overlap less." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.real,\n", + " noise_map=dataset.noise_map.real,\n", + ")\n", + "\n", + "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.imag,\n", + " noise_map=dataset.noise_map.imag,\n", + ")\n", + "\n", + "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", + "\n", + "print(\"Curvature Matrix F shape:\")\n", + "print(curvature_matrix.shape)\n", + "\n", + "plt.imshow(curvature_matrix)\n", + "plt.colorbar()\n", + "plt.title(\"Curvature Matrix F\")\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction (Positive-Negative)__\n", + "\n", + "The chi-squared minimised by the inversion is:\n", + "\n", + "$\\\\chi^2 = \\\\sum_{\\\\rm j=1}^{J} \\\\bigg[ \\\\frac{(\\\\sum_{\\\\rm i=1}^{I} s_{i} f_{ij}) - d_{j}}{\\\\sigma_{j}} \\\\bigg]^2$\n", + "\n", + "Where $s_i$ is the solved-for `intensity` of the $i$-th Gaussian. The solution is given by (equation 5\n", + "WD03):\n", + "\n", + " $s = F^{-1} D$\n", + "\n", + "For an MGE this often returns negative `intensity` values for some Gaussians \u2014 the \"ringing\" pattern\n", + "where alternating Gaussians take large positive/negative values. This is unphysical." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction_positive_negative = np.linalg.solve(curvature_matrix, data_vector)\n", + "\n", + "print(\"Reconstruction s (positive-negative solver, first 5):\")\n", + "print(reconstruction_positive_negative[:5])\n", + "print(\n", + " f\" number of negative entries: {(reconstruction_positive_negative < 0).sum()} / {total_gaussians}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction (Positive Only)__\n", + "\n", + "The linear algebra can be solved with the constraint that all `intensity` values are positive. The naive\n", + "approach is `scipy.optimize.nnls`, which is iterative and works directly on the transformed mapping matrix\n", + "\u2014 slow for many Gaussians.\n", + "\n", + "The source code therefore uses a \"fast nnls\" algorithm \u2014 an adaptation of:\n", + " https://github.com/jvendrow/fnnls\n", + "\n", + "`fnnls` uses `data_vector` $D$ and `curvature_matrix` $F$ (not the full mapping matrix), which makes it\n", + "much faster than scipy's nnls. This is essential for MGE because the basis size can be large." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = al.util.inversion.reconstruction_positive_only_from(\n", + " data_vector=data_vector,\n", + " curvature_reg_matrix=curvature_matrix, # ignore the _reg_ tag in this guide\n", + ")\n", + "\n", + "print(\"Reconstruction s (positive-only solver, first 5):\")\n", + "print(reconstruction[:5])\n", + "print(f\" number of zero entries: {(reconstruction == 0).sum()} / {total_gaussians}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visibilities Reconstruction__\n", + "\n", + "Using the reconstructed `intensity` values we can map the reconstruction back to the visibility plane via\n", + "the `transformed_mapping_matrix`, producing the model visibilities.\n", + "\n", + "This sums the per-Gaussian visibility-plane contributions weighted by their solved-for intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_visibilities = (\n", + " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " reconstruction=reconstruction,\n", + " )\n", + ")\n", + "\n", + "mapped_reconstructed_visibilities = al.Visibilities(\n", + " visibilities=mapped_reconstructed_visibilities\n", + ")\n", + "\n", + "aplt.plot_grid(\n", + " grid=mapped_reconstructed_visibilities.in_grid,\n", + " title=\"Model Visibilities (MGE source)\",\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our MGE source reconstruction.\n", + "\n", + "The likelihood function for parametric galaxy modeling, even with linear light profile bases, consists of\n", + "two terms in the visibility plane:\n", + "\n", + " $-2 \\\\mathrm{ln} \\\\, \\\\epsilon = \\\\chi^2 + \\\\sum_{\\\\rm j=1}^{J} { \\\\mathrm{ln}} \\\\left [2 \\\\pi (\\\\sigma_j)^2 \\\\right] \\\\, .$\n", + "\n", + "(Note: for a *pixelization* there are additional regularisation terms. The MGE inversion does not use\n", + "regularisation \u2014 its smoothness comes from the basis-function form, not from a regularisation matrix.)\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\\\chi^2$ statistic computed in the visibility plane:\n", + "\n", + " - `model_data` = `mapped_reconstructed_visibilities`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = `normalized_residual_map` ** 2.0\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "Visibilities are complex-valued, so we split into real and imaginary components, compute $\\\\chi^2$ for\n", + "each, and sum." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_visibilities = mapped_reconstructed_visibilities\n", + "\n", + "residual_map = dataset.data - model_visibilities\n", + "\n", + "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", + "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", + "\n", + "chi_squared_real = np.sum(chi_squared_map_real)\n", + "chi_squared_imag = np.sum(chi_squared_map_imag)\n", + "chi_squared = chi_squared_real + chi_squared_imag\n", + "\n", + "print(f\"chi_squared (real + imag) = {chi_squared}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "The likelihood function assumes the visibility data consists of independent Gaussian noise on every\n", + "visibility (real and imaginary parts treated independently)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization_real = float(\n", + " np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", + ")\n", + "noise_normalization_imag = float(\n", + " np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", + ")\n", + "noise_normalization = noise_normalization_real + noise_normalization_imag" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "Combine the two terms to compute the `log_likelihood`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(f\"log_likelihood (figure of merit) = {figure_of_merit}\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "The exact same likelihood evaluation is performed inside the `FitInterferometer` object. We construct one,\n", + "print its `figure_of_merit`, and confirm it matches the value we computed by hand above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", + "print(f\"FitInterferometer.figure_of_merit = {fit.figure_of_merit}\")\n", + "\n", + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit contains an `Inversion` object, which handles all the linear algebra we have covered in this\n", + "script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.inversion)\n", + "print(f\"data_vector shape: {fit.inversion.data_vector.shape}\")\n", + "print(f\"curvature_matrix shape: {fit.inversion.curvature_matrix.shape}\")\n", + "print(f\"reconstruction shape: {fit.inversion.reconstruction.shape}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `nautilus`\n", + "(https://github.com/johannesulf/nautilus); multiple MCMC and optimization algorithms are also supported.\n", + "\n", + "For an MGE source, the reduced number of free parameters (intensities are solved analytically, sigmas are\n", + "fixed, centres and ell_comps are shared) means that the sampler converges in fewer iterations than a\n", + "free-intensity Sersic source \u2014 even though each likelihood evaluation is slower per-iteration because of\n", + "the larger basis.\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the multi-Gaussian expansion interferometer likelihood\n", + "function. The pipeline:\n", + "\n", + " ray-trace \u2192 source-plane Gaussian images (intensity=1 each) \u2192 `mapping_matrix` (N columns)\n", + " \u2192 NUFFT \u2192 `transformed_mapping_matrix` \u2192 $D$ (length N) and $F$ (NxN) \u2192 solve $s = F^{-1} D$\n", + " \u2192 `mapped_reconstructed_visibilities` \u2192 visibility-plane $\\\\chi^2$ \u2192 log likelihood\n", + "\n", + "is the same as for the single-component linear light profile case\n", + "(`features/linear_light_profiles/likelihood_function.py`), just generalised to N basis components. The\n", + "NUFFT is the operation that nufftax made fast on the GPU, which is why this entire workflow is now\n", + "practical even for large MGE bases on visibility data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/multi_gaussian_expansion/modeling.ipynb b/notebooks/interferometer/features/multi_gaussian_expansion/modeling.ipynb index 65d8cf58c..296954b92 100644 --- a/notebooks/interferometer/features/multi_gaussian_expansion/modeling.ipynb +++ b/notebooks/interferometer/features/multi_gaussian_expansion/modeling.ipynb @@ -1,521 +1,558 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Multi Gaussian Expansion (Interferometer)\n", - "=============================================================\n", - "\n", - "A multi-Gaussian expansion (MGE) decomposes a galaxy's light into ~15-100 Gaussians, where the `intensity` of\n", - "every Gaussian is solved for via linear algebra using a process called an \"inversion\" (see the\n", - "`linear_light_profiles` feature for a full description of this).\n", - "\n", - "This script performs lens modeling of an `Interferometer` dataset using an MGE source bulge consisting of\n", - "30 Gaussians arranged in two groups of 15. Each group has its own elliptical components, so the source's\n", - "light is decomposed into two distinct elliptical components \u2014 which could be viewed as a bulge and a disk\n", - "of the source galaxy.\n", - "\n", - "**Role swap vs the imaging MGE example:** the imaging script fits the *lens* galaxy's complex morphology\n", - "with an MGE. For interferometer data the lens light is omitted (no detection in mm/sub-mm), so the MGE is\n", - "instead applied to the *source* galaxy. This is the standard convention for interferometer modeling.\n", - "\n", - "MGE fits to interferometer data were previously impractical because every likelihood evaluation has to\n", - "Fourier-transform each Gaussian basis component into the uv-plane. With `nufftax`\n", - "(https://github.com/GragasLab/nufftax) \u2014 a JAX-native NUFFT \u2014 the full basis is transformed inside the same\n", - "jit/vmap pipeline as the rest of the model, amortising the per-iteration NUFFT cost on the GPU. MGE source\n", - "fits are now routine even at ALMA-class visibility counts.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** Benefits and drawbacks of an MGE source for interferometer data.\n", - "- **NUFFT (nufftax):** Why MGE-on-visibilities is now practical thanks to nufftax.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Model:** Compose the lens model \u2014 `Isothermal` + `ExternalShear` lens mass and a 30-Gaussian MGE\n", - " source bulge. Lens light omitted (interferometer convention).\n", - "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", - "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", - "- **Over Sampling:** Interferometer modeling does not use over-sampling.\n", - "- **Search:** Configure the non-linear search (Nautilus).\n", - "- **Analysis:** Create the `AnalysisInterferometer` object.\n", - "- **VRAM:** Memory budget for a multi-component linear basis on GPU.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Advantages__\n", - "\n", - "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals when fitting the source\n", - "galaxy of a strong lens \u2014 they fail to capture irregular and asymmetric morphology (e.g. isophotal twists,\n", - "multi-component source emission). An MGE fully captures these features and can therefore much better\n", - "represent the emission of complex source galaxies.\n", - "\n", - "The MGE model is composed such that the `intensity` parameters and the `sigma` parameters controlling the\n", - "Gaussian sizes are *not* sampled by Nautilus. Centres and elliptical components are shared across each\n", - "group of Gaussians. This removes the most significant degeneracies in parameter space, making the model\n", - "much more reliable and efficient to fit than a free-intensity Sersic.\n", - "\n", - "Therefore, not only does an MGE source fit more complex source morphologies, it does so using fewer\n", - "non-linear parameters in a much simpler non-linear parameter space which has far less significant parameter\n", - "degeneracies.\n", - "\n", - "__Disadvantages__\n", - "\n", - "To fit an MGE model, the light of the ~15-100 Gaussians must be evaluated and NUFFT'd to the uv-plane per\n", - "likelihood evaluation. This is slower than evaluating a single `SersicCore`. With `nufftax` the per-NUFFT\n", - "cost is small enough that the total slow-down per likelihood evaluation is typically 2-5x for a 30-Gaussian\n", - "MGE \u2014 paid back in fewer iterations because the parameter space is simpler.\n", - "\n", - "The MGE source can also be less intuitive to interpret physically than a Sersic. The Sersic `effective_radius`\n", - "and `sersic_index` map directly to galaxy size and concentration; an MGE's solved-for Gaussian intensities\n", - "require an extra processing step to compute equivalent physical quantities.\n", - "\n", - "__NUFFT (nufftax)__\n", - "\n", - "The image-to-visibilities Fourier transform is performed by a Non-Uniform Fast Fourier Transform (NUFFT),\n", - "exposed in **PyAutoLens** as `TransformerNUFFT`. The default backend is `nufftax`, a pure-JAX NUFFT that\n", - "jit-compiles and vmap-batches like the rest of the library:\n", - "\n", - " https://github.com/GragasLab/nufftax\n", - "\n", - "Because `nufftax` is JAX-native, NUFFT-ing every Gaussian in the MGE basis happens inside the same compiled\n", - "likelihood that does the inversion, mass model ray-tracing, and chi-squared sum. There is no host round-trip\n", - "between NUFFT calls, so a model with N Gaussians costs only N forward-NUFFTs per iteration on the GPU \u2014\n", - "fast enough that MGE-on-visibilities is now routinely practical.\n", - "\n", - "If `nufftax` is not installed, install it via `pip install nufftax`.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and\n", - "negative values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This is problematic, as it means that negative surface brightnesses values can be computed to represent a\n", - "galaxy's light, which is clearly unphysical. For an MGE, this produces a positive-negative \"ringing\", where\n", - "the Gaussians alternate between large positive and negative values. This is clearly undesirable and\n", - "unphysical.\n", - "\n", - "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it\n", - "is as fast as positive-negative solvers. This ensures that all Gaussian intensities are positive and\n", - "therefore physical.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data). This is the standard\n", - " convention for interferometer modeling, as the lens galaxy's optical/IR emission is typically below the\n", - " detection threshold of mm/sub-mm interferometers.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's bulge is a multi-Gaussian expansion of 30 linear `Gaussian` profiles, arranged in\n", - " two groups of 15 (each group shares a centre and ell_comps).\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook.\n", - "\n", - "__Imaging Equivalent__\n", - "\n", - "For the CCD-imaging version of this script (MGE on the lens galaxy, not the source), see\n", - "`autolens_workspace/*/imaging/features/multi_gaussian_expansion/modeling.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit with\n", - "the lens model.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane\n", - "and compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native Non-Uniform Fast Fourier\n", - "Transform backed by `nufftax`, which is required for fast MGE modeling and scales efficiently from a few\n", - "hundred visibilities to tens of millions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "If you are familiar with using imaging data, you may have seen that a numerical technique called over\n", - "sampling is used, which evaluates light profiles on a higher resolution grid than the image data to ensure\n", - "the calculation is accurate.\n", - "\n", - "Interferometer data does not observe galaxies in a way where over sampling is necessary, therefore all\n", - "interferometer calculations are performed without over sampling.\n", - "\n", - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's bulge is 10 linear `Gaussian` profiles [6 parameters total].\n", - " - The centres and elliptical components of the Gaussians are linked together in two groups of 5.\n", - " - The `sigma` size of the Gaussians increases in log10 increments from 0.01 to the mask radius.\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=13.\n", - "\n", - "The MGE comprises 2 groups of 5 Gaussians. Each group has its own elliptical components, so the source's\n", - "light is decomposed into two distinct elliptical components which could be viewed as a bulge and a disk of\n", - "the source.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook:\n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 5\n", - "gaussian_per_basis = 2\n", - "\n", - "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.5\".\n", - "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", - "\n", - "# By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.\n", - "\n", - "centre_0 = af.UniformPrior(lower_limit=-0.3, upper_limit=0.3)\n", - "centre_1 = af.UniformPrior(lower_limit=-0.3, upper_limit=0.3)\n", - "\n", - "bulge_gaussian_list = []\n", - "\n", - "for j in range(gaussian_per_basis):\n", - " # A list of Gaussian model components whose parameters are customized below.\n", - "\n", - " gaussian_list = af.Collection(\n", - " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", - " )\n", - "\n", - " # Iterate over every Gaussian and customize its parameters.\n", - "\n", - " for i, gaussian in enumerate(gaussian_list):\n", - " gaussian.centre.centre_0 = centre_0 # All Gaussians have same y centre.\n", - " gaussian.centre.centre_1 = centre_1 # All Gaussians have same x centre.\n", - " gaussian.ell_comps = gaussian_list[\n", - " 0\n", - " ].ell_comps # All Gaussians in this group share ell_comps.\n", - " gaussian.sigma = (\n", - " 10 ** log10_sigma_list[i]\n", - " ) # All Gaussian sigmas are fixed to values above.\n", - "\n", - " bulge_gaussian_list += gaussian_list\n", - "\n", - "# The Basis object groups many light profiles together into a single model component.\n", - "source_bulge = af.Model(al.lp_basis.Basis, profile_list=bulge_gaussian_list)\n", - "\n", - "# Lens:\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", - "\n", - "# Overall Lens Model:\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", - "refer to `start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This shows every single Gaussian light profile in the model, which is a lot of parameters! However, the\n", - "vast majority of these parameters are fixed to the values we set above, so the model actually has far fewer\n", - "free parameters than it looks!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", - "full description).\n", - "\n", - "Owing to the simplicity of fitting an MGE (no intensity or sigma free parameters), we use fewer live points\n", - "than the `interferometer/modeling.py` example: 75 live points speeds up convergence." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\") / \"features\",\n", - " name=\"mge\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisInterferometer` object defining how Nautilus fits the model to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `interferometer/modeling.py` example explains how VRAM is used during GPU-based fitting and how to\n", - "print the estimated VRAM required by a model.\n", - "\n", - "For each linear `Gaussian` profile, extra VRAM is used to store its NUFFT'd mapping matrix column. For\n", - "around 30 linear Gaussians this typically requires a modest amount of VRAM (e.g. 10-50 MB per batched\n", - "likelihood). Models that use hundreds of Gaussians, especially in combination with a large batch size, may\n", - "therefore exceed GBs of VRAM and require you to adjust the batch size to fit within your GPU's VRAM.\n", - "\n", - "VRAM on interferometer datasets is driven primarily by the visibility count and the real-space mask size,\n", - "not the number of Gaussians in the MGE basis.\n", - "\n", - "__Run Time__\n", - "\n", - "The likelihood evaluation time for an MGE is slower than a single linear `SersicCore` source, because the\n", - "image of every Gaussian must be evaluated and NUFFT'd to the uv-plane. With `nufftax`, the per-NUFFT cost\n", - "is small enough that the total slow-down per likelihood is typically 2-5x for a 30-Gaussian MGE compared\n", - "to a one-component source \u2014 paid back in fewer iterations because the parameter space is simpler.\n", - "\n", - "Because the MGE has no free `intensity` or `sigma` parameters and shares centres/ell_comps across groups,\n", - "Nautilus converges significantly faster than for a free-intensity Sersic source. We also use fewer live\n", - "points (75 vs the 100 used in `interferometer/modeling.py`), further speeding up the model-fit.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the\n", - "output folder for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format (if\n", - "this does not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix\n", - "this):\n", - "\n", - "This confirms there are many `Gaussian`s in the source light model and it lists their inferred parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "In particular, checkout the results example `linear.py` which details how to extract all information about\n", - "linear light profiles from a fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "To further showcase MGE source modeling, see the SLaM pipeline at\n", - "`autolens_workspace/scripts/interferometer/features/multi_gaussian_expansion/slam.py`, which uses the\n", - "`al.model_util.mge_model_from` helper to compose an MGE source in the SOURCE LP stage before chaining\n", - "to a pixelized source reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Multi Gaussian Expansion (Interferometer)\n", + "=============================================================\n", + "\n", + "A multi-Gaussian expansion (MGE) decomposes a galaxy's light into ~15-100 Gaussians, where the `intensity` of\n", + "every Gaussian is solved for via linear algebra using a process called an \"inversion\" (see the\n", + "`linear_light_profiles` feature for a full description of this).\n", + "\n", + "This script performs lens modeling of an `Interferometer` dataset using an MGE source bulge consisting of\n", + "30 Gaussians arranged in two groups of 15. Each group has its own elliptical components, so the source's\n", + "light is decomposed into two distinct elliptical components \u2014 which could be viewed as a bulge and a disk\n", + "of the source galaxy.\n", + "\n", + "**Role swap vs the imaging MGE example:** the imaging script fits the *lens* galaxy's complex morphology\n", + "with an MGE. For interferometer data the lens light is omitted (no detection in mm/sub-mm), so the MGE is\n", + "instead applied to the *source* galaxy. This is the standard convention for interferometer modeling.\n", + "\n", + "MGE fits to interferometer data were previously impractical because every likelihood evaluation has to\n", + "Fourier-transform each Gaussian basis component into the uv-plane. With `nufftax`\n", + "(https://github.com/GragasLab/nufftax) \u2014 a JAX-native NUFFT \u2014 the full basis is transformed inside the same\n", + "jit/vmap pipeline as the rest of the model, amortising the per-iteration NUFFT cost on the GPU. MGE source\n", + "fits are now routine even at ALMA-class visibility counts.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** Benefits and drawbacks of an MGE source for interferometer data.\n", + "- **NUFFT (nufftax):** Why MGE-on-visibilities is now practical thanks to nufftax.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Model:** Compose the lens model \u2014 `Isothermal` + `ExternalShear` lens mass and a 30-Gaussian MGE\n", + " source bulge. Lens light omitted (interferometer convention).\n", + "- **Mask:** Define the `real_space_mask` which sets the grid the strong lens is evaluated on.\n", + "- **Dataset:** Load the strong lens `Interferometer` dataset using `TransformerNUFFT` (backed by `nufftax`).\n", + "- **Over Sampling:** Interferometer modeling does not use over-sampling.\n", + "- **Search:** Configure the non-linear search (Nautilus).\n", + "- **Analysis:** Create the `AnalysisInterferometer` object.\n", + "- **VRAM:** Memory budget for a multi-component linear basis on GPU.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Advantages__\n", + "\n", + "Symmetric light profiles (e.g. elliptical Sersics) may leave significant residuals when fitting the source\n", + "galaxy of a strong lens \u2014 they fail to capture irregular and asymmetric morphology (e.g. isophotal twists,\n", + "multi-component source emission). An MGE fully captures these features and can therefore much better\n", + "represent the emission of complex source galaxies.\n", + "\n", + "The MGE model is composed such that the `intensity` parameters and the `sigma` parameters controlling the\n", + "Gaussian sizes are *not* sampled by Nautilus. Centres and elliptical components are shared across each\n", + "group of Gaussians. This removes the most significant degeneracies in parameter space, making the model\n", + "much more reliable and efficient to fit than a free-intensity Sersic.\n", + "\n", + "Therefore, not only does an MGE source fit more complex source morphologies, it does so using fewer\n", + "non-linear parameters in a much simpler non-linear parameter space which has far less significant parameter\n", + "degeneracies.\n", + "\n", + "__Disadvantages__\n", + "\n", + "To fit an MGE model, the light of the ~15-100 Gaussians must be evaluated and NUFFT'd to the uv-plane per\n", + "likelihood evaluation. This is slower than evaluating a single `SersicCore`. With `nufftax` the per-NUFFT\n", + "cost is small enough that the total slow-down per likelihood evaluation is typically 2-5x for a 30-Gaussian\n", + "MGE \u2014 paid back in fewer iterations because the parameter space is simpler.\n", + "\n", + "The MGE source can also be less intuitive to interpret physically than a Sersic. The Sersic `effective_radius`\n", + "and `sersic_index` map directly to galaxy size and concentration; an MGE's solved-for Gaussian intensities\n", + "require an extra processing step to compute equivalent physical quantities.\n", + "\n", + "__NUFFT (nufftax)__\n", + "\n", + "The image-to-visibilities Fourier transform is performed by a Non-Uniform Fast Fourier Transform (NUFFT),\n", + "exposed in **PyAutoLens** as `TransformerNUFFT`. The default backend is `nufftax`, a pure-JAX NUFFT that\n", + "jit-compiles and vmap-batches like the rest of the library:\n", + "\n", + " https://github.com/GragasLab/nufftax\n", + "\n", + "Because `nufftax` is JAX-native, NUFFT-ing every Gaussian in the MGE basis happens inside the same compiled\n", + "likelihood that does the inversion, mass model ray-tracing, and chi-squared sum. There is no host round-trip\n", + "between NUFFT calls, so a model with N Gaussians costs only N forward-NUFFTs per iteration on the GPU \u2014\n", + "fast enough that MGE-on-visibilities is now routinely practical.\n", + "\n", + "If `nufftax` is not installed, install it via `pip install nufftax`.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algebra solver which allows for positive and\n", + "negative values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This is problematic, as it means that negative surface brightnesses values can be computed to represent a\n", + "galaxy's light, which is clearly unphysical. For an MGE, this produces a positive-negative \"ringing\", where\n", + "the Gaussians alternate between large positive and negative values. This is clearly undesirable and\n", + "unphysical.\n", + "\n", + "**PyAutoLens** uses a positive only linear algebra solver which has been extensively optimized to ensure it\n", + "is as fast as positive-negative solvers. This ensures that all Gaussian intensities are positive and\n", + "therefore physical.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data). This is the standard\n", + " convention for interferometer modeling, as the lens galaxy's optical/IR emission is typically below the\n", + " detection threshold of mm/sub-mm interferometers.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's bulge is a multi-Gaussian expansion of 30 linear `Gaussian` profiles, arranged in\n", + " two groups of 15 (each group shares a centre and ell_comps).\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `interferometer/start_here.ipynb` notebook.\n", + "\n", + "__Imaging Equivalent__\n", + "\n", + "For the CCD-imaging version of this script (MGE on the lens galaxy, not the source), see\n", + "`autolens_workspace/*/imaging/features/multi_gaussian_expansion/modeling.py`." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the `real_space_mask` which defines the grid the image of the strong lens is evaluated on." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit with\n", + "the lens model.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane\n", + "and compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native Non-Uniform Fast Fourier\n", + "Transform backed by `nufftax`, which is required for fast MGE modeling and scales efficiently from a few\n", + "hundred visibilities to tens of millions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "If you are familiar with using imaging data, you may have seen that a numerical technique called over\n", + "sampling is used, which evaluates light profiles on a higher resolution grid than the image data to ensure\n", + "the calculation is accurate.\n", + "\n", + "Interferometer data does not observe galaxies in a way where over sampling is necessary, therefore all\n", + "interferometer calculations are performed without over sampling.\n", + "\n", + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's bulge is 10 linear `Gaussian` profiles [6 parameters total].\n", + " - The centres and elliptical components of the Gaussians are linked together in two groups of 5.\n", + " - The `sigma` size of the Gaussians increases in log10 increments from 0.01 to the mask radius.\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=13.\n", + "\n", + "The MGE comprises 2 groups of 5 Gaussians. Each group has its own elliptical components, so the source's\n", + "light is decomposed into two distinct elliptical components which could be viewed as a bulge and a disk of\n", + "the source.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook:\n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 5\n", + "gaussian_per_basis = 2\n", + "\n", + "# The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.5\".\n", + "log10_sigma_list = np.linspace(-2, np.log10(mask_radius), total_gaussians)\n", + "\n", + "# By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.\n", + "\n", + "centre_0 = af.UniformPrior(lower_limit=-0.3, upper_limit=0.3)\n", + "centre_1 = af.UniformPrior(lower_limit=-0.3, upper_limit=0.3)\n", + "\n", + "bulge_gaussian_list = []\n", + "\n", + "for j in range(gaussian_per_basis):\n", + " # A list of Gaussian model components whose parameters are customized below.\n", + "\n", + " gaussian_list = af.Collection(\n", + " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", + " )\n", + "\n", + " # Iterate over every Gaussian and customize its parameters.\n", + "\n", + " for i, gaussian in enumerate(gaussian_list):\n", + " gaussian.centre.centre_0 = centre_0 # All Gaussians have same y centre.\n", + " gaussian.centre.centre_1 = centre_1 # All Gaussians have same x centre.\n", + " gaussian.ell_comps = gaussian_list[\n", + " 0\n", + " ].ell_comps # All Gaussians in this group share ell_comps.\n", + " gaussian.sigma = (\n", + " 10 ** log10_sigma_list[i]\n", + " ) # All Gaussian sigmas are fixed to values above.\n", + "\n", + " bulge_gaussian_list += gaussian_list\n", + "\n", + "# The Basis object groups many light profiles together into a single model component.\n", + "source_bulge = af.Model(al.lp_basis.Basis, profile_list=bulge_gaussian_list)\n", + "\n", + "# Lens:\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=source_bulge)\n", + "\n", + "# Overall Lens Model:\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen\n", + "refer to `start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This shows every single Gaussian light profile in the model, which is a lot of parameters! However, the\n", + "vast majority of these parameters are fixed to the values we set above, so the model actually has far fewer\n", + "free parameters than it looks!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start_here.py` for a\n", + "full description).\n", + "\n", + "Owing to the simplicity of fitting an MGE (no intensity or sigma free parameters), we use fewer live points\n", + "than the `interferometer/modeling.py` example: 75 live points speeds up convergence." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\") / \"features\",\n", + " name=\"mge\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisInterferometer` object defining how Nautilus fits the model to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `interferometer/modeling.py` example explains how VRAM is used during GPU-based fitting and how to\n", + "print the estimated VRAM required by a model.\n", + "\n", + "For each linear `Gaussian` profile, extra VRAM is used to store its NUFFT'd mapping matrix column. For\n", + "around 30 linear Gaussians this typically requires a modest amount of VRAM (e.g. 10-50 MB per batched\n", + "likelihood). Models that use hundreds of Gaussians, especially in combination with a large batch size, may\n", + "therefore exceed GBs of VRAM and require you to adjust the batch size to fit within your GPU's VRAM.\n", + "\n", + "VRAM on interferometer datasets is driven primarily by the visibility count and the real-space mask size,\n", + "not the number of Gaussians in the MGE basis.\n", + "\n", + "__Run Time__\n", + "\n", + "The likelihood evaluation time for an MGE is slower than a single linear `SersicCore` source, because the\n", + "image of every Gaussian must be evaluated and NUFFT'd to the uv-plane. With `nufftax`, the per-NUFFT cost\n", + "is small enough that the total slow-down per likelihood is typically 2-5x for a 30-Gaussian MGE compared\n", + "to a one-component source \u2014 paid back in fewer iterations because the parameter space is simpler.\n", + "\n", + "Because the MGE has no free `intensity` or `sigma` parameters and shares centres/ell_comps across groups,\n", + "Nautilus converges significantly faster than for a free-intensity Sersic source. We also use fewer live\n", + "points (75 vs the 100 used in `interferometer/modeling.py`), further speeding up the model-fit.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the\n", + "output folder for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format (if\n", + "this does not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix\n", + "this):\n", + "\n", + "This confirms there are many `Gaussian`s in the source light model and it lists their inferred parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "In particular, checkout the results example `linear.py` which details how to extract all information about\n", + "linear light profiles from a fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "To further showcase MGE source modeling, see the SLaM pipeline at\n", + "`autolens_workspace/scripts/interferometer/features/multi_gaussian_expansion/slam.py`, which uses the\n", + "`al.model_util.mge_model_from` helper to compose an MGE source in the SOURCE LP stage before chaining\n", + "to a pixelized source reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/multi_gaussian_expansion/slam.ipynb b/notebooks/interferometer/features/multi_gaussian_expansion/slam.ipynb index 03cc43605..5b535a66b 100644 --- a/notebooks/interferometer/features/multi_gaussian_expansion/slam.ipynb +++ b/notebooks/interferometer/features/multi_gaussian_expansion/slam.ipynb @@ -1,739 +1,776 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Multi Gaussian Expansion: SLaM (Interferometer)\n", - "================================================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for `Interferometer`\n", - "data, using a **multi-Gaussian expansion (MGE)** for the source galaxy in the SOURCE LP stage.\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that guide\n", - "before working through this example.\n", - "\n", - "The interferometer pixelization SLaM pipeline at `interferometer/features/pixelization/slam.py` is the\n", - "closest analogue \u2014 this script keeps the same pipeline structure (SOURCE LP \u2192 SOURCE PIX 1 \u2192 SOURCE PIX 2\n", - "\u2192 MASS TOTAL) and only differs in the size of the MGE source basis used in SOURCE LP. The LIGHT LP\n", - "pipeline from `slam_start_here.py` is omitted because interferometer data does not contain lens light\n", - "emission.\n", - "\n", - "The MGE captures the asymmetric morphology of typical sub-mm/radio-selected lensed sources better than a\n", - "single `Sersic`, while keeping the non-linear parameter space small (only the shared Gaussian centre and\n", - "elliptical components are free; the per-Gaussian intensities are solved analytically and the sigmas are\n", - "fixed). This makes SOURCE LP fast, robust, and a better initialiser for the pixelized stages that follow.\n", - "\n", - "The SOURCE LP stage uses `TransformerNUFFT` (backed by JAX-native `nufftax`,\n", - "https://github.com/GragasLab/nufftax). MGE fits to visibilities are practical at any visibility count\n", - "because the per-iteration NUFFT of every Gaussian basis component runs inside the same compiled likelihood\n", - "as the rest of the model.\n", - "\n", - "The SOURCE PIX and MASS TOTAL stages use the same `TransformerNUFFT` plus a pre-computed sparse\n", - "operator (via `apply_sparse_operator`) so the pixelized source curvature matrix is assembled via the\n", - "FFT-based W\u0303 precision matrix instead of the dense `transformed_mapping_matrix`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with `slam_start_here`.\n", - "- **SOURCE LP PIPELINE:** Initializes the mass and source-light model with an MGE source bulge, fitted\n", - " via `TransformerNUFFT`.\n", - "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP\n", - " result.\n", - "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included.\n", - "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with\n", - " `TransformerNUFFT` + sparse operator (source_pix onwards).\n", - "- **Sparse Operators:** Pre-compute the sparse operator for the pixelized stages.\n", - "- **Settings:** Disable the positive-only solver so the pixelized source reconstruction can have negative\n", - " pixel values.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization,\n", - " database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in `features/pixelization/modeling`, the mesh shape is fixed before\n", - " modeling.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines and how they\n", - " integrate into strong-lens modeling.\n", - "\n", - "- **Multi Gaussian Expansion** (`interferometer/features/multi_gaussian_expansion/modeling`)\n", - " How an MGE source basis is composed for interferometer data, and why nufftax makes large-N MGE bases\n", - " practical.\n", - "\n", - "- **Interferometer Pixelization SLaM** (`interferometer/features/pixelization/slam`)\n", - " The canonical interferometer SLaM pipeline. This script is essentially that pipeline with a larger MGE\n", - " source basis in SOURCE LP.\n", - "\n", - "__High Resolution Dataset__\n", - "\n", - "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large\n", - "files which are too big to include in the main `autolens_workspace` repository:\n", - "\n", - " https://github.com/PyAutoLabs/autolens_workspace_large_files\n", - "\n", - "After downloading the file, place it in the directory:\n", - "\n", - " `autolens_workspace/dataset/interferometer/alma`\n", - "\n", - "You can then perform modeling using this high-resolution dataset by uncommenting the relevant line of code\n", - "below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light using a\n", - "multi-Gaussian expansion (MGE). The lens galaxy mass and external shear are fitted at the same time.\n", - "\n", - "The MGE source is built via the `al.model_util.mge_model_from` helper, which constructs a basis of N\n", - "linear `Gaussian` profiles arranged in two groups (each sharing a centre and ell_comps; sigmas fixed to\n", - "log-spaced values). For interferometer data we use `total_gaussians=5`, which is enough to capture the\n", - "source morphology while keeping per-iteration cost manageable.\n", - "\n", - "This stage uses the `dataset_nufft` Interferometer (built with `TransformerNUFFT`, backed by `nufftax`),\n", - "which makes the per-iteration NUFFT of every Gaussian fast even for ALMA-class datasets with millions of\n", - "visibilities.\n", - "\n", - "The mass and source models from this search initialize the SOURCE PIX PIPELINE searches that follow, and\n", - "the result also provides the adapt image and position likelihood used by those later stages.\n", - "\n", - "Note that no lens light is fitted: interferometer data does not contain lens light emission, so\n", - "`lens.bulge` and `lens.disk` are kept at `None`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " # interferometer data does not contain lens light emission\n", - " bulge=None,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source galaxy.\n", - "\n", - "The first search fits a pixelization whose purpose is to generate a high-quality adapt image used in\n", - "search 2. It uses the adapt image computed from the SOURCE LP result and the position likelihood is also\n", - "derived automatically from the SOURCE LP result \u2014 no manual positions input is required.\n", - "\n", - "This stage uses the `dataset_sparse` Interferometer (built with `TransformerNUFFT` +\n", - "`apply_sparse_operator`). The NUFFT keeps the one-time dirty-image setup tractable for ALMA-scale\n", - "visibility counts, and the precomputed sparse operator makes per-likelihood curvature assembly use the\n", - "FFT-based W\u0303 precision matrix instead of the dense `transformed_mapping_matrix`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", - "pixelization and regularization.\n", - "\n", - "The LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not contain\n", - "lens light emission.\n", - "\n", - "Like SOURCE PIX PIPELINE 1, this stage uses the `dataset_sparse` Interferometer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " settings=settings,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", - "contain lens light emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " source_pix_result_2: af.Result,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_1.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=source_pix_result_1.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load the `Interferometer` data and define the real-space mask." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", - ")\n", - "\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "# dataset_name = \"alma\"\n", - "#\n", - "# if dataset_name == \"alma\":\n", - "#\n", - "# real_space_mask = al.Mask2D.circular(\n", - "# shape_native=(800, 800),\n", - "# pixel_scales=0.01,\n", - "# radius=mask_radius,\n", - "# )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Two Datasets__\n", - "\n", - "Both stages use `TransformerNUFFT` (backed by JAX-native `nufftax`), which keeps the dirty-image setup\n", - "tractable at any visibility count. The two datasets differ only in whether `apply_sparse_operator` has\n", - "been called:\n", - "\n", - "- `dataset_nufft` is plain `TransformerNUFFT` for the `source_lp` stage. With an MGE source bulge each\n", - " Gaussian's NUFFT runs inside the same compiled likelihood.\n", - "- `dataset_sparse` is the same `TransformerNUFFT` plus `apply_sparse_operator(...)` for `source_pix_1`,\n", - " `source_pix_2`, and `mass_total`. The precomputed sparse operator makes the pixelized curvature matrix\n", - " assembly use the FFT-based W\u0303 precision matrix instead of the dense `transformed_mapping_matrix`.\n", - "\n", - "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload)\n", - "differ." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_nufft = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "dataset_sparse = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", - "pixelized source modeling, especially for many visibilities.\n", - "\n", - "We use a try / except to load the pre-computed curvature preload, which is necessary to use the sparse\n", - "operator formalism. If this file does not exist it is made here.\n", - "\n", - "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", - "`source_lp` does not need it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " nufft_precision_operator = np.load(\n", - " file=dataset_path / \"nufft_precision_operator.npy\",\n", - " )\n", - "except FileNotFoundError:\n", - " nufft_precision_operator = None\n", - "\n", - "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", - " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "Disable the default positive-only linear algebra solver so the pixelized source reconstruction can have\n", - "negative pixel values. (The MGE source bulge in SOURCE LP is still constrained to positive intensities\n", - "internally because each Gaussian's intensity is a physical normalization.)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"interferometer\") / \"slam_multi_gaussian_expansion\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step.\n", - "\n", - "Both datasets use `TransformerNUFFT`. `source_lp` is passed `dataset_nufft` (no sparse operator) while\n", - "every later stage is passed `dataset_sparse` (`TransformerNUFFT` + `apply_sparse_operator`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset_nufft,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " settings=settings,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Multi Gaussian Expansion: SLaM (Interferometer)\n", + "================================================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for `Interferometer`\n", + "data, using a **multi-Gaussian expansion (MGE)** for the source galaxy in the SOURCE LP stage.\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that guide\n", + "before working through this example.\n", + "\n", + "The interferometer pixelization SLaM pipeline at `interferometer/features/pixelization/slam.py` is the\n", + "closest analogue \u2014 this script keeps the same pipeline structure (SOURCE LP \u2192 SOURCE PIX 1 \u2192 SOURCE PIX 2\n", + "\u2192 MASS TOTAL) and only differs in the size of the MGE source basis used in SOURCE LP. The LIGHT LP\n", + "pipeline from `slam_start_here.py` is omitted because interferometer data does not contain lens light\n", + "emission.\n", + "\n", + "The MGE captures the asymmetric morphology of typical sub-mm/radio-selected lensed sources better than a\n", + "single `Sersic`, while keeping the non-linear parameter space small (only the shared Gaussian centre and\n", + "elliptical components are free; the per-Gaussian intensities are solved analytically and the sigmas are\n", + "fixed). This makes SOURCE LP fast, robust, and a better initialiser for the pixelized stages that follow.\n", + "\n", + "The SOURCE LP stage uses `TransformerNUFFT` (backed by JAX-native `nufftax`,\n", + "https://github.com/GragasLab/nufftax). MGE fits to visibilities are practical at any visibility count\n", + "because the per-iteration NUFFT of every Gaussian basis component runs inside the same compiled likelihood\n", + "as the rest of the model.\n", + "\n", + "The SOURCE PIX and MASS TOTAL stages use the same `TransformerNUFFT` plus a pre-computed sparse\n", + "operator (via `apply_sparse_operator`) so the pixelized source curvature matrix is assembled via the\n", + "FFT-based W\u0303 precision matrix instead of the dense `transformed_mapping_matrix`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with `slam_start_here`.\n", + "- **SOURCE LP PIPELINE:** Initializes the mass and source-light model with an MGE source bulge, fitted\n", + " via `TransformerNUFFT`.\n", + "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP\n", + " result.\n", + "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included.\n", + "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with\n", + " `TransformerNUFFT` + sparse operator (source_pix onwards).\n", + "- **Sparse Operators:** Pre-compute the sparse operator for the pixelized stages.\n", + "- **Settings:** Disable the positive-only solver so the pixelized source reconstruction can have negative\n", + " pixel values.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization,\n", + " database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in `features/pixelization/modeling`, the mesh shape is fixed before\n", + " modeling.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines and how they\n", + " integrate into strong-lens modeling.\n", + "\n", + "- **Multi Gaussian Expansion** (`interferometer/features/multi_gaussian_expansion/modeling`)\n", + " How an MGE source basis is composed for interferometer data, and why nufftax makes large-N MGE bases\n", + " practical.\n", + "\n", + "- **Interferometer Pixelization SLaM** (`interferometer/features/pixelization/slam`)\n", + " The canonical interferometer SLaM pipeline. This script is essentially that pipeline with a larger MGE\n", + " source basis in SOURCE LP.\n", + "\n", + "__High Resolution Dataset__\n", + "\n", + "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large\n", + "files which are too big to include in the main `autolens_workspace` repository:\n", + "\n", + " https://github.com/PyAutoLabs/autolens_workspace_large_files\n", + "\n", + "After downloading the file, place it in the directory:\n", + "\n", + " `autolens_workspace/dataset/interferometer/alma`\n", + "\n", + "You can then perform modeling using this high-resolution dataset by uncommenting the relevant line of code\n", + "below." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light using a\n", + "multi-Gaussian expansion (MGE). The lens galaxy mass and external shear are fitted at the same time.\n", + "\n", + "The MGE source is built via the `al.model_util.mge_model_from` helper, which constructs a basis of N\n", + "linear `Gaussian` profiles arranged in two groups (each sharing a centre and ell_comps; sigmas fixed to\n", + "log-spaced values). For interferometer data we use `total_gaussians=5`, which is enough to capture the\n", + "source morphology while keeping per-iteration cost manageable.\n", + "\n", + "This stage uses the `dataset_nufft` Interferometer (built with `TransformerNUFFT`, backed by `nufftax`),\n", + "which makes the per-iteration NUFFT of every Gaussian fast even for ALMA-class datasets with millions of\n", + "visibilities.\n", + "\n", + "The mass and source models from this search initialize the SOURCE PIX PIPELINE searches that follow, and\n", + "the result also provides the adapt image and position likelihood used by those later stages.\n", + "\n", + "Note that no lens light is fitted: interferometer data does not contain lens light emission, so\n", + "`lens.bulge` and `lens.disk` are kept at `None`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " # interferometer data does not contain lens light emission\n", + " bulge=None,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source galaxy.\n", + "\n", + "The first search fits a pixelization whose purpose is to generate a high-quality adapt image used in\n", + "search 2. It uses the adapt image computed from the SOURCE LP result and the position likelihood is also\n", + "derived automatically from the SOURCE LP result \u2014 no manual positions input is required.\n", + "\n", + "This stage uses the `dataset_sparse` Interferometer (built with `TransformerNUFFT` +\n", + "`apply_sparse_operator`). The NUFFT keeps the one-time dirty-image setup tractable for ALMA-scale\n", + "visibility counts, and the precomputed sparse operator makes per-likelihood curvature assembly use the\n", + "FFT-based W\u0303 precision matrix instead of the dense `transformed_mapping_matrix`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", + "pixelization and regularization.\n", + "\n", + "The LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not contain\n", + "lens light emission.\n", + "\n", + "Like SOURCE PIX PIPELINE 1, this stage uses the `dataset_sparse` Interferometer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " settings=settings,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", + "contain lens light emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " source_pix_result_2: af.Result,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_1.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=source_pix_result_1.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load the `Interferometer` data and define the real-space mask." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", + ")\n", + "\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "# dataset_name = \"alma\"\n", + "#\n", + "# if dataset_name == \"alma\":\n", + "#\n", + "# real_space_mask = al.Mask2D.circular(\n", + "# shape_native=(800, 800),\n", + "# pixel_scales=0.01,\n", + "# radius=mask_radius,\n", + "# )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Two Datasets__\n", + "\n", + "Both stages use `TransformerNUFFT` (backed by JAX-native `nufftax`), which keeps the dirty-image setup\n", + "tractable at any visibility count. The two datasets differ only in whether `apply_sparse_operator` has\n", + "been called:\n", + "\n", + "- `dataset_nufft` is plain `TransformerNUFFT` for the `source_lp` stage. With an MGE source bulge each\n", + " Gaussian's NUFFT runs inside the same compiled likelihood.\n", + "- `dataset_sparse` is the same `TransformerNUFFT` plus `apply_sparse_operator(...)` for `source_pix_1`,\n", + " `source_pix_2`, and `mass_total`. The precomputed sparse operator makes the pixelized curvature matrix\n", + " assembly use the FFT-based W\u0303 precision matrix instead of the dense `transformed_mapping_matrix`.\n", + "\n", + "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload)\n", + "differ." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_nufft = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "dataset_sparse = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", + "pixelized source modeling, especially for many visibilities.\n", + "\n", + "We use a try / except to load the pre-computed curvature preload, which is necessary to use the sparse\n", + "operator formalism. If this file does not exist it is made here.\n", + "\n", + "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", + "`source_lp` does not need it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " nufft_precision_operator = np.load(\n", + " file=dataset_path / \"nufft_precision_operator.npy\",\n", + " )\n", + "except FileNotFoundError:\n", + " nufft_precision_operator = None\n", + "\n", + "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", + " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "Disable the default positive-only linear algebra solver so the pixelized source reconstruction can have\n", + "negative pixel values. (The MGE source bulge in SOURCE LP is still constrained to positive intensities\n", + "internally because each Gaussian's intensity is a physical normalization.)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"interferometer\") / \"slam_multi_gaussian_expansion\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step.\n", + "\n", + "Both datasets use `TransformerNUFFT`. `source_lp` is passed `dataset_nufft` (no sparse operator) while\n", + "every later stage is passed `dataset_sparse` (`TransformerNUFFT` + `apply_sparse_operator`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset_nufft,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " settings=settings,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/pixelization/delaunay.ipynb b/notebooks/interferometer/features/pixelization/delaunay.ipynb index 6a459926d..00e8645b0 100644 --- a/notebooks/interferometer/features/pixelization/delaunay.ipynb +++ b/notebooks/interferometer/features/pixelization/delaunay.ipynb @@ -1,1747 +1,1784 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: Delaunay\n", - "======================\n", - "\n", - "The majority of pixelized source reconstructions in the workspace use a rectangular mesh to reconstruct\n", - "the source's surface brightness.\n", - "\n", - "This example illustrates an alternative pixelization that uses a Delaunay triangulation mesh to reconstruct the\n", - "source.\n", - "\n", - "The approach is distinct from the rectangular mesh and has a number of traits which are unique to it:\n", - "\n", - "- `Adaptive Mesh`: In the source plane, the Delaunay mesh uses irregularly shaped triangles to reconstruct the\n", - " source, as opposed to uniform rectangular pixels. This allows the mesh to better adapt to irregular and\n", - " asymmetric source morphologies and change the distribution of source pixels to better match the source's\n", - " surface brightness.\n", - "\n", - "- `Image Mesh`: The vertexes of the Delaunay triangles are computed by overlaying a coarse uniform grid in the\n", - " image-plane and ray-tracing these coordinates to the source-plane. This is unlike the rectangular mesh, which\n", - " simply overlays a uniform grid in the source-plane. This again helps the Delaunay mesh to better adapt to the\n", - " source's surface brightness.\n", - "\n", - "- `Interpolation`: The Delaunay mesh uses a different interpolation scheme to the rectangular mesh, which is\n", - " barycentric interpolation within each triangle. This is different to the rectangular mesh, which uses bilinear\n", - " interpolation within each rectangular pixel.\n", - "\n", - "- `Regularization`: The Delaunay mesh provides different approaches to regularization, with the default being\n", - " one which uses the barycentric coordinates of the triangles to compute how source pixels are regularized with\n", - " their neighbors.\n", - "\n", - "Currently it is not expected that the Delaunay is better or worse than the rectangular mesh, it is simply a different\n", - "approach to pixelization that may work better for certain datasets.\n", - "\n", - "__JAX + GPU__\n", - "\n", - "Generating a Delaunay mesh supports JAX and GPU acceleration, however certain operations (e.g. generating the Delaunay\n", - "triangulation itself) do not run on the GPU because they cannot be easily converted to JAX.\n", - "\n", - "Instead, JAX sends them to a CPU, runs them there, and then sends the results back to the GPU. This process is\n", - "very efficient, because these operations run very fast on a CPU and the data being sent back and forth is small.\n", - "Current benchmarking suggests the Delaunay runs less than twice as long as the same fit using a rectangular mesh,\n", - "but scientfically offers better results in many cases.\n", - "\n", - "If you do want to run only on CPU, you can use fast CPU method described in\n", - "example `imaging/features/pixelization/cpu_fast_modeling` with the Delaunay mesh.\n", - "\n", - "\n", - "__Source Science (Magnification, Flux and More)__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", - "size of the source.\n", - "\n", - "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", - "including examples using a Delaunay source reconstruction. Once you have completed lens modeling using a Delaunay mesh,\n", - "you can jump to that example to study the source galaxy.\n", - "\n", - "__Contents__\n", - "\n", - "- **Image Mesh:** For a Delaunay mesh, the vertices of the triangles are defined by (y, x) coordinates in the.\n", - "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **VRAM:** The `pixelization/modeling` example explains how VRAM use is an important consideration for.\n", - "- **Adaptive Delaunay:** The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to.\n", - "- **SLaM Pipelines:** The API above allows you to use adaptive features yourself, and you should go ahead an explore them.\n", - "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, using an MGE for the lens and source light profiles.\n", - "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", - "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included as interferometer data.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "- **Prerequisites:** The likelihood function of pixelizations is the most complicated likelihood function.\n", - "- **Likelihood Function:** The example `interferometer/pixelization/likelihood_function.py` provides a step-by-step.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Source Galaxy Pixelization and Regularization:** We combine the pixelization into a single `Galaxy` object.\n", - "- **Source Pixel Centre Calculation:** In order to reconstruct the source galaxy using a Delaunay mesh, we need to determine the centres.\n", - "- **Ray Tracing:** Overview of ray tracing for this example.\n", - "- **Border Relocation:** Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace.\n", - "- **Delaunay Mesh:** The relocated mesh grid is used to create the `Pixelization`'s Delaunay mesh using the.\n", - "- **Interpolation:** Overview of interpolation for this example.\n", - "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "\n", - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "dataset = dataset.apply_sparse_operator(use_jax=True, show_progress=True)\n", - "\n", - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n", - "\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)\n", - "\n", - "\n", - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Image Mesh__\n", - "\n", - "For a Delaunay mesh, the vertices of the triangles are defined by (y, x) coordinates in the image-plane. These\n", - "coordinates are then ray-traced into the source-plane for each mass model sampled during the non-linear search.\n", - "This `image_plane_mesh_grid` must be computed before lens modeling.\n", - "\n", - "We compute this `image_plane_mesh_grid` using an `Overlay` image-mesh, which places a regular grid of\n", - "(y, x) points across the image-plane. This has a mild adaptive effect: regions of high lens magnification receive\n", - "more source pixels once they are ray-traced. Later in this example, we switch to a `Hilbert` image-mesh, which adapts\n", - "the pixel distribution more strongly to the source\u2019s surface brightness.\n", - "\n", - "The `Delaunay` mesh has an input number of `pixels`, which is the number of source pixels used to reconstruct the \n", - "source. The number of `pixels` must be equal to the number of coordinates in the `image_plane_mesh_grid`. \n", - "\n", - "Like for the `mesh_shape` rectangular mesh, `pixels` must be fixed for lens modeling because JAX uses the \n", - "number of `pixels` to determine static array shapes. \n", - "\n", - "To pass the `image_plane_mesh_grid` to the modeling, we use the `AdaptImages` object below, which pairs\n", - "the `image_plane_mesh_grid` to the source galaxy. For double source plane lenses, this means we can\n", - "attach an `image_plane_mesh_grid` to each source galaxy and use adaptive meshes for each source plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Edge Zeroing__\n", - "\n", - "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", - "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", - "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", - "\n", - "For a rectangular mesh, the source code computes edge pixels internally using the known pixels at the edge of the mesh,\n", - "requiring no input from the user. \n", - "\n", - "For the `Delaunay` mesh, we use the `append_with_circle_edge_points` function to manually setup the Delaunay image \n", - "mesh to include a ring of edge pixels and then input the total number into the mesh to perform zeroing. \n", - "\n", - "These points are added to the edge of the image-plane mesh, ray-traced to the source-plane during lens modeling, \n", - "included in the Delaunay triangulation but zeroed during the inversion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "edge_pixels_total = 30\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=real_space_mask.mask_centre,\n", - " radius=mask_radius + real_space_mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "In the example `interferometer/features/pixelization/fit.py`, we illustrate how to use a pixelized source\n", - "with a rectangular mesh.\n", - "\n", - "Below, we use a Delaunay mesh to perform a fit using the Delaunay source reconstruction.\n", - "\n", - "The API is nearly identical to the rectangular mesh example, noting that the inputs to the `Delaunay` \n", - "mesh are different to the rectangular mesh and use image mesh quantities computed above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh = al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - ")\n", - "regularization = al.reg.ConstantSplit(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", - "\n", - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "adapt_images = al.AdaptImages(\n", - " galaxy_image_plane_mesh_grid_dict={source: image_plane_mesh_grid}\n", - ")\n", - "\n", - "fit = al.FitInterferometer(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " adapt_images=adapt_images,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the Delaunay source does a good job at capturing the appearance of the source galaxy\n", - "using adaptive triangular pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We now perform lens modeling using the Delaunay pixelization with the Overlay image-mesh.\n", - "\n", - "The code below is a simple adaptive modeling example using the Delaunay mesh, which mirrors the\n", - "API used in other pixelization modeling examples.\n", - "\n", - "The example `interferometer/features/pixelization/adaptive.py` illustrates how to use adaptive features to\n", - "adapt the rectangular mesh and its regularization to the source's surface brightness. In particular, an image\n", - "of the lensed source is passed to the modeling via the `AdaptImages` object, in order to adapt\n", - "the mesh and regularization during the model-fit.\n", - "\n", - "The same object is used to pass the `image_plane_mesh_grid` to the modeling. Above, this image-plane mesh grid\n", - "is an `Overlay` mesh and does not specifically adapt to the source's surface brightness, thus pairing it with\n", - "the source as done below seems redundant. However, in a moment we will switch to a `Hilbert` image-mesh, which\n", - "does adapt to the source's surface brightness, meaning this pairing is necessary." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "adapt_images = al.AdaptImages(\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In the first\n", - "search our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light uses an `Overlay` image-mesh with fixed resolution 30 x 30 pixels [0 parameters].\n", - "\n", - " - The source-galaxy's light uses a `Delaunay` mesh [0 parameters].\n", - "\n", - " - This pixelization is regularized using a `Constant` scheme [1 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=8." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", - ")\n", - "\n", - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.ConstantSplit,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "search_1 = af.Nautilus(\n", - " path_prefix=Path(\"features\"),\n", - " name=\"delaunay\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - ")\n", - "\n", - "analysis_1 = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `pixelization/modeling` example explains how VRAM use is an important consideration for pixelization models\n", - "and how it depends on image resolution, number of source pixels and batch size.\n", - "\n", - "This is true for the Delaunay mesh, therefore we print out the estimated VRAM required for this model-fit.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_1.print_vram_use(model=model_1, batch_size=search_1.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit (Search 1)__\n", - "\n", - "Perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Adaptive Delaunay__\n", - "\n", - "The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to\n", - "adapt the rectangular mesh and its regularization to the source's surface brightness.\n", - "\n", - "The image-mesh has a special adaptive variant called the `Hilbert` image-mesh, which adapts the distribution \n", - "of source-pixels to the source's unlensed morphology. This means that the source's brightest regions are \n", - "reconstructed using significantly more source pixels than seen for the `Overlay` image mesh. \n", - "Conversely, the source's faintest regions are reconstructed using significantly fewer source pixels.\n", - "\n", - "Unlike the adaptive rectangular mesh, the Hilbert image-plane mesh is computed before modeling, passed\n", - "to the `AdaptImages` object, and remains fixed during the model-fit.\n", - "\n", - "It is recommend that the parameters governing these features are always fitted from using a fixed lens light and\n", - "mass model. This ensures the adaptation is performed quickly, and removes degeneracies in the lens model that\n", - "are difficult to sample. Given the Hilbert mesh is fixed, this modeling only fits for the regularization coefficients\n", - "of the adaptive regularization scheme.\n", - "\n", - "For this reason, search 2 fixes the lens galaxy's light and mass model to the best-fit model of search 1. A third\n", - "search will then fit for the lens galaxy's light and mass model using these adaptive features.\n", - "\n", - "The details of how the above features work is not provided here, but is given at the end of chapter 4 of the HowToLens\n", - "lecture series." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=result_1, use_model_images=True\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result_1)\n", - "\n", - "image_mesh = al.image_mesh.Hilbert(pixels=1000, weight_power=3.5, weight_floor=0.01)\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask, adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", - ")\n", - "\n", - "# Repeat edge zeroing set up describe above.\n", - "\n", - "edge_pixels_total = 30\n", - "\n", - "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=real_space_mask.mask_centre,\n", - " radius=mask_radius + real_space_mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model (Search 2)__\n", - "\n", - "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In \n", - "the second search our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` with fixed parameters from \n", - " search 1 [0 parameters].\n", - "\n", - " - The source galaxy's light uses a `Hilbert` image-mesh with fixed resolution 1000 pixels [2 parameters].\n", - "\n", - " - The source-galaxy's light uses a `Delaunay` mesh [0 parameters].\n", - "\n", - " - This pixelization is regularized using a `AdaptSplit` scheme [2 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=4." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.AdaptSplit,\n", - ")\n", - "\n", - "source = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " pixelization=pixelization,\n", - ")\n", - "\n", - "model_2 = af.Collection(\n", - " galaxies=af.Collection(lens=result_1.instance.galaxies.lens, source=source)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis (Search 2)__\n", - "\n", - "We now create the analysis for the second search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_2 = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search + Model-Fit (Search 2)__\n", - "\n", - "We now create the non-linear search and perform the model-fit using this model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search_2 = af.Nautilus(\n", - " path_prefix=Path(\"features\"),\n", - " name=\"delaunay_adapt\",\n", - " unique_tag=dataset_name,\n", - " n_live=75,\n", - ")\n", - "\n", - "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We could perform a third fit where we free all lens model parameters and fit them using the adaptive \n", - "image mesh and regularization.\n", - "\n", - "However, it is better to use all of these features with the Delaunay via the\n", - "SLaM pipelines, which we jump to immediately below.\n", - "\n", - "__SLaM Pipelines__\n", - "\n", - "The API above allows you to use adaptive features yourself, and you should go ahead an explore them on datasets you\n", - "are familiar with.\n", - "\n", - "However, you may also wish to use the Source, Light and Mass (SLaM) pipelines, which are pipelines that\n", - "have been carefully crafted to automate lens modeling of large samples whilst ensuring models of the highest\n", - "complexity can be reliably fitted.\n", - "\n", - "These pipelines are built around the use of adaptive features -- for example the Source pipeline comes first so that\n", - "these features are set up robustly before more complex lens light and mass models are fitted." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# %%\n", - "'''\n", - "__SOURCE LP PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, using an MGE for the lens and source light profiles.\n", - "\n", - "Note that unlike the other interferometer SLaM scripts, this Delaunay script does include a source_lp pipeline.\n", - "Its result provides adapt images for source_pix_1, which are used to initialise the Delaunay image mesh.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisInterferometer(dataset=dataset)\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " # interferometry does not support lens light\n", - " bulge=None,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", - "\n", - "The `source_pix_1` search uses an `Overlay` image-mesh to place the initial Delaunay mesh pixels, with\n", - "additional edge points added around the mask boundary to ensure full coverage.\n", - "\n", - "Adapt images from the source LP result provide the initial image-plane mesh grid via `AdaptImages`, and\n", - "positions from the source LP result constrain the mass model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_lp_result: af.Result,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result, use_model_images=True\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask,\n", - " )\n", - "\n", - " edge_pixels_total = 30\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=dataset.mask.mask_centre,\n", - " radius=mask_radius + dataset.mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0],\n", - " zeroed_pixels=edge_pixels_total,\n", - " ),\n", - " regularization=al.reg.ConstantSplit,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", - "\n", - "The `source_pix_2` search uses a `Hilbert` image-mesh to place the final Delaunay mesh pixels, which adapts\n", - "the mesh to the source morphology using the high-quality adapt images from search 1." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " image_mesh = al.image_mesh.Hilbert(pixels=1000, weight_power=3.5, weight_floor=0.01)\n", - "\n", - " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask, adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", - " )\n", - "\n", - " edge_pixels_total = 30\n", - "\n", - " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - " centre=dataset.mask.mask_centre,\n", - " radius=mask_radius + dataset.mask.pixel_scale / 2.0,\n", - " n_points=edge_pixels_total,\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - " )\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " settings=settings,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " # interferometry does not support lens light\n", - " bulge=None,\n", - " disk=None,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0],\n", - " zeroed_pixels=edge_pixels_total,\n", - " ),\n", - " regularization=al.reg.AdaptSplit,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", - "contain lens light emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " source_pix_result_2: af.Result,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(\n", - " galaxy_name_image_dict=galaxy_image_name_dict,\n", - " galaxy_name_image_plane_mesh_grid_dict=(\n", - " source_pix_result_2.analysis.adapt_images.galaxy_name_image_plane_mesh_grid_dict\n", - " ),\n", - " )\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_1.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=source_pix_result_1.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"interferometer\") / \"slam_delaunay\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_lp_result=source_lp_result,\n", - " settings=settings,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " mask_radius=mask_radius,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " settings=settings,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " settings=settings,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function: Pixelization__\n", - "\n", - "This script provides a step-by-step guide of the **PyAutoLens** `log_likelihood_function` which is used to fit\n", - "`Interferometer` data with an inversion (specifically a `Delaunay` mesh and `Constant` regularization scheme`).\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can understand the\n", - " likelihood function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - " - To make inversions in **PyAutoLens** less of a \"black-box\" to users.\n", - "\n", - "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", - "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", - "packages are called when the likelihood is evaluated.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function of pixelizations is the most complicated likelihood function.\n", - "\n", - "It is advised you read through the following two simpler likelihood functions first, which break down a number of the\n", - "concepts used in this script:\n", - "\n", - " - `interferometer/light_profile/log_likelihood_function.py` the likelihood function for a light profile.\n", - " - `interferometer/linear_light_profile/log_likelihood_function.py` the likelihood function for a linear light profile, which\n", - " introduces the linear algebra used for a pixelization but with a simpler use case.\n", - "\n", - "This script repeats all text and code examples in the above likelihood function examples. It therefore can be used to\n", - "learn about the linear light profile likelihood function without reading other likelihood scripts.\n", - "\n", - "__Likelihood Function__\n", - "\n", - "The example `interferometer/pixelization/likelihood_function.py` provides a step-by-step description of how\n", - "a likelihood evaluation is performed for interferometer data using a pixelized source reconstruction with a rectangular\n", - "mesh.\n", - "\n", - "We now give the same step-by-step description for a pixelized source reconstruction using a Delaunay mesh and\n", - "adaptive features.\n", - "\n", - "We only describe code which is specific to Delaunay meshes and adaptive features -- for all other aspects of the likelihood\n", - "evaluation, refer to rectangular mesh example.\n", - "\n", - "__Mask__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(80, 80), pixel_scales=0.05, radius=4.0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)\n", - "\n", - "aplt.plot_grid(grid=dataset.grids.pixelization, title=\"\")\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Pixelization and Regularization__\n", - "\n", - "We combine the pixelization into a single `Galaxy` object.\n", - "\n", - "The galaxy includes the Delaunay mesh and constant regularization scheme, which will ultimately be used\n", - "to reconstruct its star forming clumps.\n", - "\n", - "One of the biggest differences between a Delaunay mesh and rectangular mesh is how the centres of the mesh pixels\n", - "in the source-plane are computed. \n", - "\n", - "For the rectangular mesh, the pixel centres are computed by overlaying a uniform grid over the source-plane.\n", - "\n", - "For a Delaunay mesh, the uniform grid is instead laid over the image-plane to create a course grid of (y,x) coordinates.\n", - "These are then ray-traced to the source-plane and are used as the vertexes of the Delaunay triangles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.Delaunay(\n", - " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", - " ),\n", - " regularization=al.reg.ConstantSplit(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Pixel Centre Calculation__\n", - "\n", - "In order to reconstruct the source galaxy using a Delaunay mesh, we need to determine the centres of the Delaunay\n", - "source pixels.\n", - "\n", - "The image-mesh `Overlay` object computes the source-pixel centres in the image-plane (which are ray-traced to the\n", - "source-plane below). The source pixelization therefore adapts to the lens model magnification, because more\n", - "source pixels will congregate in higher magnification regions.\n", - "\n", - "This calculation is performed by overlaying a uniform regular grid with an `pixelization_shape_2d` over the image\n", - "mask and retaining all pixels that fall within the mask. This uses a `Grid2DSparse` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_mesh = al.image_mesh.Overlay(shape=(30, 30)) # Specific to Delaunay\n", - "\n", - "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", - " mask=dataset.mask,\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(\n", - " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid},\n", - " galaxy_name_image_plane_mesh_grid_dict={\n", - " \"('galaxies', 'source')\": image_plane_mesh_grid\n", - " },\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting this grid shows a sparse grid of (y,x) coordinates within the mask, which will form our source pixel centres." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Ray Tracing__\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The source code gets quite complex when handling grids for a pixelization, but it is all handled in\n", - "the `TracerToInversion` objects.\n", - "\n", - "The plots at the bottom of this cell show the traced grids used by the source pixelization, showing\n", - "how the Delaunay mesh and traced image pixels are constructed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(\n", - " tracer=tracer, dataset=dataset, adapt_images=adapt_images\n", - ")\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", - "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", - " grid=dataset.grids.pixelization\n", - ")[-1]\n", - "\n", - "# This functions a bit weird - it returns a list of lists of ndarrays. Best not to worry about it for now!\n", - "traced_mesh_grid = tracer_to_inversion.traced_mesh_grid_pg_list[-1][-1]\n", - "\n", - "\n", - "aplt.plot_grid(grid=traced_grid_pixelization, title=\"\")\n", - "\n", - "aplt.plot_grid(grid=traced_mesh_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Border Relocation__\n", - "\n", - "Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace to far outskirts of\n", - "the source-plane. \n", - "\n", - "Border relocation is performed on both the traced image-pixel grid and traced mesh pixels, therefore ensuring that\n", - "the vertexes of the Delaunay triangles are not at the extreme outskirts of the source-plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", - "\n", - "border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)\n", - "\n", - "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", - "\n", - "relocated_mesh_grid = border_relocator.relocated_mesh_grid_from(\n", - " grid=traced_grid_pixelization, mesh_grid=traced_mesh_grid\n", - ")\n", - "\n", - "\n", - "aplt.plot_grid(grid=relocated_grid, title=\"\")\n", - "\n", - "aplt.plot_grid(grid=relocated_mesh_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Delaunay Mesh__\n", - "\n", - "The relocated mesh grid is used to create the `Pixelization`'s Delaunay mesh using the `scipy.spatial` library." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "interpolator = al.InterpolatorDelaunay(\n", - " mesh=pixelization.mesh,\n", - " mesh_grid=relocated_mesh_grid,\n", - " data_grid=relocated_grid,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting the Delaunay mesh shows that the source-plane and been discretized into a grid of irregular Delaunay pixels.\n", - "\n", - "(To plot the Delaunay mesh, we have to convert it to a `Mapper` object, which is described in the next likelihood step).\n", - "\n", - "Below, we plot the Delaunay mesh without the traced image-grid pixels (for clarity) and with them as black dots in order\n", - "to show how each set of image-pixels fall within a Delaunay pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = al.Mapper(\n", - " interpolator=interpolator,\n", - " image_plane_mesh_grid=image_plane_mesh_grid,\n", - ")\n", - "\n", - "# mapper_plotter.figure_2d()\n", - "#\n", - "# grid=mapper.source_plane_data_grid,\n", - "# )\n", - "# mapper_plotter.figure_2d()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolation__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pix_indexes_for_sub_slim_index = mapper.pix_indexes_for_sub_slim_index\n", - "\n", - "print(pix_indexes_for_sub_slim_index[0:9])\n", - "\n", - "\n", - "aplt.plot_array(\n", - " array=dataset.dirty_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", - ")\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "pix_indexes = [[200]]\n", - "\n", - "indexes = mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)\n", - "\n", - "\n", - "aplt.plot_array(\n", - " array=dataset.dirty_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", - ")\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", - " pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,\n", - " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index, # unused for Delaunay\n", - " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index, # unused for Delaunay\n", - " pixels=mapper.pixels,\n", - " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", - " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", - " sub_fraction=mapper.over_sampler.sub_fraction,\n", - ")\n", - "\n", - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "indexes_source_pix_200 = np.nonzero(mapping_matrix[:, 200])\n", - "\n", - "print(indexes_source_pix_200[0])\n", - "\n", - "array_2d = al.Array2D(values=mapping_matrix[:, 200], mask=dataset.mask)\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")\n", - "\n", - "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", - " mapping_matrix=mapping_matrix\n", - ")\n", - "\n", - "plt.imshow(\n", - " transformed_mapping_matrix.real,\n", - " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "plt.imshow(\n", - " transformed_mapping_matrix.imag,\n", - " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "indexes_pix_200 = np.nonzero(transformed_mapping_matrix[:, 200])\n", - "\n", - "print(indexes_pix_200[0])\n", - "\n", - "visibilities = al.Visibilities(visibilities=transformed_mapping_matrix[:, 200])\n", - "\n", - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n", - "\n", - "print(f\"Mapping between visibility 0 and Delaunay pixel 2 = {mapping_matrix[0, 2]}\")\n", - "\n", - "data_vector = (\n", - " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " visibilities=dataset.data,\n", - " noise_map=dataset.noise_map,\n", - " )\n", - ")\n", - "\n", - "plt.imshow(\n", - " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "print(\"Data Vector:\")\n", - "print(data_vector)\n", - "print(data_vector.shape)\n", - "\n", - "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.real,\n", - " noise_map=dataset.noise_map.real,\n", - ")\n", - "\n", - "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.imag,\n", - " noise_map=dataset.noise_map.imag,\n", - ")\n", - "\n", - "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", - "\n", - "plt.imshow(curvature_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "source_pixel_0 = 0\n", - "source_pixel_1 = 1\n", - "\n", - "print(curvature_matrix[source_pixel_0, source_pixel_1])\n", - "\n", - "visibilities = al.Visibilities(\n", - " visibilities=transformed_mapping_matrix[:, source_pixel_0],\n", - ")\n", - "\n", - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n", - "\n", - "visibilities = al.Visibilities(\n", - " visibilities=transformed_mapping_matrix[:, source_pixel_1],\n", - ")\n", - "\n", - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n", - "\n", - "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", - " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", - " neighbors=mapper.neighbors,\n", - " neighbors_sizes=mapper.neighbors.sizes,\n", - ")\n", - "\n", - "plt.imshow(regularization_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)\n", - "\n", - "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)\n", - "\n", - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "mapped_reconstructed_visibilities = (\n", - " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " reconstruction=reconstruction,\n", - " )\n", - ")\n", - "\n", - "mapped_reconstructed_visibilities = al.Visibilities(\n", - " visibilities=mapped_reconstructed_visibilities\n", - ")\n", - "\n", - "aplt.plot_grid(grid=mapped_reconstructed_visibilities.in_grid, title=\"\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_visibilities = mapped_reconstructed_visibilities\n", - "\n", - "residual_map = dataset.data - model_visibilities\n", - "\n", - "\n", - "normalized_residual_map_real = (residual_map.real / dataset.noise_map.real).astype(\n", - " \"complex128\"\n", - ")\n", - "normalized_residual_map_imag = (residual_map.imag / dataset.noise_map.imag).astype(\n", - " \"complex128\"\n", - ")\n", - "normalized_residual_map = (\n", - " normalized_residual_map_real + 1j * normalized_residual_map_imag\n", - ")\n", - "\n", - "\n", - "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", - "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", - "chi_squared_map = chi_squared_map_real + 1j * chi_squared_map_imag\n", - "\n", - "\n", - "chi_squared_real = np.sum(chi_squared_map.real)\n", - "chi_squared_imag = np.sum(chi_squared_map.imag)\n", - "chi_squared = chi_squared_real + chi_squared_imag\n", - "\n", - "print(chi_squared)\n", - "\n", - "chi_squared_map = al.Visibilities(visibilities=chi_squared_map)\n", - "\n", - "aplt.plot_grid(grid=chi_squared_map.in_grid, title=\"\")\n", - "\n", - "regularization_term = np.matmul(\n", - " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", - ")\n", - "\n", - "print(regularization_term)\n", - "\n", - "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", - "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]\n", - "\n", - "print(log_curvature_reg_matrix_term)\n", - "print(log_regularization_matrix_term)\n", - "\n", - "noise_normalization_real = np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", - "noise_normalization_imag = np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", - "noise_normalization = noise_normalization_real + noise_normalization_imag\n", - "\n", - "log_evidence = float(\n", - " -0.5\n", - " * (\n", - " chi_squared\n", - " + regularization_term\n", - " + log_curvature_reg_matrix_term\n", - " - log_regularization_matrix_term\n", - " + noise_normalization\n", - " )\n", - ")\n", - "\n", - "print(log_evidence)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This process to perform a likelihood function evaluation performed via the `FitInterferometer` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitInterferometer(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " adapt_images=adapt_images,\n", - " settings=al.Settings(use_border_relocator=True),\n", - ")\n", - "fit_log_evidence = fit.log_evidence\n", - "print(fit_log_evidence)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "__Log Likelihood Function: Source Code Speed Up__\n", - "\n", - "The interferometer pixelization likelihood function described in this notebook performs certain calculations using\n", - "functions which are easier to understand, but are computationally slower than the actual source code implementation\n", - "(but the two produce identical results).\n", - "\n", - "We end by pointing out some of these, but we do not provide an step-by-step description of how they work.\n", - "If you are interested, you will need to dive into the source code itself.\n", - "\n", - "**Fast Chi Squared:** The `chi_squared` above is computed using the `transformed_mapping_matrix`, which requires\n", - "many NUFFT's to compute and requires large memroy store. The source code uses a trick which computes the chi-squared\n", - "but bypasses the need to ever compute the `transformed_mapping_matrix`.\n", - "\n", - "**Sparse Operator Curvature Matrix:** The `curvature_matrix` above is also computed using the `transformed_mapping_matrix`, \n", - "which again means slow run times and large memory usage. The source code can instead use sparse operators to \n", - "compute the curvature matrix in a way which again bypasses the need to compute the `transformed_mapping_matrix`.\n", - "\n", - "The two tricks in combination lead to a significant speed up in the likelihood function evaluation and mean that\n", - "the large matrix of size [source pixels, visibilities] never needs to be stored in memory. This is at the heart\n", - "of why lens modeling interferometer data with pixelized source reconstructions is so fast!\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the pixelization likelihood function.\n", - "\n", - "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", - "are described in additional notebooks found in this package. In brief, these describe:\n", - "\n", - " - **Over Sampling**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", - " paired fractionally with each Delaunay pixel.\n", - "\n", - " - **Source-plane Interpolation**: Using bilinear interpolation on the Delaunay pixelization to pair each \n", - " image (sub-)pixel to multiple Delaunay pixels with interpolation weights.\n", - "\n", - " - **Luminosity Weighted Regularization**: Using an adaptive regularization coefficient which adapts the level of \n", - " regularization applied to the source galaxy based on its luminosity." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: Delaunay\n", + "======================\n", + "\n", + "The majority of pixelized source reconstructions in the workspace use a rectangular mesh to reconstruct\n", + "the source's surface brightness.\n", + "\n", + "This example illustrates an alternative pixelization that uses a Delaunay triangulation mesh to reconstruct the\n", + "source.\n", + "\n", + "The approach is distinct from the rectangular mesh and has a number of traits which are unique to it:\n", + "\n", + "- `Adaptive Mesh`: In the source plane, the Delaunay mesh uses irregularly shaped triangles to reconstruct the\n", + " source, as opposed to uniform rectangular pixels. This allows the mesh to better adapt to irregular and\n", + " asymmetric source morphologies and change the distribution of source pixels to better match the source's\n", + " surface brightness.\n", + "\n", + "- `Image Mesh`: The vertexes of the Delaunay triangles are computed by overlaying a coarse uniform grid in the\n", + " image-plane and ray-tracing these coordinates to the source-plane. This is unlike the rectangular mesh, which\n", + " simply overlays a uniform grid in the source-plane. This again helps the Delaunay mesh to better adapt to the\n", + " source's surface brightness.\n", + "\n", + "- `Interpolation`: The Delaunay mesh uses a different interpolation scheme to the rectangular mesh, which is\n", + " barycentric interpolation within each triangle. This is different to the rectangular mesh, which uses bilinear\n", + " interpolation within each rectangular pixel.\n", + "\n", + "- `Regularization`: The Delaunay mesh provides different approaches to regularization, with the default being\n", + " one which uses the barycentric coordinates of the triangles to compute how source pixels are regularized with\n", + " their neighbors.\n", + "\n", + "Currently it is not expected that the Delaunay is better or worse than the rectangular mesh, it is simply a different\n", + "approach to pixelization that may work better for certain datasets.\n", + "\n", + "__JAX + GPU__\n", + "\n", + "Generating a Delaunay mesh supports JAX and GPU acceleration, however certain operations (e.g. generating the Delaunay\n", + "triangulation itself) do not run on the GPU because they cannot be easily converted to JAX.\n", + "\n", + "Instead, JAX sends them to a CPU, runs them there, and then sends the results back to the GPU. This process is\n", + "very efficient, because these operations run very fast on a CPU and the data being sent back and forth is small.\n", + "Current benchmarking suggests the Delaunay runs less than twice as long as the same fit using a rectangular mesh,\n", + "but scientfically offers better results in many cases.\n", + "\n", + "If you do want to run only on CPU, you can use fast CPU method described in\n", + "example `imaging/features/pixelization/cpu_fast_modeling` with the Delaunay mesh.\n", + "\n", + "\n", + "__Source Science (Magnification, Flux and More)__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", + "size of the source.\n", + "\n", + "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", + "including examples using a Delaunay source reconstruction. Once you have completed lens modeling using a Delaunay mesh,\n", + "you can jump to that example to study the source galaxy.\n", + "\n", + "__Contents__\n", + "\n", + "- **Image Mesh:** For a Delaunay mesh, the vertices of the triangles are defined by (y, x) coordinates in the.\n", + "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **VRAM:** The `pixelization/modeling` example explains how VRAM use is an important consideration for.\n", + "- **Adaptive Delaunay:** The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to.\n", + "- **SLaM Pipelines:** The API above allows you to use adaptive features yourself, and you should go ahead an explore them.\n", + "- **SOURCE LP PIPELINE:** Identical to `slam_start_here.py`, using an MGE for the lens and source light profiles.\n", + "- **SOURCE PIX PIPELINE 1:** Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", + "- **SOURCE PIX PIPELINE 2:** Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included as interferometer data.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "- **Prerequisites:** The likelihood function of pixelizations is the most complicated likelihood function.\n", + "- **Likelihood Function:** The example `interferometer/pixelization/likelihood_function.py` provides a step-by-step.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Source Galaxy Pixelization and Regularization:** We combine the pixelization into a single `Galaxy` object.\n", + "- **Source Pixel Centre Calculation:** In order to reconstruct the source galaxy using a Delaunay mesh, we need to determine the centres.\n", + "- **Ray Tracing:** Overview of ray tracing for this example.\n", + "- **Border Relocation:** Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace.\n", + "- **Delaunay Mesh:** The relocated mesh grid is used to create the `Pixelization`'s Delaunay mesh using the.\n", + "- **Interpolation:** Overview of interpolation for this example.\n", + "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "\n", + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "dataset = dataset.apply_sparse_operator(use_jax=True, show_progress=True)\n", + "\n", + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n", + "\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)\n", + "\n", + "\n", + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Image Mesh__\n", + "\n", + "For a Delaunay mesh, the vertices of the triangles are defined by (y, x) coordinates in the image-plane. These\n", + "coordinates are then ray-traced into the source-plane for each mass model sampled during the non-linear search.\n", + "This `image_plane_mesh_grid` must be computed before lens modeling.\n", + "\n", + "We compute this `image_plane_mesh_grid` using an `Overlay` image-mesh, which places a regular grid of\n", + "(y, x) points across the image-plane. This has a mild adaptive effect: regions of high lens magnification receive\n", + "more source pixels once they are ray-traced. Later in this example, we switch to a `Hilbert` image-mesh, which adapts\n", + "the pixel distribution more strongly to the source\u2019s surface brightness.\n", + "\n", + "The `Delaunay` mesh has an input number of `pixels`, which is the number of source pixels used to reconstruct the \n", + "source. The number of `pixels` must be equal to the number of coordinates in the `image_plane_mesh_grid`. \n", + "\n", + "Like for the `mesh_shape` rectangular mesh, `pixels` must be fixed for lens modeling because JAX uses the \n", + "number of `pixels` to determine static array shapes. \n", + "\n", + "To pass the `image_plane_mesh_grid` to the modeling, we use the `AdaptImages` object below, which pairs\n", + "the `image_plane_mesh_grid` to the source galaxy. For double source plane lenses, this means we can\n", + "attach an `image_plane_mesh_grid` to each source galaxy and use adaptive meshes for each source plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Edge Zeroing__\n", + "\n", + "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", + "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", + "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", + "\n", + "For a rectangular mesh, the source code computes edge pixels internally using the known pixels at the edge of the mesh,\n", + "requiring no input from the user. \n", + "\n", + "For the `Delaunay` mesh, we use the `append_with_circle_edge_points` function to manually setup the Delaunay image \n", + "mesh to include a ring of edge pixels and then input the total number into the mesh to perform zeroing. \n", + "\n", + "These points are added to the edge of the image-plane mesh, ray-traced to the source-plane during lens modeling, \n", + "included in the Delaunay triangulation but zeroed during the inversion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "edge_pixels_total = 30\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=real_space_mask.mask_centre,\n", + " radius=mask_radius + real_space_mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "In the example `interferometer/features/pixelization/fit.py`, we illustrate how to use a pixelized source\n", + "with a rectangular mesh.\n", + "\n", + "Below, we use a Delaunay mesh to perform a fit using the Delaunay source reconstruction.\n", + "\n", + "The API is nearly identical to the rectangular mesh example, noting that the inputs to the `Delaunay` \n", + "mesh are different to the rectangular mesh and use image mesh quantities computed above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh = al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + ")\n", + "regularization = al.reg.ConstantSplit(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", + "\n", + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "adapt_images = al.AdaptImages(\n", + " galaxy_image_plane_mesh_grid_dict={source: image_plane_mesh_grid}\n", + ")\n", + "\n", + "fit = al.FitInterferometer(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " adapt_images=adapt_images,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the Delaunay source does a good job at capturing the appearance of the source galaxy\n", + "using adaptive triangular pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We now perform lens modeling using the Delaunay pixelization with the Overlay image-mesh.\n", + "\n", + "The code below is a simple adaptive modeling example using the Delaunay mesh, which mirrors the\n", + "API used in other pixelization modeling examples.\n", + "\n", + "The example `interferometer/features/pixelization/adaptive.py` illustrates how to use adaptive features to\n", + "adapt the rectangular mesh and its regularization to the source's surface brightness. In particular, an image\n", + "of the lensed source is passed to the modeling via the `AdaptImages` object, in order to adapt\n", + "the mesh and regularization during the model-fit.\n", + "\n", + "The same object is used to pass the `image_plane_mesh_grid` to the modeling. Above, this image-plane mesh grid\n", + "is an `Overlay` mesh and does not specifically adapt to the source's surface brightness, thus pairing it with\n", + "the source as done below seems redundant. However, in a moment we will switch to a `Hilbert` image-mesh, which\n", + "does adapt to the source's surface brightness, meaning this pairing is necessary." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "adapt_images = al.AdaptImages(\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In the first\n", + "search our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light uses an `Overlay` image-mesh with fixed resolution 30 x 30 pixels [0 parameters].\n", + "\n", + " - The source-galaxy's light uses a `Delaunay` mesh [0 parameters].\n", + "\n", + " - This pixelization is regularized using a `Constant` scheme [1 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=8." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", + ")\n", + "\n", + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.ConstantSplit,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "model_1 = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "search_1 = af.Nautilus(\n", + " path_prefix=Path(\"features\"),\n", + " name=\"delaunay\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + ")\n", + "\n", + "analysis_1 = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[al.PositionsLH(positions=positions, threshold=0.3)],\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `pixelization/modeling` example explains how VRAM use is an important consideration for pixelization models\n", + "and how it depends on image resolution, number of source pixels and batch size.\n", + "\n", + "This is true for the Delaunay mesh, therefore we print out the estimated VRAM required for this model-fit.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_1.print_vram_use(model=model_1, batch_size=search_1.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit (Search 1)__\n", + "\n", + "Perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result_1 = search_1.fit(model=model_1, analysis=analysis_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Adaptive Delaunay__\n", + "\n", + "The example `imaging/features/pixelization/adaptive.py` illustrates how to use adaptive features to\n", + "adapt the rectangular mesh and its regularization to the source's surface brightness.\n", + "\n", + "The image-mesh has a special adaptive variant called the `Hilbert` image-mesh, which adapts the distribution \n", + "of source-pixels to the source's unlensed morphology. This means that the source's brightest regions are \n", + "reconstructed using significantly more source pixels than seen for the `Overlay` image mesh. \n", + "Conversely, the source's faintest regions are reconstructed using significantly fewer source pixels.\n", + "\n", + "Unlike the adaptive rectangular mesh, the Hilbert image-plane mesh is computed before modeling, passed\n", + "to the `AdaptImages` object, and remains fixed during the model-fit.\n", + "\n", + "It is recommend that the parameters governing these features are always fitted from using a fixed lens light and\n", + "mass model. This ensures the adaptation is performed quickly, and removes degeneracies in the lens model that\n", + "are difficult to sample. Given the Hilbert mesh is fixed, this modeling only fits for the regularization coefficients\n", + "of the adaptive regularization scheme.\n", + "\n", + "For this reason, search 2 fixes the lens galaxy's light and mass model to the best-fit model of search 1. A third\n", + "search will then fit for the lens galaxy's light and mass model using these adaptive features.\n", + "\n", + "The details of how the above features work is not provided here, but is given at the end of chapter 4 of the HowToLens\n", + "lecture series." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=result_1, use_model_images=True\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result_1)\n", + "\n", + "image_mesh = al.image_mesh.Hilbert(pixels=1000, weight_power=3.5, weight_floor=0.01)\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask, adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", + ")\n", + "\n", + "# Repeat edge zeroing set up describe above.\n", + "\n", + "edge_pixels_total = 30\n", + "\n", + "image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=real_space_mask.mask_centre,\n", + " radius=mask_radius + real_space_mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model (Search 2)__\n", + "\n", + "We therefore compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In \n", + "the second search our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` with fixed parameters from \n", + " search 1 [0 parameters].\n", + "\n", + " - The source galaxy's light uses a `Hilbert` image-mesh with fixed resolution 1000 pixels [2 parameters].\n", + "\n", + " - The source-galaxy's light uses a `Delaunay` mesh [0 parameters].\n", + "\n", + " - This pixelization is regularized using a `AdaptSplit` scheme [2 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=4." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.AdaptSplit,\n", + ")\n", + "\n", + "source = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " pixelization=pixelization,\n", + ")\n", + "\n", + "model_2 = af.Collection(\n", + " galaxies=af.Collection(lens=result_1.instance.galaxies.lens, source=source)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis (Search 2)__\n", + "\n", + "We now create the analysis for the second search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_2 = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search + Model-Fit (Search 2)__\n", + "\n", + "We now create the non-linear search and perform the model-fit using this model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search_2 = af.Nautilus(\n", + " path_prefix=Path(\"features\"),\n", + " name=\"delaunay_adapt\",\n", + " unique_tag=dataset_name,\n", + " n_live=75,\n", + ")\n", + "\n", + "result_2 = search_2.fit(model=model_2, analysis=analysis_2)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We could perform a third fit where we free all lens model parameters and fit them using the adaptive \n", + "image mesh and regularization.\n", + "\n", + "However, it is better to use all of these features with the Delaunay via the\n", + "SLaM pipelines, which we jump to immediately below.\n", + "\n", + "__SLaM Pipelines__\n", + "\n", + "The API above allows you to use adaptive features yourself, and you should go ahead an explore them on datasets you\n", + "are familiar with.\n", + "\n", + "However, you may also wish to use the Source, Light and Mass (SLaM) pipelines, which are pipelines that\n", + "have been carefully crafted to automate lens modeling of large samples whilst ensuring models of the highest\n", + "complexity can be reliably fitted.\n", + "\n", + "These pipelines are built around the use of adaptive features -- for example the Source pipeline comes first so that\n", + "these features are set up robustly before more complex lens light and mass models are fitted." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# %%\n", + "'''\n", + "__SOURCE LP PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, using an MGE for the lens and source light profiles.\n", + "\n", + "Note that unlike the other interferometer SLaM scripts, this Delaunay script does include a source_lp pipeline.\n", + "Its result provides adapt images for source_pix_1, which are used to initialise the Delaunay image mesh.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisInterferometer(dataset=dataset)\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " # interferometry does not support lens light\n", + " bulge=None,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", + "\n", + "The `source_pix_1` search uses an `Overlay` image-mesh to place the initial Delaunay mesh pixels, with\n", + "additional edge points added around the mask boundary to ensure full coverage.\n", + "\n", + "Adapt images from the source LP result provide the initial image-plane mesh grid via `AdaptImages`, and\n", + "positions from the source LP result constrain the mass model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_lp_result: af.Result,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result, use_model_images=True\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Overlay(shape=(26, 26))\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask,\n", + " )\n", + "\n", + " edge_pixels_total = 30\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=dataset.mask.mask_centre,\n", + " radius=mask_radius + dataset.mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0],\n", + " zeroed_pixels=edge_pixels_total,\n", + " ),\n", + " regularization=al.reg.ConstantSplit,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, except the source pixelization uses a Delaunay mesh.\n", + "\n", + "The `source_pix_2` search uses a `Hilbert` image-mesh to place the final Delaunay mesh pixels, which adapts\n", + "the mesh to the source morphology using the high-quality adapt images from search 1." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " image_mesh = al.image_mesh.Hilbert(pixels=1000, weight_power=3.5, weight_floor=0.01)\n", + "\n", + " image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask, adapt_data=galaxy_image_name_dict[\"('galaxies', 'source')\"]\n", + " )\n", + "\n", + " edge_pixels_total = 30\n", + "\n", + " image_plane_mesh_grid = al.image_mesh.append_with_circle_edge_points(\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + " centre=dataset.mask.mask_centre,\n", + " radius=mask_radius + dataset.mask.pixel_scale / 2.0,\n", + " n_points=edge_pixels_total,\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + " )\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " settings=settings,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " # interferometry does not support lens light\n", + " bulge=None,\n", + " disk=None,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0],\n", + " zeroed_pixels=edge_pixels_total,\n", + " ),\n", + " regularization=al.reg.AdaptSplit,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", + "contain lens light emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " source_pix_result_2: af.Result,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(\n", + " galaxy_name_image_dict=galaxy_image_name_dict,\n", + " galaxy_name_image_plane_mesh_grid_dict=(\n", + " source_pix_result_2.analysis.adapt_images.galaxy_name_image_plane_mesh_grid_dict\n", + " ),\n", + " )\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_1.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=source_pix_result_1.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"interferometer\") / \"slam_delaunay\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_lp_result=source_lp_result,\n", + " settings=settings,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " mask_radius=mask_radius,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " settings=settings,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " settings=settings,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function: Pixelization__\n", + "\n", + "This script provides a step-by-step guide of the **PyAutoLens** `log_likelihood_function` which is used to fit\n", + "`Interferometer` data with an inversion (specifically a `Delaunay` mesh and `Constant` regularization scheme`).\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can understand the\n", + " likelihood function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + " - To make inversions in **PyAutoLens** less of a \"black-box\" to users.\n", + "\n", + "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", + "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", + "packages are called when the likelihood is evaluated.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function of pixelizations is the most complicated likelihood function.\n", + "\n", + "It is advised you read through the following two simpler likelihood functions first, which break down a number of the\n", + "concepts used in this script:\n", + "\n", + " - `interferometer/light_profile/log_likelihood_function.py` the likelihood function for a light profile.\n", + " - `interferometer/linear_light_profile/log_likelihood_function.py` the likelihood function for a linear light profile, which\n", + " introduces the linear algebra used for a pixelization but with a simpler use case.\n", + "\n", + "This script repeats all text and code examples in the above likelihood function examples. It therefore can be used to\n", + "learn about the linear light profile likelihood function without reading other likelihood scripts.\n", + "\n", + "__Likelihood Function__\n", + "\n", + "The example `interferometer/pixelization/likelihood_function.py` provides a step-by-step description of how\n", + "a likelihood evaluation is performed for interferometer data using a pixelized source reconstruction with a rectangular\n", + "mesh.\n", + "\n", + "We now give the same step-by-step description for a pixelized source reconstruction using a Delaunay mesh and\n", + "adaptive features.\n", + "\n", + "We only describe code which is specific to Delaunay meshes and adaptive features -- for all other aspects of the likelihood\n", + "evaluation, refer to rectangular mesh example.\n", + "\n", + "__Mask__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(80, 80), pixel_scales=0.05, radius=4.0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)\n", + "\n", + "aplt.plot_grid(grid=dataset.grids.pixelization, title=\"\")\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Pixelization and Regularization__\n", + "\n", + "We combine the pixelization into a single `Galaxy` object.\n", + "\n", + "The galaxy includes the Delaunay mesh and constant regularization scheme, which will ultimately be used\n", + "to reconstruct its star forming clumps.\n", + "\n", + "One of the biggest differences between a Delaunay mesh and rectangular mesh is how the centres of the mesh pixels\n", + "in the source-plane are computed. \n", + "\n", + "For the rectangular mesh, the pixel centres are computed by overlaying a uniform grid over the source-plane.\n", + "\n", + "For a Delaunay mesh, the uniform grid is instead laid over the image-plane to create a course grid of (y,x) coordinates.\n", + "These are then ray-traced to the source-plane and are used as the vertexes of the Delaunay triangles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.Delaunay(\n", + " pixels=image_plane_mesh_grid.shape[0], zeroed_pixels=edge_pixels_total\n", + " ),\n", + " regularization=al.reg.ConstantSplit(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Pixel Centre Calculation__\n", + "\n", + "In order to reconstruct the source galaxy using a Delaunay mesh, we need to determine the centres of the Delaunay\n", + "source pixels.\n", + "\n", + "The image-mesh `Overlay` object computes the source-pixel centres in the image-plane (which are ray-traced to the\n", + "source-plane below). The source pixelization therefore adapts to the lens model magnification, because more\n", + "source pixels will congregate in higher magnification regions.\n", + "\n", + "This calculation is performed by overlaying a uniform regular grid with an `pixelization_shape_2d` over the image\n", + "mask and retaining all pixels that fall within the mask. This uses a `Grid2DSparse` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_mesh = al.image_mesh.Overlay(shape=(30, 30)) # Specific to Delaunay\n", + "\n", + "image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from(\n", + " mask=dataset.mask,\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(\n", + " galaxy_image_plane_mesh_grid_dict={source_galaxy: image_plane_mesh_grid},\n", + " galaxy_name_image_plane_mesh_grid_dict={\n", + " \"('galaxies', 'source')\": image_plane_mesh_grid\n", + " },\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting this grid shows a sparse grid of (y,x) coordinates within the mask, which will form our source pixel centres." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Ray Tracing__\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The source code gets quite complex when handling grids for a pixelization, but it is all handled in\n", + "the `TracerToInversion` objects.\n", + "\n", + "The plots at the bottom of this cell show the traced grids used by the source pixelization, showing\n", + "how the Delaunay mesh and traced image pixels are constructed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(\n", + " tracer=tracer, dataset=dataset, adapt_images=adapt_images\n", + ")\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", + "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", + " grid=dataset.grids.pixelization\n", + ")[-1]\n", + "\n", + "# This functions a bit weird - it returns a list of lists of ndarrays. Best not to worry about it for now!\n", + "traced_mesh_grid = tracer_to_inversion.traced_mesh_grid_pg_list[-1][-1]\n", + "\n", + "\n", + "aplt.plot_grid(grid=traced_grid_pixelization, title=\"\")\n", + "\n", + "aplt.plot_grid(grid=traced_mesh_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Border Relocation__\n", + "\n", + "Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace to far outskirts of\n", + "the source-plane. \n", + "\n", + "Border relocation is performed on both the traced image-pixel grid and traced mesh pixels, therefore ensuring that\n", + "the vertexes of the Delaunay triangles are not at the extreme outskirts of the source-plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", + "\n", + "border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)\n", + "\n", + "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", + "\n", + "relocated_mesh_grid = border_relocator.relocated_mesh_grid_from(\n", + " grid=traced_grid_pixelization, mesh_grid=traced_mesh_grid\n", + ")\n", + "\n", + "\n", + "aplt.plot_grid(grid=relocated_grid, title=\"\")\n", + "\n", + "aplt.plot_grid(grid=relocated_mesh_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Delaunay Mesh__\n", + "\n", + "The relocated mesh grid is used to create the `Pixelization`'s Delaunay mesh using the `scipy.spatial` library." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "interpolator = al.InterpolatorDelaunay(\n", + " mesh=pixelization.mesh,\n", + " mesh_grid=relocated_mesh_grid,\n", + " data_grid=relocated_grid,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting the Delaunay mesh shows that the source-plane and been discretized into a grid of irregular Delaunay pixels.\n", + "\n", + "(To plot the Delaunay mesh, we have to convert it to a `Mapper` object, which is described in the next likelihood step).\n", + "\n", + "Below, we plot the Delaunay mesh without the traced image-grid pixels (for clarity) and with them as black dots in order\n", + "to show how each set of image-pixels fall within a Delaunay pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = al.Mapper(\n", + " interpolator=interpolator,\n", + " image_plane_mesh_grid=image_plane_mesh_grid,\n", + ")\n", + "\n", + "# mapper_plotter.figure_2d()\n", + "#\n", + "# grid=mapper.source_plane_data_grid,\n", + "# )\n", + "# mapper_plotter.figure_2d()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolation__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pix_indexes_for_sub_slim_index = mapper.pix_indexes_for_sub_slim_index\n", + "\n", + "print(pix_indexes_for_sub_slim_index[0:9])\n", + "\n", + "\n", + "aplt.plot_array(\n", + " array=dataset.dirty_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", + ")\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "pix_indexes = [[200]]\n", + "\n", + "indexes = mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)\n", + "\n", + "\n", + "aplt.plot_array(\n", + " array=dataset.dirty_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", + ")\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", + " pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,\n", + " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index, # unused for Delaunay\n", + " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index, # unused for Delaunay\n", + " pixels=mapper.pixels,\n", + " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", + " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", + " sub_fraction=mapper.over_sampler.sub_fraction,\n", + ")\n", + "\n", + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "indexes_source_pix_200 = np.nonzero(mapping_matrix[:, 200])\n", + "\n", + "print(indexes_source_pix_200[0])\n", + "\n", + "array_2d = al.Array2D(values=mapping_matrix[:, 200], mask=dataset.mask)\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")\n", + "\n", + "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", + " mapping_matrix=mapping_matrix\n", + ")\n", + "\n", + "plt.imshow(\n", + " transformed_mapping_matrix.real,\n", + " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "plt.imshow(\n", + " transformed_mapping_matrix.imag,\n", + " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "indexes_pix_200 = np.nonzero(transformed_mapping_matrix[:, 200])\n", + "\n", + "print(indexes_pix_200[0])\n", + "\n", + "visibilities = al.Visibilities(visibilities=transformed_mapping_matrix[:, 200])\n", + "\n", + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n", + "\n", + "print(f\"Mapping between visibility 0 and Delaunay pixel 2 = {mapping_matrix[0, 2]}\")\n", + "\n", + "data_vector = (\n", + " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " visibilities=dataset.data,\n", + " noise_map=dataset.noise_map,\n", + " )\n", + ")\n", + "\n", + "plt.imshow(\n", + " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "print(\"Data Vector:\")\n", + "print(data_vector)\n", + "print(data_vector.shape)\n", + "\n", + "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.real,\n", + " noise_map=dataset.noise_map.real,\n", + ")\n", + "\n", + "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.imag,\n", + " noise_map=dataset.noise_map.imag,\n", + ")\n", + "\n", + "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", + "\n", + "plt.imshow(curvature_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "source_pixel_0 = 0\n", + "source_pixel_1 = 1\n", + "\n", + "print(curvature_matrix[source_pixel_0, source_pixel_1])\n", + "\n", + "visibilities = al.Visibilities(\n", + " visibilities=transformed_mapping_matrix[:, source_pixel_0],\n", + ")\n", + "\n", + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n", + "\n", + "visibilities = al.Visibilities(\n", + " visibilities=transformed_mapping_matrix[:, source_pixel_1],\n", + ")\n", + "\n", + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n", + "\n", + "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", + " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", + " neighbors=mapper.neighbors,\n", + " neighbors_sizes=mapper.neighbors.sizes,\n", + ")\n", + "\n", + "plt.imshow(regularization_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)\n", + "\n", + "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)\n", + "\n", + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "mapped_reconstructed_visibilities = (\n", + " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " reconstruction=reconstruction,\n", + " )\n", + ")\n", + "\n", + "mapped_reconstructed_visibilities = al.Visibilities(\n", + " visibilities=mapped_reconstructed_visibilities\n", + ")\n", + "\n", + "aplt.plot_grid(grid=mapped_reconstructed_visibilities.in_grid, title=\"\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_visibilities = mapped_reconstructed_visibilities\n", + "\n", + "residual_map = dataset.data - model_visibilities\n", + "\n", + "\n", + "normalized_residual_map_real = (residual_map.real / dataset.noise_map.real).astype(\n", + " \"complex128\"\n", + ")\n", + "normalized_residual_map_imag = (residual_map.imag / dataset.noise_map.imag).astype(\n", + " \"complex128\"\n", + ")\n", + "normalized_residual_map = (\n", + " normalized_residual_map_real + 1j * normalized_residual_map_imag\n", + ")\n", + "\n", + "\n", + "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", + "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", + "chi_squared_map = chi_squared_map_real + 1j * chi_squared_map_imag\n", + "\n", + "\n", + "chi_squared_real = np.sum(chi_squared_map.real)\n", + "chi_squared_imag = np.sum(chi_squared_map.imag)\n", + "chi_squared = chi_squared_real + chi_squared_imag\n", + "\n", + "print(chi_squared)\n", + "\n", + "chi_squared_map = al.Visibilities(visibilities=chi_squared_map)\n", + "\n", + "aplt.plot_grid(grid=chi_squared_map.in_grid, title=\"\")\n", + "\n", + "regularization_term = np.matmul(\n", + " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", + ")\n", + "\n", + "print(regularization_term)\n", + "\n", + "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", + "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]\n", + "\n", + "print(log_curvature_reg_matrix_term)\n", + "print(log_regularization_matrix_term)\n", + "\n", + "noise_normalization_real = np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", + "noise_normalization_imag = np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", + "noise_normalization = noise_normalization_real + noise_normalization_imag\n", + "\n", + "log_evidence = float(\n", + " -0.5\n", + " * (\n", + " chi_squared\n", + " + regularization_term\n", + " + log_curvature_reg_matrix_term\n", + " - log_regularization_matrix_term\n", + " + noise_normalization\n", + " )\n", + ")\n", + "\n", + "print(log_evidence)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This process to perform a likelihood function evaluation performed via the `FitInterferometer` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitInterferometer(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " adapt_images=adapt_images,\n", + " settings=al.Settings(use_border_relocator=True),\n", + ")\n", + "fit_log_evidence = fit.log_evidence\n", + "print(fit_log_evidence)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "__Log Likelihood Function: Source Code Speed Up__\n", + "\n", + "The interferometer pixelization likelihood function described in this notebook performs certain calculations using\n", + "functions which are easier to understand, but are computationally slower than the actual source code implementation\n", + "(but the two produce identical results).\n", + "\n", + "We end by pointing out some of these, but we do not provide an step-by-step description of how they work.\n", + "If you are interested, you will need to dive into the source code itself.\n", + "\n", + "**Fast Chi Squared:** The `chi_squared` above is computed using the `transformed_mapping_matrix`, which requires\n", + "many NUFFT's to compute and requires large memroy store. The source code uses a trick which computes the chi-squared\n", + "but bypasses the need to ever compute the `transformed_mapping_matrix`.\n", + "\n", + "**Sparse Operator Curvature Matrix:** The `curvature_matrix` above is also computed using the `transformed_mapping_matrix`, \n", + "which again means slow run times and large memory usage. The source code can instead use sparse operators to \n", + "compute the curvature matrix in a way which again bypasses the need to compute the `transformed_mapping_matrix`.\n", + "\n", + "The two tricks in combination lead to a significant speed up in the likelihood function evaluation and mean that\n", + "the large matrix of size [source pixels, visibilities] never needs to be stored in memory. This is at the heart\n", + "of why lens modeling interferometer data with pixelized source reconstructions is so fast!\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the pixelization likelihood function.\n", + "\n", + "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", + "are described in additional notebooks found in this package. In brief, these describe:\n", + "\n", + " - **Over Sampling**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", + " paired fractionally with each Delaunay pixel.\n", + "\n", + " - **Source-plane Interpolation**: Using bilinear interpolation on the Delaunay pixelization to pair each \n", + " image (sub-)pixel to multiple Delaunay pixels with interpolation weights.\n", + "\n", + " - **Luminosity Weighted Regularization**: Using an adaptive regularization coefficient which adapts the level of \n", + " regularization applied to the source galaxy based on its luminosity." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/pixelization/fit.ipynb b/notebooks/interferometer/features/pixelization/fit.ipynb index a737a9880..70bfbff74 100644 --- a/notebooks/interferometer/features/pixelization/fit.ipynb +++ b/notebooks/interferometer/features/pixelization/fit.ipynb @@ -1,824 +1,861 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Pixelization\n", - "======================\n", - "\n", - "A pixelization reconstructs the source's light using a pixel-grid, which is regularized using a prior that forces\n", - "the solution to have a degree of smoothness.\n", - "\n", - "This script fits a source galaxy model which uses a pixelization to reconstruct the source's light.\n", - "\n", - "A rectangular mesh which adapts to the lens mass model magnification and constant regularization scheme are used, which\n", - "are the simplest forms of mesh and regularization with provide computationally fast and accurate solutions.\n", - "\n", - "For simplicity, the lens galaxy's light is omitted from the model and is not present in the simulated data. It is\n", - "straightforward to include the lens galaxy's light in the model.\n", - "\n", - "Pixelizations are covered in detail in chapter 4 of the **HowToLens** lectures.\n", - "\n", - "__CPU Users__\n", - "\n", - "Matrices must be set up for a pixelized source reconstruction to speed up the linear algebra. On GPU, this takes\n", - "seconds, or at most a minute for datasets with tens of millions or more visibilities. On CPU, this can be a lot\n", - "slower, taking over an hour for very large datasets. If you are on CPU, the\n", - "`feature/pixelization/many_visibilities_preparation` example explains how this initial setup can be performed\n", - "before lens modeling and saved to hard disk for fast loading before the model fit.\n", - "\n", - "__Contents__\n", - "\n", - "- **CPU Users:** Matrices must be set up for a pixelized source reconstruction which speed up the linear algebra.\n", - "- **Advantages & Disadvantages:** Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Sparse Operators:** Pixelized source modeling requires dense linear algebra operations.\n", - "- **Settings:** As discussed above, disable the default position only linear algebra solver so the source.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", - "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", - "- **Pixelization:** We create a `Pixelization` object to perform the pixelized source reconstruction, which is made up.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "- **Linear Objects:** An `Inversion` contains all of the linear objects used to reconstruct the data in its.\n", - "- **Grids:** The role of a mapper is to map between the image-plane and source-plane.\n", - "- **Reconstruction:** The source reconstruction is also available as a 1D numpy array of values representative of the.\n", - "- **Mapped Reconstructed Images:** The source reconstruction(s) are mapped to the image-plane in order to fit the lens model.\n", - "- **Simulated Interferometer:** We load the source galaxy image from the pixelized inversion of a previous fit, which was performed.\n", - "\n", - "__Advantages__\n", - "\n", - "Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies. Such structures\n", - "cannot be well approximated by analytic light profiles such as a S\u00e9rsic profile, or even combinations of multiple\n", - "S\u00e9rsic components. pixelizations are therefore required to accurately reconstruct this irregular source-plane light.\n", - "\n", - "Even alternative basis-function approaches, such as shapelets or multi-Gaussian expansions, struggle to accurately\n", - "reconstruct sources with highly complex morphologies or multiple distinct source galaxies.\n", - "\n", - "Pixelized source models are also essential for robustly constraining detailed components of the lens mass\n", - "distribution (e.g. the mass density slope or the presence of dark matter substructure). By fitting all of the lensed\n", - "source light, they reduce degeneracies between the source and lens mass model.\n", - "\n", - "Finally, many science applications aim to study the highly magnified source galaxy itself, in order to learn about\n", - "distant and intrinsically faint galaxies. pixelizations reconstruct the unlensed source emission, enabling detailed\n", - "studies of the source-plane structure.\n", - "\n", - "For CCD imaging, a disadvantage of pixelized source reconstructions is they are the most computationally expensive\n", - "modeling approach. However, for interferometer datasets, the way that JAX and GPUs can exploit the sparsity in the\n", - "linear algebra means pixelized source reconstructions are both significantly faster than other approaches (E.g.\n", - "light profiles) and can scale to millions of visibilities.\n", - "\n", - "__Disadvantages__\n", - "\n", - "Lens modeling with pixelizations is conceptually more complex. There are additional failure modes, such as\n", - "solutions where the source is reconstructed in a highly demagnified configuration due to an unphysical lens mass\n", - "model (e.g. too little or too much mass). These issues are discussed in detail later in the workspace.\n", - "\n", - "As a result, learning to successfully fit lens models with pixelizations typically requires more time and experience\n", - "than the simpler modeling approaches introduced elsewhere in the workspace.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algabra solver which allows for positive and negative\n", - "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This could be problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical. For a pixelizaiton, this often produces negative source pixels which over-fit\n", - "the data, producing unphysical solutions.\n", - "\n", - "For CCD imaging datsets pixelized source reconstructions use a positive-only solver, meaning that every source-pixel\n", - "is only allowed to reconstruct positive flux values. This ensures that the source reconstruction is physical and\n", - "that we don't reconstruct negative flux values that don't exist in the real source galaxy (a common systematic\n", - "solution in lens analysis).\n", - "\n", - "However, for interferometer datasets this positive-only solver is often disabled, because negative pixel values\n", - "can be observed from the measurement process. All interferometer examples therefore disable the positive only solver,\n", - "but you may want to consider if using the positive-only solver is appropriate for your dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "from autoarray.inversion.plot.inversion_plots import subplot_of_mapper, subplot_mappings" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit \n", - "with the lens model.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and compare \n", - "directly to the visiblities. We use a non-uniform fast Fourier transform, which is the most efficient method for \n", - "interferometer datasets containing ~1-10 million visibilities.\n", - "\n", - "If you want to use the high resolution ALMA dataset, uncomment the relevant lines of code below after downloading\n", - "the data from the repository described in the \"High Resolution Dataset\" section above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "Pixelized source modeling requires dense linear algebra operations. These calculations are greatly accelerated\n", - "using an alternative mathematical approach called the **sparse linear algebra formalism**.\n", - "\n", - "You do not need to understand the full details of the method, but the key point is:\n", - "\n", - "- It exploits the **sparsity** of the matrices used in pixelized source reconstruction, reducing memory usage.\n", - "- This leads to a **significant speed-up on GPU or CPU**, using JAX to perform the linear algebra calculations.\n", - "\n", - "To enable this feature, we call `apply_sparse_operator()` on the dataset. This computes and stores a NUFFT operator \n", - "matrix.\n", - "\n", - "For datasets with over 100000 visibilities and many pixels in their real-space mask, this computation takes seconds\n", - "on GPU, but may take 10 minutes or hours (for the small dataset loaded above its miliseconds) on CPU. The `show_progress` \n", - "input outputs a progress bar to the terminal so you can monitor the computation, which is useful when it is slow\n", - "\n", - "When computing it is slow, it is recommend you compute it once, save it to hard-disk, and load it\n", - "before modeling. The example `pixelization/many_visibilities_preparation.py` illustrates how to do this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_sparse_operator(use_jax=True, show_progress=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "As discussed above, disable the default position only linear algebra solver so the source\n", - "reconstruction can have negative pixel values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling is used, \n", - "which evaluates light profiles on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", - "calculations are performed without over sampling.\n", - "\n", - "__Mesh Shape__\n", - "\n", - "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", - "set below to 28 x 28. \n", - "\n", - "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", - "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", - "mesh, the same number of pixels must be used in the y and x directions.\n", - "\n", - "__Edge Zeroing__\n", - "\n", - "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", - "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", - "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", - "\n", - "For a rectangular mesh, the source code computes edge pixels internally using the known\n", - "pixels at the edge of the mesh. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixelization__\n", - "\n", - "We create a `Pixelization` object to perform the pixelized source reconstruction, which is made up of three\n", - "components:\n", - "\n", - "- `mesh:` Different types of mesh can be used to perform the source reconstruction, where the mesh changes the\n", - "details of how the source is reconstructed (e.g. interpolation weights). In this example, we use a rectangular mesh,\n", - "where the centres computed by overlayiong a rectangular mesh over the source plane.\n", - "\n", - "- `regularization:` A pixelization uses many pixels to reconstructed the source, which will often lead to over fitting\n", - "of the noise in the data and an unrealistically complex and structured source. Regularization smooths the source\n", - "reconstruction solution by penalizing solutions where neighboring pixels have large flux differences." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", - "regularization = al.reg.Constant(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This is to illustrate the API for performing a fit via a pixelization using standard objects like \n", - "the `Galaxy`, `Tracer` and `FitInterferometer` \n", - "\n", - "We simply create a `Pixelization` and pass it to the source galaxy, which then gets input into the tracer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens, source])\n", - "\n", - "fit = al.FitInterferometer(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=settings,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the fit, we see that the pixelized source does a good job at capturing the appearance of the source galaxy\n", - "and fitting the data to roughly the noise level." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_interferometer(fit=fit)\n", - "aplt.subplot_fit_dirty_images(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelizations have bespoke visualizations which show more details about the source-reconstruction, image-mesh\n", - "and other quantities.\n", - "\n", - "The `subplot_of_mapper` function produces a comprehensive diagnostic subplot for the inversion. The\n", - "`subplot_mappings` overlays colored circles in the image and source planes that map to one another, thereby\n", - "allowing one to assess how the mass model ray-traces image-pixels and therefore to assess how the source\n", - "reconstruction maps to the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "\n", - "subplot_of_mapper(inversion=inversion, mapper_index=0)\n", - "subplot_mappings(inversion=inversion, pixelization_index=0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Science (Magnification, Flux and More)__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", - "size of the source.\n", - "\n", - "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", - "including examples using a pixelized source reconstruction. Now you know how to fit a pixelization, go check it out!\n", - "\n", - "__Wrap Up__\n", - "\n", - "Pixelizations are the most complex but also most powerful way to model a source galaxy.\n", - "\n", - "Whether you need to use them or not depends on the science you are doing. If you are only interested in measuring a\n", - "simple quantity like the Einstein radius of a lens, you can get away with using light profiles like a Sersic, MGE or \n", - "shapelets to model the source. Low resolution data also means that using a pixelization is not necessary, as the\n", - "complex structure of the source galaxy is not resolved anyway.\n", - "\n", - "However, fitting complex mass models (e.g. a power-law, stellar / dark model or dark matter substructure) requires \n", - "this level of complexity in the source model. Furthermore, if you are interested in studying the properties of the\n", - "source itself, you won't find a better way to do this than using a pixelization.\n", - "\n", - "__Linear Objects__\n", - "\n", - "An `Inversion` contains all of the linear objects used to reconstruct the data in its `linear_obj_list`. \n", - "\n", - "This list may include the following objects:\n", - "\n", - " - `LightProfileLinearObjFuncList`: This object contains lists of linear light profiles and the functionality used\n", - " by them to reconstruct data in an inversion. For example it may only contain a list with a single light profile\n", - " (e.g. `lp_linear.Sersic`) or many light profiles combined in a `Basis` (e.g. `lp_basis.Basis`).\n", - "\n", - "- `Mapper`: The linear objected used by a `Pixelization` to reconstruct data via an `Inversion`, where the `Mapper` \n", - "is specific to the `Pixelization`'s `Mesh` (e.g. a `RectnagularMapper` is used for a `Voronoi` mesh).\n", - "\n", - "In this example, the only linear object used to fit the data was a `Pixelization`, thus the `linear_obj_list`\n", - "contains just one entry corresponding to a `Mapper`:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.linear_obj_list)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To extract results from an inversion many quantities will come in lists or require that we specific the linear object\n", - "we with to use. \n", - "\n", - "Thus, knowing what linear objects are contained in the `linear_obj_list` and what indexes they correspond to\n", - "is important." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Mapper = {inversion.linear_obj_list[0]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grids__\n", - "\n", - "The role of a mapper is to map between the image-plane and source-plane. \n", - "\n", - "This includes mapping grids corresponding to the data grid (e.g. the centers of each image-pixel in the image and\n", - "source plane) and the pixelization grid (e.g. the centre of the Delaunay triangulation in the image-plane and \n", - "source-plane).\n", - "\n", - "All grids are available in a mapper via its `mapper` property." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = inversion.linear_obj_list[0]\n", - "\n", - "# Centre of each masked image pixel in the image-plane.\n", - "print(mapper.image_plane_data_grid)\n", - "\n", - "# Centre of each source pixel in the source-plane.\n", - "print(mapper.source_plane_data_grid)\n", - "\n", - "# Centre of each pixelization pixel in the image-plane (the `Overlay` image_mesh computes these in the image-plane\n", - "# and maps to the source-plane).\n", - "print(mapper.image_plane_mesh_grid)\n", - "\n", - "# Centre of each pixelization pixel in the source-plane.\n", - "print(mapper.source_plane_mesh_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction__\n", - "\n", - "The source reconstruction is also available as a 1D numpy array of values representative of the source pixelization\n", - "itself (in this example, the reconstructed source values at the vertexes of each Voronoi triangle)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.reconstruction)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The (y,x) grid of coordinates associated with these values is given by the `Inversion`'s `Mapper` (which are \n", - "described in chapter 4 of **HowToLens**." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = inversion.linear_obj_list[0]\n", - "print(mapper.source_plane_mesh_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The mapper also contains the (y,x) grid of coordinates that correspond to the ray-traced image sub-pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(mapper.source_plane_data_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapped Reconstructed Images__\n", - "\n", - "The source reconstruction(s) are mapped to the image-plane in order to fit the lens model.\n", - "\n", - "These mapped reconstructed images are also accessible via the `Inversion`. \n", - "\n", - "Note that any light profiles in the lens model (e.g. the `bulge` and `disk` of a lens galaxy) are not \n", - "included in this image -- it only contains the source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.mapped_reconstructed_operated_data.native)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Linear Algebra Matrices (Advanced)__\n", - "\n", - "To perform an `Inversion` a number of matrices are constructed which use linear algebra to perform the reconstruction.\n", - "\n", - "These are accessible in the inversion object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.curvature_matrix)\n", - "print(inversion.regularization_matrix)\n", - "print(inversion.curvature_reg_matrix)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Evidence Terms (Advanced)__\n", - "\n", - "In **HowToLens** and the papers below, we cover how an `Inversion` uses a Bayesian evidence to quantify the goodness\n", - "of fit:\n", - "\n", - "https://arxiv.org/abs/1708.07377\n", - "https://arxiv.org/abs/astro-ph/0601493\n", - "\n", - "This evidence balances solutions which fit the data accurately, without using an overly complex regularization source.\n", - "\n", - "The individual terms of the evidence and accessed via the following properties:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(inversion.regularization_term)\n", - "print(inversion.log_det_regularization_matrix_term)\n", - "print(inversion.log_det_curvature_reg_matrix_term)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulated Interferometer__\n", - "\n", - "We load the source galaxy image from the pixelized inversion of a previous fit, which was performed on an irregular mesh. \n", - "\n", - "Since irregular meshes cannot be directly used to simulate lensed images, we interpolate the source onto a uniform \n", - "grid with shape `interpolated_pixelized_shape`. This grid should have a high resolution (e.g., 1000 \u00d7 1000) to preserve \n", - "all resolved structure from the original mesh. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.interpolate import griddata\n", - "\n", - "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", - "\n", - "reconstruction = inversion.reconstruction\n", - "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", - "\n", - "interpolated_reconstruction = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", - ")\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_ndarray,\n", - " pixel_scales=interpolation_grid.pixel_scales,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To create the lensed image, we ray-trace image pixels to the source plane and interpolate them onto the source \n", - "galaxy image. \n", - "\n", - "This requires an image-plane grid of (y, x) coordinates. In this example, we use the real space mask grid.\n", - "\n", - "To ensure accurate ray-tracing, we apply an 8\u00d78 oversampling scheme. This means that for each pixel in the \n", - "image-plane grid, an 8\u00d78 sub-pixel grid is ray-traced. This approach fully resolves how light is distributed \n", - "across each simulated image pixel, given the source pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=real_space_mask.shape_native,\n", - " pixel_scales=real_space_mask.pixel_scales,\n", - " over_sample_size=8,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We create a tracer to generate the lensed grid onto which we overlay the interpolated source galaxy image, \n", - "producing the lensed source galaxy image. \n", - "\n", - "The source-plane requires a source galaxy with a defined `redshift` for the tracer to function. Since the source\u2019s \n", - "emission is entirely determined by the source galaxy image, this galaxy has no light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(\n", - " galaxies=[\n", - " lens,\n", - " al.Galaxy(redshift=source.redshift),\n", - " ]\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Using the tracer, we generate the lensed source galaxy image on the image-plane grid. This process incorporates \n", - "the `interpolated_reconstruction`, preserving the irregular and asymmetric morphological features captured by the source reconstruction. \n", - "\n", - "Next, we configure the grid, PSF, and simulator settings to match the signal-to-noise ratio (S/N) and noise properties \n", - "of the observed data used for sensitivity mapping. \n", - "\n", - "The `SimulatorInterferometer` takes the generated strong lens image and convolves it with the PSF before adding noise. To \n", - "prevent edge effects, the image is padded before convolution and then trimmed to restore its original `shape_native`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorInterferometer(\n", - " uv_wavelengths=dataset.uv_wavelengths,\n", - " exposure_time=300.0,\n", - " noise_sigma=1000.0,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "# dataset = simulator.via_interpolated_reconstruction_from(\n", - "# tracer=tracer, grid=grid, interpolated_reconstruction=interpolated_reconstruction\n", - "# )\n", - "#\n", - "#\n", - "#\n", - "# dataset=dataset\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Future Ideas / Contributions__\n", - "\n", - "Here are a list of things I would like to add to this tutorial but haven't found the time. If you are interested\n", - "in having a go at adding them contact me on SLACK! :)\n", - "\n", - "- More magnification calculations.\n", - "- Source gradient calculations.\n", - "- A calculation which shows differential lensing effects (e.g. magnification across the source plane)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Pixelization\n", + "======================\n", + "\n", + "A pixelization reconstructs the source's light using a pixel-grid, which is regularized using a prior that forces\n", + "the solution to have a degree of smoothness.\n", + "\n", + "This script fits a source galaxy model which uses a pixelization to reconstruct the source's light.\n", + "\n", + "A rectangular mesh which adapts to the lens mass model magnification and constant regularization scheme are used, which\n", + "are the simplest forms of mesh and regularization with provide computationally fast and accurate solutions.\n", + "\n", + "For simplicity, the lens galaxy's light is omitted from the model and is not present in the simulated data. It is\n", + "straightforward to include the lens galaxy's light in the model.\n", + "\n", + "Pixelizations are covered in detail in chapter 4 of the **HowToLens** lectures.\n", + "\n", + "__CPU Users__\n", + "\n", + "Matrices must be set up for a pixelized source reconstruction to speed up the linear algebra. On GPU, this takes\n", + "seconds, or at most a minute for datasets with tens of millions or more visibilities. On CPU, this can be a lot\n", + "slower, taking over an hour for very large datasets. If you are on CPU, the\n", + "`feature/pixelization/many_visibilities_preparation` example explains how this initial setup can be performed\n", + "before lens modeling and saved to hard disk for fast loading before the model fit.\n", + "\n", + "__Contents__\n", + "\n", + "- **CPU Users:** Matrices must be set up for a pixelized source reconstruction which speed up the linear algebra.\n", + "- **Advantages & Disadvantages:** Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Sparse Operators:** Pixelized source modeling requires dense linear algebra operations.\n", + "- **Settings:** As discussed above, disable the default position only linear algebra solver so the source.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", + "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", + "- **Pixelization:** We create a `Pixelization` object to perform the pixelized source reconstruction, which is made up.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "- **Linear Objects:** An `Inversion` contains all of the linear objects used to reconstruct the data in its.\n", + "- **Grids:** The role of a mapper is to map between the image-plane and source-plane.\n", + "- **Reconstruction:** The source reconstruction is also available as a 1D numpy array of values representative of the.\n", + "- **Mapped Reconstructed Images:** The source reconstruction(s) are mapped to the image-plane in order to fit the lens model.\n", + "- **Simulated Interferometer:** We load the source galaxy image from the pixelized inversion of a previous fit, which was performed.\n", + "\n", + "__Advantages__\n", + "\n", + "Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies. Such structures\n", + "cannot be well approximated by analytic light profiles such as a S\u00e9rsic profile, or even combinations of multiple\n", + "S\u00e9rsic components. pixelizations are therefore required to accurately reconstruct this irregular source-plane light.\n", + "\n", + "Even alternative basis-function approaches, such as shapelets or multi-Gaussian expansions, struggle to accurately\n", + "reconstruct sources with highly complex morphologies or multiple distinct source galaxies.\n", + "\n", + "Pixelized source models are also essential for robustly constraining detailed components of the lens mass\n", + "distribution (e.g. the mass density slope or the presence of dark matter substructure). By fitting all of the lensed\n", + "source light, they reduce degeneracies between the source and lens mass model.\n", + "\n", + "Finally, many science applications aim to study the highly magnified source galaxy itself, in order to learn about\n", + "distant and intrinsically faint galaxies. pixelizations reconstruct the unlensed source emission, enabling detailed\n", + "studies of the source-plane structure.\n", + "\n", + "For CCD imaging, a disadvantage of pixelized source reconstructions is they are the most computationally expensive\n", + "modeling approach. However, for interferometer datasets, the way that JAX and GPUs can exploit the sparsity in the\n", + "linear algebra means pixelized source reconstructions are both significantly faster than other approaches (E.g.\n", + "light profiles) and can scale to millions of visibilities.\n", + "\n", + "__Disadvantages__\n", + "\n", + "Lens modeling with pixelizations is conceptually more complex. There are additional failure modes, such as\n", + "solutions where the source is reconstructed in a highly demagnified configuration due to an unphysical lens mass\n", + "model (e.g. too little or too much mass). These issues are discussed in detail later in the workspace.\n", + "\n", + "As a result, learning to successfully fit lens models with pixelizations typically requires more time and experience\n", + "than the simpler modeling approaches introduced elsewhere in the workspace.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algabra solver which allows for positive and negative\n", + "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This could be problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical. For a pixelizaiton, this often produces negative source pixels which over-fit\n", + "the data, producing unphysical solutions.\n", + "\n", + "For CCD imaging datsets pixelized source reconstructions use a positive-only solver, meaning that every source-pixel\n", + "is only allowed to reconstruct positive flux values. This ensures that the source reconstruction is physical and\n", + "that we don't reconstruct negative flux values that don't exist in the real source galaxy (a common systematic\n", + "solution in lens analysis).\n", + "\n", + "However, for interferometer datasets this positive-only solver is often disabled, because negative pixel values\n", + "can be observed from the measurement process. All interferometer examples therefore disable the positive only solver,\n", + "but you may want to consider if using the positive-only solver is appropriate for your dataset." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "from autoarray.inversion.plot.inversion_plots import subplot_of_mapper, subplot_mappings" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit \n", + "with the lens model.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and compare \n", + "directly to the visiblities. We use a non-uniform fast Fourier transform, which is the most efficient method for \n", + "interferometer datasets containing ~1-10 million visibilities.\n", + "\n", + "If you want to use the high resolution ALMA dataset, uncomment the relevant lines of code below after downloading\n", + "the data from the repository described in the \"High Resolution Dataset\" section above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "Pixelized source modeling requires dense linear algebra operations. These calculations are greatly accelerated\n", + "using an alternative mathematical approach called the **sparse linear algebra formalism**.\n", + "\n", + "You do not need to understand the full details of the method, but the key point is:\n", + "\n", + "- It exploits the **sparsity** of the matrices used in pixelized source reconstruction, reducing memory usage.\n", + "- This leads to a **significant speed-up on GPU or CPU**, using JAX to perform the linear algebra calculations.\n", + "\n", + "To enable this feature, we call `apply_sparse_operator()` on the dataset. This computes and stores a NUFFT operator \n", + "matrix.\n", + "\n", + "For datasets with over 100000 visibilities and many pixels in their real-space mask, this computation takes seconds\n", + "on GPU, but may take 10 minutes or hours (for the small dataset loaded above its miliseconds) on CPU. The `show_progress` \n", + "input outputs a progress bar to the terminal so you can monitor the computation, which is useful when it is slow\n", + "\n", + "When computing it is slow, it is recommend you compute it once, save it to hard-disk, and load it\n", + "before modeling. The example `pixelization/many_visibilities_preparation.py` illustrates how to do this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_sparse_operator(use_jax=True, show_progress=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "As discussed above, disable the default position only linear algebra solver so the source\n", + "reconstruction can have negative pixel values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling is used, \n", + "which evaluates light profiles on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", + "calculations are performed without over sampling.\n", + "\n", + "__Mesh Shape__\n", + "\n", + "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", + "set below to 28 x 28. \n", + "\n", + "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", + "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", + "mesh, the same number of pixels must be used in the y and x directions.\n", + "\n", + "__Edge Zeroing__\n", + "\n", + "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", + "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", + "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", + "\n", + "For a rectangular mesh, the source code computes edge pixels internally using the known\n", + "pixels at the edge of the mesh. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixelization__\n", + "\n", + "We create a `Pixelization` object to perform the pixelized source reconstruction, which is made up of three\n", + "components:\n", + "\n", + "- `mesh:` Different types of mesh can be used to perform the source reconstruction, where the mesh changes the\n", + "details of how the source is reconstructed (e.g. interpolation weights). In this example, we use a rectangular mesh,\n", + "where the centres computed by overlayiong a rectangular mesh over the source plane.\n", + "\n", + "- `regularization:` A pixelization uses many pixels to reconstructed the source, which will often lead to over fitting\n", + "of the noise in the data and an unrealistically complex and structured source. Regularization smooths the source\n", + "reconstruction solution by penalizing solutions where neighboring pixels have large flux differences." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", + "regularization = al.reg.Constant(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This is to illustrate the API for performing a fit via a pixelization using standard objects like \n", + "the `Galaxy`, `Tracer` and `FitInterferometer` \n", + "\n", + "We simply create a `Pixelization` and pass it to the source galaxy, which then gets input into the tracer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens, source])\n", + "\n", + "fit = al.FitInterferometer(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=settings,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the fit, we see that the pixelized source does a good job at capturing the appearance of the source galaxy\n", + "and fitting the data to roughly the noise level." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_interferometer(fit=fit)\n", + "aplt.subplot_fit_dirty_images(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelizations have bespoke visualizations which show more details about the source-reconstruction, image-mesh\n", + "and other quantities.\n", + "\n", + "The `subplot_of_mapper` function produces a comprehensive diagnostic subplot for the inversion. The\n", + "`subplot_mappings` overlays colored circles in the image and source planes that map to one another, thereby\n", + "allowing one to assess how the mass model ray-traces image-pixels and therefore to assess how the source\n", + "reconstruction maps to the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "\n", + "subplot_of_mapper(inversion=inversion, mapper_index=0)\n", + "subplot_mappings(inversion=inversion, pixelization_index=0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Science (Magnification, Flux and More)__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", + "size of the source.\n", + "\n", + "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", + "including examples using a pixelized source reconstruction. Now you know how to fit a pixelization, go check it out!\n", + "\n", + "__Wrap Up__\n", + "\n", + "Pixelizations are the most complex but also most powerful way to model a source galaxy.\n", + "\n", + "Whether you need to use them or not depends on the science you are doing. If you are only interested in measuring a\n", + "simple quantity like the Einstein radius of a lens, you can get away with using light profiles like a Sersic, MGE or \n", + "shapelets to model the source. Low resolution data also means that using a pixelization is not necessary, as the\n", + "complex structure of the source galaxy is not resolved anyway.\n", + "\n", + "However, fitting complex mass models (e.g. a power-law, stellar / dark model or dark matter substructure) requires \n", + "this level of complexity in the source model. Furthermore, if you are interested in studying the properties of the\n", + "source itself, you won't find a better way to do this than using a pixelization.\n", + "\n", + "__Linear Objects__\n", + "\n", + "An `Inversion` contains all of the linear objects used to reconstruct the data in its `linear_obj_list`. \n", + "\n", + "This list may include the following objects:\n", + "\n", + " - `LightProfileLinearObjFuncList`: This object contains lists of linear light profiles and the functionality used\n", + " by them to reconstruct data in an inversion. For example it may only contain a list with a single light profile\n", + " (e.g. `lp_linear.Sersic`) or many light profiles combined in a `Basis` (e.g. `lp_basis.Basis`).\n", + "\n", + "- `Mapper`: The linear objected used by a `Pixelization` to reconstruct data via an `Inversion`, where the `Mapper` \n", + "is specific to the `Pixelization`'s `Mesh` (e.g. a `RectnagularMapper` is used for a `Voronoi` mesh).\n", + "\n", + "In this example, the only linear object used to fit the data was a `Pixelization`, thus the `linear_obj_list`\n", + "contains just one entry corresponding to a `Mapper`:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.linear_obj_list)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To extract results from an inversion many quantities will come in lists or require that we specific the linear object\n", + "we with to use. \n", + "\n", + "Thus, knowing what linear objects are contained in the `linear_obj_list` and what indexes they correspond to\n", + "is important." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Mapper = {inversion.linear_obj_list[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grids__\n", + "\n", + "The role of a mapper is to map between the image-plane and source-plane. \n", + "\n", + "This includes mapping grids corresponding to the data grid (e.g. the centers of each image-pixel in the image and\n", + "source plane) and the pixelization grid (e.g. the centre of the Delaunay triangulation in the image-plane and \n", + "source-plane).\n", + "\n", + "All grids are available in a mapper via its `mapper` property." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = inversion.linear_obj_list[0]\n", + "\n", + "# Centre of each masked image pixel in the image-plane.\n", + "print(mapper.image_plane_data_grid)\n", + "\n", + "# Centre of each source pixel in the source-plane.\n", + "print(mapper.source_plane_data_grid)\n", + "\n", + "# Centre of each pixelization pixel in the image-plane (the `Overlay` image_mesh computes these in the image-plane\n", + "# and maps to the source-plane).\n", + "print(mapper.image_plane_mesh_grid)\n", + "\n", + "# Centre of each pixelization pixel in the source-plane.\n", + "print(mapper.source_plane_mesh_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction__\n", + "\n", + "The source reconstruction is also available as a 1D numpy array of values representative of the source pixelization\n", + "itself (in this example, the reconstructed source values at the vertexes of each Voronoi triangle)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.reconstruction)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The (y,x) grid of coordinates associated with these values is given by the `Inversion`'s `Mapper` (which are \n", + "described in chapter 4 of **HowToLens**." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = inversion.linear_obj_list[0]\n", + "print(mapper.source_plane_mesh_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The mapper also contains the (y,x) grid of coordinates that correspond to the ray-traced image sub-pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(mapper.source_plane_data_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapped Reconstructed Images__\n", + "\n", + "The source reconstruction(s) are mapped to the image-plane in order to fit the lens model.\n", + "\n", + "These mapped reconstructed images are also accessible via the `Inversion`. \n", + "\n", + "Note that any light profiles in the lens model (e.g. the `bulge` and `disk` of a lens galaxy) are not \n", + "included in this image -- it only contains the source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.mapped_reconstructed_operated_data.native)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Linear Algebra Matrices (Advanced)__\n", + "\n", + "To perform an `Inversion` a number of matrices are constructed which use linear algebra to perform the reconstruction.\n", + "\n", + "These are accessible in the inversion object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.curvature_matrix)\n", + "print(inversion.regularization_matrix)\n", + "print(inversion.curvature_reg_matrix)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Evidence Terms (Advanced)__\n", + "\n", + "In **HowToLens** and the papers below, we cover how an `Inversion` uses a Bayesian evidence to quantify the goodness\n", + "of fit:\n", + "\n", + "https://arxiv.org/abs/1708.07377\n", + "https://arxiv.org/abs/astro-ph/0601493\n", + "\n", + "This evidence balances solutions which fit the data accurately, without using an overly complex regularization source.\n", + "\n", + "The individual terms of the evidence and accessed via the following properties:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(inversion.regularization_term)\n", + "print(inversion.log_det_regularization_matrix_term)\n", + "print(inversion.log_det_curvature_reg_matrix_term)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulated Interferometer__\n", + "\n", + "We load the source galaxy image from the pixelized inversion of a previous fit, which was performed on an irregular mesh. \n", + "\n", + "Since irregular meshes cannot be directly used to simulate lensed images, we interpolate the source onto a uniform \n", + "grid with shape `interpolated_pixelized_shape`. This grid should have a high resolution (e.g., 1000 \u00d7 1000) to preserve \n", + "all resolved structure from the original mesh. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.interpolate import griddata\n", + "\n", + "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", + "\n", + "reconstruction = inversion.reconstruction\n", + "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", + "\n", + "interpolated_reconstruction = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", + ")\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_ndarray,\n", + " pixel_scales=interpolation_grid.pixel_scales,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To create the lensed image, we ray-trace image pixels to the source plane and interpolate them onto the source \n", + "galaxy image. \n", + "\n", + "This requires an image-plane grid of (y, x) coordinates. In this example, we use the real space mask grid.\n", + "\n", + "To ensure accurate ray-tracing, we apply an 8\u00d78 oversampling scheme. This means that for each pixel in the \n", + "image-plane grid, an 8\u00d78 sub-pixel grid is ray-traced. This approach fully resolves how light is distributed \n", + "across each simulated image pixel, given the source pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=real_space_mask.shape_native,\n", + " pixel_scales=real_space_mask.pixel_scales,\n", + " over_sample_size=8,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We create a tracer to generate the lensed grid onto which we overlay the interpolated source galaxy image, \n", + "producing the lensed source galaxy image. \n", + "\n", + "The source-plane requires a source galaxy with a defined `redshift` for the tracer to function. Since the source\u2019s \n", + "emission is entirely determined by the source galaxy image, this galaxy has no light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(\n", + " galaxies=[\n", + " lens,\n", + " al.Galaxy(redshift=source.redshift),\n", + " ]\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Using the tracer, we generate the lensed source galaxy image on the image-plane grid. This process incorporates \n", + "the `interpolated_reconstruction`, preserving the irregular and asymmetric morphological features captured by the source reconstruction. \n", + "\n", + "Next, we configure the grid, PSF, and simulator settings to match the signal-to-noise ratio (S/N) and noise properties \n", + "of the observed data used for sensitivity mapping. \n", + "\n", + "The `SimulatorInterferometer` takes the generated strong lens image and convolves it with the PSF before adding noise. To \n", + "prevent edge effects, the image is padded before convolution and then trimmed to restore its original `shape_native`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorInterferometer(\n", + " uv_wavelengths=dataset.uv_wavelengths,\n", + " exposure_time=300.0,\n", + " noise_sigma=1000.0,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "# dataset = simulator.via_interpolated_reconstruction_from(\n", + "# tracer=tracer, grid=grid, interpolated_reconstruction=interpolated_reconstruction\n", + "# )\n", + "#\n", + "#\n", + "#\n", + "# dataset=dataset\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Future Ideas / Contributions__\n", + "\n", + "Here are a list of things I would like to add to this tutorial but haven't found the time. If you are interested\n", + "in having a go at adding them contact me on SLACK! :)\n", + "\n", + "- More magnification calculations.\n", + "- Source gradient calculations.\n", + "- A calculation which shows differential lensing effects (e.g. magnification across the source plane)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/pixelization/likelihood_function.ipynb b/notebooks/interferometer/features/pixelization/likelihood_function.ipynb index 84a4e474f..f813fff8e 100644 --- a/notebooks/interferometer/features/pixelization/likelihood_function.ipynb +++ b/notebooks/interferometer/features/pixelization/likelihood_function.ipynb @@ -1,1583 +1,1620 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Pixelization__\n", - "\n", - "This script provides a step-by-step guide of the **PyAutoLens** `log_likelihood_function` which is used to fit\n", - "`Interferometer` data with an inversion (specifically a `RectangularUniform` mesh and `Constant` regularization scheme`).\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can understand the\n", - " likelihood function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - " - To make inversions in **PyAutoLens** less of a \"black-box\" to users.\n", - "\n", - "__Contents__\n", - "\n", - "- **Simplifications:** This example uses a `RectangularUniform` mesh, where all rectangular source pixels have the same.\n", - "- **Prerequisites:** The likelihood function of pixelizations is the most complicated likelihood function.\n", - "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", - "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Masked Image Grid:** To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", - "- **Lens Galaxy:** We set up a lens galaxy with the lens light and mass, which we will use to demonstrate a pixelized.\n", - "- **Source Galaxy Pixelization and Regularization:** We combine the pixelization into a single `Galaxy` object.\n", - "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the.\n", - "- **Border Relocation:** Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace.\n", - "- **Source Pixel Centre Calculation:** In order to reconstruct the source galaxy using a mesh, we need to determine the centres of the.\n", - "- **Interpolation:** We now combine grids computed above to create an `Interpolator`, which describes how image grid.\n", - "- **Mapper:** We now use the interpolator to create a `Mapper`, which describes the mapping between every image.\n", - "- **Alternative Meshes:** We can briefly consider how this step differs for other mesh types.\n", - "- **Mapping Matrix:** The `mapping_matrix` represents the image-pixel to source-pixel mappings above in a 2D matrix.\n", - "- **Visibilities Reconstruction:** Using the reconstructed pixel fluxes we can map the reconstruction back to the image plane (via the.\n", - "- **Likelihood Function:** We now quantify the goodness-of-fit of our pixelization source galaxy reconstruction.\n", - "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", - "- **Regularization Term:** The second term, $s^{T} H s$, corresponds to the $\\lambda $G_{\\rm L}$ regularization term we added.\n", - "- **Complexity Terms:** Up to this point, it is unclear why we chose a value of `regularization_coefficient=1.0`.\n", - "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", - "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the model, by combining the five terms.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Simplifications__\n", - "\n", - "This example uses a `RectangularUniform` mesh, where all rectangular source pixels have the same size. Most\n", - "pixelization examples use a `RectangularAdaptDensity` mesh, which adapts the size of source pixels to the\n", - "density of points in the source-plane (e.g. the caustic).\n", - "\n", - "The `RectangularUniform` mesh is used here because it is simpler to explain the likelihood function\n", - "and illustrate the key steps in the calculation. The same principles apply to other mesh types, which this\n", - "example will explain where relevant.\n", - "\n", - "__Prerequisites__\n", - "\n", - "The likelihood function of pixelizations is the most complicated likelihood function.\n", - "\n", - "It is advised you read through the following two simpler likelihood functions first, which break down a number of the\n", - "concepts used in this script:\n", - "\n", - " - `interferometer/light_profile/log_likelihood_function.py` the likelihood function for a light profile.\n", - " - `imaging/linear_light_profile/log_likelihood_function.py` the likelihood function for a linear light profile, which\n", - " introduces the linear algebra used for a pixelization but with a simpler use case.\n", - "\n", - "This script repeats all text and code examples in the above likelihood function examples. It therefore can be used to\n", - "learn about the linear light profile likelihood function without reading other likelihood scripts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", - "set below to 28 x 28. \n", - "\n", - "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", - "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", - "mesh, the same number of pixels must be used in the y and x directions.\n", - "\n", - "__Edge Zeroing__\n", - "\n", - "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", - "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", - "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", - "\n", - "For a rectangular mesh, the source code computes edge pixels internally using the known\n", - "pixels at the edge of the mesh. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the \u2018real_space_mask\u2019 which defines the grid the image the galaxy is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(80, 80), pixel_scales=0.05, radius=4.0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the galaxy `Interferometer` dataset `simple` from .fits files, which we will fit \n", - "with the model.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the galaxy to the uv-plane and compare \n", - "directly to the visibilities. We use a non-uniform fast Fourier transform, which is the most efficient method for \n", - "interferometer datasets containing ~1-10 million visibilities. We will discuss how the calculation of the likelihood\n", - "function changes for different methods of Fourier transforming in this guide." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This guide uses in-built visualization tools for plotting. \n", - "\n", - "For example, using the `aplt.subplot_interferometer_dirty_images` the dataset we perform a likelihood evaluation on is plotted.\n", - "\n", - "The `subplot_dataset` displays the visibilities in the uv-plane, which are the raw data of the interferometer\n", - "dataset. These are what will ultimately be directly fitted in the Fourier space.\n", - "\n", - "The `subplot_dirty_images` displays the dirty images of the dataset, which are the reconstructed images of visibilities\n", - "using an inverse Fourier transform to convert these to real-space. These dirty images are not the images we fit, but\n", - "visualization of the dirty images are often used in radio interferometry to show the data in a way that is more\n", - "interpretable to the human eye." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Over sampling evaluates a light profile using multiple samples of its intensity per image-pixel.\n", - "\n", - "For interferometer datasets, over sampling is not used in the pixelization (or for light profiles)\n", - "therefore it is implicitly set to 1 and can be ignored hereafter.\n", - "\n", - "For CCD imaging datasets, over sampling is normally used and slightly changes the likelihood function.\n", - "\n", - "__Masked Image Grid__\n", - "\n", - "To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", - "\n", - "For light profiles these are given by `dataset.lp`, which is a uniform grid of (y,x) Cartesian coordinates\n", - "which have had the 3.0\" circular mask applied.\n", - "\n", - "A pixelization uses a separate grid of (y,x) coordinates, called `dataset.grids.pixelization`, which is\n", - "identical to the light profile grid but may of had a different over-sampling scale applied (but in this example\n", - "does not).\n", - "\n", - "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", - "used to construct a pixelization there is a straight forward mapping between the image data and pixelization pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=dataset.grids.pixelization, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "We set up a lens galaxy with the lens light and mass, which we will use to demonstrate a pixelized source\n", - "reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Pixelization and Regularization__\n", - "\n", - "We combine the pixelization into a single `Galaxy` object.\n", - "\n", - "The galaxy includes the rectangular mesh and constant regularization scheme, which will ultimately be used\n", - "to reconstruct its star forming clumps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixelization = al.Pixelization(\n", - " mesh=al.mesh.RectangularUniform(mesh_shape),\n", - " regularization=al.reg.Constant(coefficient=1.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the image-plane to its (y,x) \n", - "source-plane coordinate $\\beta$ using the summed deflection angles $\\alpha$ of the mass profiles:\n", - "\n", - " $\\beta = \\theta - \\alpha(\\theta)$\n", - "\n", - "The likelihood function of a pixelized source reconstruction ray-traces two grids from the image-plane to the source-plane:\n", - "\n", - " 1) A 2D grid of (y,x) coordinates aligned with the imaging data's image-pixels.\n", - "\n", - " 2) The sparse 2D grid of (y,x) coordinates above which form the centres of the rectangular source pixels.\n", - "\n", - "The function below computes the 2D deflection angles of the tracer's lens galaxies and subtracts them from the \n", - "image-plane 2D (y,x) coordinates $\\theta$ of each grid, thus ray-tracing their coordinates to the source plane to \n", - "compute their $\\beta$ values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The source code gets quite complex when handling grids for a pixelization, but it is all handled in\n", - "the `TracerToInversion` objects.\n", - "\n", - "The plots at the bottom of this cell show the traced grids used by the source pixelization, showing\n", - "how the rectangular mesh and traced image pixels are constructed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_to_inversion = al.TracerToInversion(tracer=tracer, dataset=dataset)\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", - "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", - " grid=dataset.grids.pixelization\n", - ")[-1]\n", - "\n", - "\n", - "aplt.plot_grid(grid=traced_grid_pixelization, title=\"\")\n", - "\n", - "# grid_plotter.figure_2d()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Border Relocation__\n", - "\n", - "Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace to far outskirts of\n", - "the source-plane. \n", - "\n", - "We relocate these pixels (for both grids above) to the edge of the source-plane border (defined via the border of the \n", - "image-plane mask). This is detailed in **HowToLens chapter 4 tutorial 5** and figure 2 of https://arxiv.org/abs/1708.07377." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", - "\n", - "border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)\n", - "\n", - "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", - "\n", - "\n", - "aplt.plot_grid(grid=relocated_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Pixel Centre Calculation__\n", - "\n", - "In order to reconstruct the source galaxy using a mesh, we need to determine the centres of the rectangular mesh's \n", - "source pixels.\n", - "\n", - "We do this by overlying a rectangular grid on the relocated traced image-plane grid computed above.\n", - "\n", - "This distributes the rectangular mesh so it fully overlaps the region of the source-plane containing the traced \n", - "image-pixels without having edge pixels that extend beyond this region." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoarray.inversion.mesh.mesh.rectangular_adapt_density import overlay_grid_from\n", - "\n", - "mesh_grid = overlay_grid_from(\n", - " shape_native=mesh_shape, grid=al.Grid2DIrregular(relocated_grid)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolation__\n", - "\n", - "We now combine grids computed above to create an `Interpolator`, which describes how image grid pixel maps to\n", - "every rectangular mesh pixel. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "interpolator = pixelization.mesh.interpolator_from(\n", - " source_plane_data_grid=relocated_grid,\n", - " source_plane_mesh_grid=mesh_grid,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mapper__\n", - "\n", - "We now use the interpolator to create a `Mapper`, which describes the mapping between every image pixel and every \n", - "rectangular pixel, based on the interpolation scheme above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = al.Mapper(interpolator=interpolator)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For the rectangular mesh, the interpolation scheme is called bilinear interpolation, which means that every image \n", - "pixel maps to the rectangular pixel it lands in and the three neighboring rectangular pixels. \n", - "\n", - "The weight of each mapping is determined by the bilinear interpolation scheme, which is a function of how close the \n", - "image pixel is to the centre of the rectangular pixel it lands in and the three neighboring rectangular pixels.\n", - "\n", - "We can print the mappings and weights, for example of the first image pixel, to confirm this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(interpolator.mappings[0])\n", - "print(interpolator.weights[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting the rectangular mesh shows that the source-plane and been discretized into a uniform grid of source pixels.\n", - "\n", - "(To plot the RectangularUniform mesh, we have to convert it to a `Mapper` object, which is described in the next \n", - "likelihood step).\n", - "\n", - "Below, we plot the rectangular mesh without the traced image-grid pixels (for clarity) and with them as \n", - "black dots in order to show how each set of image-pixels fall within a source pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Mapper` contains:\n", - "\n", - " 1) `source_plane_data_grid`: the traced grid of (y,x) image-pixel coordinate centres (`relocated_grid`).\n", - " 2) `source_plane_mesh_grid`: The rectangular mesh of traced (y,x) source-pixel coordinates (`grid_Rectangular`).\n", - "\n", - "We have therefore discretized the source-plane into a rectangular mesh, and can pair every traced image-pixel coordinate\n", - "with the corresponding source pixel it lands in.\n", - "\n", - "This pairing is contained in the ndarray `pix_indexes_for_sub_slim_index` which maps every image-pixel index to \n", - "every source-pixel index.\n", - "\n", - "In the API, the `pix_indexes` refers to the source pixel indexes (e.g. source pixel 0, 1, 2 etc.) and `sub_slim_index` \n", - "refers to the index of an image pixel (e.g. image-pixel 0, 1, 2 etc.). \n", - "\n", - "For example, printing the first ten entries of `pix_indexes_for_sub_slim_index` shows the first ten source-pixel\n", - "indexes these image sub-pixels map too." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pix_indexes_for_sub_slim_index = mapper.pix_indexes_for_sub_slim_index\n", - "\n", - "print(pix_indexes_for_sub_slim_index[0:9])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This array can be used to visualize how an input list of image-pixel indexes map to the source-plane.\n", - "\n", - "It also shows that image-pixel indexing begins from the top-left and goes rightwards and downwards, accounting for \n", - "all image-pixels which are not masked." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(\n", - " array=dataset.dirty_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", - ")\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The reverse mappings of source-pixels to image-pixels can also be used.\n", - "\n", - "If we choose the right source-pixel index, we can see that multiple imaging occur whereby image-pixels in different\n", - "regions of the image-plane are grouped into the same source-pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pix_indexes = [[200]]\n", - "\n", - "indexes = mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)\n", - "\n", - "\n", - "aplt.plot_array(\n", - " array=dataset.dirty_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", - ")\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolation__\n", - "\n", - "The right hand plot shows more laying over source pixel 200 than its retangular black lines. Pixels further \n", - "out than the pixel appear to be mapped to this source pixel. \n", - "\n", - "This is because the mesh uses an interpolation mapping scheme whereby each image pixels is paired with four source \n", - "pixels. For a rectangular mesh, this scheme is called bilinear interpolation, and it means that every pixel maps\n", - "not only to the rectangular source pixel it lands in, but also the three neighbouring source pixels. Interpolation is \n", - "key to ensuring that the pixelization can reconstruct smooth source morphologies.\n", - "\n", - "We can confirm that every image pixel maps to four source pixels by printing \n", - "the `pix_sizes_for_sub_slim_index`, which gives the number of mapped source pixels for every image pixel.\n", - "\n", - "We can also confirm that the interpolation introduces weights to each mapping by printing the \n", - "`pix_weights_for_sub_slim_index`, which gives the weight of each mapping for every image pixel." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(mapper.pix_sizes_for_sub_slim_index[0:9])\n", - "print(mapper.pix_weights_for_sub_slim_index[0:9])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "Lets quickly think about what happens when we use over sampling in the pixelization (e.g. `sub_size>1`). For\n", - "the `sub_size=1` case above, each image pixel maps to 4 source pixels (due to bilinear interpolation)\n", - "with a weight determined from the bilinear interpolation scheme.\n", - "\n", - "However, the default over sampling for a pixelization is `sub_size=4`, meaning each image pixel is divided\n", - "into a 4x4 grid of sub-pixels (16 sub-pixels in total). Each of these sub-pixels maps to 4 source pixels\n", - "(due to bilinear interpolation), where the weight of each mapping is determined by the bilinear interpolation\n", - "scheme divided by 16 (because there are 16 sub-pixels).\n", - "\n", - "This example therefore used a `sub_size=1` to keep the explanation of the likelihood, visualization of the\n", - "arrays above and understanding of the mapping scheme as simple as possible. You can manually increase the\n", - "`sub_size` above and re-run the notebook to see how this changes the mapping scheme.\n", - "\n", - "__Alternative Meshes__\n", - "\n", - "We can briefly consider how this step differs for other mesh types. Above, we simply overlaid a uniform rectangular\n", - "grid to define the source pixel centres and then mapped image pixels to these source pixels.\n", - "\n", - "The `RectangularAdaptDensity` mesh pretty much works exactly the same, its just that a calculation (which we don't\n", - "describe here) works out how to make a grid of rectangular pixels that adapt to the source-plane density and thus\n", - "vary in size. \n", - "\n", - "There is also a `RectangularAdaptImage` mesh which uses the image of the lensed source to adapt\n", - "the rectangular pixel sizes. This often puts even smaller pixels in the brightest regions of the source,\n", - "even if it lies offset or away from the caustic.\n", - "\n", - "There is also a `Delaunay` mesh which uses a Delaunay triangulation to define an irregular grid of source pixels.\n", - "This is described fully in the `delaunay` example including a likelihood function guide.\n", - "\n", - "__Mapping Matrix__\n", - "\n", - "The `mapping_matrix` represents the image-pixel to source-pixel mappings above in a 2D matrix. \n", - "\n", - "It has dimensions `(total_image_pixels, total_source_pixels)`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", - " pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,\n", - " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,\n", - " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,\n", - " pixels=mapper.pixels,\n", - " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", - " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", - " sub_fraction=mapper.over_sampler.sub_fraction,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `mapping_matrix` shows of all image-source pixel mappings.\n", - "\n", - "No row of pixels has more than one non-zero entry. It is not possible for two image pixels to map to the same source \n", - "pixel (meaning that there are no correlated pixels in the mapping matrix)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Each column of the `mapping_matrix` can therefore be used to show all image-pixels it maps too. \n", - "\n", - "For example, above, we plotted all image-pixels of source-pixel 200 (as well as 202 and 204). We can extract all\n", - "image-pixel indexes of source pixels 200 using the `mapping_matrix` and use them to plot the image of this\n", - "source-pixel (which corresponds to only values of zeros or ones)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "indexes_source_pix_200 = np.nonzero(mapping_matrix[:, 200])\n", - "\n", - "print(indexes_source_pix_200[0])\n", - "\n", - "array_2d = al.Array2D(values=mapping_matrix[:, 200], mask=dataset.mask)\n", - "\n", - "aplt.plot_array(array=array_2d, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Transformed Mapping Matrix ($f$)__\n", - "\n", - "Each pixelization pixel can therefore be thought of as an image (where all entries of this image are zeros and ones). \n", - "\n", - "However, for interferometer datasets we want to fit the visibilities in the uv-plane, not the image-plane. Therefore,\n", - "each image in the `mapping_matrix` must be transformed to the uv-plane via a Fourier transform, such that each\n", - "column in the `transformed_mapping_matrix` represents the visibilities in the uv-plane of each pixelization pixel.\n", - "\n", - "This operation changes the dimensions of the mapping matrix, meaning the `transformed_mapping_matrix` has\n", - "dimensions `(total_image_pixels, total_visibilities)`. \n", - "\n", - "If the number of visibilities is large (e.g. 10^6) this matrix becomes extremely large and computationally expensive to \n", - "store memory, meaning the sparse operator likelihood function must be used instead.\n", - "\n", - "The `transformed_mapping_matrix` is also complex, storing all entries of the visibilities after the NUFFT as real\n", - "and complex values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", - " mapping_matrix=mapping_matrix\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A 2D plot of the `transformed_mapping_matrix` shows all visibility-source pixel mappings.\n", - "\n", - "Note how, unlike for the `mapping_matrix`, every row of image-pixels fully consists of non-zero entries. This\n", - "means the matrix is fully dense, making it even more difficult to store in memory for large datasets.\n", - "\n", - "Below, we plot the real and imaginary components of the `transformed_mapping_matrix` separately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " transformed_mapping_matrix.real,\n", - " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()\n", - "\n", - "plt.imshow(\n", - " transformed_mapping_matrix.imag,\n", - " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Each column of the `transformed_mapping_matrix` shows all visibilities it maps to after the NUFFT." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "indexes_pix_200 = np.nonzero(transformed_mapping_matrix[:, 200])\n", - "\n", - "print(indexes_pix_200[0])\n", - "\n", - "visibilities = al.Visibilities(visibilities=transformed_mapping_matrix[:, 200])\n", - "\n", - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) the `transformed_mapping_matrix` is denoted $f_{ij}$\n", - "where $i$ maps over all $I$ source pixels and $j$ maps over all $J$ visibilities. \n", - "\n", - "For example: \n", - "\n", - " - $f_{0, 2} = 0.3$ indicates that visibility number $2$ maps to pixelization pixel $0$ with a weight of $0.3$ after the NUFFT.\n", - "\n", - "The indexing of the `mapping_matrix` is reversed compared to the notation of WD03 (e.g. visibilities\n", - "are the first entry of `mapping_matrix` whereas for $f$ they are the second index)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " f\"Mapping between visibility 0 and RectangularUniform pixel 2 = {mapping_matrix[0, 2]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Data Vector (D)__\n", - "\n", - "To solve for the source pixel fluxes we now pose the problem as a linear inversion.\n", - "\n", - "This requires us to convert the `transformed_mapping_matrix` and our `data` and `noise map` into matrices of certain dimensions. \n", - "\n", - "The `data_vector`, $D$, is the first matrix and it has dimensions `(total_Rectangular_pixels,)`.\n", - "\n", - "In WD03 (https://arxiv.org/abs/astro-ph/0302587) and N15 (https://arxiv.org/abs/1412.7436) the data vector \n", - "is give by: \n", - "\n", - " $\\vec{D}_{i} = \\sum_{\\rm j=1}^{J}f_{ij}(d_{j})/\\sigma_{j}^2 \\, \\, .$\n", - "\n", - "Where:\n", - "\n", - " - $d_{\\rm j}$ are the image-pixel data flux values.\n", - " - $\\sigma{\\rm _j}^2$ are the statistical uncertainties of each image-pixel value.\n", - "\n", - "$i$ maps over all $I$ source pixels and $j$ maps over all $J$ image pixels. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data_vector = (\n", - " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " visibilities=dataset.data,\n", - " noise_map=dataset.noise_map,\n", - " )\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$D$ describes which source pixels trace to which visibilities, with associated weights, after the NUFFT. This \n", - "ensures the reconstruction fully accounts for the NUFFT when fitting the data.\n", - "\n", - "We can plot $D$ as a column vector:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(\n", - " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", - ")\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dimensions of $D$ are the number of source pixels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Data Vector:\")\n", - "print(data_vector)\n", - "print(data_vector.shape)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Matrix (F)__\n", - "\n", - "The `curvature_matrix` $F$ is the second matrix and it has \n", - "dimensions `(total_Rectangular_pixels, total_Rectangular_pixels)`.\n", - "\n", - "In WD03 / N15 (https://arxiv.org/abs/astro-ph/0302587) the curvature matrix is a 2D matrix given by:\n", - "\n", - " ${F}_{ik} = \\sum_{\\rm j=1}^{J}f_{ij}f_{kj}/\\sigma_{j}^2 \\, \\, .$\n", - "\n", - "NOTE: this notation implicitly assumes a summation over $K$, where $k$ runs over all pixelization pixel indexes $K$.\n", - "\n", - "Note how summation over $J$ runs over $f$ twice, such that every entry of $F$ is the sum of the multiplication\n", - "between all values in every two columns of $f$.\n", - "\n", - "For example, $F_{0,1}$ is the sum of all visibility values in $f$ of source pixel 0 multiplied by\n", - "all visibility values of source pixel 1.\n", - "\n", - "Visibilities are both real and complex values, and the `curvature_matrix` is computed separately for the real and\n", - "imaginary components of the visibilities and then summed together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.real,\n", - " noise_map=dataset.noise_map.real,\n", - ")\n", - "\n", - "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", - " mapping_matrix=transformed_mapping_matrix.imag,\n", - " noise_map=dataset.noise_map.imag,\n", - ")\n", - "\n", - "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", - "\n", - "plt.imshow(curvature_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For $F_{ik}$ to be non-zero, this requires that the images of source pixels $i$ and $k$ share at least one\n", - "image-pixel, which for visibilities after the NUFFT is always true for all $i$ and $k$.\n", - "\n", - "For example, we can see a non-zero entry for $F_{100,101}$ and plotting their images\n", - "show overlap." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_pixel_0 = 0\n", - "source_pixel_1 = 1\n", - "\n", - "print(curvature_matrix[source_pixel_0, source_pixel_1])\n", - "\n", - "visibilities = al.Visibilities(\n", - " visibilities=transformed_mapping_matrix[:, source_pixel_0],\n", - ")\n", - "\n", - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n", - "\n", - "visibilities = al.Visibilities(\n", - " visibilities=transformed_mapping_matrix[:, source_pixel_1],\n", - ")\n", - "\n", - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The following chi-squared is minimized when we perform the inversion and reconstruct the source_galaxy:\n", - "\n", - "$\\chi^2 = \\sum_{\\rm j=1}^{J} \\bigg[ \\frac{(\\sum_{\\rm i=1}^{I} s_{i} f_{ij}) - d_{j}}{\\sigma_{j}} \\bigg]$\n", - "\n", - "Where $s$ is the reconstructed pixel fluxes in all $I$ source pixels.\n", - "\n", - "The solution for $s$ is therefore given by (equation 5 WD03):\n", - "\n", - " $s = F^{-1} D$\n", - "\n", - "We can compute this using NumPy linear algebra:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# Because we are no using regularizartion (see below) it is common for the curvature matrix to be singular and lead\n", - "# to a LinAlgException. The loop below mitigates this -- you can ignore it as it is not important for understanding\n", - "# the PyAutoLens likelihood function.\n", - "\n", - "for i in range(curvature_matrix.shape[0]):\n", - " curvature_matrix[i, i] += 1e-8\n", - "\n", - "reconstruction = np.linalg.solve(curvature_matrix, data_vector)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot this reconstruction -- it looks like a mess.\n", - "\n", - "The pixelization pixels have noisy and unsmooth values, and it is hard to make out if a source galaxy is even being \n", - "reconstructed. \n", - "\n", - "In fact, the linear inversion is (over-)fitting noise in the image data, meaning this system of equations is \n", - "ill-posed. We need to apply some form of smoothing on the reconstruction to avoid over fitting noise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Regularization Matrix (H)__\n", - "\n", - "Regularization adds a linear regularization term $G_{\\rm L}$ to the $\\chi^2$ we solve for giving us a new merit \n", - "function $G$ (equation 11 WD03):\n", - "\n", - " $G = \\chi^2 + \\lambda \\, G_{\\rm L}$\n", - "\n", - "where $\\lambda$ is the `regularization_coefficient` which describes the magnitude of smoothness that is applied. A \n", - "higher $\\lambda$ will regularize the source more, leading to a smoother source galaxy reconstruction.\n", - "\n", - "Different forms for $G_{\\rm L}$ can be defined which regularize the reconstruction in different ways. The \n", - "`Constant` regularization scheme used in this example applies gradient regularization (equation 14 WD03):\n", - "\n", - " $G_{\\rm L} = \\sum_{\\rm i}^{I} \\sum_{\\rm n=1}^{N} [s_{i} - s_{i, v}]$\n", - "\n", - "This regularization scheme is easier to express in words -- the summation goes to each rectangular pixelization pixel,\n", - "determines all rectangular pixels with which it shares a direct vertex (e.g. its neighbors) and penalizes solutions \n", - "where the difference in reconstructed flux of these two neighboring pixels is large.\n", - "\n", - "The summation does this for all source pixels, thus it favours solutions where neighboring source\n", - "pixels reconstruct similar values to one another (e.g. it favours a smooth source galaxy reconstruction).\n", - "\n", - "We now define the `regularization matrix`, $H$, which allows us to include this smoothing when we solve for $s$. $H$\n", - "has dimensions `(total_Rectangular_pixels, total_Rectangular_pixels)`.\n", - "\n", - "This relates to $G_{\\rm L}$ as (equation 13 WD03):\n", - "\n", - " $H_{ik} = \\frac{1}{2} \\frac{\\partial G_{\\rm L}}{\\partial s_{i} \\partial s_{k}}$\n", - "\n", - "$H$ has the `regularization_coefficient` $\\lambda$ folded into it such $\\lambda$'s control on the degree of smoothing\n", - "is accounted for." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", - " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", - " neighbors=mapper.neighbors,\n", - " neighbors_sizes=mapper.neighbors.sizes,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can plot the regularization matrix and note that:\n", - "\n", - " - non-zero entries indicate that two source pixels are neighbors and therefore are regularized with one another.\n", - "\n", - " - Zeros indicate the two source pixels do not neighbor one another.\n", - "\n", - "The majority of entries are zero, because the majority of source pixels are not neighbors with one another." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "plt.imshow(regularization_matrix)\n", - "plt.colorbar()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__F + Lamdba H__\n", - "\n", - "$H$ enters the linear algebra system we solve for as follows (WD03 equation (12)):\n", - "\n", - " $s = [F + H]^{-1} D$" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Galaxy Reconstruction (s)__\n", - "\n", - "We can now solve the linear system above using NumPy linear algebra. \n", - "\n", - "Note that the for loop used above to prevent a LinAlgException is no longer required." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting this source galaxy reconstruction we can see that regularization has lead us to reconstruct a smoother \n", - "source galaxy, which actually looks like the star forming clumps in the imaging data! \n", - "\n", - "This also implies we are not over-fitting the noise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visibilities Reconstruction__\n", - "\n", - "Using the reconstructed pixel fluxes we can map the reconstruction back to the image plane (via \n", - "the `blurred mapping_matrix`) and produce a reconstruction of the image data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_visibilities = (\n", - " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", - " transformed_mapping_matrix=transformed_mapping_matrix,\n", - " reconstruction=reconstruction,\n", - " )\n", - ")\n", - "\n", - "mapped_reconstructed_visibilities = al.Visibilities(\n", - " visibilities=mapped_reconstructed_visibilities\n", - ")\n", - "\n", - "aplt.plot_grid(grid=mapped_reconstructed_visibilities.in_grid, title=\"\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our pixelization source galaxy reconstruction. \n", - "\n", - "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", - "\n", - "The likelihood function for source galaxy modeling consists of five terms:\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + s^{T} H s + \\mathrm{ln} \\, \\left[ \\mathrm{det} (F + H) \\right] - { \\mathrm{ln}} \\, \\left[ \\mathrm{det} (H) \\right] + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "This expression was first derived by Suyu 2006 (https://arxiv.org/abs/astro-ph/0601493) and is given by equation (19).\n", - "It was derived into **PyAutoLens** notation in Dye 2008 (https://arxiv.org/abs/0804.4002) equation (5).\n", - "\n", - "We now explain what each of these terms mean.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", - "\n", - " - `model_data` = `mapped_reconstructed_visibilities`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", - "\n", - "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", - "for, corresponding to a fit with a lower likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_visibilities = mapped_reconstructed_visibilities\n", - "\n", - "residual_map = dataset.data - model_visibilities\n", - "\n", - "\n", - "normalized_residual_map_real = (residual_map.real / dataset.noise_map.real).astype(\n", - " \"complex128\"\n", - ")\n", - "normalized_residual_map_imag = (residual_map.imag / dataset.noise_map.imag).astype(\n", - " \"complex128\"\n", - ")\n", - "normalized_residual_map = (\n", - " normalized_residual_map_real + 1j * normalized_residual_map_imag\n", - ")\n", - "\n", - "\n", - "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", - "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", - "chi_squared_map = chi_squared_map_real + 1j * chi_squared_map_imag\n", - "\n", - "\n", - "chi_squared_real = np.sum(chi_squared_map.real)\n", - "chi_squared_imag = np.sum(chi_squared_map.imag)\n", - "chi_squared = chi_squared_real + chi_squared_imag\n", - "\n", - "print(chi_squared)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_map = al.Visibilities(visibilities=chi_squared_map)\n", - "\n", - "aplt.plot_grid(grid=chi_squared_map.in_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Regularization Term__\n", - "\n", - "The second term, $s^{T} H s$, corresponds to the $\\lambda $G_{\\rm L}$ regularization term we added to our merit \n", - "function above.\n", - "\n", - "This is the term which sums up the difference in flux of all reconstructed source pixels, and reduces the \n", - "likelihood of solutions where there are large differences in flux (e.g. the source galaxy is less smooth and more \n", - "likely to be overfitting noise).\n", - "\n", - "We compute it below via matrix multiplication, noting that the `regularization_coefficient`, $\\lambda$, is built into \n", - "the `regularization_matrix` already." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "regularization_term = np.matmul(\n", - " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", - ")\n", - "\n", - "print(regularization_term)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Complexity Terms__\n", - "\n", - "Up to this point, it is unclear why we chose a value of `regularization_coefficient=1.0`. \n", - "\n", - "We cannot rely on the `chi_squared` and `regularization_term` above to optimally choose its value, because increasing \n", - "the `regularization_coefficient` smooths the solution more and therefore:\n", - "\n", - " - Decreases `chi_squared` by fitting the data worse, producing a lower `log_likelihood`.\n", - "\n", - " - Increases the `regularization_term` by penalizing the differences between source pixel fluxes more, again reducing\n", - " the inferred `log_likelihood`.\n", - "\n", - "If we set the regularization coefficient based purely on these two terms, we would set a value of 0.0 and be back where\n", - "we started over-fitting noise!\n", - "\n", - "The terms $\\left[ \\mathrm{det} (F + H) \\right]$ and $ - { \\mathrm{ln}} \\, \\left[ \\mathrm{det} (H) \\right]$ address \n", - "this problem. \n", - "\n", - "They quantify how complex the reconstruction is, and penalize solutions where *it is more complex*. Reducing \n", - "the `regularization_coefficient` makes the source galaxy reconstruction more complex (because a galaxy that is \n", - "smoothed less uses more flexibility to fit the data better).\n", - "\n", - "These two terms therefore counteract the `chi_squared` and `regularization_term`, so as to attribute a higher\n", - "`log_likelihood` to solutions which fit the data with a more smoothed and less complex source (e.g. one with a higher \n", - "`regularization_coefficient`).\n", - "\n", - "In **HowToGalaxy** -> `chapter 4` -> `tutorial_4_bayesian_regularization` we expand on this further and give a more\n", - "detailed description of how these different terms impact the `log_likelihood_function`. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", - "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]\n", - "\n", - "print(log_curvature_reg_matrix_term)\n", - "print(log_regularization_matrix_term)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", - "\n", - "The final term ins the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", - "of the log of every noise-map value squared. \n", - "\n", - "Given the `noise_map` is fixed, this term does not change during the lens modeling process and has no impact on the \n", - "model we infer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization_real = np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", - "noise_normalization_imag = np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", - "noise_normalization = noise_normalization_real + noise_normalization_imag" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "We can now, finally, compute the `log_likelihood` of the model, by combining the five terms computed above using\n", - "the likelihood function defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "log_evidence = float(\n", - " -0.5\n", - " * (\n", - " chi_squared\n", - " + regularization_term\n", - " + log_curvature_reg_matrix_term\n", - " - log_regularization_matrix_term\n", - " + noise_normalization\n", - " )\n", - ")\n", - "\n", - "print(log_evidence)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This process to perform a likelihood function evaluation performed via the `FitInterferometer` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitInterferometer(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=al.Settings(use_border_relocator=True),\n", - ")\n", - "fit_log_evidence = fit.log_evidence\n", - "print(fit_log_evidence)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "__Log Likelihood Function: Source Code Speed Up__\n", - "\n", - "The interferometer pixelization likelihood function described in this notebook performs certain calculations using\n", - "functions which are easier to understand, but are computationally slower than the actual source code implementation\n", - "(but the two produce identical results).\n", - "\n", - "We end by pointing out some of these, but we do not provide an step-by-step description of how they work.\n", - "If you are interested, you will need to dive into the source code itself.\n", - "\n", - "**Fast Chi Squared:** The `chi_squared` above is computed using the `transformed_mapping_matrix`, which requires\n", - "many NUFFT's to compute and requires large memroy store. The source code uses a trick which computes the chi-squared\n", - "but bypasses the need to ever compute the `transformed_mapping_matrix`.\n", - "\n", - "**Sparse Operator Curvature Matrix:** The `curvature_matrix` above is also computed using the `transformed_mapping_matrix`, \n", - "which again means slow run times and large memory usage. The source code can instead use sparse operators to \n", - "compute the curvature matrix in a way which again bypasses the need to compute the `transformed_mapping_matrix`.\n", - "\n", - "The two tricks in combination lead to a significant speed up in the likelihood function evaluation and mean that\n", - "the large matrix of size [source pixels, visibilities] never needs to be stored in memory. This is at the heart\n", - "of why lens modeling interferometer data with pixelized source reconstructions is so fast!\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the pixelization likelihood function.\n", - "\n", - "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", - "are described in additional notebooks found in this package. In brief, these describe:\n", - "\n", - " - **Over Sampling**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", - " paired fractionally with each `RectangularAdaptDensity` pixel.\n", - "\n", - " - **Source-plane Interpolation**: Using bilinear interpolation on the `RectangularAdaptDensity` pixelization to pair \n", - " each image (sub-)pixel to multiple `RectangularAdaptDensity` pixels with interpolation weights.\n", - "\n", - " - **Source Morphology Pixelization Adaption**: Adapting the pixelization such that is congregates source pixels around\n", - " the source's brightest regions, as opposed to the magnification-based pixelization used here.\n", - "\n", - " - **Luminosity Weighted Regularization**: Using an adaptive regularization coefficient which adapts the level of \n", - " regularization applied to the source galaxy based on its luminosity." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Pixelization__\n", + "\n", + "This script provides a step-by-step guide of the **PyAutoLens** `log_likelihood_function` which is used to fit\n", + "`Interferometer` data with an inversion (specifically a `RectangularUniform` mesh and `Constant` regularization scheme`).\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers using **PyAutoLens**, so that readers can understand the\n", + " likelihood function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + " - To make inversions in **PyAutoLens** less of a \"black-box\" to users.\n", + "\n", + "__Contents__\n", + "\n", + "- **Simplifications:** This example uses a `RectangularUniform` mesh, where all rectangular source pixels have the same.\n", + "- **Prerequisites:** The likelihood function of pixelizations is the most complicated likelihood function.\n", + "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", + "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Masked Image Grid:** To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", + "- **Lens Galaxy:** We set up a lens galaxy with the lens light and mass, which we will use to demonstrate a pixelized.\n", + "- **Source Galaxy Pixelization and Regularization:** We combine the pixelization into a single `Galaxy` object.\n", + "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the.\n", + "- **Border Relocation:** Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace.\n", + "- **Source Pixel Centre Calculation:** In order to reconstruct the source galaxy using a mesh, we need to determine the centres of the.\n", + "- **Interpolation:** We now combine grids computed above to create an `Interpolator`, which describes how image grid.\n", + "- **Mapper:** We now use the interpolator to create a `Mapper`, which describes the mapping between every image.\n", + "- **Alternative Meshes:** We can briefly consider how this step differs for other mesh types.\n", + "- **Mapping Matrix:** The `mapping_matrix` represents the image-pixel to source-pixel mappings above in a 2D matrix.\n", + "- **Visibilities Reconstruction:** Using the reconstructed pixel fluxes we can map the reconstruction back to the image plane (via the.\n", + "- **Likelihood Function:** We now quantify the goodness-of-fit of our pixelization source galaxy reconstruction.\n", + "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", + "- **Regularization Term:** The second term, $s^{T} H s$, corresponds to the $\\lambda $G_{\\rm L}$ regularization term we added.\n", + "- **Complexity Terms:** Up to this point, it is unclear why we chose a value of `regularization_coefficient=1.0`.\n", + "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", + "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the model, by combining the five terms.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Simplifications__\n", + "\n", + "This example uses a `RectangularUniform` mesh, where all rectangular source pixels have the same size. Most\n", + "pixelization examples use a `RectangularAdaptDensity` mesh, which adapts the size of source pixels to the\n", + "density of points in the source-plane (e.g. the caustic).\n", + "\n", + "The `RectangularUniform` mesh is used here because it is simpler to explain the likelihood function\n", + "and illustrate the key steps in the calculation. The same principles apply to other mesh types, which this\n", + "example will explain where relevant.\n", + "\n", + "__Prerequisites__\n", + "\n", + "The likelihood function of pixelizations is the most complicated likelihood function.\n", + "\n", + "It is advised you read through the following two simpler likelihood functions first, which break down a number of the\n", + "concepts used in this script:\n", + "\n", + " - `interferometer/light_profile/log_likelihood_function.py` the likelihood function for a light profile.\n", + " - `imaging/linear_light_profile/log_likelihood_function.py` the likelihood function for a linear light profile, which\n", + " introduces the linear algebra used for a pixelization but with a simpler use case.\n", + "\n", + "This script repeats all text and code examples in the above likelihood function examples. It therefore can be used to\n", + "learn about the linear light profile likelihood function without reading other likelihood scripts." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", + "set below to 28 x 28. \n", + "\n", + "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", + "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", + "mesh, the same number of pixels must be used in the y and x directions.\n", + "\n", + "__Edge Zeroing__\n", + "\n", + "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", + "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", + "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", + "\n", + "For a rectangular mesh, the source code computes edge pixels internally using the known\n", + "pixels at the edge of the mesh. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the \u2018real_space_mask\u2019 which defines the grid the image the galaxy is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(80, 80), pixel_scales=0.05, radius=4.0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the galaxy `Interferometer` dataset `simple` from .fits files, which we will fit \n", + "with the model.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the galaxy to the uv-plane and compare \n", + "directly to the visibilities. We use a non-uniform fast Fourier transform, which is the most efficient method for \n", + "interferometer datasets containing ~1-10 million visibilities. We will discuss how the calculation of the likelihood\n", + "function changes for different methods of Fourier transforming in this guide." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This guide uses in-built visualization tools for plotting. \n", + "\n", + "For example, using the `aplt.subplot_interferometer_dirty_images` the dataset we perform a likelihood evaluation on is plotted.\n", + "\n", + "The `subplot_dataset` displays the visibilities in the uv-plane, which are the raw data of the interferometer\n", + "dataset. These are what will ultimately be directly fitted in the Fourier space.\n", + "\n", + "The `subplot_dirty_images` displays the dirty images of the dataset, which are the reconstructed images of visibilities\n", + "using an inverse Fourier transform to convert these to real-space. These dirty images are not the images we fit, but\n", + "visualization of the dirty images are often used in radio interferometry to show the data in a way that is more\n", + "interpretable to the human eye." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Over sampling evaluates a light profile using multiple samples of its intensity per image-pixel.\n", + "\n", + "For interferometer datasets, over sampling is not used in the pixelization (or for light profiles)\n", + "therefore it is implicitly set to 1 and can be ignored hereafter.\n", + "\n", + "For CCD imaging datasets, over sampling is normally used and slightly changes the likelihood function.\n", + "\n", + "__Masked Image Grid__\n", + "\n", + "To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", + "\n", + "For light profiles these are given by `dataset.lp`, which is a uniform grid of (y,x) Cartesian coordinates\n", + "which have had the 3.0\" circular mask applied.\n", + "\n", + "A pixelization uses a separate grid of (y,x) coordinates, called `dataset.grids.pixelization`, which is\n", + "identical to the light profile grid but may of had a different over-sampling scale applied (but in this example\n", + "does not).\n", + "\n", + "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", + "used to construct a pixelization there is a straight forward mapping between the image data and pixelization pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=dataset.grids.pixelization, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "We set up a lens galaxy with the lens light and mass, which we will use to demonstrate a pixelized source\n", + "reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Pixelization and Regularization__\n", + "\n", + "We combine the pixelization into a single `Galaxy` object.\n", + "\n", + "The galaxy includes the rectangular mesh and constant regularization scheme, which will ultimately be used\n", + "to reconstruct its star forming clumps." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixelization = al.Pixelization(\n", + " mesh=al.mesh.RectangularUniform(mesh_shape),\n", + " regularization=al.reg.Constant(coefficient=1.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the image-plane to its (y,x) \n", + "source-plane coordinate $\\beta$ using the summed deflection angles $\\alpha$ of the mass profiles:\n", + "\n", + " $\\beta = \\theta - \\alpha(\\theta)$\n", + "\n", + "The likelihood function of a pixelized source reconstruction ray-traces two grids from the image-plane to the source-plane:\n", + "\n", + " 1) A 2D grid of (y,x) coordinates aligned with the imaging data's image-pixels.\n", + "\n", + " 2) The sparse 2D grid of (y,x) coordinates above which form the centres of the rectangular source pixels.\n", + "\n", + "The function below computes the 2D deflection angles of the tracer's lens galaxies and subtracts them from the \n", + "image-plane 2D (y,x) coordinates $\\theta$ of each grid, thus ray-tracing their coordinates to the source plane to \n", + "compute their $\\beta$ values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The source code gets quite complex when handling grids for a pixelization, but it is all handled in\n", + "the `TracerToInversion` objects.\n", + "\n", + "The plots at the bottom of this cell show the traced grids used by the source pixelization, showing\n", + "how the rectangular mesh and traced image pixels are constructed." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_to_inversion = al.TracerToInversion(tracer=tracer, dataset=dataset)\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", + "traced_grid_pixelization = tracer.traced_grid_2d_list_from(\n", + " grid=dataset.grids.pixelization\n", + ")[-1]\n", + "\n", + "\n", + "aplt.plot_grid(grid=traced_grid_pixelization, title=\"\")\n", + "\n", + "# grid_plotter.figure_2d()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Border Relocation__\n", + "\n", + "Coordinates that are ray-traced near the mass profile centres are heavily demagnified and may trace to far outskirts of\n", + "the source-plane. \n", + "\n", + "We relocate these pixels (for both grids above) to the edge of the source-plane border (defined via the border of the \n", + "image-plane mask). This is detailed in **HowToLens chapter 4 tutorial 5** and figure 2 of https://arxiv.org/abs/1708.07377." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoarray.inversion.mesh.border_relocator import BorderRelocator\n", + "\n", + "border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1)\n", + "\n", + "relocated_grid = border_relocator.relocated_grid_from(grid=traced_grid_pixelization)\n", + "\n", + "\n", + "aplt.plot_grid(grid=relocated_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Pixel Centre Calculation__\n", + "\n", + "In order to reconstruct the source galaxy using a mesh, we need to determine the centres of the rectangular mesh's \n", + "source pixels.\n", + "\n", + "We do this by overlying a rectangular grid on the relocated traced image-plane grid computed above.\n", + "\n", + "This distributes the rectangular mesh so it fully overlaps the region of the source-plane containing the traced \n", + "image-pixels without having edge pixels that extend beyond this region." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoarray.inversion.mesh.mesh.rectangular_adapt_density import overlay_grid_from\n", + "\n", + "mesh_grid = overlay_grid_from(\n", + " shape_native=mesh_shape, grid=al.Grid2DIrregular(relocated_grid)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolation__\n", + "\n", + "We now combine grids computed above to create an `Interpolator`, which describes how image grid pixel maps to\n", + "every rectangular mesh pixel. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "interpolator = pixelization.mesh.interpolator_from(\n", + " source_plane_data_grid=relocated_grid,\n", + " source_plane_mesh_grid=mesh_grid,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mapper__\n", + "\n", + "We now use the interpolator to create a `Mapper`, which describes the mapping between every image pixel and every \n", + "rectangular pixel, based on the interpolation scheme above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = al.Mapper(interpolator=interpolator)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For the rectangular mesh, the interpolation scheme is called bilinear interpolation, which means that every image \n", + "pixel maps to the rectangular pixel it lands in and the three neighboring rectangular pixels. \n", + "\n", + "The weight of each mapping is determined by the bilinear interpolation scheme, which is a function of how close the \n", + "image pixel is to the centre of the rectangular pixel it lands in and the three neighboring rectangular pixels.\n", + "\n", + "We can print the mappings and weights, for example of the first image pixel, to confirm this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(interpolator.mappings[0])\n", + "print(interpolator.weights[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting the rectangular mesh shows that the source-plane and been discretized into a uniform grid of source pixels.\n", + "\n", + "(To plot the RectangularUniform mesh, we have to convert it to a `Mapper` object, which is described in the next \n", + "likelihood step).\n", + "\n", + "Below, we plot the rectangular mesh without the traced image-grid pixels (for clarity) and with them as \n", + "black dots in order to show how each set of image-pixels fall within a source pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")\n", + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Mapper` contains:\n", + "\n", + " 1) `source_plane_data_grid`: the traced grid of (y,x) image-pixel coordinate centres (`relocated_grid`).\n", + " 2) `source_plane_mesh_grid`: The rectangular mesh of traced (y,x) source-pixel coordinates (`grid_Rectangular`).\n", + "\n", + "We have therefore discretized the source-plane into a rectangular mesh, and can pair every traced image-pixel coordinate\n", + "with the corresponding source pixel it lands in.\n", + "\n", + "This pairing is contained in the ndarray `pix_indexes_for_sub_slim_index` which maps every image-pixel index to \n", + "every source-pixel index.\n", + "\n", + "In the API, the `pix_indexes` refers to the source pixel indexes (e.g. source pixel 0, 1, 2 etc.) and `sub_slim_index` \n", + "refers to the index of an image pixel (e.g. image-pixel 0, 1, 2 etc.). \n", + "\n", + "For example, printing the first ten entries of `pix_indexes_for_sub_slim_index` shows the first ten source-pixel\n", + "indexes these image sub-pixels map too." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pix_indexes_for_sub_slim_index = mapper.pix_indexes_for_sub_slim_index\n", + "\n", + "print(pix_indexes_for_sub_slim_index[0:9])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This array can be used to visualize how an input list of image-pixel indexes map to the source-plane.\n", + "\n", + "It also shows that image-pixel indexing begins from the top-left and goes rightwards and downwards, accounting for \n", + "all image-pixels which are not masked." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(\n", + " array=dataset.dirty_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", + ")\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The reverse mappings of source-pixels to image-pixels can also be used.\n", + "\n", + "If we choose the right source-pixel index, we can see that multiple imaging occur whereby image-pixels in different\n", + "regions of the image-plane are grouped into the same source-pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pix_indexes = [[200]]\n", + "\n", + "indexes = mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)\n", + "\n", + "\n", + "aplt.plot_array(\n", + " array=dataset.dirty_image, title=\"Image\", positions=mapper.image_plane_data_grid\n", + ")\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolation__\n", + "\n", + "The right hand plot shows more laying over source pixel 200 than its retangular black lines. Pixels further \n", + "out than the pixel appear to be mapped to this source pixel. \n", + "\n", + "This is because the mesh uses an interpolation mapping scheme whereby each image pixels is paired with four source \n", + "pixels. For a rectangular mesh, this scheme is called bilinear interpolation, and it means that every pixel maps\n", + "not only to the rectangular source pixel it lands in, but also the three neighbouring source pixels. Interpolation is \n", + "key to ensuring that the pixelization can reconstruct smooth source morphologies.\n", + "\n", + "We can confirm that every image pixel maps to four source pixels by printing \n", + "the `pix_sizes_for_sub_slim_index`, which gives the number of mapped source pixels for every image pixel.\n", + "\n", + "We can also confirm that the interpolation introduces weights to each mapping by printing the \n", + "`pix_weights_for_sub_slim_index`, which gives the weight of each mapping for every image pixel." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(mapper.pix_sizes_for_sub_slim_index[0:9])\n", + "print(mapper.pix_weights_for_sub_slim_index[0:9])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "Lets quickly think about what happens when we use over sampling in the pixelization (e.g. `sub_size>1`). For\n", + "the `sub_size=1` case above, each image pixel maps to 4 source pixels (due to bilinear interpolation)\n", + "with a weight determined from the bilinear interpolation scheme.\n", + "\n", + "However, the default over sampling for a pixelization is `sub_size=4`, meaning each image pixel is divided\n", + "into a 4x4 grid of sub-pixels (16 sub-pixels in total). Each of these sub-pixels maps to 4 source pixels\n", + "(due to bilinear interpolation), where the weight of each mapping is determined by the bilinear interpolation\n", + "scheme divided by 16 (because there are 16 sub-pixels).\n", + "\n", + "This example therefore used a `sub_size=1` to keep the explanation of the likelihood, visualization of the\n", + "arrays above and understanding of the mapping scheme as simple as possible. You can manually increase the\n", + "`sub_size` above and re-run the notebook to see how this changes the mapping scheme.\n", + "\n", + "__Alternative Meshes__\n", + "\n", + "We can briefly consider how this step differs for other mesh types. Above, we simply overlaid a uniform rectangular\n", + "grid to define the source pixel centres and then mapped image pixels to these source pixels.\n", + "\n", + "The `RectangularAdaptDensity` mesh pretty much works exactly the same, its just that a calculation (which we don't\n", + "describe here) works out how to make a grid of rectangular pixels that adapt to the source-plane density and thus\n", + "vary in size. \n", + "\n", + "There is also a `RectangularAdaptImage` mesh which uses the image of the lensed source to adapt\n", + "the rectangular pixel sizes. This often puts even smaller pixels in the brightest regions of the source,\n", + "even if it lies offset or away from the caustic.\n", + "\n", + "There is also a `Delaunay` mesh which uses a Delaunay triangulation to define an irregular grid of source pixels.\n", + "This is described fully in the `delaunay` example including a likelihood function guide.\n", + "\n", + "__Mapping Matrix__\n", + "\n", + "The `mapping_matrix` represents the image-pixel to source-pixel mappings above in a 2D matrix. \n", + "\n", + "It has dimensions `(total_image_pixels, total_source_pixels)`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapping_matrix = al.util.mapper.mapping_matrix_from(\n", + " pix_indexes_for_sub_slim_index=pix_indexes_for_sub_slim_index,\n", + " pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index,\n", + " pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index,\n", + " pixels=mapper.pixels,\n", + " total_mask_pixels=mapper.source_plane_data_grid.mask.pixels_in_mask,\n", + " slim_index_for_sub_slim_index=mapper.slim_index_for_sub_slim_index,\n", + " sub_fraction=mapper.over_sampler.sub_fraction,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `mapping_matrix` shows of all image-source pixel mappings.\n", + "\n", + "No row of pixels has more than one non-zero entry. It is not possible for two image pixels to map to the same source \n", + "pixel (meaning that there are no correlated pixels in the mapping matrix)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(mapping_matrix, aspect=(mapping_matrix.shape[1] / mapping_matrix.shape[0]))\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each column of the `mapping_matrix` can therefore be used to show all image-pixels it maps too. \n", + "\n", + "For example, above, we plotted all image-pixels of source-pixel 200 (as well as 202 and 204). We can extract all\n", + "image-pixel indexes of source pixels 200 using the `mapping_matrix` and use them to plot the image of this\n", + "source-pixel (which corresponds to only values of zeros or ones)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "indexes_source_pix_200 = np.nonzero(mapping_matrix[:, 200])\n", + "\n", + "print(indexes_source_pix_200[0])\n", + "\n", + "array_2d = al.Array2D(values=mapping_matrix[:, 200], mask=dataset.mask)\n", + "\n", + "aplt.plot_array(array=array_2d, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Transformed Mapping Matrix ($f$)__\n", + "\n", + "Each pixelization pixel can therefore be thought of as an image (where all entries of this image are zeros and ones). \n", + "\n", + "However, for interferometer datasets we want to fit the visibilities in the uv-plane, not the image-plane. Therefore,\n", + "each image in the `mapping_matrix` must be transformed to the uv-plane via a Fourier transform, such that each\n", + "column in the `transformed_mapping_matrix` represents the visibilities in the uv-plane of each pixelization pixel.\n", + "\n", + "This operation changes the dimensions of the mapping matrix, meaning the `transformed_mapping_matrix` has\n", + "dimensions `(total_image_pixels, total_visibilities)`. \n", + "\n", + "If the number of visibilities is large (e.g. 10^6) this matrix becomes extremely large and computationally expensive to \n", + "store memory, meaning the sparse operator likelihood function must be used instead.\n", + "\n", + "The `transformed_mapping_matrix` is also complex, storing all entries of the visibilities after the NUFFT as real\n", + "and complex values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "transformed_mapping_matrix = dataset.transformer.transform_mapping_matrix(\n", + " mapping_matrix=mapping_matrix\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A 2D plot of the `transformed_mapping_matrix` shows all visibility-source pixel mappings.\n", + "\n", + "Note how, unlike for the `mapping_matrix`, every row of image-pixels fully consists of non-zero entries. This\n", + "means the matrix is fully dense, making it even more difficult to store in memory for large datasets.\n", + "\n", + "Below, we plot the real and imaginary components of the `transformed_mapping_matrix` separately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " transformed_mapping_matrix.real,\n", + " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()\n", + "\n", + "plt.imshow(\n", + " transformed_mapping_matrix.imag,\n", + " aspect=(transformed_mapping_matrix.shape[1] / transformed_mapping_matrix.shape[0]),\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each column of the `transformed_mapping_matrix` shows all visibilities it maps to after the NUFFT." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "indexes_pix_200 = np.nonzero(transformed_mapping_matrix[:, 200])\n", + "\n", + "print(indexes_pix_200[0])\n", + "\n", + "visibilities = al.Visibilities(visibilities=transformed_mapping_matrix[:, 200])\n", + "\n", + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In Warren & Dye 2003 (https://arxiv.org/abs/astro-ph/0302587) the `transformed_mapping_matrix` is denoted $f_{ij}$\n", + "where $i$ maps over all $I$ source pixels and $j$ maps over all $J$ visibilities. \n", + "\n", + "For example: \n", + "\n", + " - $f_{0, 2} = 0.3$ indicates that visibility number $2$ maps to pixelization pixel $0$ with a weight of $0.3$ after the NUFFT.\n", + "\n", + "The indexing of the `mapping_matrix` is reversed compared to the notation of WD03 (e.g. visibilities\n", + "are the first entry of `mapping_matrix` whereas for $f$ they are the second index)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " f\"Mapping between visibility 0 and RectangularUniform pixel 2 = {mapping_matrix[0, 2]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Data Vector (D)__\n", + "\n", + "To solve for the source pixel fluxes we now pose the problem as a linear inversion.\n", + "\n", + "This requires us to convert the `transformed_mapping_matrix` and our `data` and `noise map` into matrices of certain dimensions. \n", + "\n", + "The `data_vector`, $D$, is the first matrix and it has dimensions `(total_Rectangular_pixels,)`.\n", + "\n", + "In WD03 (https://arxiv.org/abs/astro-ph/0302587) and N15 (https://arxiv.org/abs/1412.7436) the data vector \n", + "is give by: \n", + "\n", + " $\\vec{D}_{i} = \\sum_{\\rm j=1}^{J}f_{ij}(d_{j})/\\sigma_{j}^2 \\, \\, .$\n", + "\n", + "Where:\n", + "\n", + " - $d_{\\rm j}$ are the image-pixel data flux values.\n", + " - $\\sigma{\\rm _j}^2$ are the statistical uncertainties of each image-pixel value.\n", + "\n", + "$i$ maps over all $I$ source pixels and $j$ maps over all $J$ image pixels. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data_vector = (\n", + " al.util.inversion_interferometer.data_vector_via_transformed_mapping_matrix_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " visibilities=dataset.data,\n", + " noise_map=dataset.noise_map,\n", + " )\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$D$ describes which source pixels trace to which visibilities, with associated weights, after the NUFFT. This \n", + "ensures the reconstruction fully accounts for the NUFFT when fitting the data.\n", + "\n", + "We can plot $D$ as a column vector:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(\n", + " data_vector.reshape(data_vector.shape[0], 1), aspect=10.0 / data_vector.shape[0]\n", + ")\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dimensions of $D$ are the number of source pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Data Vector:\")\n", + "print(data_vector)\n", + "print(data_vector.shape)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Matrix (F)__\n", + "\n", + "The `curvature_matrix` $F$ is the second matrix and it has \n", + "dimensions `(total_Rectangular_pixels, total_Rectangular_pixels)`.\n", + "\n", + "In WD03 / N15 (https://arxiv.org/abs/astro-ph/0302587) the curvature matrix is a 2D matrix given by:\n", + "\n", + " ${F}_{ik} = \\sum_{\\rm j=1}^{J}f_{ij}f_{kj}/\\sigma_{j}^2 \\, \\, .$\n", + "\n", + "NOTE: this notation implicitly assumes a summation over $K$, where $k$ runs over all pixelization pixel indexes $K$.\n", + "\n", + "Note how summation over $J$ runs over $f$ twice, such that every entry of $F$ is the sum of the multiplication\n", + "between all values in every two columns of $f$.\n", + "\n", + "For example, $F_{0,1}$ is the sum of all visibility values in $f$ of source pixel 0 multiplied by\n", + "all visibility values of source pixel 1.\n", + "\n", + "Visibilities are both real and complex values, and the `curvature_matrix` is computed separately for the real and\n", + "imaginary components of the visibilities and then summed together." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "real_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.real,\n", + " noise_map=dataset.noise_map.real,\n", + ")\n", + "\n", + "imag_curvature_matrix = al.util.inversion.curvature_matrix_via_mapping_matrix_from(\n", + " mapping_matrix=transformed_mapping_matrix.imag,\n", + " noise_map=dataset.noise_map.imag,\n", + ")\n", + "\n", + "curvature_matrix = np.add(real_curvature_matrix, imag_curvature_matrix)\n", + "\n", + "plt.imshow(curvature_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For $F_{ik}$ to be non-zero, this requires that the images of source pixels $i$ and $k$ share at least one\n", + "image-pixel, which for visibilities after the NUFFT is always true for all $i$ and $k$.\n", + "\n", + "For example, we can see a non-zero entry for $F_{100,101}$ and plotting their images\n", + "show overlap." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_pixel_0 = 0\n", + "source_pixel_1 = 1\n", + "\n", + "print(curvature_matrix[source_pixel_0, source_pixel_1])\n", + "\n", + "visibilities = al.Visibilities(\n", + " visibilities=transformed_mapping_matrix[:, source_pixel_0],\n", + ")\n", + "\n", + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n", + "\n", + "visibilities = al.Visibilities(\n", + " visibilities=transformed_mapping_matrix[:, source_pixel_1],\n", + ")\n", + "\n", + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The following chi-squared is minimized when we perform the inversion and reconstruct the source_galaxy:\n", + "\n", + "$\\chi^2 = \\sum_{\\rm j=1}^{J} \\bigg[ \\frac{(\\sum_{\\rm i=1}^{I} s_{i} f_{ij}) - d_{j}}{\\sigma_{j}} \\bigg]$\n", + "\n", + "Where $s$ is the reconstructed pixel fluxes in all $I$ source pixels.\n", + "\n", + "The solution for $s$ is therefore given by (equation 5 WD03):\n", + "\n", + " $s = F^{-1} D$\n", + "\n", + "We can compute this using NumPy linear algebra:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# Because we are no using regularizartion (see below) it is common for the curvature matrix to be singular and lead\n", + "# to a LinAlgException. The loop below mitigates this -- you can ignore it as it is not important for understanding\n", + "# the PyAutoLens likelihood function.\n", + "\n", + "for i in range(curvature_matrix.shape[0]):\n", + " curvature_matrix[i, i] += 1e-8\n", + "\n", + "reconstruction = np.linalg.solve(curvature_matrix, data_vector)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot this reconstruction -- it looks like a mess.\n", + "\n", + "The pixelization pixels have noisy and unsmooth values, and it is hard to make out if a source galaxy is even being \n", + "reconstructed. \n", + "\n", + "In fact, the linear inversion is (over-)fitting noise in the image data, meaning this system of equations is \n", + "ill-posed. We need to apply some form of smoothing on the reconstruction to avoid over fitting noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Regularization Matrix (H)__\n", + "\n", + "Regularization adds a linear regularization term $G_{\\rm L}$ to the $\\chi^2$ we solve for giving us a new merit \n", + "function $G$ (equation 11 WD03):\n", + "\n", + " $G = \\chi^2 + \\lambda \\, G_{\\rm L}$\n", + "\n", + "where $\\lambda$ is the `regularization_coefficient` which describes the magnitude of smoothness that is applied. A \n", + "higher $\\lambda$ will regularize the source more, leading to a smoother source galaxy reconstruction.\n", + "\n", + "Different forms for $G_{\\rm L}$ can be defined which regularize the reconstruction in different ways. The \n", + "`Constant` regularization scheme used in this example applies gradient regularization (equation 14 WD03):\n", + "\n", + " $G_{\\rm L} = \\sum_{\\rm i}^{I} \\sum_{\\rm n=1}^{N} [s_{i} - s_{i, v}]$\n", + "\n", + "This regularization scheme is easier to express in words -- the summation goes to each rectangular pixelization pixel,\n", + "determines all rectangular pixels with which it shares a direct vertex (e.g. its neighbors) and penalizes solutions \n", + "where the difference in reconstructed flux of these two neighboring pixels is large.\n", + "\n", + "The summation does this for all source pixels, thus it favours solutions where neighboring source\n", + "pixels reconstruct similar values to one another (e.g. it favours a smooth source galaxy reconstruction).\n", + "\n", + "We now define the `regularization matrix`, $H$, which allows us to include this smoothing when we solve for $s$. $H$\n", + "has dimensions `(total_Rectangular_pixels, total_Rectangular_pixels)`.\n", + "\n", + "This relates to $G_{\\rm L}$ as (equation 13 WD03):\n", + "\n", + " $H_{ik} = \\frac{1}{2} \\frac{\\partial G_{\\rm L}}{\\partial s_{i} \\partial s_{k}}$\n", + "\n", + "$H$ has the `regularization_coefficient` $\\lambda$ folded into it such $\\lambda$'s control on the degree of smoothing\n", + "is accounted for." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "regularization_matrix = al.util.regularization.constant_regularization_matrix_from(\n", + " coefficient=source_galaxy.pixelization.regularization.coefficient,\n", + " neighbors=mapper.neighbors,\n", + " neighbors_sizes=mapper.neighbors.sizes,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot the regularization matrix and note that:\n", + "\n", + " - non-zero entries indicate that two source pixels are neighbors and therefore are regularized with one another.\n", + "\n", + " - Zeros indicate the two source pixels do not neighbor one another.\n", + "\n", + "The majority of entries are zero, because the majority of source pixels are not neighbors with one another." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "plt.imshow(regularization_matrix)\n", + "plt.colorbar()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__F + Lamdba H__\n", + "\n", + "$H$ enters the linear algebra system we solve for as follows (WD03 equation (12)):\n", + "\n", + " $s = [F + H]^{-1} D$" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "curvature_reg_matrix = np.add(curvature_matrix, regularization_matrix)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Galaxy Reconstruction (s)__\n", + "\n", + "We can now solve the linear system above using NumPy linear algebra. \n", + "\n", + "Note that the for loop used above to prevent a LinAlgException is no longer required." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = np.linalg.solve(curvature_reg_matrix, data_vector)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting this source galaxy reconstruction we can see that regularization has lead us to reconstruct a smoother \n", + "source galaxy, which actually looks like the star forming clumps in the imaging data! \n", + "\n", + "This also implies we are not over-fitting the noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_grid(grid=mapper.source_plane_mesh_grid, title=\"Source-Plane Mesh Grid\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visibilities Reconstruction__\n", + "\n", + "Using the reconstructed pixel fluxes we can map the reconstruction back to the image plane (via \n", + "the `blurred mapping_matrix`) and produce a reconstruction of the image data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_visibilities = (\n", + " al.util.inversion_interferometer.mapped_reconstructed_visibilities_from(\n", + " transformed_mapping_matrix=transformed_mapping_matrix,\n", + " reconstruction=reconstruction,\n", + " )\n", + ")\n", + "\n", + "mapped_reconstructed_visibilities = al.Visibilities(\n", + " visibilities=mapped_reconstructed_visibilities\n", + ")\n", + "\n", + "aplt.plot_grid(grid=mapped_reconstructed_visibilities.in_grid, title=\"\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our pixelization source galaxy reconstruction. \n", + "\n", + "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", + "\n", + "The likelihood function for source galaxy modeling consists of five terms:\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + s^{T} H s + \\mathrm{ln} \\, \\left[ \\mathrm{det} (F + H) \\right] - { \\mathrm{ln}} \\, \\left[ \\mathrm{det} (H) \\right] + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "This expression was first derived by Suyu 2006 (https://arxiv.org/abs/astro-ph/0601493) and is given by equation (19).\n", + "It was derived into **PyAutoLens** notation in Dye 2008 (https://arxiv.org/abs/0804.4002) equation (5).\n", + "\n", + "We now explain what each of these terms mean.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", + "\n", + " - `model_data` = `mapped_reconstructed_visibilities`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", + "\n", + "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", + "for, corresponding to a fit with a lower likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_visibilities = mapped_reconstructed_visibilities\n", + "\n", + "residual_map = dataset.data - model_visibilities\n", + "\n", + "\n", + "normalized_residual_map_real = (residual_map.real / dataset.noise_map.real).astype(\n", + " \"complex128\"\n", + ")\n", + "normalized_residual_map_imag = (residual_map.imag / dataset.noise_map.imag).astype(\n", + " \"complex128\"\n", + ")\n", + "normalized_residual_map = (\n", + " normalized_residual_map_real + 1j * normalized_residual_map_imag\n", + ")\n", + "\n", + "\n", + "chi_squared_map_real = (residual_map.real / dataset.noise_map.real) ** 2\n", + "chi_squared_map_imag = (residual_map.imag / dataset.noise_map.imag) ** 2\n", + "chi_squared_map = chi_squared_map_real + 1j * chi_squared_map_imag\n", + "\n", + "\n", + "chi_squared_real = np.sum(chi_squared_map.real)\n", + "chi_squared_imag = np.sum(chi_squared_map.imag)\n", + "chi_squared = chi_squared_real + chi_squared_imag\n", + "\n", + "print(chi_squared)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_map = al.Visibilities(visibilities=chi_squared_map)\n", + "\n", + "aplt.plot_grid(grid=chi_squared_map.in_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Regularization Term__\n", + "\n", + "The second term, $s^{T} H s$, corresponds to the $\\lambda $G_{\\rm L}$ regularization term we added to our merit \n", + "function above.\n", + "\n", + "This is the term which sums up the difference in flux of all reconstructed source pixels, and reduces the \n", + "likelihood of solutions where there are large differences in flux (e.g. the source galaxy is less smooth and more \n", + "likely to be overfitting noise).\n", + "\n", + "We compute it below via matrix multiplication, noting that the `regularization_coefficient`, $\\lambda$, is built into \n", + "the `regularization_matrix` already." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "regularization_term = np.matmul(\n", + " reconstruction.T, np.matmul(regularization_matrix, reconstruction)\n", + ")\n", + "\n", + "print(regularization_term)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Complexity Terms__\n", + "\n", + "Up to this point, it is unclear why we chose a value of `regularization_coefficient=1.0`. \n", + "\n", + "We cannot rely on the `chi_squared` and `regularization_term` above to optimally choose its value, because increasing \n", + "the `regularization_coefficient` smooths the solution more and therefore:\n", + "\n", + " - Decreases `chi_squared` by fitting the data worse, producing a lower `log_likelihood`.\n", + "\n", + " - Increases the `regularization_term` by penalizing the differences between source pixel fluxes more, again reducing\n", + " the inferred `log_likelihood`.\n", + "\n", + "If we set the regularization coefficient based purely on these two terms, we would set a value of 0.0 and be back where\n", + "we started over-fitting noise!\n", + "\n", + "The terms $\\left[ \\mathrm{det} (F + H) \\right]$ and $ - { \\mathrm{ln}} \\, \\left[ \\mathrm{det} (H) \\right]$ address \n", + "this problem. \n", + "\n", + "They quantify how complex the reconstruction is, and penalize solutions where *it is more complex*. Reducing \n", + "the `regularization_coefficient` makes the source galaxy reconstruction more complex (because a galaxy that is \n", + "smoothed less uses more flexibility to fit the data better).\n", + "\n", + "These two terms therefore counteract the `chi_squared` and `regularization_term`, so as to attribute a higher\n", + "`log_likelihood` to solutions which fit the data with a more smoothed and less complex source (e.g. one with a higher \n", + "`regularization_coefficient`).\n", + "\n", + "In **HowToGalaxy** -> `chapter 4` -> `tutorial_4_bayesian_regularization` we expand on this further and give a more\n", + "detailed description of how these different terms impact the `log_likelihood_function`. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_curvature_reg_matrix_term = np.linalg.slogdet(curvature_reg_matrix)[1]\n", + "log_regularization_matrix_term = np.linalg.slogdet(regularization_matrix)[1]\n", + "\n", + "print(log_curvature_reg_matrix_term)\n", + "print(log_regularization_matrix_term)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", + "\n", + "The final term ins the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", + "of the log of every noise-map value squared. \n", + "\n", + "Given the `noise_map` is fixed, this term does not change during the lens modeling process and has no impact on the \n", + "model we infer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization_real = np.sum(np.log(2 * np.pi * dataset.noise_map.real**2.0))\n", + "noise_normalization_imag = np.sum(np.log(2 * np.pi * dataset.noise_map.imag**2.0))\n", + "noise_normalization = noise_normalization_real + noise_normalization_imag" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "We can now, finally, compute the `log_likelihood` of the model, by combining the five terms computed above using\n", + "the likelihood function defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_evidence = float(\n", + " -0.5\n", + " * (\n", + " chi_squared\n", + " + regularization_term\n", + " + log_curvature_reg_matrix_term\n", + " - log_regularization_matrix_term\n", + " + noise_normalization\n", + " )\n", + ")\n", + "\n", + "print(log_evidence)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This process to perform a likelihood function evaluation performed via the `FitInterferometer` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitInterferometer(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=al.Settings(use_border_relocator=True),\n", + ")\n", + "fit_log_evidence = fit.log_evidence\n", + "print(fit_log_evidence)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "__Log Likelihood Function: Source Code Speed Up__\n", + "\n", + "The interferometer pixelization likelihood function described in this notebook performs certain calculations using\n", + "functions which are easier to understand, but are computationally slower than the actual source code implementation\n", + "(but the two produce identical results).\n", + "\n", + "We end by pointing out some of these, but we do not provide an step-by-step description of how they work.\n", + "If you are interested, you will need to dive into the source code itself.\n", + "\n", + "**Fast Chi Squared:** The `chi_squared` above is computed using the `transformed_mapping_matrix`, which requires\n", + "many NUFFT's to compute and requires large memroy store. The source code uses a trick which computes the chi-squared\n", + "but bypasses the need to ever compute the `transformed_mapping_matrix`.\n", + "\n", + "**Sparse Operator Curvature Matrix:** The `curvature_matrix` above is also computed using the `transformed_mapping_matrix`, \n", + "which again means slow run times and large memory usage. The source code can instead use sparse operators to \n", + "compute the curvature matrix in a way which again bypasses the need to compute the `transformed_mapping_matrix`.\n", + "\n", + "The two tricks in combination lead to a significant speed up in the likelihood function evaluation and mean that\n", + "the large matrix of size [source pixels, visibilities] never needs to be stored in memory. This is at the heart\n", + "of why lens modeling interferometer data with pixelized source reconstructions is so fast!\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the pixelization likelihood function.\n", + "\n", + "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", + "are described in additional notebooks found in this package. In brief, these describe:\n", + "\n", + " - **Over Sampling**: Oversampling the image grid into a finer grid of sub-pixels, which are all individually \n", + " paired fractionally with each `RectangularAdaptDensity` pixel.\n", + "\n", + " - **Source-plane Interpolation**: Using bilinear interpolation on the `RectangularAdaptDensity` pixelization to pair \n", + " each image (sub-)pixel to multiple `RectangularAdaptDensity` pixels with interpolation weights.\n", + "\n", + " - **Source Morphology Pixelization Adaption**: Adapting the pixelization such that is congregates source pixels around\n", + " the source's brightest regions, as opposed to the magnification-based pixelization used here.\n", + "\n", + " - **Luminosity Weighted Regularization**: Using an adaptive regularization coefficient which adapts the level of \n", + " regularization applied to the source galaxy based on its luminosity." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/pixelization/many_visibilities_preparation.ipynb b/notebooks/interferometer/features/pixelization/many_visibilities_preparation.ipynb index bd9599ffd..92cac8ff1 100644 --- a/notebooks/interferometer/features/pixelization/many_visibilities_preparation.ipynb +++ b/notebooks/interferometer/features/pixelization/many_visibilities_preparation.ipynb @@ -1,334 +1,371 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: Many Visibilities Preparation\n", - "===========================================\n", - "\n", - "To perform many visibility modeling, a matrix called `nufft_precision_operator` is created and used, which encodes information\n", - "and symmetries into the Fourier transform operation performed when modeling interferometer datasets, in a way that\n", - "exploits the sparsity of the pixelized source reconstructions and means a very small amount of memory or VRAM is used.\n", - "\n", - "The details can be found in the source code, but you do not need to know them to do science with the code,\n", - "nevertheless this ultimately means datasets exceeding millions of visibilities can be modeled in under an hour on a GPU.\n", - "\n", - "The time to compute this matrix can vary between seconds and hours, depending on the number of visibilities in the\n", - "dataset, the number of image pixels in the real space mask and if a CPU or GPU is used. If this matrix is not saved\n", - "and loaded from hard disk, this recalculation would need to be performed before every model-fit, which if for your\n", - "setup and hardware takes hours would be prohibitive.\n", - "\n", - "On HPC GPUs via JAX, this computation is fast even for large datasets with many visibilities, with profiling\n", - "of high resolution datasets with over 1 million visibilities showing that computation takes under 20 seconds. For\n", - "10s or 100s of millions of visibilities computation on a GPU may stretch to minutes, but this is still very fast.\n", - "If you are lucky enough to have a modern enough GPU, you can therefore compute this matrix on-the-fly during modeling.\n", - "\n", - "On consumer laptop GPUs or CPU, for datasets with over 100000 visibilities and many pixels in their real-space mask, this\n", - "computation may take 10 minutes or hours (for the small dataset loaded above its miliseconds). Computing it once,\n", - "in this script, saving it to hard-disk, and loading it for modeling is therefore recommended.\n", - "\n", - "On CPU, the `show_progress` input outputs a progress bar to the terminal so you can monitor the computation,\n", - "which is useful when it is slow.\n", - "\n", - "This example therefore creates the `nufft_precision_operator` matrix using independent Python code and saves it to hard-disk\n", - "for modeling. The `cpu_fast_modeling` example loads this matrix from hard-disk if it is available,\n", - "and computes it from scratch if not.\n", - "\n", - "__Contents__\n", - "\n", - "- **High Resolution Dataset:** A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts.\n", - "- **Profiling Dataset:** The code above loads a dataset with very few visibilities and a low resolution real space mask, so.\n", - "- **Curvature Preload:** Pixelized source modeling requires dense linear algebra operations.\n", - "- **Curvature Preload Output:** We now output the `nufft_precision_operator` object to hard-disk, so it can be loaded quickly in.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__High Resolution Dataset__\n", - "\n", - "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large files which\n", - "are too big to include in the main `autolens_workspace` repository:\n", - "\n", - "https://github.com/PyAutoLabs/autolens_workspace_large_files\n", - "\n", - "After downloading the file, place it in the directory:\n", - "\n", - "`autolens_workspace/dataset/interferometer/alma`\n", - "\n", - "You can then compute the `nufft_precision_operator` matrix for this dataset by uncommenting\n", - "the line `dataset_name = \"alma\"` below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import subprocess\n", - "import sys\n", - "import time\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__ \n", - "\n", - "Load the `Interferometer` data, define the visibility and real-space masks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", - ")\n", - "\n", - "dataset_name = \"simple\"\n", - "# dataset_name = \"alma\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not (dataset_path / \"data.fits\").exists():\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Profiling Dataset__\n", - "\n", - "The code above loads a dataset with very few visibilities and a low resolution real space mask, so the \n", - "`nufft_precision_operator` computation is fast.\n", - "\n", - "Real datasets often have 100,000+ visibilities, and a high resolution real space mask, which makes the \n", - "`nufft_precision_operator` computation much slower.\n", - "\n", - "It may therefore be useful to profile the run times for different dataset sizes using the code below, which overwrites \n", - "the dataset above. This will allow you to plan ahead how long the `nufft_precision_operator` computation will take for your \n", - "dataset, and whether doing it on a HPC is necessary.\n", - "\n", - "This code is commented out by default, so your dataset is used instead, but you can uncomment it to run the profiling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# ### Key run time parameters ###\n", - "#\n", - "# mask_radius = 3.0\n", - "# total_visibilities = 1000000\n", - "#\n", - "# ### Setup Data ###\n", - "#\n", - "# real_space_mask = al.Mask2D.circular(\n", - "# shape_native=(800, 800), pixel_scales=0.05, radius=mask_radius\n", - "# )\n", - "#\n", - "# data = al.Visibilities(np.random.normal(loc=0.0, scale=1.0, size=total_visibilities) + 1j * np.random.normal(\n", - "# loc=0.0, scale=1.0, size=total_visibilities\n", - "# ))\n", - "#\n", - "# noise_map = al.VisibilitiesNoiseMap(np.ones(total_visibilities) + 1j * np.ones(total_visibilities))\n", - "#\n", - "# uv_wavelengths = np.random.uniform(\n", - "# low=-300.0, high=300.0, size=(total_visibilities, 2)\n", - "# )\n", - "#\n", - "# dataset = al.Interferometer(\n", - "# data=data,\n", - "# noise_map=noise_map,\n", - "# uv_wavelengths=uv_wavelengths,\n", - "# real_space_mask=real_space_mask,\n", - "# transformer_class=al.TransformerNUFFT,\n", - "# )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Preload__\n", - "\n", - "Pixelized source modeling requires dense linear algebra operations. These calculations are greatly accelerated\n", - "using an alternative mathematical approach called the **sparse linear algebra formalism**.\n", - "\n", - "You do not need to understand the full details of the method, but the key point is:\n", - "\n", - "- It exploits the **sparsity** of the matrices used in pixelized source reconstruction.\n", - "- This leads to a **significant speed-up on GPU or CPU**, using JAX to perform the linear algebra calculations.\n", - "\n", - "To enable this feature, we call `apply_sparse_operator()` on the dataset. This computes and stores a\n", - "preload matrix, which reused in all subsequent pixelized source fits.\n", - "\n", - "As discussed above, the computation of this matrix can take a long time for datasets with many visibilities\n", - "and high resolution real-space masks, unless a modern GPU is used.\n", - "\n", - "We comment out the sparse linear algebra calculation below as we are going to illustrate how you can compute it on CPU.\n", - "\n", - "The code has the following inputs:\n", - "\n", - "- `chunk_k`: The chunk size of visibilities to process at a time. Decreasing this value decreases the memory\n", - " requirements of the computation, but increases the run time. You should set this as high as your system's\n", - " memory allows.\n", - " \n", - "- `show_progress`: Whether to output a progress bar to the terminal, which is on here as for runs which take over\n", - " an hour this is useful to monitor.\n", - "\n", - "- `show_memory`: Whether to output memory usage to the terminal, which is useful to ensure your system has enough\n", - " memory to complete the computation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_sparse_operator(\n", - " use_jax=True,\n", - " chunk_k=2048,\n", - " show_progress=True,\n", - " show_memory=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Curvature Preload Output__\n", - "\n", - "We now output the `nufft_precision_operator` object to hard-disk, so it can be loaded quickly in the \n", - "`cpu_fast_modeling` example.\n", - "\n", - "We save it using a numpy `npy` file, which compresses the data to save hard-disk space, and put it in the \n", - "dataset folder so it can be easily found. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "nufft_precision_operator = dataset.psf_precision_operator_from(\n", - " use_jax=True,\n", - " chunk_k=2048,\n", - " show_progress=True,\n", - " show_memory=True,\n", - ")\n", - "\n", - "np.save(\n", - " file=dataset_path / f\"nufft_precision_operator_{mask_radius}.npy\",\n", - " arr=nufft_precision_operator,\n", - " allow_pickle=False,\n", - ")\n", - "# %%\n", - "'''\n", - "To load the `nufft_precision_operator` matrix from hard-disk in your model-fit, you can use the code:\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "nufft_precision_operator = np.load(\n", - " file=dataset_path / f\"nufft_precision_operator_{mask_radius}.npy\",\n", - " allow_pickle=False,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This example has demonstrated how to set up the linear algebra to perform fast pixelized source modeling on\n", - "interferometer datasets with many visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: Many Visibilities Preparation\n", + "===========================================\n", + "\n", + "To perform many visibility modeling, a matrix called `nufft_precision_operator` is created and used, which encodes information\n", + "and symmetries into the Fourier transform operation performed when modeling interferometer datasets, in a way that\n", + "exploits the sparsity of the pixelized source reconstructions and means a very small amount of memory or VRAM is used.\n", + "\n", + "The details can be found in the source code, but you do not need to know them to do science with the code,\n", + "nevertheless this ultimately means datasets exceeding millions of visibilities can be modeled in under an hour on a GPU.\n", + "\n", + "The time to compute this matrix can vary between seconds and hours, depending on the number of visibilities in the\n", + "dataset, the number of image pixels in the real space mask and if a CPU or GPU is used. If this matrix is not saved\n", + "and loaded from hard disk, this recalculation would need to be performed before every model-fit, which if for your\n", + "setup and hardware takes hours would be prohibitive.\n", + "\n", + "On HPC GPUs via JAX, this computation is fast even for large datasets with many visibilities, with profiling\n", + "of high resolution datasets with over 1 million visibilities showing that computation takes under 20 seconds. For\n", + "10s or 100s of millions of visibilities computation on a GPU may stretch to minutes, but this is still very fast.\n", + "If you are lucky enough to have a modern enough GPU, you can therefore compute this matrix on-the-fly during modeling.\n", + "\n", + "On consumer laptop GPUs or CPU, for datasets with over 100000 visibilities and many pixels in their real-space mask, this\n", + "computation may take 10 minutes or hours (for the small dataset loaded above its miliseconds). Computing it once,\n", + "in this script, saving it to hard-disk, and loading it for modeling is therefore recommended.\n", + "\n", + "On CPU, the `show_progress` input outputs a progress bar to the terminal so you can monitor the computation,\n", + "which is useful when it is slow.\n", + "\n", + "This example therefore creates the `nufft_precision_operator` matrix using independent Python code and saves it to hard-disk\n", + "for modeling. The `cpu_fast_modeling` example loads this matrix from hard-disk if it is available,\n", + "and computes it from scratch if not.\n", + "\n", + "__Contents__\n", + "\n", + "- **High Resolution Dataset:** A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts.\n", + "- **Profiling Dataset:** The code above loads a dataset with very few visibilities and a low resolution real space mask, so.\n", + "- **Curvature Preload:** Pixelized source modeling requires dense linear algebra operations.\n", + "- **Curvature Preload Output:** We now output the `nufft_precision_operator` object to hard-disk, so it can be loaded quickly in.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__High Resolution Dataset__\n", + "\n", + "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large files which\n", + "are too big to include in the main `autolens_workspace` repository:\n", + "\n", + "https://github.com/PyAutoLabs/autolens_workspace_large_files\n", + "\n", + "After downloading the file, place it in the directory:\n", + "\n", + "`autolens_workspace/dataset/interferometer/alma`\n", + "\n", + "You can then compute the `nufft_precision_operator` matrix for this dataset by uncommenting\n", + "the line `dataset_name = \"alma\"` below." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import subprocess\n", + "import sys\n", + "import time\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__ \n", + "\n", + "Load the `Interferometer` data, define the visibility and real-space masks." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", + ")\n", + "\n", + "dataset_name = \"simple\"\n", + "# dataset_name = \"alma\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not (dataset_path / \"data.fits\").exists():\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Profiling Dataset__\n", + "\n", + "The code above loads a dataset with very few visibilities and a low resolution real space mask, so the \n", + "`nufft_precision_operator` computation is fast.\n", + "\n", + "Real datasets often have 100,000+ visibilities, and a high resolution real space mask, which makes the \n", + "`nufft_precision_operator` computation much slower.\n", + "\n", + "It may therefore be useful to profile the run times for different dataset sizes using the code below, which overwrites \n", + "the dataset above. This will allow you to plan ahead how long the `nufft_precision_operator` computation will take for your \n", + "dataset, and whether doing it on a HPC is necessary.\n", + "\n", + "This code is commented out by default, so your dataset is used instead, but you can uncomment it to run the profiling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# ### Key run time parameters ###\n", + "#\n", + "# mask_radius = 3.0\n", + "# total_visibilities = 1000000\n", + "#\n", + "# ### Setup Data ###\n", + "#\n", + "# real_space_mask = al.Mask2D.circular(\n", + "# shape_native=(800, 800), pixel_scales=0.05, radius=mask_radius\n", + "# )\n", + "#\n", + "# data = al.Visibilities(np.random.normal(loc=0.0, scale=1.0, size=total_visibilities) + 1j * np.random.normal(\n", + "# loc=0.0, scale=1.0, size=total_visibilities\n", + "# ))\n", + "#\n", + "# noise_map = al.VisibilitiesNoiseMap(np.ones(total_visibilities) + 1j * np.ones(total_visibilities))\n", + "#\n", + "# uv_wavelengths = np.random.uniform(\n", + "# low=-300.0, high=300.0, size=(total_visibilities, 2)\n", + "# )\n", + "#\n", + "# dataset = al.Interferometer(\n", + "# data=data,\n", + "# noise_map=noise_map,\n", + "# uv_wavelengths=uv_wavelengths,\n", + "# real_space_mask=real_space_mask,\n", + "# transformer_class=al.TransformerNUFFT,\n", + "# )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Preload__\n", + "\n", + "Pixelized source modeling requires dense linear algebra operations. These calculations are greatly accelerated\n", + "using an alternative mathematical approach called the **sparse linear algebra formalism**.\n", + "\n", + "You do not need to understand the full details of the method, but the key point is:\n", + "\n", + "- It exploits the **sparsity** of the matrices used in pixelized source reconstruction.\n", + "- This leads to a **significant speed-up on GPU or CPU**, using JAX to perform the linear algebra calculations.\n", + "\n", + "To enable this feature, we call `apply_sparse_operator()` on the dataset. This computes and stores a\n", + "preload matrix, which reused in all subsequent pixelized source fits.\n", + "\n", + "As discussed above, the computation of this matrix can take a long time for datasets with many visibilities\n", + "and high resolution real-space masks, unless a modern GPU is used.\n", + "\n", + "We comment out the sparse linear algebra calculation below as we are going to illustrate how you can compute it on CPU.\n", + "\n", + "The code has the following inputs:\n", + "\n", + "- `chunk_k`: The chunk size of visibilities to process at a time. Decreasing this value decreases the memory\n", + " requirements of the computation, but increases the run time. You should set this as high as your system's\n", + " memory allows.\n", + " \n", + "- `show_progress`: Whether to output a progress bar to the terminal, which is on here as for runs which take over\n", + " an hour this is useful to monitor.\n", + "\n", + "- `show_memory`: Whether to output memory usage to the terminal, which is useful to ensure your system has enough\n", + " memory to complete the computation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_sparse_operator(\n", + " use_jax=True,\n", + " chunk_k=2048,\n", + " show_progress=True,\n", + " show_memory=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Curvature Preload Output__\n", + "\n", + "We now output the `nufft_precision_operator` object to hard-disk, so it can be loaded quickly in the \n", + "`cpu_fast_modeling` example.\n", + "\n", + "We save it using a numpy `npy` file, which compresses the data to save hard-disk space, and put it in the \n", + "dataset folder so it can be easily found. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "nufft_precision_operator = dataset.psf_precision_operator_from(\n", + " use_jax=True,\n", + " chunk_k=2048,\n", + " show_progress=True,\n", + " show_memory=True,\n", + ")\n", + "\n", + "np.save(\n", + " file=dataset_path / f\"nufft_precision_operator_{mask_radius}.npy\",\n", + " arr=nufft_precision_operator,\n", + " allow_pickle=False,\n", + ")\n", + "# %%\n", + "'''\n", + "To load the `nufft_precision_operator` matrix from hard-disk in your model-fit, you can use the code:\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "nufft_precision_operator = np.load(\n", + " file=dataset_path / f\"nufft_precision_operator_{mask_radius}.npy\",\n", + " allow_pickle=False,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This example has demonstrated how to set up the linear algebra to perform fast pixelized source modeling on\n", + "interferometer datasets with many visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/pixelization/modeling.ipynb b/notebooks/interferometer/features/pixelization/modeling.ipynb index 5959be879..5b82f3cc9 100644 --- a/notebooks/interferometer/features/pixelization/modeling.ipynb +++ b/notebooks/interferometer/features/pixelization/modeling.ipynb @@ -1,742 +1,779 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Features: Pixelization\n", - "======================\n", - "\n", - "A pixelization reconstructs the source galaxy\u2019s light on a grid of pixels, which is regularized using a prior that\n", - "enforces a degree of smoothness in the solution.\n", - "\n", - "This script fits a source galaxy model that uses a pixelization to reconstruct the source\u2019s light. It employs a\n", - "rectangular mesh with a constant regularization scheme, which together form the simplest pixelization and\n", - "regularization choices available. Despite their simplicity, these choices provide fast and accurate solutions.\n", - "\n", - "For simplicity, the lens galaxy\u2019s light is omitted from both the simulated data and the model. For interferometer\n", - "datasets, the lens light is rarely present and this is the common scenario.\n", - "\n", - "You may wish to first read the pixelization/fit.py example, which demonstrates how a pixelized source reconstruction\n", - "is applied to a single dataset.\n", - "\n", - "Pixelizations are covered in detail in Chapter 4 of the HowToLens lecture series.\n", - "\n", - "__CPU Users__\n", - "\n", - "Matrices must be set up for a pixelized source reconstruction to speed up the linear algebra. On GPU, this takes\n", - "seconds, or at most a minute for datasets with tens of millions or more visibilities. On CPU, this can be a lot\n", - "slower, taking over an hour for very large datasets. If you are on CPU, the\n", - "`feature/pixelization/many_visibilities_preparation` example explains how this initial setup can be performed\n", - "before lens modeling and saved to hard disk for fast loading before the model fit.\n", - "\n", - "__Contents__\n", - "\n", - "- **CPU Users:** Matrices must be set up for a pixelized source reconstruction which speed up the linear algebra.\n", - "- **Advantages & Disadvantages:** Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies.\n", - "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **High Resolution Dataset:** A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Sparse Operators:** Pixelized source modeling requires dense linear algebra operations.\n", - "- **Settings:** As discussed above, disable the default position only linear algebra solver so the source.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", - "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Position Likelihood:** We add a penalty term ot the likelihood function, which penalizes models where the brightest.\n", - "- **Brief Description:** Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH`.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Result Use:** There are many things you can do with the result of a pixelixaiton, including analysing the.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "- **Chaining:** Modeling using a pixelization can be more efficient, robust and automated using the non-linear.\n", - "- **HowToLens:** A full description of how pixelizations work, which comes down to a lot of linear algebra, Bayesian.\n", - "\n", - "__Advantages__\n", - "\n", - "Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies. Such structures\n", - "cannot be well approximated by analytic light profiles such as a S\u00e9rsic profile, or even combinations of multiple\n", - "S\u00e9rsic components. pixelizations are therefore required to accurately reconstruct this irregular source-plane light.\n", - "\n", - "Even alternative basis-function approaches, such as shapelets or multi-Gaussian expansions, struggle to accurately\n", - "reconstruct sources with highly complex morphologies or multiple distinct source galaxies.\n", - "\n", - "Pixelized source models are also essential for robustly constraining detailed components of the lens mass\n", - "distribution (e.g. the mass density slope or the presence of dark matter substructure). By fitting all of the lensed\n", - "source light, they reduce degeneracies between the source and lens mass model.\n", - "\n", - "Finally, many science applications aim to study the highly magnified source galaxy itself, in order to learn about\n", - "distant and intrinsically faint galaxies. pixelizations reconstruct the unlensed source emission, enabling detailed\n", - "studies of the source-plane structure.\n", - "\n", - "For CCD imaging, a disadvantage of pixelized source reconstructions is that they are the most computationally\n", - "expensive modeling approach. For interferometer datasets, the situation is more nuanced: light-profile modeling\n", - "via the JAX-native `TransformerNUFFT` (backed by `nufftax`) is now also fast at large visibility counts, so\n", - "pixelizations are no longer required purely for performance. Their continuing strengths are morphological\n", - "flexibility (irregular sources) and VRAM efficiency on very large real-space masks, where their sparsity-aware\n", - "linear algebra still wins.\n", - "\n", - "__Disadvantages__\n", - "\n", - "Lens modeling with pixelizations is conceptually more complex. There are additional failure modes, such as\n", - "solutions where the source is reconstructed in a highly demagnified configuration due to an unphysical lens mass\n", - "model (e.g. too little or too much mass). These issues are discussed in detail later in the workspace.\n", - "\n", - "As a result, learning to successfully fit lens models with pixelizations typically requires more time and experience\n", - "than the simpler modeling approaches introduced elsewhere in the workspace.\n", - "\n", - "__Positive Only Solver__\n", - "\n", - "Many codes which use linear algebra typically rely on a linear algabra solver which allows for positive and negative\n", - "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", - "\n", - "This could be problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", - "light, which is clearly unphysical. For a pixelizaiton, this often produces negative source pixels which over-fit\n", - "the data, producing unphysical solutions.\n", - "\n", - "For CCD imaging datsets pixelized source reconstructions use a positive-only solver, meaning that every source-pixel\n", - "is only allowed to reconstruct positive flux values. This ensures that the source reconstruction is physical and\n", - "that we don't reconstruct negative flux values that don't exist in the real source galaxy (a common systematic\n", - "solution in lens analysis).\n", - "\n", - "However, for interferometer datasets this positive-only solver is often disabled, because negative pixel values\n", - "can be observed from the measurement process. All interferometer examples therefore disable the positive only solver,\n", - "but you may want to consider if using the positive-only solver is appropriate for your dataset.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's surface-brightness is reconstructed using a `RectangularAdaptDensity` mesh\n", - " and `Constant` regularization scheme.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook.\n", - "\n", - "__High Resolution Dataset__\n", - "\n", - "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large files which\n", - "are too big to include in the main `autolens_workspace` repository:\n", - "\n", - "https://github.com/PyAutoLabs/autolens_workspace_large_files\n", - "\n", - "After downloading the file, place it in the directory:\n", - "\n", - "`autolens_workspace/dataset/interferometer/alma`\n", - "\n", - "You can then perform modeling using this high-resolution dataset by uncommenting the relevant line of code\n", - "below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit \n", - "with the lens model.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and compare \n", - "directly to the visiblities. We use a non-uniform fast Fourier transform, which is the most efficient method for \n", - "interferometer datasets containing ~1-10 million visibilities.\n", - "\n", - "If you want to use the high resolution ALMA dataset, uncomment the relevant lines of code below after downloading\n", - "the data from the repository described in the \"High Resolution Dataset\" section above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "Pixelized source modeling requires dense linear algebra operations. These calculations are greatly accelerated\n", - "using an alternative mathematical approach called the **sparse linear algebra formalism**.\n", - "\n", - "You do not need to understand the full details of the method, but the key point is:\n", - "\n", - "- It exploits the **sparsity** of the matrices used in pixelized source reconstruction, reducing memory usage.\n", - "- This leads to a **significant speed-up on GPU or CPU**, using JAX to perform the linear algebra calculations.\n", - "\n", - "To enable this feature, we call `apply_sparse_operator()` on the dataset. This computes and stores a matrix\n", - "which represents how the NUFFT is applied to the noise-map, which reused in all subsequent pixelized source fits.\n", - "\n", - "On GPU via JAX, this computation is fast even for large datasets with many visibilities, with profiling\n", - "of high resolution datasets with over 1 million visibilities showing that computation takes under 20 seconds. For\n", - "10s or 100s of millions of visibilities computation on a GPU may stretch to minutes, but this is still very fast.\n", - "\n", - "On CPU, for datasets with over 100000 visibilities and many pixels in their real-space mask, this computation\n", - "can take 10 minutes or hours (for the small dataset loaded above its miliseconds). The `show_progress` input outputs \n", - "a progress bar to the terminal so you can monitor the computation, which is useful when it is slow.\n", - "\n", - "When computing it is slow, it is recommend you compute it once, save it to hard-disk, and load it\n", - "before modeling. The example `pixelization/many_visibilities_preparation.py` illustrates how to do this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = dataset.apply_sparse_operator(use_jax=True, show_progress=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "As discussed above, disable the default position only linear algebra solver so the source\n", - "reconstruction can have negative pixel values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling is used, \n", - "which evaluates light profiles on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", - "calculations are performed without over sampling.\n", - "\n", - "__Mesh Shape__\n", - "\n", - "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", - "set below to 28 x 28. \n", - "\n", - "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", - "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", - "mesh, the same number of pixels must be used in the y and x directions.\n", - "\n", - "__Edge Zeroing__\n", - "\n", - "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", - "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", - "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", - "\n", - "For a rectangular mesh, the source code computes edge pixels internally using the known\n", - "pixels at the edge of the mesh. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In this \n", - "example fits a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source-galaxy's light uses a 20 x 20 `RectangularAdaptDensity` mesh [0 parameters].\n", - "\n", - " - This pixelization is regularized using a `Constant` scheme which smooths every source pixel equally [1 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=6. \n", - "\n", - "It is worth noting the pixelization fits the source using significantly fewer parameters (1 parameter for \n", - "regularization) than fitting the source using light profiles or an MGE (4+ parameters). \n", - "\n", - "The lens model therefore includes a mesh and regularization scheme, which are used together to create the \n", - "pixelization. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.PowerLaw)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", - "regularization = af.Model(al.reg.Constant)\n", - "\n", - "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the source galaxy's has a mesh and regularization scheme, which are combined into a pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\"),\n", - " name=\"pixelization\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=50000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Position Likelihood__\n", - "\n", - "We add a penalty term ot the likelihood function, which penalizes models where the brightest multiple images of\n", - "the lensed source galaxy do not trace close to one another in the source plane. This removes \"demagnified source\n", - "solutions\" from the source pixelization, which one is likely to infer without this penalty.\n", - "\n", - "A comprehensive description of why we do this is given at the following readthedocs page. I strongly recommend you \n", - "read this page in full if you are not familiar with the positions likelihood penalty and demagnified source \n", - "reconstructions:\n", - "\n", - " https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", - "\n", - "__Brief Description__\n", - "\n", - "Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH` object, which\n", - "includes the positions we loaded above, alongside a `threshold`.\n", - "\n", - "This is because `Inversion`'s suffer a bias whereby they fit unphysical lens models where the source galaxy is \n", - "reconstructed as a demagnified version of the lensed source. \n", - "\n", - "To prevent these solutions biasing the model-fit we specify a `position_threshold` of 0.5\", which requires that a \n", - "mass model traces the four (y,x) coordinates specified by our positions (that correspond to the brightest regions of the \n", - "lensed source) within 0.5\" of one another in the source-plane. If this criteria is not met, a large penalty term is\n", - "added to likelihood that massively reduces the overall likelihood. This penalty is larger if the ``positions``\n", - "trace further from one another.\n", - "\n", - "This ensures the unphysical solutions that bias a pixelization have a lower likelihood that the physical solutions\n", - "we desire. Furthermore, the penalty term reduces as the image-plane multiple image positions trace closer in the \n", - "source-plane, ensuring Nautilus converges towards an accurate mass model. It does this very fast, as \n", - "ray-tracing just a few multiple image positions is computationally cheap. \n", - "\n", - "The threshold of 0.3\" is large. For an accurate lens model we would anticipate the positions trace within < 0.01\" of\n", - "one another. The high threshold ensures only the initial mass models at the start of the fit are penalized.\n", - "\n", - "Position thresholding is described in more detail in the \n", - "script `autolens_workspace/*/guides/modeling/customize`\n", - "\n", - "The arc-second positions of the multiply imaged lensed source galaxy were drawn onto the\n", - "image via the GUI described in the file `autolens_workspace/*/imaging/data_preparation/gui/positions.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n", - "\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisInterferometer` object defining how the via Nautilus the model is fitted to the data. \n", - "\n", - "The `positions_likelihood_list` is passed to the analysis, which applies the likelihood penalty described above\n", - "for everyone lens mass model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " settings=settings,\n", - " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM__\n", - "\n", - "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", - "required by a model.\n", - "\n", - "Pixelizations use a lot less VRAM than light profile-only models, provided the sparse operator \n", - "formalism is used (as it is above). In this mode, datasets with tens of millions of visibilities and real space\n", - "masks with pixel scales below 0.05\" can be stored in just GB's of VRAM, which is remarkable given how much\n", - "data they contain.\n", - "\n", - "In sparse operator mode, the **amount of VRAM used is independent of the number of visibilities in the dataset**. \n", - "This is because the sparse operator compresses all the visibility information into a matrix whose size depends only on\n", - "the number of pixels in the real-space mask. VRAM use is therefore mostly driven by how many pixels are in the real space mask.\n", - "\n", - "VRAM does scale with batch size though, and for high resoluiton datasets may require you to reduce from the value of\n", - "20 set above if your GPU does not have too much VRAM (e.g. < 4GB).\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "The run time of a pixelization are fast provided that the GPU VRAM exceeds the amount of memory required to perform\n", - "a likelihood evaluation.\n", - "\n", - "The **run times of a pixelization are independent of the number of visibilities in the dataset**. This is again \n", - "because the sparse operator method compresses all the visibility information into the `nufft_precision_operator` matrix, whose size \n", - "depends only on the number of pixels in the real-space mask.\n", - "\n", - "Therefore, like VRAM, the main driver of trun time is the number of pixels in the real-space mask,\n", - "not the number of visibilities in the dataset. The calculation also runs the same speed irrespective of whether\n", - "the real space mask is circular, or irregularly shaped, therefore using a circlular mask is recommended as it is\n", - "simpler to set up.\n", - "\n", - "Assuming the use of a 20 x 20 mesh grid above means this is the case, the run times of this model-fit on a GPU\n", - "should take under 10 minutes. Increasing the batch size will speed up the fit, provided VRAM allows it.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this\n", - "does not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix this):\n", - "\n", - "This confirms that the source galaxy's has a mesh and regularization scheme, which are combined into a pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "The end of this example provides a detailed description of all result options for a pixelization." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The example `pixelization/fit` provides a full description of the different calculations that can be performed\n", - "with the result of a pixelization model-fit.\n", - "\n", - "__Result Use__\n", - "\n", - "There are many things you can do with the result of a pixelixaiton, including analysing the reconstructing source, \n", - "magnification calculations of the source and much more.\n", - "\n", - "These are documented in the `fit.py` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = result.max_log_likelihood_fit.inversion" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Science (Magnification, Flux and More)__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", - "size of the source.\n", - "\n", - "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", - "including examples using a pixelized source reconstruction. \n", - "\n", - "If you want to study the source galaxy after modeling has reconstructed its unlensed, then check out this example.\n", - "\n", - "__Wrap Up__\n", - "\n", - "pixelizations are the most complex but also most powerful way to model a source galaxy.\n", - "\n", - "Whether you need to use them or not depends on the science you are doing. If you are only interested in measuring a\n", - "simple quantity like the Einstein radius of a lens, you can get away with using light profiles like a Sersic, MGE or \n", - "shapelets to model the source. Low resolution data also means that using a pixelization is not necessary, as the\n", - "complex structure of the source galaxy is not resolved anyway.\n", - "\n", - "However, fitting complex mass models (e.g. a power-law, stellar / dark model or dark matter substructure) requires \n", - "this level of complexity in the source model. Furthermore, if you are interested in studying the properties of the\n", - "source itself, you won't find a better way to do this than using a pixelization.\n", - "\n", - "__Chaining__\n", - "\n", - "Modeling using a pixelization can be more efficient, robust and automated using the non-linear chaining feature to \n", - "compose a pipeline which begins by fitting a simpler model using a parametric source.\n", - "\n", - "More information on chaining is provided in the `autolens_workspace/notebooks/guides/modeling/chaining` folder,\n", - "chapter 3 of the **HowToLens** lectures.\n", - "\n", - "__HowToLens__\n", - "\n", - "A full description of how pixelizations work, which comes down to a lot of linear algebra, Bayesian statistics and\n", - "2D geometry, is provided in chapter 4 of the **HowToLens** lectures.\n", - "\n", - "__Future Ideas / Contributions__\n", - "\n", - "Here are a list of things I would like to add to this tutorial but haven't found the time. If you are interested\n", - "in having a go at adding them contact me on SLACK! :)\n", - "\n", - "- More magnification calculations.\n", - "- Source gradient calculations.\n", - "- A calculation which shows differential lensing effects (e.g. magnification across the source plane)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Features: Pixelization\n", + "======================\n", + "\n", + "A pixelization reconstructs the source galaxy\u2019s light on a grid of pixels, which is regularized using a prior that\n", + "enforces a degree of smoothness in the solution.\n", + "\n", + "This script fits a source galaxy model that uses a pixelization to reconstruct the source\u2019s light. It employs a\n", + "rectangular mesh with a constant regularization scheme, which together form the simplest pixelization and\n", + "regularization choices available. Despite their simplicity, these choices provide fast and accurate solutions.\n", + "\n", + "For simplicity, the lens galaxy\u2019s light is omitted from both the simulated data and the model. For interferometer\n", + "datasets, the lens light is rarely present and this is the common scenario.\n", + "\n", + "You may wish to first read the pixelization/fit.py example, which demonstrates how a pixelized source reconstruction\n", + "is applied to a single dataset.\n", + "\n", + "Pixelizations are covered in detail in Chapter 4 of the HowToLens lecture series.\n", + "\n", + "__CPU Users__\n", + "\n", + "Matrices must be set up for a pixelized source reconstruction to speed up the linear algebra. On GPU, this takes\n", + "seconds, or at most a minute for datasets with tens of millions or more visibilities. On CPU, this can be a lot\n", + "slower, taking over an hour for very large datasets. If you are on CPU, the\n", + "`feature/pixelization/many_visibilities_preparation` example explains how this initial setup can be performed\n", + "before lens modeling and saved to hard disk for fast loading before the model fit.\n", + "\n", + "__Contents__\n", + "\n", + "- **CPU Users:** Matrices must be set up for a pixelized source reconstruction which speed up the linear algebra.\n", + "- **Advantages & Disadvantages:** Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies.\n", + "- **Positive Only Solver:** Ensuring positive-only solutions for linear light profile intensities.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **High Resolution Dataset:** A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Sparse Operators:** Pixelized source modeling requires dense linear algebra operations.\n", + "- **Settings:** As discussed above, disable the default position only linear algebra solver so the source.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Mesh Shape:** The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the.\n", + "- **Edge Zeroing:** By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Position Likelihood:** We add a penalty term ot the likelihood function, which penalizes models where the brightest.\n", + "- **Brief Description:** Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH`.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **VRAM:** The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Result Use:** There are many things you can do with the result of a pixelixaiton, including analysing the.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "- **Chaining:** Modeling using a pixelization can be more efficient, robust and automated using the non-linear.\n", + "- **HowToLens:** A full description of how pixelizations work, which comes down to a lot of linear algebra, Bayesian.\n", + "\n", + "__Advantages__\n", + "\n", + "Many strongly lensed source galaxies exhibit complex, asymmetric, and irregular morphologies. Such structures\n", + "cannot be well approximated by analytic light profiles such as a S\u00e9rsic profile, or even combinations of multiple\n", + "S\u00e9rsic components. pixelizations are therefore required to accurately reconstruct this irregular source-plane light.\n", + "\n", + "Even alternative basis-function approaches, such as shapelets or multi-Gaussian expansions, struggle to accurately\n", + "reconstruct sources with highly complex morphologies or multiple distinct source galaxies.\n", + "\n", + "Pixelized source models are also essential for robustly constraining detailed components of the lens mass\n", + "distribution (e.g. the mass density slope or the presence of dark matter substructure). By fitting all of the lensed\n", + "source light, they reduce degeneracies between the source and lens mass model.\n", + "\n", + "Finally, many science applications aim to study the highly magnified source galaxy itself, in order to learn about\n", + "distant and intrinsically faint galaxies. pixelizations reconstruct the unlensed source emission, enabling detailed\n", + "studies of the source-plane structure.\n", + "\n", + "For CCD imaging, a disadvantage of pixelized source reconstructions is that they are the most computationally\n", + "expensive modeling approach. For interferometer datasets, the situation is more nuanced: light-profile modeling\n", + "via the JAX-native `TransformerNUFFT` (backed by `nufftax`) is now also fast at large visibility counts, so\n", + "pixelizations are no longer required purely for performance. Their continuing strengths are morphological\n", + "flexibility (irregular sources) and VRAM efficiency on very large real-space masks, where their sparsity-aware\n", + "linear algebra still wins.\n", + "\n", + "__Disadvantages__\n", + "\n", + "Lens modeling with pixelizations is conceptually more complex. There are additional failure modes, such as\n", + "solutions where the source is reconstructed in a highly demagnified configuration due to an unphysical lens mass\n", + "model (e.g. too little or too much mass). These issues are discussed in detail later in the workspace.\n", + "\n", + "As a result, learning to successfully fit lens models with pixelizations typically requires more time and experience\n", + "than the simpler modeling approaches introduced elsewhere in the workspace.\n", + "\n", + "__Positive Only Solver__\n", + "\n", + "Many codes which use linear algebra typically rely on a linear algabra solver which allows for positive and negative\n", + "values of the solution (e.g. `np.linalg.solve`), because they are computationally fast.\n", + "\n", + "This could be problematic, as it means that negative surface brightnesses values can be computed to represent a galaxy's\n", + "light, which is clearly unphysical. For a pixelizaiton, this often produces negative source pixels which over-fit\n", + "the data, producing unphysical solutions.\n", + "\n", + "For CCD imaging datsets pixelized source reconstructions use a positive-only solver, meaning that every source-pixel\n", + "is only allowed to reconstruct positive flux values. This ensures that the source reconstruction is physical and\n", + "that we don't reconstruct negative flux values that don't exist in the real source galaxy (a common systematic\n", + "solution in lens analysis).\n", + "\n", + "However, for interferometer datasets this positive-only solver is often disabled, because negative pixel values\n", + "can be observed from the measurement process. All interferometer examples therefore disable the positive only solver,\n", + "but you may want to consider if using the positive-only solver is appropriate for your dataset.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's surface-brightness is reconstructed using a `RectangularAdaptDensity` mesh\n", + " and `Constant` regularization scheme.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook.\n", + "\n", + "__High Resolution Dataset__\n", + "\n", + "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large files which\n", + "are too big to include in the main `autolens_workspace` repository:\n", + "\n", + "https://github.com/PyAutoLabs/autolens_workspace_large_files\n", + "\n", + "After downloading the file, place it in the directory:\n", + "\n", + "`autolens_workspace/dataset/interferometer/alma`\n", + "\n", + "You can then perform modeling using this high-resolution dataset by uncommenting the relevant line of code\n", + "below." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit \n", + "with the lens model.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and compare \n", + "directly to the visiblities. We use a non-uniform fast Fourier transform, which is the most efficient method for \n", + "interferometer datasets containing ~1-10 million visibilities.\n", + "\n", + "If you want to use the high resolution ALMA dataset, uncomment the relevant lines of code below after downloading\n", + "the data from the repository described in the \"High Resolution Dataset\" section above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "Pixelized source modeling requires dense linear algebra operations. These calculations are greatly accelerated\n", + "using an alternative mathematical approach called the **sparse linear algebra formalism**.\n", + "\n", + "You do not need to understand the full details of the method, but the key point is:\n", + "\n", + "- It exploits the **sparsity** of the matrices used in pixelized source reconstruction, reducing memory usage.\n", + "- This leads to a **significant speed-up on GPU or CPU**, using JAX to perform the linear algebra calculations.\n", + "\n", + "To enable this feature, we call `apply_sparse_operator()` on the dataset. This computes and stores a matrix\n", + "which represents how the NUFFT is applied to the noise-map, which reused in all subsequent pixelized source fits.\n", + "\n", + "On GPU via JAX, this computation is fast even for large datasets with many visibilities, with profiling\n", + "of high resolution datasets with over 1 million visibilities showing that computation takes under 20 seconds. For\n", + "10s or 100s of millions of visibilities computation on a GPU may stretch to minutes, but this is still very fast.\n", + "\n", + "On CPU, for datasets with over 100000 visibilities and many pixels in their real-space mask, this computation\n", + "can take 10 minutes or hours (for the small dataset loaded above its miliseconds). The `show_progress` input outputs \n", + "a progress bar to the terminal so you can monitor the computation, which is useful when it is slow.\n", + "\n", + "When computing it is slow, it is recommend you compute it once, save it to hard-disk, and load it\n", + "before modeling. The example `pixelization/many_visibilities_preparation.py` illustrates how to do this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = dataset.apply_sparse_operator(use_jax=True, show_progress=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "As discussed above, disable the default position only linear algebra solver so the source\n", + "reconstruction can have negative pixel values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling is used, \n", + "which evaluates light profiles on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", + "calculations are performed without over sampling.\n", + "\n", + "__Mesh Shape__\n", + "\n", + "The `mesh_shape` parameter defines number of pixels used by the rectangular mesh to reconstruct the source,\n", + "set below to 28 x 28. \n", + "\n", + "The `mesh_shape` must be fixed before modeling and cannot be a free parameter of the model, because JAX uses the\n", + "mesh shape to define static shaped arrays which use the mesh to reconstruct the source. For a rectangular\n", + "mesh, the same number of pixels must be used in the y and x directions.\n", + "\n", + "__Edge Zeroing__\n", + "\n", + "By default, all pixels at the edge of the mesh in the source-plane are forced to solutions of zero brightness by \n", + "the linear algebra solver. This prevents unphysical solutions where pixels at the edge of the mesh reconstruct \n", + "bright surface brightnesses, often because they fit residuals from the lens light subtraction.\n", + "\n", + "For a rectangular mesh, the source code computes edge pixels internally using the known\n", + "pixels at the edge of the mesh. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In this \n", + "example fits a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source-galaxy's light uses a 20 x 20 `RectangularAdaptDensity` mesh [0 parameters].\n", + "\n", + " - This pixelization is regularized using a `Constant` scheme which smooths every source pixel equally [1 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=6. \n", + "\n", + "It is worth noting the pixelization fits the source using significantly fewer parameters (1 parameter for \n", + "regularization) than fitting the source using light profiles or an MGE (4+ parameters). \n", + "\n", + "The lens model therefore includes a mesh and regularization scheme, which are used together to create the \n", + "pixelization. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.PowerLaw)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", + "regularization = af.Model(al.reg.Constant)\n", + "\n", + "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the source galaxy's has a mesh and regularization scheme, which are combined into a pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\"),\n", + " name=\"pixelization\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=50000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Position Likelihood__\n", + "\n", + "We add a penalty term ot the likelihood function, which penalizes models where the brightest multiple images of\n", + "the lensed source galaxy do not trace close to one another in the source plane. This removes \"demagnified source\n", + "solutions\" from the source pixelization, which one is likely to infer without this penalty.\n", + "\n", + "A comprehensive description of why we do this is given at the following readthedocs page. I strongly recommend you \n", + "read this page in full if you are not familiar with the positions likelihood penalty and demagnified source \n", + "reconstructions:\n", + "\n", + " https://pyautolens.readthedocs.io/en/latest/general/demagnified_solutions.html\n", + "\n", + "__Brief Description__\n", + "\n", + "Unlike other example scripts, we also pass the `AnalysisImaging` object below a `PositionsLH` object, which\n", + "includes the positions we loaded above, alongside a `threshold`.\n", + "\n", + "This is because `Inversion`'s suffer a bias whereby they fit unphysical lens models where the source galaxy is \n", + "reconstructed as a demagnified version of the lensed source. \n", + "\n", + "To prevent these solutions biasing the model-fit we specify a `position_threshold` of 0.5\", which requires that a \n", + "mass model traces the four (y,x) coordinates specified by our positions (that correspond to the brightest regions of the \n", + "lensed source) within 0.5\" of one another in the source-plane. If this criteria is not met, a large penalty term is\n", + "added to likelihood that massively reduces the overall likelihood. This penalty is larger if the ``positions``\n", + "trace further from one another.\n", + "\n", + "This ensures the unphysical solutions that bias a pixelization have a lower likelihood that the physical solutions\n", + "we desire. Furthermore, the penalty term reduces as the image-plane multiple image positions trace closer in the \n", + "source-plane, ensuring Nautilus converges towards an accurate mass model. It does this very fast, as \n", + "ray-tracing just a few multiple image positions is computationally cheap. \n", + "\n", + "The threshold of 0.3\" is large. For an accurate lens model we would anticipate the positions trace within < 0.01\" of\n", + "one another. The high threshold ensures only the initial mass models at the start of the fit are penalized.\n", + "\n", + "Position thresholding is described in more detail in the \n", + "script `autolens_workspace/*/guides/modeling/customize`\n", + "\n", + "The arc-second positions of the multiply imaged lensed source galaxy were drawn onto the\n", + "image via the GUI described in the file `autolens_workspace/*/imaging/data_preparation/gui/positions.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n", + "\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisInterferometer` object defining how the via Nautilus the model is fitted to the data. \n", + "\n", + "The `positions_likelihood_list` is passed to the analysis, which applies the likelihood penalty described above\n", + "for everyone lens mass model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " settings=settings,\n", + " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM__\n", + "\n", + "The `modeling` example explains how VRAM is used during GPU-based fitting and how to print the estimated VRAM \n", + "required by a model.\n", + "\n", + "Pixelizations use a lot less VRAM than light profile-only models, provided the sparse operator \n", + "formalism is used (as it is above). In this mode, datasets with tens of millions of visibilities and real space\n", + "masks with pixel scales below 0.05\" can be stored in just GB's of VRAM, which is remarkable given how much\n", + "data they contain.\n", + "\n", + "In sparse operator mode, the **amount of VRAM used is independent of the number of visibilities in the dataset**. \n", + "This is because the sparse operator compresses all the visibility information into a matrix whose size depends only on\n", + "the number of pixels in the real-space mask. VRAM use is therefore mostly driven by how many pixels are in the real space mask.\n", + "\n", + "VRAM does scale with batch size though, and for high resoluiton datasets may require you to reduce from the value of\n", + "20 set above if your GPU does not have too much VRAM (e.g. < 4GB).\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "The run time of a pixelization are fast provided that the GPU VRAM exceeds the amount of memory required to perform\n", + "a likelihood evaluation.\n", + "\n", + "The **run times of a pixelization are independent of the number of visibilities in the dataset**. This is again \n", + "because the sparse operator method compresses all the visibility information into the `nufft_precision_operator` matrix, whose size \n", + "depends only on the number of pixels in the real-space mask.\n", + "\n", + "Therefore, like VRAM, the main driver of trun time is the number of pixels in the real-space mask,\n", + "not the number of visibilities in the dataset. The calculation also runs the same speed irrespective of whether\n", + "the real space mask is circular, or irregularly shaped, therefore using a circlular mask is recommended as it is\n", + "simpler to set up.\n", + "\n", + "Assuming the use of a 20 x 20 mesh grid above means this is the case, the run times of this model-fit on a GPU\n", + "should take under 10 minutes. Increasing the batch size will speed up the fit, provided VRAM allows it.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format (if this\n", + "does not display clearly on your screen refer to `start_here.ipynb` for a description of how to fix this):\n", + "\n", + "This confirms that the source galaxy's has a mesh and regularization scheme, which are combined into a pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "The end of this example provides a detailed description of all result options for a pixelization." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The example `pixelization/fit` provides a full description of the different calculations that can be performed\n", + "with the result of a pixelization model-fit.\n", + "\n", + "__Result Use__\n", + "\n", + "There are many things you can do with the result of a pixelixaiton, including analysing the reconstructing source, \n", + "magnification calculations of the source and much more.\n", + "\n", + "These are documented in the `fit.py` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = result.max_log_likelihood_fit.inversion" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Science (Magnification, Flux and More)__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", + "size of the source.\n", + "\n", + "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", + "including examples using a pixelized source reconstruction. \n", + "\n", + "If you want to study the source galaxy after modeling has reconstructed its unlensed, then check out this example.\n", + "\n", + "__Wrap Up__\n", + "\n", + "pixelizations are the most complex but also most powerful way to model a source galaxy.\n", + "\n", + "Whether you need to use them or not depends on the science you are doing. If you are only interested in measuring a\n", + "simple quantity like the Einstein radius of a lens, you can get away with using light profiles like a Sersic, MGE or \n", + "shapelets to model the source. Low resolution data also means that using a pixelization is not necessary, as the\n", + "complex structure of the source galaxy is not resolved anyway.\n", + "\n", + "However, fitting complex mass models (e.g. a power-law, stellar / dark model or dark matter substructure) requires \n", + "this level of complexity in the source model. Furthermore, if you are interested in studying the properties of the\n", + "source itself, you won't find a better way to do this than using a pixelization.\n", + "\n", + "__Chaining__\n", + "\n", + "Modeling using a pixelization can be more efficient, robust and automated using the non-linear chaining feature to \n", + "compose a pipeline which begins by fitting a simpler model using a parametric source.\n", + "\n", + "More information on chaining is provided in the `autolens_workspace/notebooks/guides/modeling/chaining` folder,\n", + "chapter 3 of the **HowToLens** lectures.\n", + "\n", + "__HowToLens__\n", + "\n", + "A full description of how pixelizations work, which comes down to a lot of linear algebra, Bayesian statistics and\n", + "2D geometry, is provided in chapter 4 of the **HowToLens** lectures.\n", + "\n", + "__Future Ideas / Contributions__\n", + "\n", + "Here are a list of things I would like to add to this tutorial but haven't found the time. If you are interested\n", + "in having a go at adding them contact me on SLACK! :)\n", + "\n", + "- More magnification calculations.\n", + "- Source gradient calculations.\n", + "- A calculation which shows differential lensing effects (e.g. magnification across the source plane)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/pixelization/slam.ipynb b/notebooks/interferometer/features/pixelization/slam.ipynb index 93cfe501e..d54a47e9b 100644 --- a/notebooks/interferometer/features/pixelization/slam.ipynb +++ b/notebooks/interferometer/features/pixelization/slam.ipynb @@ -1,732 +1,769 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: SLaM\n", - "==================\n", - "\n", - "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for pixelized source modeling.\n", - "\n", - "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", - "guide before working through this example.\n", - "\n", - "Because the SLaM pipelines are designed around pixelized source modeling, the example `slam_start_here` fully\n", - "describes all design choices and modeling decisions made in this script. This script therefore does not repeat\n", - "that documentation, therefore `slam_start_here` should be read first.\n", - "\n", - "The interferometer SLaM pipeline closely mirrors the imaging SLaM pipeline, with one notable difference:\n", - "\n", - "- It uses **two datasets with different transformers**. The `source_lp` stage uses `TransformerNUFFT` (backed\n", - " by the JAX-native `nufftax`, https://github.com/GragasLab/nufftax) so light-profile fitting runs at full GPU\n", - " speed even on ALMA-class datasets with millions of visibilities. The `source_pix` and `mass` stages switch to\n", - " `TransformerNUFFT` combined with the pre-computed sparse operator, because pixelized source reconstructions\n", - " exploit sparsity rather than the NUFFT path.\n", - "\n", - "The interferometer SLaM pipeline still omits the `light_lp` stage, because interferometer data does not contain\n", - "lens light emission. The lens galaxy therefore has no `bulge`/`disk` light components anywhere in the pipeline.\n", - "\n", - "Now that `source_lp` is included, the position likelihood and adapt image needed by the pixelized pipelines are\n", - "derived automatically from the `source_lp` result (mirroring the imaging pipeline). Manual position input is no\n", - "longer required.\n", - "\n", - "__Contents__\n", - "\n", - "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **Interferometer SLaM Description:** The `slam_start_here` notebook provides a detailed description of the SLaM pipelines, but it does.\n", - "- **High Resolution Dataset:** A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts.\n", - "- **SOURCE LP PIPELINE:** Initializes the mass and source-light model with MGE light profiles, fitted via `TransformerNUFFT`.\n", - "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP result.\n", - "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included as interferometer data.\n", - "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with `TransformerNUFFT` + sparse operator (source_pix onwards).\n", - "- **Sparse Operators:** The `pixelization/modeling` example describes how the sparse operator formalism speeds up.\n", - "- **Settings:** Disable the default position only linear algebra solver so the source reconstruction can have.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__Prerequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "You can still run the script without fully understanding the guide, but reviewing it later will\n", - "make the structure and choices of the SLaM workflow clearer.\n", - "\n", - "__Interferometer SLaM Description__\n", - "\n", - "The `slam_start_here` notebook provides a detailed description of the SLaM pipelines, but it does this using CCD\n", - "imaging data.\n", - "\n", - "There is no dedicated example which provides full descriptions of the SLaM pipelines using interferometer data, however,\n", - "the concepts and API described in the `slam_start_here` are identical to what is required for interferometer data.\n", - "\n", - "Therefore, by reading the `slam_start_here` example you will fully understand everything required to use this\n", - "interferometer SLaM script.\n", - "\n", - "__High Resolution Dataset__\n", - "\n", - "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large files which\n", - "are too big to include in the main `autolens_workspace` repository:\n", - "\n", - "https://github.com/PyAutoLabs/autolens_workspace_large_files\n", - "\n", - "After downloading the file, place it in the directory:\n", - "\n", - "`autolens_workspace/dataset/interferometer/alma`\n", - "\n", - "You can then perform modeling using this high-resolution dataset by uncommenting the relevant line of code\n", - "below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light using a Multi\n", - "Gaussian Expansion (MGE). The lens galaxy mass and external shear are fitted at the same time.\n", - "\n", - "This stage uses the `dataset_nufft` Interferometer (built with `TransformerNUFFT`, backed by `nufftax`), which\n", - "makes light-profile fitting fast even for ALMA-class datasets with millions of visibilities.\n", - "\n", - "The mass and source models from this search initialize the SOURCE PIX PIPELINE searches that follow, and the\n", - "result also provides the adapt image and position likelihood used by those later stages.\n", - "\n", - "Note that no lens light is fitted: interferometer data does not contain lens light emission, so `lens.bulge` and\n", - "`lens.disk` are kept at `None`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " # interferometer data does not contain lens light emission\n", - " bulge=None,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source galaxy.\n", - "\n", - "The first search fits a pixelization whose purpose is to generate a high-quality adapt image used in\n", - "search 2. It uses the adapt image computed from the SOURCE LP result (passed via `source_lp_result`) and the\n", - "position likelihood is also derived automatically from the SOURCE LP result via\n", - "`source_lp_result.positions_likelihood_from(...)` \u2014 no manual positions input is required.\n", - "\n", - "This stage uses the `dataset_sparse` Interferometer (built with `TransformerNUFFT` + `apply_sparse_operator`).\n", - "The NUFFT keeps the one-time dirty-image setup tractable at ALMA-scale visibility counts, and the precomputed\n", - "sparse operator makes per-likelihood curvature assembly use the FFT-based W\u0303 precision matrix instead of the\n", - "dense `transformed_mapping_matrix`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", - "pixelization and regularization.\n", - "\n", - "Note that the LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not\n", - "contain lens light emission.\n", - "\n", - "Like SOURCE PIX PIPELINE 1, this stage uses the `dataset_sparse` Interferometer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " settings=settings,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " # interferometry does not support lens light\n", - " bulge=None,\n", - " disk=None,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", - "contain lens light emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " source_pix_result_2: af.Result,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_1.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=source_pix_result_1.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load the `Interferometer` data, define the visibility and real-space masks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", - ")\n", - "\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "# dataset_name = \"alma\"\n", - "\n", - "# if dataset_name == \"alma\":\n", - "#\n", - "# real_space_mask = al.Mask2D.circular(\n", - "# shape_native=(800, 800),\n", - "# pixel_scales=0.01,\n", - "# radius=mask_radius,\n", - "# )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Two Datasets__\n", - "\n", - "The SLaM pipeline runs in two phases that prefer different transformers:\n", - "\n", - "- `dataset_nufft` uses `TransformerNUFFT` (backed by JAX-native `nufftax`) for the `source_lp` stage. With\n", - " light profiles this is the fast path at any visibility count, including ALMA-class datasets.\n", - "- `dataset_sparse` uses `TransformerNUFFT` combined with `apply_sparse_operator(...)` for `source_pix_1`,\n", - " `source_pix_2` and `mass_total`. Pixelized source reconstructions exploit sparsity in the linear inversion\n", - " rather than the NUFFT, so this combination is the right choice for the pixelized stages.\n", - "\n", - "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload) differ." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_nufft = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "dataset_sparse = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", - "pixelized source modeling, especially for many visibilities.\n", - "\n", - "We use a try / except to load the pre-computed curvature preload, which is necessary to use\n", - "the sparse operator formalism. If this file does not exist (e.g. you have not made it manually via\n", - "the `many_visibilities_preparation` example) it is made here.\n", - "\n", - "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", - "`source_lp` does not need it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " nufft_precision_operator = np.load(\n", - " file=dataset_path / \"nufft_precision_operator.npy\",\n", - " )\n", - "except FileNotFoundError:\n", - " nufft_precision_operator = None\n", - "\n", - "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", - " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "Disable the default position only linear algebra solver so the source reconstruction can have\n", - "negative pixel values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"interferometer\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step.\n", - "\n", - "Note the transformer split: `source_lp` is passed `dataset_nufft` (TransformerNUFFT), while every later stage\n", - "is passed `dataset_sparse` (TransformerNUFFT + sparse operator)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset_nufft,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " settings=settings,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: SLaM\n", + "==================\n", + "\n", + "This script provides an example of the Source, (Lens) Light, and Mass (SLaM) pipelines for pixelized source modeling.\n", + "\n", + "A full overview of SLaM is provided in `guides/modeling/slam_start_here`. You should read that\n", + "guide before working through this example.\n", + "\n", + "Because the SLaM pipelines are designed around pixelized source modeling, the example `slam_start_here` fully\n", + "describes all design choices and modeling decisions made in this script. This script therefore does not repeat\n", + "that documentation, therefore `slam_start_here` should be read first.\n", + "\n", + "The interferometer SLaM pipeline closely mirrors the imaging SLaM pipeline, with one notable difference:\n", + "\n", + "- It uses **two datasets with different transformers**. The `source_lp` stage uses `TransformerNUFFT` (backed\n", + " by the JAX-native `nufftax`, https://github.com/GragasLab/nufftax) so light-profile fitting runs at full GPU\n", + " speed even on ALMA-class datasets with millions of visibilities. The `source_pix` and `mass` stages switch to\n", + " `TransformerNUFFT` combined with the pre-computed sparse operator, because pixelized source reconstructions\n", + " exploit sparsity rather than the NUFFT path.\n", + "\n", + "The interferometer SLaM pipeline still omits the `light_lp` stage, because interferometer data does not contain\n", + "lens light emission. The lens galaxy therefore has no `bulge`/`disk` light components anywhere in the pipeline.\n", + "\n", + "Now that `source_lp` is included, the position likelihood and adapt image needed by the pixelized pipelines are\n", + "derived automatically from the `source_lp` result (mirroring the imaging pipeline). Manual position input is no\n", + "longer required.\n", + "\n", + "__Contents__\n", + "\n", + "- **Prerequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **Interferometer SLaM Description:** The `slam_start_here` notebook provides a detailed description of the SLaM pipelines, but it does.\n", + "- **High Resolution Dataset:** A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts.\n", + "- **SOURCE LP PIPELINE:** Initializes the mass and source-light model with MGE light profiles, fitted via `TransformerNUFFT`.\n", + "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP result.\n", + "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included as interferometer data.\n", + "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with `TransformerNUFFT` + sparse operator (source_pix onwards).\n", + "- **Sparse Operators:** The `pixelization/modeling` example describes how the sparse operator formalism speeds up.\n", + "- **Settings:** Disable the default position only linear algebra solver so the source reconstruction can have.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__Prerequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "You can still run the script without fully understanding the guide, but reviewing it later will\n", + "make the structure and choices of the SLaM workflow clearer.\n", + "\n", + "__Interferometer SLaM Description__\n", + "\n", + "The `slam_start_here` notebook provides a detailed description of the SLaM pipelines, but it does this using CCD\n", + "imaging data.\n", + "\n", + "There is no dedicated example which provides full descriptions of the SLaM pipelines using interferometer data, however,\n", + "the concepts and API described in the `slam_start_here` are identical to what is required for interferometer data.\n", + "\n", + "Therefore, by reading the `slam_start_here` example you will fully understand everything required to use this\n", + "interferometer SLaM script.\n", + "\n", + "__High Resolution Dataset__\n", + "\n", + "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large files which\n", + "are too big to include in the main `autolens_workspace` repository:\n", + "\n", + "https://github.com/PyAutoLabs/autolens_workspace_large_files\n", + "\n", + "After downloading the file, place it in the directory:\n", + "\n", + "`autolens_workspace/dataset/interferometer/alma`\n", + "\n", + "You can then perform modeling using this high-resolution dataset by uncommenting the relevant line of code\n", + "below." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "The SOURCE LP PIPELINE uses one search to initialize a robust model for the source galaxy's light using a Multi\n", + "Gaussian Expansion (MGE). The lens galaxy mass and external shear are fitted at the same time.\n", + "\n", + "This stage uses the `dataset_nufft` Interferometer (built with `TransformerNUFFT`, backed by `nufftax`), which\n", + "makes light-profile fitting fast even for ALMA-class datasets with millions of visibilities.\n", + "\n", + "The mass and source models from this search initialize the SOURCE PIX PIPELINE searches that follow, and the\n", + "result also provides the adapt image and position likelihood used by those later stages.\n", + "\n", + "Note that no lens light is fitted: interferometer data does not contain lens light emission, so `lens.bulge` and\n", + "`lens.disk` are kept at `None`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " # interferometer data does not contain lens light emission\n", + " bulge=None,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "The SOURCE PIX PIPELINE uses two searches to initialize a robust pixelized model of the source galaxy.\n", + "\n", + "The first search fits a pixelization whose purpose is to generate a high-quality adapt image used in\n", + "search 2. It uses the adapt image computed from the SOURCE LP result (passed via `source_lp_result`) and the\n", + "position likelihood is also derived automatically from the SOURCE LP result via\n", + "`source_lp_result.positions_likelihood_from(...)` \u2014 no manual positions input is required.\n", + "\n", + "This stage uses the `dataset_sparse` Interferometer (built with `TransformerNUFFT` + `apply_sparse_operator`).\n", + "The NUFFT keeps the one-time dirty-image setup tractable at ALMA-scale visibility counts, and the precomputed\n", + "sparse operator makes per-likelihood curvature assembly use the FFT-based W\u0303 precision matrix instead of the\n", + "dense `transformed_mapping_matrix`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", + "pixelization and regularization.\n", + "\n", + "Note that the LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not\n", + "contain lens light emission.\n", + "\n", + "Like SOURCE PIX PIPELINE 1, this stage uses the `dataset_sparse` Interferometer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " settings=settings,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " # interferometry does not support lens light\n", + " bulge=None,\n", + " disk=None,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", + "contain lens light emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " source_pix_result_2: af.Result,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_1.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=source_pix_result_1.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load the `Interferometer` data, define the visibility and real-space masks." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", + ")\n", + "\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "# dataset_name = \"alma\"\n", + "\n", + "# if dataset_name == \"alma\":\n", + "#\n", + "# real_space_mask = al.Mask2D.circular(\n", + "# shape_native=(800, 800),\n", + "# pixel_scales=0.01,\n", + "# radius=mask_radius,\n", + "# )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Two Datasets__\n", + "\n", + "The SLaM pipeline runs in two phases that prefer different transformers:\n", + "\n", + "- `dataset_nufft` uses `TransformerNUFFT` (backed by JAX-native `nufftax`) for the `source_lp` stage. With\n", + " light profiles this is the fast path at any visibility count, including ALMA-class datasets.\n", + "- `dataset_sparse` uses `TransformerNUFFT` combined with `apply_sparse_operator(...)` for `source_pix_1`,\n", + " `source_pix_2` and `mass_total`. Pixelized source reconstructions exploit sparsity in the linear inversion\n", + " rather than the NUFFT, so this combination is the right choice for the pixelized stages.\n", + "\n", + "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload) differ." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_nufft = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "dataset_sparse = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", + "pixelized source modeling, especially for many visibilities.\n", + "\n", + "We use a try / except to load the pre-computed curvature preload, which is necessary to use\n", + "the sparse operator formalism. If this file does not exist (e.g. you have not made it manually via\n", + "the `many_visibilities_preparation` example) it is made here.\n", + "\n", + "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", + "`source_lp` does not need it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " nufft_precision_operator = np.load(\n", + " file=dataset_path / \"nufft_precision_operator.npy\",\n", + " )\n", + "except FileNotFoundError:\n", + " nufft_precision_operator = None\n", + "\n", + "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", + " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "Disable the default position only linear algebra solver so the source reconstruction can have\n", + "negative pixel values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"interferometer\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step.\n", + "\n", + "Note the transformer split: `source_lp` is passed `dataset_nufft` (TransformerNUFFT), while every later stage\n", + "is passed `dataset_sparse` (TransformerNUFFT + sparse operator)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset_nufft,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " settings=settings,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/pixelization/source_science.ipynb b/notebooks/interferometer/features/pixelization/source_science.ipynb index 49b0ac455..fab29a962 100644 --- a/notebooks/interferometer/features/pixelization/source_science.ipynb +++ b/notebooks/interferometer/features/pixelization/source_science.ipynb @@ -1,756 +1,793 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pixelization: Source Reconstruction\n", - "===================================\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source pixelization, we can compute key quantities such as the magnification, total flux, and\n", - "intrinsic size of the source.\n", - "\n", - "For pixelized source reconstructions, these calculations can be quite involved as they required speciifc code to\n", - "handle irregular mesh pixels and other quantities. We illustrate how to perform these calculations below.\n", - "\n", - "However, this does make the source reconstructions different to share with other people, as it would mean they need\n", - "to understand how to manipulate irregular meshes. The end of this example shows how a .csv source reconstruction file\n", - "is output by a pixelization model-fit, which allows anyone to easy interpolate the source reconstruction on to a uniform grid\n", - "for analysis without the need for PyAutoLens.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model Fit:** Perform the model-fit using the search and analysis.\n", - "- **Interpolated Source:** The simplest way to perform source science calculations on a pixelized source reconstruction is to.\n", - "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", - "- **Zoom:** The interpolation grid above was large in extent (-3.0\" to 3.0\" in both the y and x directions).\n", - "- **Errors:** The interpolated errors on the source reconstruction can also be computed, which will allow you to.\n", - "- **Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", - "- **Masking:** Reconstructions can be imperfect, for example having faint source flux in pixels at the edge of the.\n", - "- **Magnification via Mesh:** The calculations above used an interpolation of the source-plane reconstruction to a 2D grid of.\n", - "- **Reconstruction CSV:** In the results `image` folder there is a .csv file called `source_plane_reconstruction_0.csv` which." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model Fit__\n", - "\n", - "The code below is identical to the pixelizaiton `modeling` example, crucially creating a model-fit which\n", - "outputs the pixelization source reconstruction to a .csv file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "dataset = dataset.apply_sparse_operator(use_jax=True, show_progress=True)\n", - "\n", - "settings = al.Settings(use_positive_only_solver=False)\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", - "\n", - "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", - "regularization = al.reg.Constant(coefficient=1.0)\n", - "\n", - "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitInterferometer(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " settings=settings,\n", - ")\n", - "\n", - "inversion = fit.inversion\n", - "\n", - "mapper = inversion.cls_list_from(cls=al.Mapper)[\n", - " 0\n", - "] # Extract the mapper from the inversion\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the fit, confirming that the pixelized source reconstruction provides a good fit to the data.\n", - "\n", - "Note how the pixelized source reconstruction is performed on an irregular adaptive grid of rectangular pixels,\n", - "which is denser in regions of high magnification. This non-uniform distribution of pixels means we need to be care\n", - "when performing source science calculations, especially a quantity like the magnification which depends on area." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "All information about the pixelized source reconstruction is contained in the `Inversion` object, which can be\n", - "accessed via `fit.inversion`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "inversion = fit.inversion\n", - "print(f\"Inversion Object: {inversion}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For example, the reconstructed source pixel flux values are stored in the `reconstruction` attribute of the inversion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction = inversion.reconstruction\n", - "\n", - "print(f\"Reconstructed Source Pixel Fluxes: {reconstruction}\")\n", - "\n", - "total_flux = np.sum(reconstruction)\n", - "\n", - "print(f\"Total Source Flux via Pixelization: {total_flux} mJy beam^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In order to perform source science calculations we need to know which flux value corresponds to which pixel in the \n", - "source-plane.\n", - "\n", - "This information is available in the inversion, below we print the (y,x) centre of each source pixel corresponding to \n", - "the `reconstruction` values printed above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapper = inversion.cls_list_from(cls=al.Mapper)[\n", - " 0\n", - "] # Extract the mapper from the inversion\n", - "\n", - "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", - "\n", - "print(f\"Source Plane Mesh Grid Coordinates: {source_plane_mesh_grid}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The image-plane reconstruction can also be computed from the inversion, which is called the `mapped_reconstructed_data` \n", - "and as seen above is needed to compute the magnification." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mapped_reconstructed_data = inversion.mapped_reconstructed_data\n", - "\n", - "print(f\"Mapped Reconstructed Image: {mapped_reconstructed_data}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interpolated Source__\n", - "\n", - "The simplest way to perform source science calculations on a pixelized source reconstruction is to interpolate\n", - "its values to a uniform 2D grid of pixels, which can therefore be stored using a `Array2D` object,\n", - "which is basically just a 2D numpy array (see the `Data Structure` section at the top of this example).\n", - "\n", - "We interpolate the rectangular pixelized source reconstruction to a new uniform grid we call the `interpolation_grid`.\n", - "This calculation can be quite slow, so to make this example run fast we use a relatively small grid, but in practice\n", - "you may wish to use a larger grid (e.g. 100x1000 pixels or larger for actual science calculations)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.interpolate import griddata\n", - "\n", - "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", - "\n", - "interpolated_reconstruction = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", - ")\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_ndarray,\n", - " pixel_scales=interpolation_grid.pixel_scales,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By printing the interpolated array, we confirm it is a 2D array and can see the pixel values of the source \n", - "reconstruction.\n", - "\n", - "We also plot the interpolated source reconstruction using an `aplt.plot_array`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(interpolated_reconstruction.native)\n", - "\n", - "aplt.plot_array(array=interpolated_reconstruction, title=\"\")\n", - "aplt.plot_array(array=interpolated_reconstruction, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Flux__\n", - "\n", - "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", - "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", - "\n", - "The total flux of the source reconstruction can now be computed by summing the interpolated array.\n", - "\n", - "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", - "we will assume everything is in miliJansky / beam (`mJy beam^-1`), which is typical for ALMA data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_source_flux = np.sum(interpolated_reconstruction)\n", - "\n", - "print(\n", - " f\"Total Source Flux via Interpolated Pixelization: {total_source_flux} mJy beam^-1\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Zoom__\n", - "\n", - "The interpolation grid above was large in extent (-3.0\" to 3.0\" in both the y and x directions), meaning that\n", - "the source was a small flux was a small region of this grid.\n", - "\n", - "By changing the `extent` of the interpolation grid, we can performed the interpolation zoomed in on only the\n", - "regions of the source-plane where the source reconstruction has non-negligible flux. This\n", - "makes the interpolation more accurate, as the interpolation ican use more pixels in the region of interest,\n", - "and also makes visualizing the source reconstruction easier." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extent = (-1.0, 1.0, -1.0, 1.0)\n", - "shape_native = (401, 401)\n", - "\n", - "interpolation_grid_zoom = al.Grid2D.from_extent(\n", - " extent=extent,\n", - " shape_native=shape_native,\n", - ")\n", - "\n", - "interpolated_reconstruction = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid_zoom\n", - ")\n", - "\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", - " interpolation_grid_zoom.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_ndarray,\n", - " pixel_scales=interpolation_grid_zoom.pixel_scales,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Errors__\n", - "\n", - "The interpolated errors on the source reconstruction can also be computed, which will allow you to perform\n", - "model-fitting of the source reconstruction." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "reconstruction_noise_map = inversion.reconstruction_noise_map\n", - "\n", - "interpolated_noise_map = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", - ")\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_noise_map_ndarray = interpolated_noise_map.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_noise_map = al.Array2D.no_mask(\n", - " values=interpolated_noise_map_ndarray, pixel_scales=interpolation_grid.pixel_scales\n", - ")\n", - "\n", - "aplt.plot_array(array=interpolated_noise_map, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Magnification__\n", - "\n", - "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", - "total surface brightness in the source-plane.\n", - "\n", - "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", - "We therefore explicitly mention how area folds into the calculation below.\n", - "\n", - "The interpolated source reconstruction above has different sized pixels in the image-plane and source-plane, so \n", - "we need to explicitly account for area when computing the magnification.\n", - "\n", - "The `pixel_area` attribute of the `Array2D` object gives us the area of each pixel in arcseconds squared, which we\n", - "can use to compute the magnification below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnification = np.sum(\n", - " mapped_reconstructed_data * mapped_reconstructed_data.pixel_area\n", - ") / np.sum(interpolated_reconstruction * interpolated_reconstruction.pixel_area)\n", - "\n", - "print(f\"Magnification via Interpolated Source: {magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Masking__\n", - "\n", - "Reconstructions can be imperfect, for example having faint source flux in pixels at the edge of the\n", - "source-plane that through comparison to the data are not a genuine part of the source. This can impact\n", - "the calculation of the source flux and magnification.\n", - "\n", - "If you want to be extra careful, you can use a mask to zero the source-plane pixels that you do not trust and use\n", - "this to remove pixels from source science calculations.\n", - "\n", - "Another approach, which we use below, is we create a source-plane signal-to-noise map and use this to create a mask \n", - "that removes all pixels with a signal-to-noise < 5.0." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "signal_to_noise_map = reconstruction / reconstruction_noise_map\n", - "\n", - "mesh_pixel_mask = signal_to_noise_map < 5.0\n", - "\n", - "reconstruction_masked = reconstruction.copy()\n", - "reconstruction_masked[mesh_pixel_mask] = 0.0\n", - "\n", - "interpolated_reconstruction_masked = griddata(\n", - " points=source_plane_mesh_grid, values=reconstruction_masked, xi=interpolation_grid\n", - ")\n", - "\n", - "# As a pure 2D numpy array in case its useful for calculations\n", - "interpolated_reconstruction_masked_ndarray = interpolated_reconstruction_masked.reshape(\n", - " interpolation_grid.shape_native\n", - ")\n", - "\n", - "interpolated_reconstruction_masked = al.Array2D.no_mask(\n", - " values=interpolated_reconstruction_masked_ndarray,\n", - " pixel_scales=interpolation_grid.pixel_scales,\n", - ")\n", - "\n", - "aplt.plot_array(array=interpolated_reconstruction_masked, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Magnification via Mesh__\n", - "\n", - "The calculations above used an interpolation of the source-plane reconstruction to a 2D grid of 1000 x 1000\n", - "pixels.\n", - "\n", - "However, we can use directly the irregular rectangular mesh of the pixelized source reconstruction to compute\n", - "quantities. This is more accurate as it does not introduce interpolation errors, but requires more care as the\n", - "pixels are irregularly spaced and have different areas. \n", - "\n", - "We have already computed the total source flux using the mesh above, but we can also compute the magnification.\n", - "\n", - "Computed the areas of every pixel in the irregular rectangular mesh is a bit involved, therefore the values can be\n", - "accessed from the source code via the `mesh_areas` attribute of the `Mapper` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_areas = mapper.mesh_geometry.areas_for_magnification\n", - "\n", - "magnification = np.sum(\n", - " mapped_reconstructed_data * mapped_reconstructed_data.pixel_area\n", - ") / np.sum(reconstruction * mesh_areas)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction CSV__\n", - "\n", - "In the results `image` folder there is a .csv file called `source_plane_reconstruction_0.csv` which contains the\n", - "y and x coordinates of the pixelization mesh, the reconstruct values and the noise map of these values.\n", - "\n", - "This file is provides all information on the source reconstruction in a format that does not depend autolens\n", - "and therefore be easily loaded to create images of the source or shared collaborations who do not have PyAutoLens\n", - "installed.\n", - "\n", - "We now perform a lens model fit, which will create this .csv file in the modeling output folder.\n", - "\n", - "First, lets load `source_plane_reconstruction_0.csv` as a dictionary, using basic `csv` functionality in Python." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.PowerLaw)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", - "regularization = af.Model(al.reg.Constant)\n", - "\n", - "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\"),\n", - " name=\"pixelization\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=50000,\n", - ")\n", - "\n", - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n", - "\n", - "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)\n", - "\n", - "analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " settings=settings,\n", - " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Reconstruction CSV__\n", - "\n", - "The file ``source_plane_reconstruction_0.csv` provides all information on the source reconstruction in a format that \n", - "does not depend autolens and therefore be easily loaded to create images of the source or shared collaborations who \n", - "do not have PyAutoLens installed.\n", - "\n", - "First, lets load `source_plane_reconstruction_0.csv` as a dictionary, using basic `csv` functionality in Python.\n", - "\n", - "NOTE: If the .csv file does not exist, we create a dictionary with the same format but with dummy values so the rest of\n", - "the script can be run." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import csv\n", - "\n", - "try:\n", - "\n", - " with open(\n", - " search.paths.image_path / \"source_plane_reconstruction_0.csv\", mode=\"r\"\n", - " ) as file:\n", - " reader = csv.reader(file)\n", - " header_list = next(reader) # ['y', 'x', 'reconstruction', 'noise_map']\n", - "\n", - " reconstruction_dict = {header: [] for header in header_list}\n", - "\n", - " for row in reader:\n", - " for key, value in zip(header_list, row):\n", - " reconstruction_dict[key].append(float(value))\n", - "\n", - " # Convert lists to NumPy arrays\n", - " for key in reconstruction_dict:\n", - " reconstruction_dict[key] = np.array(reconstruction_dict[key])\n", - "\n", - "except FileNotFoundError:\n", - "\n", - " print(\"`source_plane_reconstruction_0.csv` not found. Using dummy data instead.\")\n", - "\n", - " x = np.array([-1.0, 0.0, 1.0, -1.0, 0.0, 1.0, -1.0, 0.0, 1.0])\n", - " y = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0, -1.0, -1.0, -1.0])\n", - " reconstruction = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])\n", - " noise_map = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])\n", - "\n", - " reconstruction_dict = {\n", - " \"x\": x,\n", - " \"y\": y,\n", - " \"reconstruction\": reconstruction,\n", - " \"noise_map\": noise_map,\n", - " }" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can now use standard libraries to performed calculations with the reconstruction on the mesh, again avoiding\n", - "the need to use autolens.\n", - "\n", - "For example, we can create a RectangularAdaptDensity mesh using the scipy.spatial library, which is a triangulation\n", - "of the y and x coordinates of the pixelization mesh. This is useful for visualizing the pixelization\n", - "and performing calculations on the mesh." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "import scipy\n", - "\n", - "points = np.stack(arrays=(reconstruction_dict[\"x\"], reconstruction_dict[\"y\"]), axis=-1)\n", - "\n", - "mesh = scipy.spatial.Delaunay(points)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Interpolating the result to a uniform grid is also possible using the scipy.interpolate library, which means the result\n", - "can be turned into a uniform 2D image which can be useful to analyse the source with tools which require an uniform grid.\n", - "\n", - "Below, we interpolate the result onto a 201 x 201 grid of pixels with the extent spanning -1.0\" to 1.0\", which\n", - "capture the majority of the source reconstruction without being too high resolution." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from scipy.interpolate import griddata\n", - "\n", - "values = reconstruction_dict[\"reconstruction\"]\n", - "\n", - "interpolation_grid = al.Grid2D.from_extent(\n", - " extent=(-1.0, 1.0, -1.0, 1.0), shape_native=(201, 201)\n", - ")\n", - "\n", - "interpolated_array = griddata(points=points, values=values, xi=interpolation_grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pixelization: Source Reconstruction\n", + "===================================\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source pixelization, we can compute key quantities such as the magnification, total flux, and\n", + "intrinsic size of the source.\n", + "\n", + "For pixelized source reconstructions, these calculations can be quite involved as they required speciifc code to\n", + "handle irregular mesh pixels and other quantities. We illustrate how to perform these calculations below.\n", + "\n", + "However, this does make the source reconstructions different to share with other people, as it would mean they need\n", + "to understand how to manipulate irregular meshes. The end of this example shows how a .csv source reconstruction file\n", + "is output by a pixelization model-fit, which allows anyone to easy interpolate the source reconstruction on to a uniform grid\n", + "for analysis without the need for PyAutoLens.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model Fit:** Perform the model-fit using the search and analysis.\n", + "- **Interpolated Source:** The simplest way to perform source science calculations on a pixelized source reconstruction is to.\n", + "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", + "- **Zoom:** The interpolation grid above was large in extent (-3.0\" to 3.0\" in both the y and x directions).\n", + "- **Errors:** The interpolated errors on the source reconstruction can also be computed, which will allow you to.\n", + "- **Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", + "- **Masking:** Reconstructions can be imperfect, for example having faint source flux in pixels at the edge of the.\n", + "- **Magnification via Mesh:** The calculations above used an interpolation of the source-plane reconstruction to a 2D grid of.\n", + "- **Reconstruction CSV:** In the results `image` folder there is a .csv file called `source_plane_reconstruction_0.csv` which." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model Fit__\n", + "\n", + "The code below is identical to the pixelizaiton `modeling` example, crucially creating a model-fit which\n", + "outputs the pixelization source reconstruction to a .csv file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "dataset = dataset.apply_sparse_operator(use_jax=True, show_progress=True)\n", + "\n", + "settings = al.Settings(use_positive_only_solver=False)\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", + "\n", + "mesh = al.mesh.RectangularAdaptDensity(shape=mesh_shape)\n", + "regularization = al.reg.Constant(coefficient=1.0)\n", + "\n", + "pixelization = al.Pixelization(mesh=mesh, regularization=regularization)\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, pixelization=pixelization)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitInterferometer(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " settings=settings,\n", + ")\n", + "\n", + "inversion = fit.inversion\n", + "\n", + "mapper = inversion.cls_list_from(cls=al.Mapper)[\n", + " 0\n", + "] # Extract the mapper from the inversion\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the fit, confirming that the pixelized source reconstruction provides a good fit to the data.\n", + "\n", + "Note how the pixelized source reconstruction is performed on an irregular adaptive grid of rectangular pixels,\n", + "which is denser in regions of high magnification. This non-uniform distribution of pixels means we need to be care\n", + "when performing source science calculations, especially a quantity like the magnification which depends on area." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "All information about the pixelized source reconstruction is contained in the `Inversion` object, which can be\n", + "accessed via `fit.inversion`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "inversion = fit.inversion\n", + "print(f\"Inversion Object: {inversion}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For example, the reconstructed source pixel flux values are stored in the `reconstruction` attribute of the inversion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction = inversion.reconstruction\n", + "\n", + "print(f\"Reconstructed Source Pixel Fluxes: {reconstruction}\")\n", + "\n", + "total_flux = np.sum(reconstruction)\n", + "\n", + "print(f\"Total Source Flux via Pixelization: {total_flux} mJy beam^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In order to perform source science calculations we need to know which flux value corresponds to which pixel in the \n", + "source-plane.\n", + "\n", + "This information is available in the inversion, below we print the (y,x) centre of each source pixel corresponding to \n", + "the `reconstruction` values printed above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapper = inversion.cls_list_from(cls=al.Mapper)[\n", + " 0\n", + "] # Extract the mapper from the inversion\n", + "\n", + "source_plane_mesh_grid = mapper.source_plane_mesh_grid\n", + "\n", + "print(f\"Source Plane Mesh Grid Coordinates: {source_plane_mesh_grid}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The image-plane reconstruction can also be computed from the inversion, which is called the `mapped_reconstructed_data` \n", + "and as seen above is needed to compute the magnification." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mapped_reconstructed_data = inversion.mapped_reconstructed_data\n", + "\n", + "print(f\"Mapped Reconstructed Image: {mapped_reconstructed_data}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interpolated Source__\n", + "\n", + "The simplest way to perform source science calculations on a pixelized source reconstruction is to interpolate\n", + "its values to a uniform 2D grid of pixels, which can therefore be stored using a `Array2D` object,\n", + "which is basically just a 2D numpy array (see the `Data Structure` section at the top of this example).\n", + "\n", + "We interpolate the rectangular pixelized source reconstruction to a new uniform grid we call the `interpolation_grid`.\n", + "This calculation can be quite slow, so to make this example run fast we use a relatively small grid, but in practice\n", + "you may wish to use a larger grid (e.g. 100x1000 pixels or larger for actual science calculations)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.interpolate import griddata\n", + "\n", + "interpolation_grid = al.Grid2D.uniform(shape_native=(200, 200), pixel_scales=0.05)\n", + "\n", + "interpolated_reconstruction = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", + ")\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_ndarray,\n", + " pixel_scales=interpolation_grid.pixel_scales,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By printing the interpolated array, we confirm it is a 2D array and can see the pixel values of the source \n", + "reconstruction.\n", + "\n", + "We also plot the interpolated source reconstruction using an `aplt.plot_array`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(interpolated_reconstruction.native)\n", + "\n", + "aplt.plot_array(array=interpolated_reconstruction, title=\"\")\n", + "aplt.plot_array(array=interpolated_reconstruction, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Flux__\n", + "\n", + "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", + "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", + "\n", + "The total flux of the source reconstruction can now be computed by summing the interpolated array.\n", + "\n", + "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", + "we will assume everything is in miliJansky / beam (`mJy beam^-1`), which is typical for ALMA data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_source_flux = np.sum(interpolated_reconstruction)\n", + "\n", + "print(\n", + " f\"Total Source Flux via Interpolated Pixelization: {total_source_flux} mJy beam^-1\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Zoom__\n", + "\n", + "The interpolation grid above was large in extent (-3.0\" to 3.0\" in both the y and x directions), meaning that\n", + "the source was a small flux was a small region of this grid.\n", + "\n", + "By changing the `extent` of the interpolation grid, we can performed the interpolation zoomed in on only the\n", + "regions of the source-plane where the source reconstruction has non-negligible flux. This\n", + "makes the interpolation more accurate, as the interpolation ican use more pixels in the region of interest,\n", + "and also makes visualizing the source reconstruction easier." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extent = (-1.0, 1.0, -1.0, 1.0)\n", + "shape_native = (401, 401)\n", + "\n", + "interpolation_grid_zoom = al.Grid2D.from_extent(\n", + " extent=extent,\n", + " shape_native=shape_native,\n", + ")\n", + "\n", + "interpolated_reconstruction = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid_zoom\n", + ")\n", + "\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_reconstruction_ndarray = interpolated_reconstruction.reshape(\n", + " interpolation_grid_zoom.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_ndarray,\n", + " pixel_scales=interpolation_grid_zoom.pixel_scales,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Errors__\n", + "\n", + "The interpolated errors on the source reconstruction can also be computed, which will allow you to perform\n", + "model-fitting of the source reconstruction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "reconstruction_noise_map = inversion.reconstruction_noise_map\n", + "\n", + "interpolated_noise_map = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction, xi=interpolation_grid\n", + ")\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_noise_map_ndarray = interpolated_noise_map.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_noise_map = al.Array2D.no_mask(\n", + " values=interpolated_noise_map_ndarray, pixel_scales=interpolation_grid.pixel_scales\n", + ")\n", + "\n", + "aplt.plot_array(array=interpolated_noise_map, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Magnification__\n", + "\n", + "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", + "total surface brightness in the source-plane.\n", + "\n", + "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", + "We therefore explicitly mention how area folds into the calculation below.\n", + "\n", + "The interpolated source reconstruction above has different sized pixels in the image-plane and source-plane, so \n", + "we need to explicitly account for area when computing the magnification.\n", + "\n", + "The `pixel_area` attribute of the `Array2D` object gives us the area of each pixel in arcseconds squared, which we\n", + "can use to compute the magnification below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnification = np.sum(\n", + " mapped_reconstructed_data * mapped_reconstructed_data.pixel_area\n", + ") / np.sum(interpolated_reconstruction * interpolated_reconstruction.pixel_area)\n", + "\n", + "print(f\"Magnification via Interpolated Source: {magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Masking__\n", + "\n", + "Reconstructions can be imperfect, for example having faint source flux in pixels at the edge of the\n", + "source-plane that through comparison to the data are not a genuine part of the source. This can impact\n", + "the calculation of the source flux and magnification.\n", + "\n", + "If you want to be extra careful, you can use a mask to zero the source-plane pixels that you do not trust and use\n", + "this to remove pixels from source science calculations.\n", + "\n", + "Another approach, which we use below, is we create a source-plane signal-to-noise map and use this to create a mask \n", + "that removes all pixels with a signal-to-noise < 5.0." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "signal_to_noise_map = reconstruction / reconstruction_noise_map\n", + "\n", + "mesh_pixel_mask = signal_to_noise_map < 5.0\n", + "\n", + "reconstruction_masked = reconstruction.copy()\n", + "reconstruction_masked[mesh_pixel_mask] = 0.0\n", + "\n", + "interpolated_reconstruction_masked = griddata(\n", + " points=source_plane_mesh_grid, values=reconstruction_masked, xi=interpolation_grid\n", + ")\n", + "\n", + "# As a pure 2D numpy array in case its useful for calculations\n", + "interpolated_reconstruction_masked_ndarray = interpolated_reconstruction_masked.reshape(\n", + " interpolation_grid.shape_native\n", + ")\n", + "\n", + "interpolated_reconstruction_masked = al.Array2D.no_mask(\n", + " values=interpolated_reconstruction_masked_ndarray,\n", + " pixel_scales=interpolation_grid.pixel_scales,\n", + ")\n", + "\n", + "aplt.plot_array(array=interpolated_reconstruction_masked, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Magnification via Mesh__\n", + "\n", + "The calculations above used an interpolation of the source-plane reconstruction to a 2D grid of 1000 x 1000\n", + "pixels.\n", + "\n", + "However, we can use directly the irregular rectangular mesh of the pixelized source reconstruction to compute\n", + "quantities. This is more accurate as it does not introduce interpolation errors, but requires more care as the\n", + "pixels are irregularly spaced and have different areas. \n", + "\n", + "We have already computed the total source flux using the mesh above, but we can also compute the magnification.\n", + "\n", + "Computed the areas of every pixel in the irregular rectangular mesh is a bit involved, therefore the values can be\n", + "accessed from the source code via the `mesh_areas` attribute of the `Mapper` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_areas = mapper.mesh_geometry.areas_for_magnification\n", + "\n", + "magnification = np.sum(\n", + " mapped_reconstructed_data * mapped_reconstructed_data.pixel_area\n", + ") / np.sum(reconstruction * mesh_areas)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction CSV__\n", + "\n", + "In the results `image` folder there is a .csv file called `source_plane_reconstruction_0.csv` which contains the\n", + "y and x coordinates of the pixelization mesh, the reconstruct values and the noise map of these values.\n", + "\n", + "This file is provides all information on the source reconstruction in a format that does not depend autolens\n", + "and therefore be easily loaded to create images of the source or shared collaborations who do not have PyAutoLens\n", + "installed.\n", + "\n", + "We now perform a lens model fit, which will create this .csv file in the modeling output folder.\n", + "\n", + "First, lets load `source_plane_reconstruction_0.csv` as a dictionary, using basic `csv` functionality in Python." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.PowerLaw)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "mesh = af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape)\n", + "regularization = af.Model(al.reg.Constant)\n", + "\n", + "pixelization = af.Model(al.Pixelization, mesh=mesh, regularization=regularization)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\"),\n", + " name=\"pixelization\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=20, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=50000,\n", + ")\n", + "\n", + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n", + "\n", + "positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)\n", + "\n", + "analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " settings=settings,\n", + " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Reconstruction CSV__\n", + "\n", + "The file ``source_plane_reconstruction_0.csv` provides all information on the source reconstruction in a format that \n", + "does not depend autolens and therefore be easily loaded to create images of the source or shared collaborations who \n", + "do not have PyAutoLens installed.\n", + "\n", + "First, lets load `source_plane_reconstruction_0.csv` as a dictionary, using basic `csv` functionality in Python.\n", + "\n", + "NOTE: If the .csv file does not exist, we create a dictionary with the same format but with dummy values so the rest of\n", + "the script can be run." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import csv\n", + "\n", + "try:\n", + "\n", + " with open(\n", + " search.paths.image_path / \"source_plane_reconstruction_0.csv\", mode=\"r\"\n", + " ) as file:\n", + " reader = csv.reader(file)\n", + " header_list = next(reader) # ['y', 'x', 'reconstruction', 'noise_map']\n", + "\n", + " reconstruction_dict = {header: [] for header in header_list}\n", + "\n", + " for row in reader:\n", + " for key, value in zip(header_list, row):\n", + " reconstruction_dict[key].append(float(value))\n", + "\n", + " # Convert lists to NumPy arrays\n", + " for key in reconstruction_dict:\n", + " reconstruction_dict[key] = np.array(reconstruction_dict[key])\n", + "\n", + "except FileNotFoundError:\n", + "\n", + " print(\"`source_plane_reconstruction_0.csv` not found. Using dummy data instead.\")\n", + "\n", + " x = np.array([-1.0, 0.0, 1.0, -1.0, 0.0, 1.0, -1.0, 0.0, 1.0])\n", + " y = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0, -1.0, -1.0, -1.0])\n", + " reconstruction = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])\n", + " noise_map = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])\n", + "\n", + " reconstruction_dict = {\n", + " \"x\": x,\n", + " \"y\": y,\n", + " \"reconstruction\": reconstruction,\n", + " \"noise_map\": noise_map,\n", + " }" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can now use standard libraries to performed calculations with the reconstruction on the mesh, again avoiding\n", + "the need to use autolens.\n", + "\n", + "For example, we can create a RectangularAdaptDensity mesh using the scipy.spatial library, which is a triangulation\n", + "of the y and x coordinates of the pixelization mesh. This is useful for visualizing the pixelization\n", + "and performing calculations on the mesh." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import scipy\n", + "\n", + "points = np.stack(arrays=(reconstruction_dict[\"x\"], reconstruction_dict[\"y\"]), axis=-1)\n", + "\n", + "mesh = scipy.spatial.Delaunay(points)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Interpolating the result to a uniform grid is also possible using the scipy.interpolate library, which means the result\n", + "can be turned into a uniform 2D image which can be useful to analyse the source with tools which require an uniform grid.\n", + "\n", + "Below, we interpolate the result onto a 201 x 201 grid of pixels with the extent spanning -1.0\" to 1.0\", which\n", + "capture the majority of the source reconstruction without being too high resolution." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from scipy.interpolate import griddata\n", + "\n", + "values = reconstruction_dict[\"reconstruction\"]\n", + "\n", + "interpolation_grid = al.Grid2D.from_extent(\n", + " extent=(-1.0, 1.0, -1.0, 1.0), shape_native=(201, 201)\n", + ")\n", + "\n", + "interpolated_array = griddata(points=points, values=values, xi=interpolation_grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/subhalo/detect/start_here.ipynb b/notebooks/interferometer/features/subhalo/detect/start_here.ipynb index f3c67ac3c..e036ede1e 100644 --- a/notebooks/interferometer/features/subhalo/detect/start_here.ipynb +++ b/notebooks/interferometer/features/subhalo/detect/start_here.ipynb @@ -1,932 +1,969 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Subhalo Detection: Start Here\n", - "=============================\n", - "\n", - "Strong gravitational lenses can be used to detect the presence of small-scale dark matter (DM) subhalos. This occurs\n", - "when the DM subhalo overlaps the lensed source emission, and therefore gravitationally perturbs the observed image of\n", - "the lensed source galaxy.\n", - "\n", - "When a DM subhalo is not included in the lens model, residuals will be present in the fit to the data in the lensed\n", - "source regions near the subhalo. By adding a DM subhalo to the lens model, these residuals can be reduced. Bayesian\n", - "model comparison can then be used to quantify whether or not the improvement to the fit is significant enough to\n", - "claim the detection of a DM subhalo.\n", - "\n", - "The example illustrates DM subhalo detection with **PyAutoLens**.\n", - "\n", - "__Contents__\n", - "\n", - "- **SLaM Pipelines:** The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which.\n", - "- **Grid Search:** The second stage of the SUBHALO PIPELINE uses a grid-search of non-linear searches to determine the.\n", - "- **Pixelized Source:** Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **SOURCE LP PIPELINE:** Initializes the mass model + source-light using MGE light profiles via `TransformerNUFFT`.\n", - "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP result.\n", - "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", - "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included as interferometer data.\n", - "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with `TransformerNUFFT` + sparse operator (source_pix onwards).\n", - "- **Sparse Operators:** The `pixelization/modeling` example describes how the sparse operator formalism speeds up.\n", - "- **Settings:** Disable the default position only linear algebra solver so the source reconstruction can have.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", - "\n", - "__SLaM Pipelines__\n", - "\n", - "The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which automate the fitting\n", - "of complex lens models. The SLaM pipelines are used for all DM subhalo detection analyses. Therefore\n", - "you should be familiar with the SLaM pipelines before performing DM subhalo detection yourself. If you are unfamiliar\n", - "with the SLaM pipelines, checkout the\n", - "example `autolens_workspace/notebooks/guides/modeling/slam_start_here`.\n", - "\n", - "Dark matter subhalo detection runs the standard SLaM pipelines, and then extends them with a SUBHALO PIPELINE which\n", - "performs the following three chained non-linear searches:\n", - "\n", - " 1) Fits the lens model fitted in the MASS PIPELINE again, without a DM subhalo, to estimate the Bayesian evidence\n", - " of the model without a DM subhalo.\n", - "\n", - " 2) Performs a grid-search of non-linear searches, where each grid cell includes a DM subhalo whose (y,x) centre is\n", - " confined to a small 2D section of the image plane via uniform priors (we explain this in more detail below).\n", - "\n", - " 3) Fit the lens model again, including a DM subhalo whose (y,x) centre is initialized from the highest log evidence\n", - " grid cell of the grid-search. The Bayesian evidence estimated in this model-fit is compared to the model-fit\n", - " which did not include a DM subhalo, to determine whether or not a DM subhalo was detected.\n", - "\n", - "__Grid Search__\n", - "\n", - "The second stage of the SUBHALO PIPELINE uses a grid-search of non-linear searches to determine the highest log\n", - "evidence model with a DM subhalo. This grid search confines each DM subhalo in the lens model to a small 2D section\n", - "of the image plane via priors on its (y,x) centre. The reasons for this are as follows:\n", - "\n", - " - Lens models including a DM subhalo often have a multi-model parameter space. This means there are multiple lens\n", - " models with high likelihood solutions, each of which place the DM subhalo in different (y,x) image-plane location.\n", - " Multi-modal parameter spaces are synonomously difficult for non-linear searches to fit, and often produce\n", - " incorrect or inefficient fitting. The grid search breaks the multi-modal parameter space into many single-peaked\n", - " parameter spaces, making the model-fitting faster and more reliable.\n", - "\n", - " - By inferring how placing a DM subhalo at different locations in the image-plane changes the Bayesian evidence, we\n", - " map out spatial information on where a DM subhalo is detected. This can help our interpretation of the DM subhalo\n", - " detection.\n", - "\n", - "__Pixelized Source__\n", - "\n", - "Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of the source's light\n", - "are at a level where the subhalo's perturbing lensing effects can be detected.\n", - "\n", - "This requires the source reconstruction to be performed using a pixelized source, as this provides a more detailed\n", - "reconstruction of the source's light than fits using light profiles.\n", - "\n", - "This example therefore using a pixelized source and the corresponding SLaM pipelines.\n", - "\n", - "The `subhalo/detection/examples` folder contains an example using light profile sources, if you have a use-case where\n", - "using light profile source is feasible (e.g. fitting simple simulated datasets).\n", - "\n", - "__Model__\n", - "\n", - "Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE, MASS TOTAL PIPELINE and SUBHALO PIPELINE this SLaM script\n", - "fits `Interferometer` of a strong lens system, where in the final model:\n", - "\n", - " - The lens galaxy's light is an MGE bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - A dark matter subhalo near The lens galaxy mass is included as a`NFWMCRLudlowSph`.\n", - " - The source galaxy is an `Inversion`.\n", - "\n", - "This uses the SLaM pipelines:\n", - "\n", - " `source_lp`\n", - " `source_pix`\n", - " `light_lp`\n", - " `mass_total`\n", - " `subhalo/detection`\n", - "\n", - "Check them out for a full description of the analysis!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE LP PIPELINE__\n", - "\n", - "Initializes a robust mass + source-light model using a Multi Gaussian Expansion (MGE), fit with light profiles.\n", - "\n", - "This stage uses `dataset_nufft` (built with `TransformerNUFFT`, backed by JAX-native `nufftax`), which makes\n", - "light-profile fitting fast even on ALMA-class datasets with millions of visibilities. The result provides the\n", - "adapt image and position likelihood threaded into the pixelized pipelines that follow.\n", - "\n", - "No lens light is fitted: interferometer data does not contain lens light emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " mask_radius: float,\n", - " redshift_lens: float,\n", - " redshift_source: float,\n", - " n_batch: int = 50,\n", - ") -> af.Result:\n", - " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", - "\n", - " source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=af.Model(al.mp.Isothermal),\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 1__\n", - "\n", - "The first search of the SOURCE PIX PIPELINE fits a pixelization whose purpose is to generate a high-quality\n", - "adapt image used in search 2. It uses the adapt image computed from the SOURCE LP result, with the position\n", - "likelihood derived automatically via `source_lp_result.positions_likelihood_from(...)`.\n", - "\n", - "This stage uses `dataset_sparse` (built with `TransformerNUFFT` + `apply_sparse_operator`). Pixelizations\n", - "exploit sparsity in the linear inversion rather than the NUFFT path." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_1(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " mesh_init,\n", - " regularization_init,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - " shear = source_lp_result.model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh_init,\n", - " regularization=regularization_init,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SOURCE PIX PIPELINE 2__\n", - "\n", - "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", - "pixelization and regularization.\n", - "\n", - "Note that the LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not\n", - "contain lens light emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_pix_2(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_lp_result: af.Result,\n", - " source_pix_result_1: af.Result,\n", - " mesh,\n", - " regularization,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " settings=settings,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " # interferometry does not support lens light\n", - " bulge=None,\n", - " disk=None,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=mesh,\n", - " regularization=regularization,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__MASS TOTAL PIPELINE__\n", - "\n", - "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", - "contain lens light emission." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def mass_total(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " source_pix_result_2: af.Result,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_1.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - " )\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_pix_result_2)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", - " bulge=None,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=source_pix_result_1.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (grid search)__\n", - "\n", - "The second search of the SUBHALO PIPELINE performs a [number_of_steps x number_of_steps] grid search of\n", - "non-linear searches. Each grid cell includes a DM subhalo whose (y,x) centre is confined to a small 2D section\n", - "of the image plane via uniform priors.\n", - "\n", - "This grid search maps out where in the image plane including a DM subhalo provides a better fit to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def subhalo_grid_search(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " mass_result: af.Result,\n", - " subhalo_mass: af.Model,\n", - " settings,\n", - " grid_dimension_arcsec: float = 3.0,\n", - " number_of_steps: int = 2,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " subhalo = af.Model(al.Galaxy, mass=subhalo_mass)\n", - "\n", - " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", - " subhalo.mass.centre_0 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - " )\n", - " subhalo.mass.centre_1 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - " )\n", - "\n", - " subhalo.redshift = mass_result.instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_object = mass_result.instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_source = mass_result.instance.galaxies.source.redshift\n", - "\n", - " lens = mass_result.model.galaxies.lens\n", - " source = al.util.chaining.source_from(result=mass_result)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, subhalo=subhalo, source=source),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[1]_[search_lens_plane]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " subhalo_grid_search = af.SearchGridSearch(\n", - " search=search,\n", - " number_of_steps=number_of_steps,\n", - " )\n", - "\n", - " return subhalo_grid_search.fit(\n", - " model=model,\n", - " analysis=analysis,\n", - " grid_priors=[\n", - " model.galaxies.subhalo.mass.centre_1,\n", - " model.galaxies.subhalo.mass.centre_0,\n", - " ],\n", - " info=settings_search.info,\n", - " )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SUBHALO PIPELINE (refine)__\n", - "\n", - "The third search of the SUBHALO PIPELINE refits the lens model including a DM subhalo, initializing the\n", - "subhalo centre from the highest log evidence grid cell of the grid search.\n", - "\n", - "The Bayesian evidence from this fit is compared to the no-subhalo fit to determine whether a DM subhalo\n", - "was detected." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def subhalo_refine(\n", - " settings_search: af.SettingsSearch,\n", - " dataset,\n", - " source_pix_result_1: af.Result,\n", - " mass_result: af.Result,\n", - " subhalo_grid_search_result: af.Result,\n", - " subhalo_mass: af.Model,\n", - " settings,\n", - " n_batch: int = 20,\n", - ") -> af.Result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1, use_model_images=True\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " settings=settings,\n", - " )\n", - "\n", - " subhalo = af.Model(\n", - " al.Galaxy,\n", - " redshift=mass_result.instance.galaxies.lens.redshift,\n", - " mass=subhalo_mass,\n", - " )\n", - "\n", - " subhalo.redshift = mass_result.instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_object = mass_result.instance.galaxies.lens.redshift\n", - " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", - " subhalo.mass.centre = subhalo_grid_search_result.model_centred_absolute(\n", - " a=1.0\n", - " ).galaxies.subhalo.mass.centre\n", - "\n", - " subhalo.redshift = subhalo_grid_search_result.model.galaxies.subhalo.redshift\n", - " subhalo.mass.redshift_object = subhalo.redshift\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=subhalo_grid_search_result.model.galaxies.lens,\n", - " subhalo=subhalo,\n", - " source=subhalo_grid_search_result.model.galaxies.source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[2]_[single_plane_refine]\",\n", - " **settings_search.search_dict,\n", - " n_live=600,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset + Masking__\n", - "\n", - "Load the `Interferometer` data, define the visibility and real-space masks." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", - ")\n", - "\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "# dataset_name = \"alma\"\n", - "\n", - "# if dataset_name == \"alma\":\n", - "#\n", - "# real_space_mask = al.Mask2D.circular(\n", - "# shape_native=(800, 800),\n", - "# pixel_scales=0.01,\n", - "# radius=mask_radius,\n", - "# )\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Two Datasets__\n", - "\n", - "The SLaM pipeline runs in two phases that prefer different transformers:\n", - "\n", - "- `dataset_nufft` uses `TransformerNUFFT` (backed by JAX-native `nufftax`) for the `source_lp` stage.\n", - "- `dataset_sparse` uses `TransformerNUFFT` + `apply_sparse_operator(...)` for `source_pix_1`,\n", - " `source_pix_2`, `mass_total`, and every search of the SUBHALO PIPELINE.\n", - "\n", - "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload) differ." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_nufft = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "dataset_sparse = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sparse Operators__\n", - "\n", - "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", - "pixelized source modeling, especially for many visibilities.\n", - "\n", - "We use a try / except to load the pre-computed curvature preload, which is necessary to use\n", - "the sparse operator formalism. If this file does not exist (e.g. you have not made it manually via\n", - "the `many_visibilities_preparation` example) it is made here.\n", - "\n", - "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", - "`source_lp` does not need it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "try:\n", - " nufft_precision_operator = np.load(\n", - " file=dataset_path / \"nufft_precision_operator.npy\",\n", - " )\n", - "except FileNotFoundError:\n", - " nufft_precision_operator = None\n", - "\n", - "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", - " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings__\n", - "\n", - "Disable the default position only linear algebra solver so the source reconstruction can have\n", - "negative pixel values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings = al.Settings(use_positive_only_solver=False)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"interferometer\") / \"slam\",\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__\n", - "\n", - "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", - "a description of each pipeline step.\n", - "\n", - "Note the transformer split: `source_lp` is passed `dataset_nufft` (TransformerNUFFT), while every later stage\n", - "\u2014 including the SUBHALO PIPELINE \u2014 is passed `dataset_sparse` (TransformerNUFFT + sparse operator)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " dataset=dataset_nufft,\n", - " mask_radius=mask_radius,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization_init=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " settings=settings,\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_pix_result_1=source_pix_result_1,\n", - " source_pix_result_2=source_pix_result_2,\n", - " settings=settings,\n", - ")\n", - "\n", - "result_subhalo_grid_search = subhalo_grid_search(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mass_result=mass_result,\n", - " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", - " settings=settings,\n", - " grid_dimension_arcsec=3.0,\n", - " number_of_steps=2,\n", - ")\n", - "\n", - "result_with_subhalo = subhalo_refine(\n", - " settings_search=settings_search,\n", - " dataset=dataset_sparse,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mass_result=mass_result,\n", - " subhalo_grid_search_result=result_subhalo_grid_search,\n", - " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", - " settings=settings,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Subhalo Detection: Start Here\n", + "=============================\n", + "\n", + "Strong gravitational lenses can be used to detect the presence of small-scale dark matter (DM) subhalos. This occurs\n", + "when the DM subhalo overlaps the lensed source emission, and therefore gravitationally perturbs the observed image of\n", + "the lensed source galaxy.\n", + "\n", + "When a DM subhalo is not included in the lens model, residuals will be present in the fit to the data in the lensed\n", + "source regions near the subhalo. By adding a DM subhalo to the lens model, these residuals can be reduced. Bayesian\n", + "model comparison can then be used to quantify whether or not the improvement to the fit is significant enough to\n", + "claim the detection of a DM subhalo.\n", + "\n", + "The example illustrates DM subhalo detection with **PyAutoLens**.\n", + "\n", + "__Contents__\n", + "\n", + "- **SLaM Pipelines:** The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which.\n", + "- **Grid Search:** The second stage of the SUBHALO PIPELINE uses a grid-search of non-linear searches to determine the.\n", + "- **Pixelized Source:** Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **SOURCE LP PIPELINE:** Initializes the mass model + source-light using MGE light profiles via `TransformerNUFFT`.\n", + "- **SOURCE PIX PIPELINE 1:** Initializes a pixelized source using the adapt image from the SOURCE LP result.\n", + "- **SOURCE PIX PIPELINE 2:** Improves the pixelized source using adapt images from `source_pix_result_1`.\n", + "- **MASS TOTAL PIPELINE:** Identical to `slam_start_here.py`, except no lens light model is included as interferometer data.\n", + "- **Two Datasets:** Build one Interferometer with `TransformerNUFFT` (source_lp) and one with `TransformerNUFFT` + sparse operator (source_pix onwards).\n", + "- **Sparse Operators:** The `pixelization/modeling` example describes how the sparse operator formalism speeds up.\n", + "- **Settings:** Disable the default position only linear algebra solver so the source reconstruction can have.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **SLaM Pipeline:** The code below calls the full SLaM PIPELINE.\n", + "\n", + "__SLaM Pipelines__\n", + "\n", + "The Source, (lens) Light and Mass (SLaM) pipelines are advanced lens modeling pipelines which automate the fitting\n", + "of complex lens models. The SLaM pipelines are used for all DM subhalo detection analyses. Therefore\n", + "you should be familiar with the SLaM pipelines before performing DM subhalo detection yourself. If you are unfamiliar\n", + "with the SLaM pipelines, checkout the\n", + "example `autolens_workspace/notebooks/guides/modeling/slam_start_here`.\n", + "\n", + "Dark matter subhalo detection runs the standard SLaM pipelines, and then extends them with a SUBHALO PIPELINE which\n", + "performs the following three chained non-linear searches:\n", + "\n", + " 1) Fits the lens model fitted in the MASS PIPELINE again, without a DM subhalo, to estimate the Bayesian evidence\n", + " of the model without a DM subhalo.\n", + "\n", + " 2) Performs a grid-search of non-linear searches, where each grid cell includes a DM subhalo whose (y,x) centre is\n", + " confined to a small 2D section of the image plane via uniform priors (we explain this in more detail below).\n", + "\n", + " 3) Fit the lens model again, including a DM subhalo whose (y,x) centre is initialized from the highest log evidence\n", + " grid cell of the grid-search. The Bayesian evidence estimated in this model-fit is compared to the model-fit\n", + " which did not include a DM subhalo, to determine whether or not a DM subhalo was detected.\n", + "\n", + "__Grid Search__\n", + "\n", + "The second stage of the SUBHALO PIPELINE uses a grid-search of non-linear searches to determine the highest log\n", + "evidence model with a DM subhalo. This grid search confines each DM subhalo in the lens model to a small 2D section\n", + "of the image plane via priors on its (y,x) centre. The reasons for this are as follows:\n", + "\n", + " - Lens models including a DM subhalo often have a multi-model parameter space. This means there are multiple lens\n", + " models with high likelihood solutions, each of which place the DM subhalo in different (y,x) image-plane location.\n", + " Multi-modal parameter spaces are synonomously difficult for non-linear searches to fit, and often produce\n", + " incorrect or inefficient fitting. The grid search breaks the multi-modal parameter space into many single-peaked\n", + " parameter spaces, making the model-fitting faster and more reliable.\n", + "\n", + " - By inferring how placing a DM subhalo at different locations in the image-plane changes the Bayesian evidence, we\n", + " map out spatial information on where a DM subhalo is detected. This can help our interpretation of the DM subhalo\n", + " detection.\n", + "\n", + "__Pixelized Source__\n", + "\n", + "Detecting a DM subhalo requires the lens model to be sufficiently accurate that the residuals of the source's light\n", + "are at a level where the subhalo's perturbing lensing effects can be detected.\n", + "\n", + "This requires the source reconstruction to be performed using a pixelized source, as this provides a more detailed\n", + "reconstruction of the source's light than fits using light profiles.\n", + "\n", + "This example therefore using a pixelized source and the corresponding SLaM pipelines.\n", + "\n", + "The `subhalo/detection/examples` folder contains an example using light profile sources, if you have a use-case where\n", + "using light profile source is feasible (e.g. fitting simple simulated datasets).\n", + "\n", + "__Model__\n", + "\n", + "Using a SOURCE LP PIPELINE, LIGHT LP PIPELINE, MASS TOTAL PIPELINE and SUBHALO PIPELINE this SLaM script\n", + "fits `Interferometer` of a strong lens system, where in the final model:\n", + "\n", + " - The lens galaxy's light is an MGE bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - A dark matter subhalo near The lens galaxy mass is included as a`NFWMCRLudlowSph`.\n", + " - The source galaxy is an `Inversion`.\n", + "\n", + "This uses the SLaM pipelines:\n", + "\n", + " `source_lp`\n", + " `source_pix`\n", + " `light_lp`\n", + " `mass_total`\n", + " `subhalo/detection`\n", + "\n", + "Check them out for a full description of the analysis!" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE LP PIPELINE__\n", + "\n", + "Initializes a robust mass + source-light model using a Multi Gaussian Expansion (MGE), fit with light profiles.\n", + "\n", + "This stage uses `dataset_nufft` (built with `TransformerNUFFT`, backed by JAX-native `nufftax`), which makes\n", + "light-profile fitting fast even on ALMA-class datasets with millions of visibilities. The result provides the\n", + "adapt image and position likelihood threaded into the pixelized pipelines that follow.\n", + "\n", + "No lens light is fitted: interferometer data does not contain lens light emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " mask_radius: float,\n", + " redshift_lens: float,\n", + " redshift_source: float,\n", + " n_batch: int = 50,\n", + ") -> af.Result:\n", + " analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", + "\n", + " source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=af.Model(al.mp.Isothermal),\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 1__\n", + "\n", + "The first search of the SOURCE PIX PIPELINE fits a pixelization whose purpose is to generate a high-quality\n", + "adapt image used in search 2. It uses the adapt image computed from the SOURCE LP result, with the position\n", + "likelihood derived automatically via `source_lp_result.positions_likelihood_from(...)`.\n", + "\n", + "This stage uses `dataset_sparse` (built with `TransformerNUFFT` + `apply_sparse_operator`). Pixelizations\n", + "exploit sparsity in the linear inversion rather than the NUFFT path." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_1(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " mesh_init,\n", + " regularization_init,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + " shear = source_lp_result.model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh_init,\n", + " regularization=regularization_init,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SOURCE PIX PIPELINE 2__\n", + "\n", + "Identical to `slam_start_here.py`, using adapt images from `source_pix_result_1` to improve the source\n", + "pixelization and regularization.\n", + "\n", + "Note that the LIGHT LP PIPELINE from `slam_start_here.py` is omitted here, as interferometer data does not\n", + "contain lens light emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_pix_2(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_lp_result: af.Result,\n", + " source_pix_result_1: af.Result,\n", + " mesh,\n", + " regularization,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " settings=settings,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " # interferometry does not support lens light\n", + " bulge=None,\n", + " disk=None,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=mesh,\n", + " regularization=regularization,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__MASS TOTAL PIPELINE__\n", + "\n", + "Identical to `slam_start_here.py`, except no lens light model is included as interferometer data does not\n", + "contain lens light emission." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def mass_total(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " source_pix_result_2: af.Result,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_1.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + " )\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_pix_result_1.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_pix_result_2)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_pix_result_1.instance.galaxies.lens.redshift,\n", + " bulge=None,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=source_pix_result_1.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (grid search)__\n", + "\n", + "The second search of the SUBHALO PIPELINE performs a [number_of_steps x number_of_steps] grid search of\n", + "non-linear searches. Each grid cell includes a DM subhalo whose (y,x) centre is confined to a small 2D section\n", + "of the image plane via uniform priors.\n", + "\n", + "This grid search maps out where in the image plane including a DM subhalo provides a better fit to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def subhalo_grid_search(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " mass_result: af.Result,\n", + " subhalo_mass: af.Model,\n", + " settings,\n", + " grid_dimension_arcsec: float = 3.0,\n", + " number_of_steps: int = 2,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " subhalo = af.Model(al.Galaxy, mass=subhalo_mass)\n", + "\n", + " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", + " subhalo.mass.centre_0 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + " )\n", + " subhalo.mass.centre_1 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + " )\n", + "\n", + " subhalo.redshift = mass_result.instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_object = mass_result.instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_source = mass_result.instance.galaxies.source.redshift\n", + "\n", + " lens = mass_result.model.galaxies.lens\n", + " source = al.util.chaining.source_from(result=mass_result)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, subhalo=subhalo, source=source),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[1]_[search_lens_plane]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " subhalo_grid_search = af.SearchGridSearch(\n", + " search=search,\n", + " number_of_steps=number_of_steps,\n", + " )\n", + "\n", + " return subhalo_grid_search.fit(\n", + " model=model,\n", + " analysis=analysis,\n", + " grid_priors=[\n", + " model.galaxies.subhalo.mass.centre_1,\n", + " model.galaxies.subhalo.mass.centre_0,\n", + " ],\n", + " info=settings_search.info,\n", + " )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SUBHALO PIPELINE (refine)__\n", + "\n", + "The third search of the SUBHALO PIPELINE refits the lens model including a DM subhalo, initializing the\n", + "subhalo centre from the highest log evidence grid cell of the grid search.\n", + "\n", + "The Bayesian evidence from this fit is compared to the no-subhalo fit to determine whether a DM subhalo\n", + "was detected." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def subhalo_refine(\n", + " settings_search: af.SettingsSearch,\n", + " dataset,\n", + " source_pix_result_1: af.Result,\n", + " mass_result: af.Result,\n", + " subhalo_grid_search_result: af.Result,\n", + " subhalo_mass: af.Model,\n", + " settings,\n", + " n_batch: int = 20,\n", + ") -> af.Result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1, use_model_images=True\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " mass_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " settings=settings,\n", + " )\n", + "\n", + " subhalo = af.Model(\n", + " al.Galaxy,\n", + " redshift=mass_result.instance.galaxies.lens.redshift,\n", + " mass=subhalo_mass,\n", + " )\n", + "\n", + " subhalo.redshift = mass_result.instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_object = mass_result.instance.galaxies.lens.redshift\n", + " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", + " subhalo.mass.centre = subhalo_grid_search_result.model_centred_absolute(\n", + " a=1.0\n", + " ).galaxies.subhalo.mass.centre\n", + "\n", + " subhalo.redshift = subhalo_grid_search_result.model.galaxies.subhalo.redshift\n", + " subhalo.mass.redshift_object = subhalo.redshift\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=subhalo_grid_search_result.model.galaxies.lens,\n", + " subhalo=subhalo,\n", + " source=subhalo_grid_search_result.model.galaxies.source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[2]_[single_plane_refine]\",\n", + " **settings_search.search_dict,\n", + " n_live=600,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset + Masking__\n", + "\n", + "Load the `Interferometer` data, define the visibility and real-space masks." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256), pixel_scales=0.1, radius=mask_radius\n", + ")\n", + "\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "# dataset_name = \"alma\"\n", + "\n", + "# if dataset_name == \"alma\":\n", + "#\n", + "# real_space_mask = al.Mask2D.circular(\n", + "# shape_native=(800, 800),\n", + "# pixel_scales=0.01,\n", + "# radius=mask_radius,\n", + "# )\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Two Datasets__\n", + "\n", + "The SLaM pipeline runs in two phases that prefer different transformers:\n", + "\n", + "- `dataset_nufft` uses `TransformerNUFFT` (backed by JAX-native `nufftax`) for the `source_lp` stage.\n", + "- `dataset_sparse` uses `TransformerNUFFT` + `apply_sparse_operator(...)` for `source_pix_1`,\n", + " `source_pix_2`, `mass_total`, and every search of the SUBHALO PIPELINE.\n", + "\n", + "Both datasets are built from the same FITS files; only the transformer (and sparse-operator preload) differ." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_nufft = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "dataset_sparse = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sparse Operators__\n", + "\n", + "The `pixelization/modeling` example describes how the sparse operator formalism speeds up interferometer\n", + "pixelized source modeling, especially for many visibilities.\n", + "\n", + "We use a try / except to load the pre-computed curvature preload, which is necessary to use\n", + "the sparse operator formalism. If this file does not exist (e.g. you have not made it manually via\n", + "the `many_visibilities_preparation` example) it is made here.\n", + "\n", + "The sparse operator is applied only to `dataset_sparse` \u2014 the NUFFT-backed `dataset_nufft` used by\n", + "`source_lp` does not need it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " nufft_precision_operator = np.load(\n", + " file=dataset_path / \"nufft_precision_operator.npy\",\n", + " )\n", + "except FileNotFoundError:\n", + " nufft_precision_operator = None\n", + "\n", + "dataset_sparse = dataset_sparse.apply_sparse_operator(\n", + " nufft_precision_operator=nufft_precision_operator, use_jax=True, show_progress=True\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings__\n", + "\n", + "Disable the default position only linear algebra solver so the source reconstruction can have\n", + "negative pixel values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings = al.Settings(use_positive_only_solver=False)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"interferometer\") / \"slam\",\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__\n", + "\n", + "The code below calls the full SLaM PIPELINE. See the documentation string above each Python function for\n", + "a description of each pipeline step.\n", + "\n", + "Note the transformer split: `source_lp` is passed `dataset_nufft` (TransformerNUFFT), while every later stage\n", + "\u2014 including the SUBHALO PIPELINE \u2014 is passed `dataset_sparse` (TransformerNUFFT + sparse operator)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " dataset=dataset_nufft,\n", + " mask_radius=mask_radius,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " mesh_init=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization_init=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " settings=settings,\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_pix_result_1=source_pix_result_1,\n", + " source_pix_result_2=source_pix_result_2,\n", + " settings=settings,\n", + ")\n", + "\n", + "result_subhalo_grid_search = subhalo_grid_search(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mass_result=mass_result,\n", + " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", + " settings=settings,\n", + " grid_dimension_arcsec=3.0,\n", + " number_of_steps=2,\n", + ")\n", + "\n", + "result_with_subhalo = subhalo_refine(\n", + " settings_search=settings_search,\n", + " dataset=dataset_sparse,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mass_result=mass_result,\n", + " subhalo_grid_search_result=result_subhalo_grid_search,\n", + " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", + " settings=settings,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/subhalo/sensitivity/start_here.ipynb b/notebooks/interferometer/features/subhalo/sensitivity/start_here.ipynb index 806a28c66..cc6a7e6d8 100644 --- a/notebooks/interferometer/features/subhalo/sensitivity/start_here.ipynb +++ b/notebooks/interferometer/features/subhalo/sensitivity/start_here.ipynb @@ -1,340 +1,377 @@ { - "cells": [ - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# \"\"\"\n", - "# Feature: Sensitivity Mapping\n", - "# ============================\n", - "#\n", - "# Bayesian model comparison allows us to take a dataset, fit it with multiple models and use the Bayesian evidence to\n", - "# quantify which model objectively gives the best-fit following the principles of Occam's Razor.\n", - "#\n", - "# However, a complex model may not be favoured by model comparison not because it is the 'wrong' model, but simply\n", - "# because the dataset being fitted is not of a sufficient quality for the more complex model to be favoured. Sensitivity\n", - "# mapping addresses what quality of data would be needed for the more complex model to be favoured.\n", - "#\n", - "# In order to do this, sensitivity mapping involves us writing a function that uses the model(s) to simulate a dataset.\n", - "# We then use this function to simulate many datasets, for many different models, and fit each dataset using the same\n", - "# model-fitting procedure we used to perform Bayesian model comparison. This allows us to infer how much of a Bayesian\n", - "# evidence increase we should expect for datasets of varying quality and / or models with different parameters.\n", - "#\n", - "# For strong lensing, this process is crucial for dark matter substructure detection, as discussed in the following paper:\n", - "#\n", - "# https://arxiv.org/abs/0903.4752\n", - "#\n", - "# In substructure detection, we scan a strong lens dark matter subhalos by fitting a lens models which include a subhalo.\n", - "# This tells us whether we successfully did detect a subhalo, but it does not tell us where a subhalo has to be located\n", - "# (in relation to the source light) to be detectable, nor does to what masses of subhalo we could actually have made a\n", - "# detection.\n", - "#\n", - "# To answer these questions, we must perform sensitivity mapping, where we simulate many thousands of datasets,\n", - "# each of which include a dark matter subhalo at a given (y,x) coordinate at a given mass. We then fit each dataset twice,\n", - "# once with a lens model which does not include a subhalo and once with a lens model that does. If the Bayesian evidence\n", - "# of the model which includes a subhalo is higher than that which does not, then it means a subhalo was detectable!\n", - "# \"\"\"\n", - "# from autoconf import jax_wrapper # Ensures JAX environment variables are set before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "#\n", - "# import numpy as np\n", - "# from pathlib import Path\n", - "# import autofit as af\n", - "# import autolens as al\n", - "# import autolens.plot as aplt\n", - "#\n", - "# \"\"\"\n", - "# __Dataset + Masking__\n", - "#\n", - "# Load the `Interferometer` data, define the visibility and real-space masks and plot them.\n", - "# \"\"\"\n", - "# real_space_mask = al.Mask2D.circular(\n", - "# shape_native=(151, 151), pixel_scales=0.05, radius=3.0\n", - "# )\n", - "#\n", - "# dataset_name = \"simple\"\n", - "# dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name\n", - "#\n", - "# \"\"\"\n", - "# __Dataset Auto-Simulation__\n", - "#\n", - "# If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "# simulator script. This ensures that all example scripts can be run without manually simulating data first.\n", - "# \"\"\"\n", - "# if not dataset_path.exists():\n", - "# import subprocess\n", - "# import sys\n", - "# subprocess.run(\n", - "# [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - "# check=True,\n", - "# )\n", - "#\n", - "# dataset = al.Interferometer.from_fits(\n", - "# data_path=dataset_path / \"data.fits\",\n", - "# noise_map_path=dataset_path / \"noise_map.fits\",\n", - "# uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - "# real_space_mask=real_space_mask,\n", - "# )\n", - "# dataset = dataset.apply_settings(\n", - "# settings=al.SettingsInterferometer(transformer_class=al.TransformerDFT)\n", - "# )\n", - "#\n", - "#\n", - "#\n", - "# \"\"\"\n", - "# We are performing sensitivity mapping to determine where a subhalo is detectable. This will require us to simulate\n", - "# many realizations of our dataset with a lens model, called the `simulation_instance`. To get this model, we therefore\n", - "# fit the data before performing sensitivity mapping so that we can set the `simulation_instance` as the maximum\n", - "# likelihood model.\n", - "#\n", - "# We perform this fit using the lens model we will use to perform sensitivity mapping, which we call the `base_model`.\n", - "# \"\"\"\n", - "# base_model = af.Collection(\n", - "# galaxies=af.Collection(\n", - "# lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", - "# source=af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore),\n", - "# )\n", - "# )\n", - "#\n", - "# search_base = af.Nautilus(\n", - "# path_prefix=Path(\"interferometer\", \"misc\"),\n", - "# name=\"sensitivity_mapping_base\",\n", - "# unique_tag=dataset_name,\n", - "# n_live=100,\n", - "# )\n", - "#\n", - "# analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", - "#\n", - "# result = search_base.fit(model=base_model, analysis=analysis)\n", - "#\n", - "# \"\"\"\n", - "# We now define the `base_model` that we use to perform sensitivity mapping. This is the lens model that is fitted to\n", - "# every simulated strong lens without a subhalo, giving us the Bayesian evidence which we compare to the model which\n", - "# includes one!).\n", - "#\n", - "# For this model, we can use the `base_model` above, however we will use the result of fitting this model to the dataset\n", - "# before sensitivity mapping. This ensures the priors associated with each parameter are initialized so as to speed up\n", - "# each non-linear search performed during sensitivity mapping.\n", - "# \"\"\"\n", - "# base_model = result.model\n", - "#\n", - "# \"\"\"\n", - "# We now define the `perturb_model`, which is the model component whose parameters we iterate over to perform\n", - "# sensitivity mapping. In this case, this model is a `NFWMCRLudlowSph` model and we will iterate over its\n", - "# `centre` and `mass_at_200`. We set it up as a `Model` so it has an associated redshift and can be directly\n", - "# passed to the tracer in the simulate function below.\n", - "#\n", - "# Many instances of the `perturb_model` are created and used to simulate the many strong lens datasets that we fit.\n", - "# However, it is only included in half of the model-fits; corresponding to the lens models which include a dark matter\n", - "# subhalo and whose Bayesian evidence we compare to the simpler model-fits consisting of just the `base_model` to\n", - "# determine if the subhalo was detectable.\n", - "#\n", - "# By fitting both models to every simulated lens, we therefore infer the Bayesian evidence of every model to every\n", - "# dataset. Sensitivity mapping therefore maps out for what values of `centre` and `mass_at_200` in the dark matter\n", - "# subhalo the model-fit including a subhalo provide higher values of Bayesian evidence than the simpler model-fit (and\n", - "# therefore when it is detectable!).\n", - "# \"\"\"\n", - "# perturb_model = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.NFWMCRLudlowSph)\n", - "#\n", - "# \"\"\"\n", - "# Sensitivity mapping is typically performed over a large range of parameters. However, to make this demonstration quick\n", - "# and clear we are going to fix the `centre` of the subhalo to a value near the Einstein ring of (1.6, 0.0). We will\n", - "# iterate over just two `mass_at_200` values corresponding to subhalos of mass 1e6 and 1e13, of which only the latter\n", - "# will be shown to be detectable.\n", - "# \"\"\"\n", - "# perturb_model.mass.centre.centre_0 = 1.6\n", - "# perturb_model.mass.centre.centre_1 = 0.0\n", - "# perturb_model.mass.redshift_object = 0.5\n", - "# perturb_model.mass.redshift_source = 1.0\n", - "# perturb_model.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1e6, upper_limit=1e13)\n", - "#\n", - "# \"\"\"\n", - "# We are performing sensitivity mapping to determine where a subhalo is detectable. This will require us to\n", - "# simulate many realizations of our dataset with a lens model, called the `simulation_instance`. This model uses the\n", - "# result of the fit above.\n", - "# \"\"\"\n", - "# simulation_instance = result.instance\n", - "#\n", - "# \"\"\"\n", - "# We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above) and uses it to\n", - "# simulate a dataset which is subsequently fitted.\n", - "#\n", - "# Note that when this dataset is simulated, the quantity `instance.perturb` is used in the `simulate_cls`.\n", - "# This is an instance of the `NFWMCRLudlowSph`, and it is different every time the `simulate_cls` is called\n", - "# based on the value of sensitivity being computed.\n", - "#\n", - "# In this example, this `instance.perturb` corresponds to two different subhalos with values of `mass_at_200` of\n", - "# 1e6 MSun and 1e13 MSun.\n", - "# \"\"\"\n", - "#\n", - "#\n", - "# def __call__(instance, simulate_path):\n", - "# \"\"\"\n", - "# Set up the `Tracer` which is used to simulate the strong lens interferometer, which may include the subhalo in\n", - "# addition to the lens and source galaxy.\n", - "# \"\"\"\n", - "# tracer = al.Tracer(\n", - "# galaxies=[\n", - "# instance.galaxies.lens,\n", - "# instance.perturb,\n", - "# instance.galaxies.source,\n", - "# ]\n", - "# )\n", - "#\n", - "# \"\"\"\n", - "# Set up the grid, uv_wavelengths and simulator settings used to simulate interferometer dataset of the strong lens.\n", - "# These should be tuned to match the S/N and noise properties of the observed data you are performing sensitivity\n", - "# mapping on.\n", - "# \"\"\"\n", - "# grid = al.Grid2D.uniform(\n", - "# shape_native=real_space_mask.shape_native,\n", - "# pixel_scales=real_space_mask.pixel_scales,\n", - "# over_sampling = al.OverSamplingIterate(\n", - "# fractional_accuracy=0.9999,\n", - "# sub_steps=[2, 4, 8, 16]\n", - "# )\n", - "# )\n", - "#\n", - "# simulator = al.SimulatorInterferometer(\n", - "# uv_wavelengths=dataset.uv_wavelengths,\n", - "# exposure_time=300.0,\n", - "# noise_sigma=0.1,\n", - "# transformer_class=al.TransformerNUFFT,\n", - "# )\n", - "#\n", - "# simulated_dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "#\n", - "# \"\"\"\n", - "# The data generated by the simulate function is that which is fitted, so we should apply the mask for the analysis\n", - "# here before we return the simulated data.\n", - "# \"\"\"\n", - "# return al.Interferometer(\n", - "# data=simulated_dataset.visibilities,\n", - "# noise_map=simulated_dataset.noise_map,\n", - "# uv_wavelengths=simulated_dataset.uv_wavelengths,\n", - "# real_space_mask=real_space_mask,\n", - "# )\n", - "#\n", - "#\n", - "# \"\"\"\n", - "# Each model-fit performed by sensitivity mapping creates a new instance of an `Analysis` class, which contains the\n", - "# data simulated by the `simulate_cls` for that model.\n", - "#\n", - "# This requires us to write a wrapper around the PyAutoLens `Analysis` class.\n", - "# \"\"\"\n", - "#\n", - "#\n", - "# class AnalysisInterferometerSensitivity(al.AnalysisInterferometer):\n", - "# def __init__(self, dataset):\n", - "# super().__init__(dataset=dataset)\n", - "#\n", - "#\n", - "# \"\"\"\n", - "# We next specify the search used to perform each model fit by the sensitivity mapper.\n", - "# \"\"\"\n", - "# search = af.Nautilus(\n", - "# path_prefix=Path(\"interferometer\", \"misc\"),\n", - "# name=\"sensitivity_mapping\",\n", - "# unique_tag=dataset_name,\n", - "# n_live=100,\n", - "# force_x1_cpu=True, # ensures parallelizing over grid search works.\n", - "# )\n", - "#\n", - "# \"\"\"\n", - "# We can now combine all of the objects created above and perform sensitivity mapping. The inputs to the `Sensitivity`\n", - "# object below are:\n", - "#\n", - "# - `simulation_instance`: This is an instance of the model used to simulate every dataset that is fitted. In this example\n", - "# it is a lens model that does not include a subhalo, which was inferred by fitting the dataset we perform sensitivity\n", - "# mapping on.\n", - "#\n", - "# - `base_model`: This is the lens model that is fitted to every simulated dataset, which does not include a subhalo. In\n", - "# this example is composed of an `Isothermal` lens and MGE source.\n", - "#\n", - "# - `perturb_model`: This is the extra model component that alongside the `base_model` is fitted to every simulated\n", - "# dataset. In this example it is a `NFWMCRLudlowSph` dark matter subhalo.\n", - "#\n", - "# - `simulate_cls`: This is the function that uses the `simulation_instance` and many instances of the `perturb_model`\n", - "# to simulate many datasets that are fitted with the `base_model` and `base_model` + `perturb_model`.\n", - "#\n", - "# - `analysis_class`: The wrapper `Analysis` class that passes each simulated dataset to the `Analysis` class that fits\n", - "# the data.\n", - "#\n", - "# - `number_of_steps`: The number of steps over which the parameters in the `perturb_model` are iterated. In this\n", - "# example, `mass_at_200` has a `LogUniformPrior` with lower limit 1e6 and upper limit 1e13, therefore\n", - "# the `number_of_steps` of 2 will simulate and fit just 2 datasets where the `mass_at_200` is between 1e6 and 1e13.\n", - "#\n", - "# - `number_of_cores`: The number of cores over which the sensitivity mapping is performed, enabling parallel processing\n", - "# if set above 1.\n", - "# \"\"\"\n", - "# from autofit.non_linear.grid import sensitivity as s\n", - "#\n", - "# sensitivity = s.Sensitivity(\n", - "# search=search,\n", - "# simulation_instance=simulation_instance,\n", - "# base_model=base_model,\n", - "# perturb_model=perturb_model,\n", - "# simulate_cls=simulate_cls,\n", - "# analysis_class=AnalysisInterferometerSensitivity,\n", - "# number_of_cores=2,\n", - "# number_of_steps=2,\n", - "# )\n", - "#\n", - "# sensitivity_result = sensitivity.run()\n", - "#\n", - "# \"\"\"\n", - "# You should now look at the results of the sensitivity mapping in the folder `output/features/sensitivity_mapping`.\n", - "#\n", - "# You will note the following 4 model-fits have been performed:\n", - "#\n", - "# - The `base_model` is fitted to a simulated dataset where a subhalo with `mass_at_200=1e6` is included.\n", - "#\n", - "# - The `base_model` + `perturb_model` is fitted to a simulated dataset where a subhalo with `mass_at_200=1e6`\n", - "# is included.\n", - "#\n", - "# - The `base_model` is fitted to a simulated dataset where a subhalo with `mass_at_200=1e13` is included.\n", - "#\n", - "# - The `base_model` + `perturb_model` is fitted to a simulated dataset where a subhalo with `mass_at_200=1e13` is\n", - "# included.\n", - "#\n", - "# The fit produces a `sensitivity_result`.\n", - "#\n", - "# We are still developing the `SensitivityResult` class to provide a data structure that better streamlines the analysis\n", - "# of results. If you intend to use sensitivity mapping, the best way to interpret the resutls is currently via\n", - "# **PyAutoFit**'s database and `Aggregator` tools.\n", - "# \"\"\"\n", - "# print(sensitivity_result.results[0].result.samples.log_evidence)\n", - "# print(sensitivity_result.results[1].result.samples.log_evidence)\n", - "#\n", - "# \"\"\"\n", - "# Finish.\n", - "# \"\"\"\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# \"\"\"\n", + "# Feature: Sensitivity Mapping\n", + "# ============================\n", + "#\n", + "# Bayesian model comparison allows us to take a dataset, fit it with multiple models and use the Bayesian evidence to\n", + "# quantify which model objectively gives the best-fit following the principles of Occam's Razor.\n", + "#\n", + "# However, a complex model may not be favoured by model comparison not because it is the 'wrong' model, but simply\n", + "# because the dataset being fitted is not of a sufficient quality for the more complex model to be favoured. Sensitivity\n", + "# mapping addresses what quality of data would be needed for the more complex model to be favoured.\n", + "#\n", + "# In order to do this, sensitivity mapping involves us writing a function that uses the model(s) to simulate a dataset.\n", + "# We then use this function to simulate many datasets, for many different models, and fit each dataset using the same\n", + "# model-fitting procedure we used to perform Bayesian model comparison. This allows us to infer how much of a Bayesian\n", + "# evidence increase we should expect for datasets of varying quality and / or models with different parameters.\n", + "#\n", + "# For strong lensing, this process is crucial for dark matter substructure detection, as discussed in the following paper:\n", + "#\n", + "# https://arxiv.org/abs/0903.4752\n", + "#\n", + "# In substructure detection, we scan a strong lens dark matter subhalos by fitting a lens models which include a subhalo.\n", + "# This tells us whether we successfully did detect a subhalo, but it does not tell us where a subhalo has to be located\n", + "# (in relation to the source light) to be detectable, nor does to what masses of subhalo we could actually have made a\n", + "# detection.\n", + "#\n", + "# To answer these questions, we must perform sensitivity mapping, where we simulate many thousands of datasets,\n", + "# each of which include a dark matter subhalo at a given (y,x) coordinate at a given mass. We then fit each dataset twice,\n", + "# once with a lens model which does not include a subhalo and once with a lens model that does. If the Bayesian evidence\n", + "# of the model which includes a subhalo is higher than that which does not, then it means a subhalo was detectable!\n", + "# \"\"\"\n", + "# from autoconf import jax_wrapper # Ensures JAX environment variables are set before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "#\n", + "# import numpy as np\n", + "# from pathlib import Path\n", + "# import autofit as af\n", + "# import autolens as al\n", + "# import autolens.plot as aplt\n", + "#\n", + "# \"\"\"\n", + "# __Dataset + Masking__\n", + "#\n", + "# Load the `Interferometer` data, define the visibility and real-space masks and plot them.\n", + "# \"\"\"\n", + "# real_space_mask = al.Mask2D.circular(\n", + "# shape_native=(151, 151), pixel_scales=0.05, radius=3.0\n", + "# )\n", + "#\n", + "# dataset_name = \"simple\"\n", + "# dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name\n", + "#\n", + "# \"\"\"\n", + "# __Dataset Auto-Simulation__\n", + "#\n", + "# If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "# simulator script. This ensures that all example scripts can be run without manually simulating data first.\n", + "# \"\"\"\n", + "# if not dataset_path.exists():\n", + "# import subprocess\n", + "# import sys\n", + "# subprocess.run(\n", + "# [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + "# check=True,\n", + "# )\n", + "#\n", + "# dataset = al.Interferometer.from_fits(\n", + "# data_path=dataset_path / \"data.fits\",\n", + "# noise_map_path=dataset_path / \"noise_map.fits\",\n", + "# uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + "# real_space_mask=real_space_mask,\n", + "# )\n", + "# dataset = dataset.apply_settings(\n", + "# settings=al.SettingsInterferometer(transformer_class=al.TransformerDFT)\n", + "# )\n", + "#\n", + "#\n", + "#\n", + "# \"\"\"\n", + "# We are performing sensitivity mapping to determine where a subhalo is detectable. This will require us to simulate\n", + "# many realizations of our dataset with a lens model, called the `simulation_instance`. To get this model, we therefore\n", + "# fit the data before performing sensitivity mapping so that we can set the `simulation_instance` as the maximum\n", + "# likelihood model.\n", + "#\n", + "# We perform this fit using the lens model we will use to perform sensitivity mapping, which we call the `base_model`.\n", + "# \"\"\"\n", + "# base_model = af.Collection(\n", + "# galaxies=af.Collection(\n", + "# lens=af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal),\n", + "# source=af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore),\n", + "# )\n", + "# )\n", + "#\n", + "# search_base = af.Nautilus(\n", + "# path_prefix=Path(\"interferometer\", \"misc\"),\n", + "# name=\"sensitivity_mapping_base\",\n", + "# unique_tag=dataset_name,\n", + "# n_live=100,\n", + "# )\n", + "#\n", + "# analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=True)\n", + "#\n", + "# result = search_base.fit(model=base_model, analysis=analysis)\n", + "#\n", + "# \"\"\"\n", + "# We now define the `base_model` that we use to perform sensitivity mapping. This is the lens model that is fitted to\n", + "# every simulated strong lens without a subhalo, giving us the Bayesian evidence which we compare to the model which\n", + "# includes one!).\n", + "#\n", + "# For this model, we can use the `base_model` above, however we will use the result of fitting this model to the dataset\n", + "# before sensitivity mapping. This ensures the priors associated with each parameter are initialized so as to speed up\n", + "# each non-linear search performed during sensitivity mapping.\n", + "# \"\"\"\n", + "# base_model = result.model\n", + "#\n", + "# \"\"\"\n", + "# We now define the `perturb_model`, which is the model component whose parameters we iterate over to perform\n", + "# sensitivity mapping. In this case, this model is a `NFWMCRLudlowSph` model and we will iterate over its\n", + "# `centre` and `mass_at_200`. We set it up as a `Model` so it has an associated redshift and can be directly\n", + "# passed to the tracer in the simulate function below.\n", + "#\n", + "# Many instances of the `perturb_model` are created and used to simulate the many strong lens datasets that we fit.\n", + "# However, it is only included in half of the model-fits; corresponding to the lens models which include a dark matter\n", + "# subhalo and whose Bayesian evidence we compare to the simpler model-fits consisting of just the `base_model` to\n", + "# determine if the subhalo was detectable.\n", + "#\n", + "# By fitting both models to every simulated lens, we therefore infer the Bayesian evidence of every model to every\n", + "# dataset. Sensitivity mapping therefore maps out for what values of `centre` and `mass_at_200` in the dark matter\n", + "# subhalo the model-fit including a subhalo provide higher values of Bayesian evidence than the simpler model-fit (and\n", + "# therefore when it is detectable!).\n", + "# \"\"\"\n", + "# perturb_model = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.NFWMCRLudlowSph)\n", + "#\n", + "# \"\"\"\n", + "# Sensitivity mapping is typically performed over a large range of parameters. However, to make this demonstration quick\n", + "# and clear we are going to fix the `centre` of the subhalo to a value near the Einstein ring of (1.6, 0.0). We will\n", + "# iterate over just two `mass_at_200` values corresponding to subhalos of mass 1e6 and 1e13, of which only the latter\n", + "# will be shown to be detectable.\n", + "# \"\"\"\n", + "# perturb_model.mass.centre.centre_0 = 1.6\n", + "# perturb_model.mass.centre.centre_1 = 0.0\n", + "# perturb_model.mass.redshift_object = 0.5\n", + "# perturb_model.mass.redshift_source = 1.0\n", + "# perturb_model.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1e6, upper_limit=1e13)\n", + "#\n", + "# \"\"\"\n", + "# We are performing sensitivity mapping to determine where a subhalo is detectable. This will require us to\n", + "# simulate many realizations of our dataset with a lens model, called the `simulation_instance`. This model uses the\n", + "# result of the fit above.\n", + "# \"\"\"\n", + "# simulation_instance = result.instance\n", + "#\n", + "# \"\"\"\n", + "# We now write the `simulate_cls`, which takes the `simulation_instance` of our model (defined above) and uses it to\n", + "# simulate a dataset which is subsequently fitted.\n", + "#\n", + "# Note that when this dataset is simulated, the quantity `instance.perturb` is used in the `simulate_cls`.\n", + "# This is an instance of the `NFWMCRLudlowSph`, and it is different every time the `simulate_cls` is called\n", + "# based on the value of sensitivity being computed.\n", + "#\n", + "# In this example, this `instance.perturb` corresponds to two different subhalos with values of `mass_at_200` of\n", + "# 1e6 MSun and 1e13 MSun.\n", + "# \"\"\"\n", + "#\n", + "#\n", + "# def __call__(instance, simulate_path):\n", + "# \"\"\"\n", + "# Set up the `Tracer` which is used to simulate the strong lens interferometer, which may include the subhalo in\n", + "# addition to the lens and source galaxy.\n", + "# \"\"\"\n", + "# tracer = al.Tracer(\n", + "# galaxies=[\n", + "# instance.galaxies.lens,\n", + "# instance.perturb,\n", + "# instance.galaxies.source,\n", + "# ]\n", + "# )\n", + "#\n", + "# \"\"\"\n", + "# Set up the grid, uv_wavelengths and simulator settings used to simulate interferometer dataset of the strong lens.\n", + "# These should be tuned to match the S/N and noise properties of the observed data you are performing sensitivity\n", + "# mapping on.\n", + "# \"\"\"\n", + "# grid = al.Grid2D.uniform(\n", + "# shape_native=real_space_mask.shape_native,\n", + "# pixel_scales=real_space_mask.pixel_scales,\n", + "# over_sampling = al.OverSamplingIterate(\n", + "# fractional_accuracy=0.9999,\n", + "# sub_steps=[2, 4, 8, 16]\n", + "# )\n", + "# )\n", + "#\n", + "# simulator = al.SimulatorInterferometer(\n", + "# uv_wavelengths=dataset.uv_wavelengths,\n", + "# exposure_time=300.0,\n", + "# noise_sigma=0.1,\n", + "# transformer_class=al.TransformerNUFFT,\n", + "# )\n", + "#\n", + "# simulated_dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "#\n", + "# \"\"\"\n", + "# The data generated by the simulate function is that which is fitted, so we should apply the mask for the analysis\n", + "# here before we return the simulated data.\n", + "# \"\"\"\n", + "# return al.Interferometer(\n", + "# data=simulated_dataset.visibilities,\n", + "# noise_map=simulated_dataset.noise_map,\n", + "# uv_wavelengths=simulated_dataset.uv_wavelengths,\n", + "# real_space_mask=real_space_mask,\n", + "# )\n", + "#\n", + "#\n", + "# \"\"\"\n", + "# Each model-fit performed by sensitivity mapping creates a new instance of an `Analysis` class, which contains the\n", + "# data simulated by the `simulate_cls` for that model.\n", + "#\n", + "# This requires us to write a wrapper around the PyAutoLens `Analysis` class.\n", + "# \"\"\"\n", + "#\n", + "#\n", + "# class AnalysisInterferometerSensitivity(al.AnalysisInterferometer):\n", + "# def __init__(self, dataset):\n", + "# super().__init__(dataset=dataset)\n", + "#\n", + "#\n", + "# \"\"\"\n", + "# We next specify the search used to perform each model fit by the sensitivity mapper.\n", + "# \"\"\"\n", + "# search = af.Nautilus(\n", + "# path_prefix=Path(\"interferometer\", \"misc\"),\n", + "# name=\"sensitivity_mapping\",\n", + "# unique_tag=dataset_name,\n", + "# n_live=100,\n", + "# force_x1_cpu=True, # ensures parallelizing over grid search works.\n", + "# )\n", + "#\n", + "# \"\"\"\n", + "# We can now combine all of the objects created above and perform sensitivity mapping. The inputs to the `Sensitivity`\n", + "# object below are:\n", + "#\n", + "# - `simulation_instance`: This is an instance of the model used to simulate every dataset that is fitted. In this example\n", + "# it is a lens model that does not include a subhalo, which was inferred by fitting the dataset we perform sensitivity\n", + "# mapping on.\n", + "#\n", + "# - `base_model`: This is the lens model that is fitted to every simulated dataset, which does not include a subhalo. In\n", + "# this example is composed of an `Isothermal` lens and MGE source.\n", + "#\n", + "# - `perturb_model`: This is the extra model component that alongside the `base_model` is fitted to every simulated\n", + "# dataset. In this example it is a `NFWMCRLudlowSph` dark matter subhalo.\n", + "#\n", + "# - `simulate_cls`: This is the function that uses the `simulation_instance` and many instances of the `perturb_model`\n", + "# to simulate many datasets that are fitted with the `base_model` and `base_model` + `perturb_model`.\n", + "#\n", + "# - `analysis_class`: The wrapper `Analysis` class that passes each simulated dataset to the `Analysis` class that fits\n", + "# the data.\n", + "#\n", + "# - `number_of_steps`: The number of steps over which the parameters in the `perturb_model` are iterated. In this\n", + "# example, `mass_at_200` has a `LogUniformPrior` with lower limit 1e6 and upper limit 1e13, therefore\n", + "# the `number_of_steps` of 2 will simulate and fit just 2 datasets where the `mass_at_200` is between 1e6 and 1e13.\n", + "#\n", + "# - `number_of_cores`: The number of cores over which the sensitivity mapping is performed, enabling parallel processing\n", + "# if set above 1.\n", + "# \"\"\"\n", + "# from autofit.non_linear.grid import sensitivity as s\n", + "#\n", + "# sensitivity = s.Sensitivity(\n", + "# search=search,\n", + "# simulation_instance=simulation_instance,\n", + "# base_model=base_model,\n", + "# perturb_model=perturb_model,\n", + "# simulate_cls=simulate_cls,\n", + "# analysis_class=AnalysisInterferometerSensitivity,\n", + "# number_of_cores=2,\n", + "# number_of_steps=2,\n", + "# )\n", + "#\n", + "# sensitivity_result = sensitivity.run()\n", + "#\n", + "# \"\"\"\n", + "# You should now look at the results of the sensitivity mapping in the folder `output/features/sensitivity_mapping`.\n", + "#\n", + "# You will note the following 4 model-fits have been performed:\n", + "#\n", + "# - The `base_model` is fitted to a simulated dataset where a subhalo with `mass_at_200=1e6` is included.\n", + "#\n", + "# - The `base_model` + `perturb_model` is fitted to a simulated dataset where a subhalo with `mass_at_200=1e6`\n", + "# is included.\n", + "#\n", + "# - The `base_model` is fitted to a simulated dataset where a subhalo with `mass_at_200=1e13` is included.\n", + "#\n", + "# - The `base_model` + `perturb_model` is fitted to a simulated dataset where a subhalo with `mass_at_200=1e13` is\n", + "# included.\n", + "#\n", + "# The fit produces a `sensitivity_result`.\n", + "#\n", + "# We are still developing the `SensitivityResult` class to provide a data structure that better streamlines the analysis\n", + "# of results. If you intend to use sensitivity mapping, the best way to interpret the resutls is currently via\n", + "# **PyAutoFit**'s database and `Aggregator` tools.\n", + "# \"\"\"\n", + "# print(sensitivity_result.results[0].result.samples.log_evidence)\n", + "# print(sensitivity_result.results[1].result.samples.log_evidence)\n", + "#\n", + "# \"\"\"\n", + "# Finish.\n", + "# \"\"\"\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/features/subhalo/simulator.ipynb b/notebooks/interferometer/features/subhalo/simulator.ipynb index 25394b1de..af3a7ad5f 100644 --- a/notebooks/interferometer/features/subhalo/simulator.ipynb +++ b/notebooks/interferometer/features/subhalo/simulator.ipynb @@ -1,372 +1,409 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Subhalo\n", - "==================\n", - "\n", - "This script simulates `Interferometer` data of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The subhalo`s `MassProfile` is a `NFWSph`.\n", - " - The source galaxy's light is an `Sersic`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated (in this case, `Interferometer` data).\n", - "- **Simulate:** For simulating interferometer data of a strong lens, we recommend using a Grid2D object with a.\n", - "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear), subhalo (NFW) and source galaxy light (elliptical Sersic).\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated (in this case, `Interferometer` data) and `dataset_name` \n", - "gives it a descriptive name. \n", - "\n", - " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/image.fits`.\n", - " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.fits`.\n", - " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/psf.fits`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"interferometer\"\n", - "dataset_name = \"dark_matter_subhalo\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "For simulating interferometer data of a strong lens, we recommend using a Grid2D object with a `sub_size` of 1. This\n", - "simplifies the generation of the strong lens image in real space before it is transformed to Fourier space." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.05)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To perform the Fourier transform we need the wavelengths of the baselines, which we'll load from the fits file below.\n", - "\n", - "By default we use baselines from the Square Mile Array (SMA), which produces low resolution interferometer data that\n", - "can be fitted extremely efficiently. The `autolens_workspace` includes ALMA uv_wavelengths files for simulating\n", - "much high resolution datasets (which can be performed by replacing \"sma.fits\" below with \"alma.fits\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "uv_wavelengths_path = Path(\"dataset\", dataset_type, \"uv_wavelengths\")\n", - "uv_wavelengths = al.ndarray_via_fits_from(\n", - " file_path=Path(uv_wavelengths_path, \"sma.fits\"), hdu=0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the interferometer dataset we first create a simulator, which defines the shape, resolution and pixel-scale \n", - "of the visibilities that are simulated, as well as its exposure time, noise levels and uv-wavelengths." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorInterferometer(\n", - " uv_wavelengths=uv_wavelengths,\n", - " exposure_time=300.0,\n", - " noise_sigma=1000.0,\n", - " transformer_class=al.TransformerDFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's mass (SIE+Shear), subhalo (NFW) and source galaxy light (elliptical Sersic) for this \n", - "simulated lens.\n", - "\n", - "For lens modeling, defining ellipticity in terms of the `ell_comps` improves the model-fitting procedure.\n", - "\n", - "However, for simulating a strong lens you may find it more intuitive to define the elliptical geometry using the \n", - "axis-ratio of the profile (axis_ratio = semi-major axis / semi-minor axis = b/a) and position angle, where angle is\n", - "in degrees and defined counter clockwise from the positive x-axis.\n", - "\n", - "We can use the `convert` module to determine the elliptical components from the axis-ratio and angle." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " subhalo=al.mp.NFWTruncatedMCRLudlowSph(centre=(1.601, 0.0), mass_at_200=1.0e10),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.0),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "# %%\n", - "'''\n", - "Use these galaxies to setup a tracer, which will generate the image for the simulated interferometer dataset.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now pass this simulator a tracer, which creates the ray-traced image plotted above and simulates it as an\n", - "interferometer dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets plot the simulated interferometer dataset before we output it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_interferometer(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_interferometer_dirty_images(\n", - " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Subhalo\n", + "==================\n", + "\n", + "This script simulates `Interferometer` data of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The subhalo`s `MassProfile` is a `NFWSph`.\n", + " - The source galaxy's light is an `Sersic`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated (in this case, `Interferometer` data).\n", + "- **Simulate:** For simulating interferometer data of a strong lens, we recommend using a Grid2D object with a.\n", + "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear), subhalo (NFW) and source galaxy light (elliptical Sersic).\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated (in this case, `Interferometer` data) and `dataset_name` \n", + "gives it a descriptive name. \n", + "\n", + " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/image.fits`.\n", + " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.fits`.\n", + " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/psf.fits`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"interferometer\"\n", + "dataset_name = \"dark_matter_subhalo\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "For simulating interferometer data of a strong lens, we recommend using a Grid2D object with a `sub_size` of 1. This\n", + "simplifies the generation of the strong lens image in real space before it is transformed to Fourier space." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.05)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To perform the Fourier transform we need the wavelengths of the baselines, which we'll load from the fits file below.\n", + "\n", + "By default we use baselines from the Square Mile Array (SMA), which produces low resolution interferometer data that\n", + "can be fitted extremely efficiently. The `autolens_workspace` includes ALMA uv_wavelengths files for simulating\n", + "much high resolution datasets (which can be performed by replacing \"sma.fits\" below with \"alma.fits\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "uv_wavelengths_path = Path(\"dataset\", dataset_type, \"uv_wavelengths\")\n", + "uv_wavelengths = al.ndarray_via_fits_from(\n", + " file_path=Path(uv_wavelengths_path, \"sma.fits\"), hdu=0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the interferometer dataset we first create a simulator, which defines the shape, resolution and pixel-scale \n", + "of the visibilities that are simulated, as well as its exposure time, noise levels and uv-wavelengths." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorInterferometer(\n", + " uv_wavelengths=uv_wavelengths,\n", + " exposure_time=300.0,\n", + " noise_sigma=1000.0,\n", + " transformer_class=al.TransformerDFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's mass (SIE+Shear), subhalo (NFW) and source galaxy light (elliptical Sersic) for this \n", + "simulated lens.\n", + "\n", + "For lens modeling, defining ellipticity in terms of the `ell_comps` improves the model-fitting procedure.\n", + "\n", + "However, for simulating a strong lens you may find it more intuitive to define the elliptical geometry using the \n", + "axis-ratio of the profile (axis_ratio = semi-major axis / semi-minor axis = b/a) and position angle, where angle is\n", + "in degrees and defined counter clockwise from the positive x-axis.\n", + "\n", + "We can use the `convert` module to determine the elliptical components from the axis-ratio and angle." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " subhalo=al.mp.NFWTruncatedMCRLudlowSph(centre=(1.601, 0.0), mass_at_200=1.0e10),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.0),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "# %%\n", + "'''\n", + "Use these galaxies to setup a tracer, which will generate the image for the simulated interferometer dataset.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now pass this simulator a tracer, which creates the ray-traced image plotted above and simulates it as an\n", + "interferometer dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets plot the simulated interferometer dataset before we output it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_interferometer(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_interferometer_dirty_images(\n", + " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/fit.ipynb b/notebooks/interferometer/fit.ipynb index 43c78f9e1..9105daade 100644 --- a/notebooks/interferometer/fit.ipynb +++ b/notebooks/interferometer/fit.ipynb @@ -1,663 +1,700 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fits\n", - "====\n", - "\n", - "This guide shows how to fit data using the `FitInterferometer` object, including visualizing and interpreting its results.\n", - "\n", - "References\n", - "----------\n", - "\n", - "This example uses functionality described fully in other examples in the `guides` package:\n", - "\n", - "- `guides/plot`: Using the plotting API (`aplt.plot_array`, `aplt.subplot_fit_interferometer`, etc.) to visualize figures.\n", - "- `guides/units`: The source code unit conventions (e.g. arc seconds for distances and how to convert to physical units).\n", - "- `guides/data_structures`: The bespoke data structures used to store 1D and 2d arrays.\n", - "\n", - "__Contents__\n", - "\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Loading Data:** We we begin by loading the strong lens dataset `simple` from .fits files, which is the dataset we.\n", - "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", - "- **Bad Fit:** A bad lens model will show features in the residual-map and chi-squared map.\n", - "- **Fit Quantities:** The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", - "- **Figures of Merit:** There are single valued floats which quantify the goodness of fit.\n", - "- **Plane Quantities:** The `FitInterferometer` object has specific quantities which break down each image of each plane.\n", - "- **Outputting Results:** You may wish to output certain results to .fits files for later inspection.\n", - "\n", - "__JAX__\n", - "\n", - "`FitInterferometer` runs on either NumPy or JAX. For the standard\n", - "analysis-driven path \u2014 where `AnalysisInterferometer` auto-enables\n", - "`use_jax=True` and the search driver handles the JIT \u2014 see `start_here.py`\n", - "/ `modeling.py`. For the JIT-it-yourself path around individual library\n", - "methods, see `scripts/guides/lens_calc.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Data__\n", - "\n", - "We we begin by loading the strong lens dataset `simple` from .fits files, which is the dataset \n", - "we will use to demonstrate fitting.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and\n", - "compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native Non-Uniform Fast Fourier Transform\n", - "backed by `nufftax`, which scales efficiently from a few hundred visibilities to tens of millions.\n", - "\n", - "This dataset was simulated using the `interferometer/simulator` example, read through that to understand how\n", - "the data this example fits was generated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `aplt.subplot_interferometer_dirty_images` contains a subplot which plots all the key properties of the dataset simultaneously.\n", - "\n", - "This includes the observed visibility data, RMS noise map and other information." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Visibility data is in uv space, making it hard to interpret by eye.\n", - "\n", - "The dirty images of the interferometer dataset can plotted, which use the transformer of the interferometer \n", - "to map the visibilities, noise-map or other quantity to a real-space image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Fitting__\n", - "\n", - "Following the previous overview example, we can make a tracer from a collection of light profiles, mass profiles\n", - "and galaxies.\n", - "\n", - "The combination of light and mass profiels below is the same as those used to generate the simulated \n", - "dataset we loaded above.\n", - "\n", - "It therefore produces a tracer whose image looks exactly like the dataset.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=1.0,\n", - " sersic_index=2.5,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Because the tracer's light and mass profiles are the same used to make the dataset, its image is nearly the same as the\n", - "observed image.\n", - "\n", - "We can plot the image of the tracer to confirm this, noting that for a tracer its images are always in real space\n", - "(not Fourier space like the interferometer dataset) and therefore they can be directly visualized." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "However, the tracer's image is not what we observe in the interferometer dataset, because we observe the image as\n", - "visibilities in the uv-plane. \n", - "\n", - "To compare directly to the data, we therefore need to Fourier transform the tracer's image to the uv-plane. \n", - "\n", - "We do this by creating a `FitInterferometer` object, which performs this Fourier transform as part of the fitting \n", - "procedure.\n", - "\n", - "The code plots the result of this, by using the `model_data` of the fit, which performs this Fourier transform \n", - "on the tracer image above and plots the result visibilities in uv-space." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The visibilities are again hard to interpret by eye, so we can plot the dirty image of the fit's model data. This \n", - "dirty image is the Fourier transform of the fit's model data (therefore the Fourier transform of the tracer's image) and\n", - "can be compared directly to the image of the tracer above (albeit it still has the interferometer's PSF/dirty beam\n", - "convolved with it)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit does a lot more than just Fourier transform the tracer's image it also creates the following:\n", - "\n", - " - The `residual_map`: The `model_data` visibilities subtracted from the observed dataset`s `data` visibilities.\n", - " - The `normalized_residual_map`: The `residual_map `divided by the observed dataset's `noise_map`.\n", - " - The `chi_squared_map`: The `normalized_residual_map` squared.\n", - "\n", - "For a good lens model where the model and tracer are representative of the strong lens system the\n", - "residuals, normalized residuals and chi-squareds are minimized:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "A subplot can be plotted which contains all of the above quantities, as well as other information contained in the\n", - "tracer such as the source-plane image, a zoom in of the source-plane and a normalized residual map where the colorbar\n", - "goes from 1.0 sigma to -1.0 sigma, to highlight regions where the fit is poor.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_interferometer(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Once again, dirty images are often easier to interpret, so we can plot a subplot of the dirty images of the data, model\n", - "data, residuals and chi-squared." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_dirty_images(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit also provides us with a ``log_likelihood``, a single value quantifying how good the tracer fitted the dataset.\n", - "\n", - "Lens modeling, describe in the next overview example, effectively tries to maximize this log likelihood value." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Bad Fit__\n", - "\n", - "A bad lens model will show features in the residual-map and chi-squared map.\n", - "\n", - "We can produce such an image by creating a tracer with different lens and source galaxies. In the example below, we \n", - "change the centre of the source galaxy from (0.0, 0.0) to (0.05, 0.05), which leads to residuals appearing\n", - "in the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.1, 0.1),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.1, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A new fit using this plane shows residuals, normalized residuals and chi-squared which are non-zero. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=fit)\n", - "aplt.subplot_fit_dirty_images(fit=fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We also note that its likelihood decreases." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit Quantities__\n", - "\n", - "The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", - "\n", - "There is a `model_data`, which is the image-plane visibilities of the tracer.\n", - "\n", - "This is the image that is fitted to the data in order to compute the log likelihood and therefore quantify the \n", - "goodness-of-fit.\n", - "\n", - "If you are unclear on what `slim` means, refer to the section `Data Structure` at the top of this example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_data)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There are numerous ndarrays showing the goodness of fit: \n", - "\n", - " - `residual_map`: Residuals = (Data - Model_Data).\n", - " - `normalized_residual_map`: Normalized_Residual = (Data - Model_Data) / Noise\n", - " - `chi_squared_map`: Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.residual_map.slim)\n", - "print(fit.normalized_residual_map.slim)\n", - "print(fit.chi_squared_map.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There are `dirty` variants of the above maps, which transform the visibilities, residual-map, chi squared and other\n", - "values to to real-space images using the interferometer's transformer.\n", - "\n", - "These real space images can be mapped between their `slim` and `native` representations (see the\n", - "`guides/data_structures` example for more information on these terms)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.dirty_image.slim) # Data\n", - "print(fit.dirty_model_image.slim)\n", - "print(fit.dirty_residual_map.slim)\n", - "print(fit.dirty_normalized_residual_map.slim)\n", - "print(fit.dirty_chi_squared_map.slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Figures of Merit__\n", - "\n", - "There are single valued floats which quantify the goodness of fit:\n", - "\n", - " - `chi_squared`: The sum of the `chi_squared_map`.\n", - "\n", - " - `noise_normalization`: The normalizing noise term in the likelihood function \n", - " where [Noise_Term] = sum(log(2*pi*[Noise]**2.0)).\n", - "\n", - " - `log_likelihood`: The log likelihood value of the fit where [LogLikelihood] = -0.5*[Chi_Squared_Term + Noise_Term].\n", - " \n", - "These sum other both the real and imaginary components of the visibilities to give a single value for each quantity." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.chi_squared)\n", - "print(fit.noise_normalization)\n", - "print(fit.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Plane Quantities__\n", - "\n", - "The `FitInterferometer` object has specific quantities which break down each image of each plane:\n", - "\n", - " - `model_visibilities_of_planes_list`: Model-images of each individual plane, which in this example is a model image of the \n", - " lens galaxy and model image of the lensed source galaxy, both corresponding to dirty images.\n", - "\n", - " - `subtracted_images_of_planes_list`: Subtracted images of each individual plane, which are the data's image with\n", - " all other plane's model-images subtracted. For example, the first subtracted image has the source galaxy's model image\n", - " subtracted and therefore is of only the lens galaxy's emission. The second subtracted image is of the lensed source,\n", - " with the lens galaxy's light removed.\n", - "\n", - "For multi-plane lens systems these lists will be extended to provide information on every individual plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.model_visibilities_of_planes_list[1].slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There is also a `galaxy_model_visibilities_dict` which maps each galaxy in the tracer to its model visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.galaxy_model_visibilities_dict[source_galaxy].slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A dictionary which maps the model images of each galaxy is also available.\n", - "\n", - "These are not the dirty images, but instead the images of each galaxy that come from the tracer object\n", - "(e.g. simply evaluating the tracer's image on the interferometer's real-space grid)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.galaxy_image_dict[source_galaxy].slim)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Outputting Results__\n", - "\n", - "You may wish to output certain results to .fits files for later inspection. \n", - "\n", - "For example, one could output the lens light subtracted image of the lensed source galaxy to a .fits file such that\n", - "we could fit this source-only image again with an independent pipeline." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_model_image = fit.galaxy_image_dict[source_galaxy]\n", - "aplt.fits_array(\n", - " array=source_model_image,\n", - " file_path=dataset_path / \"source_model_image.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fin." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fits\n", + "====\n", + "\n", + "This guide shows how to fit data using the `FitInterferometer` object, including visualizing and interpreting its results.\n", + "\n", + "References\n", + "----------\n", + "\n", + "This example uses functionality described fully in other examples in the `guides` package:\n", + "\n", + "- `guides/plot`: Using the plotting API (`aplt.plot_array`, `aplt.subplot_fit_interferometer`, etc.) to visualize figures.\n", + "- `guides/units`: The source code unit conventions (e.g. arc seconds for distances and how to convert to physical units).\n", + "- `guides/data_structures`: The bespoke data structures used to store 1D and 2d arrays.\n", + "\n", + "__Contents__\n", + "\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Loading Data:** We we begin by loading the strong lens dataset `simple` from .fits files, which is the dataset we.\n", + "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", + "- **Bad Fit:** A bad lens model will show features in the residual-map and chi-squared map.\n", + "- **Fit Quantities:** The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", + "- **Figures of Merit:** There are single valued floats which quantify the goodness of fit.\n", + "- **Plane Quantities:** The `FitInterferometer` object has specific quantities which break down each image of each plane.\n", + "- **Outputting Results:** You may wish to output certain results to .fits files for later inspection.\n", + "\n", + "__JAX__\n", + "\n", + "`FitInterferometer` runs on either NumPy or JAX. For the standard\n", + "analysis-driven path \u2014 where `AnalysisInterferometer` auto-enables\n", + "`use_jax=True` and the search driver handles the JIT \u2014 see `start_here.py`\n", + "/ `modeling.py`. For the JIT-it-yourself path around individual library\n", + "methods, see `scripts/guides/lens_calc.py`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Data__\n", + "\n", + "We we begin by loading the strong lens dataset `simple` from .fits files, which is the dataset \n", + "we will use to demonstrate fitting.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and\n", + "compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native Non-Uniform Fast Fourier Transform\n", + "backed by `nufftax`, which scales efficiently from a few hundred visibilities to tens of millions.\n", + "\n", + "This dataset was simulated using the `interferometer/simulator` example, read through that to understand how\n", + "the data this example fits was generated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `aplt.subplot_interferometer_dirty_images` contains a subplot which plots all the key properties of the dataset simultaneously.\n", + "\n", + "This includes the observed visibility data, RMS noise map and other information." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Visibility data is in uv space, making it hard to interpret by eye.\n", + "\n", + "The dirty images of the interferometer dataset can plotted, which use the transformer of the interferometer \n", + "to map the visibilities, noise-map or other quantity to a real-space image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Fitting__\n", + "\n", + "Following the previous overview example, we can make a tracer from a collection of light profiles, mass profiles\n", + "and galaxies.\n", + "\n", + "The combination of light and mass profiels below is the same as those used to generate the simulated \n", + "dataset we loaded above.\n", + "\n", + "It therefore produces a tracer whose image looks exactly like the dataset.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=1.0,\n", + " sersic_index=2.5,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Because the tracer's light and mass profiles are the same used to make the dataset, its image is nearly the same as the\n", + "observed image.\n", + "\n", + "We can plot the image of the tracer to confirm this, noting that for a tracer its images are always in real space\n", + "(not Fourier space like the interferometer dataset) and therefore they can be directly visualized." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=dataset.grid), title=\"Tracer Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "However, the tracer's image is not what we observe in the interferometer dataset, because we observe the image as\n", + "visibilities in the uv-plane. \n", + "\n", + "To compare directly to the data, we therefore need to Fourier transform the tracer's image to the uv-plane. \n", + "\n", + "We do this by creating a `FitInterferometer` object, which performs this Fourier transform as part of the fitting \n", + "procedure.\n", + "\n", + "The code plots the result of this, by using the `model_data` of the fit, which performs this Fourier transform \n", + "on the tracer image above and plots the result visibilities in uv-space." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The visibilities are again hard to interpret by eye, so we can plot the dirty image of the fit's model data. This \n", + "dirty image is the Fourier transform of the fit's model data (therefore the Fourier transform of the tracer's image) and\n", + "can be compared directly to the image of the tracer above (albeit it still has the interferometer's PSF/dirty beam\n", + "convolved with it)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit does a lot more than just Fourier transform the tracer's image it also creates the following:\n", + "\n", + " - The `residual_map`: The `model_data` visibilities subtracted from the observed dataset`s `data` visibilities.\n", + " - The `normalized_residual_map`: The `residual_map `divided by the observed dataset's `noise_map`.\n", + " - The `chi_squared_map`: The `normalized_residual_map` squared.\n", + "\n", + "For a good lens model where the model and tracer are representative of the strong lens system the\n", + "residuals, normalized residuals and chi-squareds are minimized:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "A subplot can be plotted which contains all of the above quantities, as well as other information contained in the\n", + "tracer such as the source-plane image, a zoom in of the source-plane and a normalized residual map where the colorbar\n", + "goes from 1.0 sigma to -1.0 sigma, to highlight regions where the fit is poor.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_interferometer(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Once again, dirty images are often easier to interpret, so we can plot a subplot of the dirty images of the data, model\n", + "data, residuals and chi-squared." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_dirty_images(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit also provides us with a ``log_likelihood``, a single value quantifying how good the tracer fitted the dataset.\n", + "\n", + "Lens modeling, describe in the next overview example, effectively tries to maximize this log likelihood value." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Bad Fit__\n", + "\n", + "A bad lens model will show features in the residual-map and chi-squared map.\n", + "\n", + "We can produce such an image by creating a tracer with different lens and source galaxies. In the example below, we \n", + "change the centre of the source galaxy from (0.0, 0.0) to (0.05, 0.05), which leads to residuals appearing\n", + "in the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.1, 0.1),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.1, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A new fit using this plane shows residuals, normalized residuals and chi-squared which are non-zero. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=fit)\n", + "aplt.subplot_fit_dirty_images(fit=fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We also note that its likelihood decreases." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit Quantities__\n", + "\n", + "The maximum log likelihood fit contains many 1D and 2D arrays showing the fit.\n", + "\n", + "There is a `model_data`, which is the image-plane visibilities of the tracer.\n", + "\n", + "This is the image that is fitted to the data in order to compute the log likelihood and therefore quantify the \n", + "goodness-of-fit.\n", + "\n", + "If you are unclear on what `slim` means, refer to the section `Data Structure` at the top of this example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_data)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There are numerous ndarrays showing the goodness of fit: \n", + "\n", + " - `residual_map`: Residuals = (Data - Model_Data).\n", + " - `normalized_residual_map`: Normalized_Residual = (Data - Model_Data) / Noise\n", + " - `chi_squared_map`: Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.residual_map.slim)\n", + "print(fit.normalized_residual_map.slim)\n", + "print(fit.chi_squared_map.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There are `dirty` variants of the above maps, which transform the visibilities, residual-map, chi squared and other\n", + "values to to real-space images using the interferometer's transformer.\n", + "\n", + "These real space images can be mapped between their `slim` and `native` representations (see the\n", + "`guides/data_structures` example for more information on these terms)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.dirty_image.slim) # Data\n", + "print(fit.dirty_model_image.slim)\n", + "print(fit.dirty_residual_map.slim)\n", + "print(fit.dirty_normalized_residual_map.slim)\n", + "print(fit.dirty_chi_squared_map.slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Figures of Merit__\n", + "\n", + "There are single valued floats which quantify the goodness of fit:\n", + "\n", + " - `chi_squared`: The sum of the `chi_squared_map`.\n", + "\n", + " - `noise_normalization`: The normalizing noise term in the likelihood function \n", + " where [Noise_Term] = sum(log(2*pi*[Noise]**2.0)).\n", + "\n", + " - `log_likelihood`: The log likelihood value of the fit where [LogLikelihood] = -0.5*[Chi_Squared_Term + Noise_Term].\n", + " \n", + "These sum other both the real and imaginary components of the visibilities to give a single value for each quantity." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.chi_squared)\n", + "print(fit.noise_normalization)\n", + "print(fit.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Plane Quantities__\n", + "\n", + "The `FitInterferometer` object has specific quantities which break down each image of each plane:\n", + "\n", + " - `model_visibilities_of_planes_list`: Model-images of each individual plane, which in this example is a model image of the \n", + " lens galaxy and model image of the lensed source galaxy, both corresponding to dirty images.\n", + "\n", + " - `subtracted_images_of_planes_list`: Subtracted images of each individual plane, which are the data's image with\n", + " all other plane's model-images subtracted. For example, the first subtracted image has the source galaxy's model image\n", + " subtracted and therefore is of only the lens galaxy's emission. The second subtracted image is of the lensed source,\n", + " with the lens galaxy's light removed.\n", + "\n", + "For multi-plane lens systems these lists will be extended to provide information on every individual plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.model_visibilities_of_planes_list[1].slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There is also a `galaxy_model_visibilities_dict` which maps each galaxy in the tracer to its model visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.galaxy_model_visibilities_dict[source_galaxy].slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A dictionary which maps the model images of each galaxy is also available.\n", + "\n", + "These are not the dirty images, but instead the images of each galaxy that come from the tracer object\n", + "(e.g. simply evaluating the tracer's image on the interferometer's real-space grid)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.galaxy_image_dict[source_galaxy].slim)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Outputting Results__\n", + "\n", + "You may wish to output certain results to .fits files for later inspection. \n", + "\n", + "For example, one could output the lens light subtracted image of the lensed source galaxy to a .fits file such that\n", + "we could fit this source-only image again with an independent pipeline." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_model_image = fit.galaxy_image_dict[source_galaxy]\n", + "aplt.fits_array(\n", + " array=source_model_image,\n", + " file_path=dataset_path / \"source_model_image.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fin." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/likelihood_function.ipynb b/notebooks/interferometer/likelihood_function.ipynb index 16263d2c8..48d2f5389 100644 --- a/notebooks/interferometer/likelihood_function.ipynb +++ b/notebooks/interferometer/likelihood_function.ipynb @@ -1,850 +1,887 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Log Likelihood Function: Parametric__\n", - "\n", - "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Interferometer` data\n", - "with a lens light profile and source light profile (e.g. an elliptical Sersic lens and source).\n", - "\n", - "This script has the following aims:\n", - "\n", - " - To provide a resource that authors can include in papers, so that readers can understand the likelihood\n", - " function (including references to the previous literature from which it is defined) without having to\n", - " write large quantities of text and equations.\n", - "\n", - "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", - "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", - "packages are called when the likelihood is evaluated.\n", - "\n", - "__Contents__\n", - "\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Masked Image Grid:** To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", - "- **Lens Galaxy Mass:** We next define the mass profiles which represents the lens galaxy's mass, which will be used to.\n", - "- **Lens Galaxy:** We now combine the light and mass profiles into a single `Galaxy` object for the lens galaxy.\n", - "- **Source Galaxy Light Profile:** The source galaxy is fitted using another analytic light profile, in this example another.\n", - "- **Lens Light:** Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the an.\n", - "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the.\n", - "- **Source Image:** We pass the traced grid of coordinates to the source galaxy to evaluate its 2D image.\n", - "- **Fourier Transform:** Fourier Transform the 2D image of the galaxy above using the Non Uniform Fast Fourier Transform.\n", - "- **Likelihood Function:** We now quantify the goodness-of-fit of our galaxy model.\n", - "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", - "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", - "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms.\n", - "- **Fit:** Fit the lens model to the dataset.\n", - "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autoarray as aa\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the \u2018real_space_mask\u2019 which defines the grid the image the galaxy is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 4.0\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(800, 800), pixel_scales=0.05, radius=mask_radius\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the galaxy `Interferometer` dataset `simple` from .fits files, which we will fit \n", - "with the model.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the galaxy to the uv-plane and compare \n", - "directly to the visibilities. We use a non-uniform fast Fourier transform, which is the most efficient method for \n", - "interferometer datasets containing ~1-10 million visibilities. We will discuss how the calculation of the likelihood\n", - "function changes for different methods of Fourier transforming in this guide." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This guide uses in-built visualization tools for plotting. \n", - "\n", - "For example, using the `aplt.subplot_interferometer_dirty_images` the dataset we perform a likelihood evaluation on is plotted.\n", - "\n", - "The `subplot_dataset` displays the visibilities in the uv-plane, which are the raw data of the interferometer\n", - "dataset. These are what will ultimately be directly fitted in the Fourier space.\n", - "\n", - "The `subplot_dirty_images` displays the dirty images of the dataset, which are the reconstructed images of visibilities\n", - "using an inverse Fourier transform to convert these to real-space. These dirty images are not the images we fit, but\n", - "visualization of the dirty images are often used in radio interferometry to show the data in a way that is more\n", - "interpretable to the human eye." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling is used, \n", - "which evaluates light profiles on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", - "calculations are performed without over sampling.\n", - "\n", - "__Masked Image Grid__\n", - "\n", - "To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", - "\n", - "The dataset is defined in real-space, and is Fourier transformed to the uv-plane for the model-fit. The grid is\n", - "therefore paired to the `real_space_mask`.\n", - "\n", - "The coordinates are given by `dataset.grids.lp`, which we can plot and see is a uniform grid of (y,x) Cartesian \n", - "coordinates which have had the 3.0\" circular mask applied.\n", - "\n", - "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", - "used to evaluate a light profile the intensity of the profile at the centre of each image-pixel is computed, making\n", - "it straight forward to compute the light profile's image to the image data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=dataset.grids.lp, title=\"\")\n", - "\n", - "print(f\"(y,x) coordinates of first ten unmasked image-pixels {dataset.grid[0:9]}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To perform lensing calculations we convert this 2D (y,x) grid of coordinates to elliptical coordinates:\n", - "\n", - " $\\eta = \\sqrt{(x - x_c)^2 + (y - y_c)^2/q^2}$\n", - "\n", - "Where:\n", - "\n", - " - $y$ and $x$ are the (y,x) arc-second coordinates of each unmasked image-pixel, given by `dataset.grids.lp`.\n", - " - $y_c$ and $x_c$ are the (y,x) arc-second `centre` of the light or mass profile used to perform lensing calculations.\n", - " - $q$ is the axis-ratio of the elliptical light or mass profile (`axis_ratio=1.0` for spherical profiles).\n", - " - The elliptical coordinates is rotated by position angle $\\phi$, defined counter-clockwise from the positive \n", - " x-axis.\n", - "\n", - "$q$ and $\\phi$ are not used to parameterize a light profile but expresses these as \"elliptical components\", \n", - "or `ell_comps` for short:\n", - "\n", - "$\\epsilon_{1} =\\frac{1-q}{1+q} \\sin 2\\phi, \\,\\,$\n", - "$\\epsilon_{2} =\\frac{1-q}{1+q} \\cos 2\\phi.$\n", - "\n", - "Note that `Ell` is used as shorthand for elliptical and `Sph` for spherical." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "profile = al.EllProfile(centre=(0.1, 0.2), ell_comps=(0.1, 0.2))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Transform `dataset.grids.lp` to the centre of profile and rotate it using its angle `phi`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "transformed_grid = profile.transformed_to_reference_frame_grid_from(\n", - " grid=dataset.grids.lp\n", - ")\n", - "\n", - "aplt.plot_grid(grid=transformed_grid, title=\"\")\n", - "print(\n", - " f\"transformed coordinates of first ten unmasked image-pixels {transformed_grid[0:9]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Using these transformed (y',x') values we compute the elliptical coordinates $\\eta = \\sqrt{(x')^2 + (y')^2/q^2}$" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "elliptical_radii = profile.elliptical_radii_grid_from(grid=transformed_grid)\n", - "\n", - "print(\n", - " f\"elliptical coordinates of first ten unmasked image-pixels {elliptical_radii[0:9]}\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Light Profiles (Setup)__\n", - "\n", - "To perform a likelihood evaluation we now compose our lens model.\n", - "\n", - "We first define the light profiles which represents the lens galaxy's light, which will be used to fit the lens \n", - "light.\n", - "\n", - "A light profile is defined by its intensity $I (\\eta_{\\rm l}) $, for example the Sersic profile:\n", - "\n", - "$I_{\\rm Ser} (\\eta_{\\rm l}) = I \\exp \\bigg\\{ -k \\bigg[ \\bigg( \\frac{\\eta}{R} \\bigg)^{\\frac{1}{n}} - 1 \\bigg] \\bigg\\}$\n", - "\n", - "Where:\n", - "\n", - " - $\\eta$ are the elliptical coordinates (see above) or the masked image-grid.\n", - " - $I$ is the `intensity`, which controls the overall brightness of the Sersic profile.\n", - " - $n$ is the ``sersic_index``, which via $k$ controls the steepness of the inner profile.\n", - " - $R$ is the `effective_radius`, which defines the arc-second radius of a circle containing half the light.\n", - "\n", - "In this example, we assume our lens is composed of one light profile, an elliptical Sersic which represent the \n", - "bulge of the lens. \n", - "\n", - "It is uncommon for a lens galaxy observed with interferometer data to have luminous emission, but we show this example\n", - "to illustrate how the likelihood function works." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=4.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Using the masked 2D grid defined above, we can calculate and plot images of each light profile component in real space.\n", - "\n", - "(The transformation to elliptical coordinates above are built into the `image_2d_from` function and performed \n", - "implicitly)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image_2d_bulge = bulge.image_2d_from(grid=dataset.grid)\n", - "\n", - "aplt.plot_array(array=bulge.image_2d_from(grid=dataset.grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy Mass__\n", - "\n", - "We next define the mass profiles which represents the lens galaxy's mass, which will be used to ray-trace the \n", - "image-plane 2D grid of (y,x) coordinates to the source-plane so that the source model can be evaluated.\n", - "\n", - "In this example, we assume our lens is composed of an elliptical isothermal mass distribution and external shear.\n", - "\n", - "A mass profile is defined by its convergence $\\kappa (\\eta)$, which is related to\n", - "the surface density of the mass distribution as\n", - "\n", - "$\\kappa(\\eta)=\\frac{\\Sigma(\\eta)}{\\Sigma_\\mathrm{crit}},$\n", - "\n", - "where\n", - "\n", - "$\\Sigma_\\mathrm{crit}=\\frac{{\\rm c}^2}{4{\\rm \\pi} {\\rm G}}\\frac{D_{\\rm s}}{D_{\\rm l} D_{\\rm ls}},$\n", - "\n", - "and\n", - "\n", - " - `c` is the speed of light.\n", - " - $D_{\\rm l}$, $D_{\\rm s}$, and $D_{\\rm ls}$ are respectively the angular diameter distances to the lens, to the \n", - " source, and from the lens to the source.\n", - "\n", - "For readers less familiar with lensing, we can think of $\\kappa(\\eta)$ as a convenient and\n", - "dimensionless way to describe how light is gravitationally lensed after assuming a cosmology.\n", - "\n", - "For the for the isothermal profile:\n", - "\n", - "$\\kappa(\\eta) = \\frac{1.0}{1 + q} \\bigg( \\frac{\\theta_{\\rm E}}{\\eta} \\bigg)$\n", - "\n", - "Where:\n", - "\n", - " - $\\theta_{\\rm E}$ is the `einstein_radius` (which is rescaled compared to other einstein radius\n", - " definitions)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", - "\n", - "aplt.plot_array(array=mass.convergence_2d_from(grid=dataset.grid), title=\"Convergence\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From each mass profile we can compute its deflection angles, which describe how due to gravitational lensing\n", - "image-pixels are ray-traced to the source plane.\n", - "\n", - "The deflection angles are computed by integrating $\\kappa$: \n", - "\n", - "$\\vec{{\\alpha}}_{\\rm x,y} (\\vec{x}) = \\frac{1}{\\pi} \\int \\frac{\\vec{x} - \\vec{x'}}{\\left | \\vec{x} - \\vec{x'} \\right |^2} \\kappa(\\vec{x'}) d\\vec{x'} \\, ,$" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "deflections_yx_2d = mass.deflections_yx_2d_from(grid=dataset.grid)\n", - "\n", - "deflections = mass.deflections_yx_2d_from(grid=dataset.grid)\n", - "deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=dataset.grid.mask)\n", - "aplt.plot_array(array=deflections_y, title=\"Deflections Y\")\n", - "deflections = mass.deflections_yx_2d_from(grid=dataset.grid)\n", - "deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=dataset.grid.mask)\n", - "aplt.plot_array(array=deflections_x, title=\"Deflections X\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Galaxy__\n", - "\n", - "We now combine the light and mass profiles into a single `Galaxy` object for the lens galaxy.\n", - "\n", - "When computing quantities for the light and mass profiles from this object, it computes each individual quantity and \n", - "adds them together. \n", - "\n", - "For example, for the `bulge`, when it computes their 2D images it computes each individually and then adds\n", - "them together." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Galaxy Light Profile__\n", - "\n", - "The source galaxy is fitted using another analytic light profile, in this example another elliptical Sersic." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Light__\n", - "\n", - "Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the an MGE \n", - "bulge). \n", - "\n", - "This computes the `lens_image_2d` of each `LightProfile` and adds them together. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_image_2d = lens_galaxy.image_2d_from(grid=dataset.grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the image-plane to its (y,x) \n", - "source-plane coordinate $\\beta$ using the summed deflection angles $\\alpha$ of the mass profiles:\n", - "\n", - " $\\beta = \\theta - \\alpha(\\theta)$\n", - "\n", - "The likelihood function of a source light profile ray-traces two grids from the image-plane to the source-plane:\n", - "\n", - " 1) A 2D grid of (y,x) coordinates aligned with the imaging data's image-pixels.\n", - " \n", - " 2) The 2D blurring grid (used for the lens light above) which accounts for pixels at the edge of the mask whose\n", - " light blurs into the mask.\n", - " \n", - "The function below computes the 2D deflection angles of the tracer's lens galaxies and subtracts them from the \n", - "image-plane 2D (y,x) coordinates $\\theta$ of each grid, thus ray-tracing their coordinates to the source plane to \n", - "compute their $\\beta$ values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", - "traced_grid = tracer.traced_grid_2d_list_from(grid=dataset.grid)[-1]\n", - "\n", - "\n", - "aplt.plot_grid(grid=traced_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Image__\n", - "\n", - "We pass the traced grid of coordinates to the source galaxy to evaluate its 2D image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_image_2d = source_galaxy.image_2d_from(grid=traced_grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens + Source Light Addition__\n", - "\n", - "We add the lens and source galaxy images together, to create an overall image of the strong lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = lens_image_2d + source_image_2d\n", - "\n", - "aplt.plot_array(array=image, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If you are familiar with imaging data, you may have seen that a `blurring_image` of pixels surrounding the mask,\n", - "whose light is convolved into the masked, is also computed at this point.\n", - "\n", - "For interferometer data, this is not necessary as the Fourier transform of the real-space image to the uv-plane \n", - "does not require that the emission from outside the mask is accounted for.\n", - "\n", - "__Fourier Transform__\n", - "\n", - "Fourier Transform the 2D image of the galaxy above using the Non Uniform Fast Fourier Transform (NUFFT)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "visibilities = dataset.transformer.visibilities_from(\n", - " image=image,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The Fourier Transform converts the galaxy image from real-space, which is the observed 2D image of the galaxy we \n", - "see with our eyes, to the uv-plane, where the visibilities are measured.\n", - "\n", - "The visibilities are a grid of 2D values representing the real and imaginary components of the visibilities at each\n", - "uv-plane coordinate.\n", - "\n", - "If you are not familiar with interferometer data and the uv-plane, you will need to read up on interferometry to\n", - "fully understand how this likelihood function works." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Likelihood Function__\n", - "\n", - "We now quantify the goodness-of-fit of our galaxy model.\n", - "\n", - "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", - "\n", - "The likelihood function for parametric galaxy modeling consists of two terms:\n", - "\n", - " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", - "\n", - "We now explain what each of these terms mean.\n", - "\n", - "__Chi Squared__\n", - "\n", - "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", - "\n", - " - `model_data` = `visibilities`\n", - " - `residual_map` = (`data` - `model_data`)\n", - " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", - " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", - " - `chi_squared` = sum(`chi_squared_map`)\n", - "\n", - "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", - "\n", - "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", - "for, corresponding to a fit with a lower likelihood." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_data = visibilities\n", - "\n", - "residual_map = dataset.data - model_data\n", - "normalized_residual_map = residual_map / dataset.noise_map\n", - "chi_squared_map = normalized_residual_map**2.0\n", - "\n", - "chi_squared = np.sum(chi_squared_map)\n", - "\n", - "print(chi_squared)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "chi_squared_map = al.Visibilities(visibilities=chi_squared_map)\n", - "\n", - "aplt.plot_grid(grid=chi_squared_map.in_grid, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Noise Normalization Term__\n", - "\n", - "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", - "\n", - "The final term in the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", - "of the log of every noise-map value squared. \n", - "\n", - "Given the `noise_map` is fixed, this term does not change during the galaxy modeling process and has no impact on the \n", - "model we infer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "noise_normalization = float(np.sum(np.log(2 * np.pi * dataset.noise_map**2.0)))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Calculate The Log Likelihood__\n", - "\n", - "We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms computed above using\n", - "the likelihood function defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", - "\n", - "print(figure_of_merit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "This process to perform a likelihood function evaluation performed via the `FitInterferometer` object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", - "fit_figure_of_merit = fit.figure_of_merit\n", - "print(fit_figure_of_merit)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=fit)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lens Modeling__\n", - "\n", - "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", - "non-linear search algorithm.\n", - "\n", - "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", - "multiple MCMC and optimization algorithms are supported.\n", - "\n", - "__Wrap Up__\n", - "\n", - "We have presented a visual step-by-step guide to the parametric likelihood function, which uses \n", - "analytic light profiles to fit the galaxy light.\n", - "\n", - "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", - "are described in additional notebooks found in the `guides` package:\n", - "\n", - " - `over_sampling`: Oversampling the image grid into a finer grid of sub-pixels, which are all individually\n", - " ray-traced to the source-plane and used to evaluate the light profile more accurately.\n", - "\n", - "__JAX__\n", - "\n", - "The step-by-step interferometer likelihood you've just walked through\n", - "can be JAX-accelerated by wrapping construction in `@jax.jit`. Pattern\n", - "mirrors the `imaging/likelihood_function.py` walkthrough:\n", - "\n", - "```python\n", - "import jax\n", - "import jax.numpy as jnp\n", - "from autolens.jax import register_tracer_classes\n", - "\n", - "register_tracer_classes(tracer) # one-time pytree setup\n", - "\n", - "@jax.jit\n", - "def my_log_likelihood(instance):\n", - " tracer = al.Tracer(galaxies=instance.galaxies)\n", - " fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", - " return fit.log_likelihood\n", - "```\n", - "\n", - "Use `TransformerDFT` (the default) under JAX \u2014 `TransformerNUFFT` is not\n", - "JAX-traceable. To validate the JAX log-likelihood matches the NumPy\n", - "chi-squared you derived above, use `Fitness._vmap(jnp.array([parameters]))`\n", - "(production validation pattern \u2014 single `jax.jit(fn)(concrete)` would\n", - "hide un-threaded `xp` sites).\n", - "\n", - "For the canonical Analysis-driven path (zero JAX code on your side),\n", - "see `start_here.py` / `modeling.py`. For JIT-ing library methods\n", - "directly, see `scripts/guides/lens_calc.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Parametric__\n", + "\n", + "This script provides a step-by-step guide of the `log_likelihood_function` which is used to fit `Interferometer` data\n", + "with a lens light profile and source light profile (e.g. an elliptical Sersic lens and source).\n", + "\n", + "This script has the following aims:\n", + "\n", + " - To provide a resource that authors can include in papers, so that readers can understand the likelihood\n", + " function (including references to the previous literature from which it is defined) without having to\n", + " write large quantities of text and equations.\n", + "\n", + "Accompanying this script is the `contributor_guide.py` which provides URL's to every part of the source-code that\n", + "is illustrated in this guide. This gives contributors a sequential run through of what source-code functions, modules and\n", + "packages are called when the likelihood is evaluated.\n", + "\n", + "__Contents__\n", + "\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Masked Image Grid:** To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", + "- **Lens Galaxy Mass:** We next define the mass profiles which represents the lens galaxy's mass, which will be used to.\n", + "- **Lens Galaxy:** We now combine the light and mass profiles into a single `Galaxy` object for the lens galaxy.\n", + "- **Source Galaxy Light Profile:** The source galaxy is fitted using another analytic light profile, in this example another.\n", + "- **Lens Light:** Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the an.\n", + "- **Ray Tracing:** To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the.\n", + "- **Source Image:** We pass the traced grid of coordinates to the source galaxy to evaluate its 2D image.\n", + "- **Fourier Transform:** Fourier Transform the 2D image of the galaxy above using the Non Uniform Fast Fourier Transform.\n", + "- **Likelihood Function:** We now quantify the goodness-of-fit of our galaxy model.\n", + "- **Chi Squared:** The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is.\n", + "- **Noise Normalization Term:** Our likelihood function assumes the imaging data consists of independent Gaussian noise in every.\n", + "- **Calculate The Log Likelihood:** We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms.\n", + "- **Fit:** Fit the lens model to the dataset.\n", + "- **Lens Modeling:** To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autoarray as aa\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the \u2018real_space_mask\u2019 which defines the grid the image the galaxy is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 4.0\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(800, 800), pixel_scales=0.05, radius=mask_radius\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the galaxy `Interferometer` dataset `simple` from .fits files, which we will fit \n", + "with the model.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the galaxy to the uv-plane and compare \n", + "directly to the visibilities. We use a non-uniform fast Fourier transform, which is the most efficient method for \n", + "interferometer datasets containing ~1-10 million visibilities. We will discuss how the calculation of the likelihood\n", + "function changes for different methods of Fourier transforming in this guide." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This guide uses in-built visualization tools for plotting. \n", + "\n", + "For example, using the `aplt.subplot_interferometer_dirty_images` the dataset we perform a likelihood evaluation on is plotted.\n", + "\n", + "The `subplot_dataset` displays the visibilities in the uv-plane, which are the raw data of the interferometer\n", + "dataset. These are what will ultimately be directly fitted in the Fourier space.\n", + "\n", + "The `subplot_dirty_images` displays the dirty images of the dataset, which are the reconstructed images of visibilities\n", + "using an inverse Fourier transform to convert these to real-space. These dirty images are not the images we fit, but\n", + "visualization of the dirty images are often used in radio interferometry to show the data in a way that is more\n", + "interpretable to the human eye." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling is used, \n", + "which evaluates light profiles on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", + "calculations are performed without over sampling.\n", + "\n", + "__Masked Image Grid__\n", + "\n", + "To perform galaxy calculations we define a 2D image-plane grid of (y,x) coordinates.\n", + "\n", + "The dataset is defined in real-space, and is Fourier transformed to the uv-plane for the model-fit. The grid is\n", + "therefore paired to the `real_space_mask`.\n", + "\n", + "The coordinates are given by `dataset.grids.lp`, which we can plot and see is a uniform grid of (y,x) Cartesian \n", + "coordinates which have had the 3.0\" circular mask applied.\n", + "\n", + "Each (y,x) coordinate coordinates to the centre of each image-pixel in the dataset, meaning that when this grid is\n", + "used to evaluate a light profile the intensity of the profile at the centre of each image-pixel is computed, making\n", + "it straight forward to compute the light profile's image to the image data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=dataset.grids.lp, title=\"\")\n", + "\n", + "print(f\"(y,x) coordinates of first ten unmasked image-pixels {dataset.grid[0:9]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To perform lensing calculations we convert this 2D (y,x) grid of coordinates to elliptical coordinates:\n", + "\n", + " $\\eta = \\sqrt{(x - x_c)^2 + (y - y_c)^2/q^2}$\n", + "\n", + "Where:\n", + "\n", + " - $y$ and $x$ are the (y,x) arc-second coordinates of each unmasked image-pixel, given by `dataset.grids.lp`.\n", + " - $y_c$ and $x_c$ are the (y,x) arc-second `centre` of the light or mass profile used to perform lensing calculations.\n", + " - $q$ is the axis-ratio of the elliptical light or mass profile (`axis_ratio=1.0` for spherical profiles).\n", + " - The elliptical coordinates is rotated by position angle $\\phi$, defined counter-clockwise from the positive \n", + " x-axis.\n", + "\n", + "$q$ and $\\phi$ are not used to parameterize a light profile but expresses these as \"elliptical components\", \n", + "or `ell_comps` for short:\n", + "\n", + "$\\epsilon_{1} =\\frac{1-q}{1+q} \\sin 2\\phi, \\,\\,$\n", + "$\\epsilon_{2} =\\frac{1-q}{1+q} \\cos 2\\phi.$\n", + "\n", + "Note that `Ell` is used as shorthand for elliptical and `Sph` for spherical." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "profile = al.EllProfile(centre=(0.1, 0.2), ell_comps=(0.1, 0.2))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Transform `dataset.grids.lp` to the centre of profile and rotate it using its angle `phi`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "transformed_grid = profile.transformed_to_reference_frame_grid_from(\n", + " grid=dataset.grids.lp\n", + ")\n", + "\n", + "aplt.plot_grid(grid=transformed_grid, title=\"\")\n", + "print(\n", + " f\"transformed coordinates of first ten unmasked image-pixels {transformed_grid[0:9]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Using these transformed (y',x') values we compute the elliptical coordinates $\\eta = \\sqrt{(x')^2 + (y')^2/q^2}$" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "elliptical_radii = profile.elliptical_radii_grid_from(grid=transformed_grid)\n", + "\n", + "print(\n", + " f\"elliptical coordinates of first ten unmasked image-pixels {elliptical_radii[0:9]}\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Light Profiles (Setup)__\n", + "\n", + "To perform a likelihood evaluation we now compose our lens model.\n", + "\n", + "We first define the light profiles which represents the lens galaxy's light, which will be used to fit the lens \n", + "light.\n", + "\n", + "A light profile is defined by its intensity $I (\\eta_{\\rm l}) $, for example the Sersic profile:\n", + "\n", + "$I_{\\rm Ser} (\\eta_{\\rm l}) = I \\exp \\bigg\\{ -k \\bigg[ \\bigg( \\frac{\\eta}{R} \\bigg)^{\\frac{1}{n}} - 1 \\bigg] \\bigg\\}$\n", + "\n", + "Where:\n", + "\n", + " - $\\eta$ are the elliptical coordinates (see above) or the masked image-grid.\n", + " - $I$ is the `intensity`, which controls the overall brightness of the Sersic profile.\n", + " - $n$ is the ``sersic_index``, which via $k$ controls the steepness of the inner profile.\n", + " - $R$ is the `effective_radius`, which defines the arc-second radius of a circle containing half the light.\n", + "\n", + "In this example, we assume our lens is composed of one light profile, an elliptical Sersic which represent the \n", + "bulge of the lens. \n", + "\n", + "It is uncommon for a lens galaxy observed with interferometer data to have luminous emission, but we show this example\n", + "to illustrate how the likelihood function works." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=4.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Using the masked 2D grid defined above, we can calculate and plot images of each light profile component in real space.\n", + "\n", + "(The transformation to elliptical coordinates above are built into the `image_2d_from` function and performed \n", + "implicitly)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image_2d_bulge = bulge.image_2d_from(grid=dataset.grid)\n", + "\n", + "aplt.plot_array(array=bulge.image_2d_from(grid=dataset.grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy Mass__\n", + "\n", + "We next define the mass profiles which represents the lens galaxy's mass, which will be used to ray-trace the \n", + "image-plane 2D grid of (y,x) coordinates to the source-plane so that the source model can be evaluated.\n", + "\n", + "In this example, we assume our lens is composed of an elliptical isothermal mass distribution and external shear.\n", + "\n", + "A mass profile is defined by its convergence $\\kappa (\\eta)$, which is related to\n", + "the surface density of the mass distribution as\n", + "\n", + "$\\kappa(\\eta)=\\frac{\\Sigma(\\eta)}{\\Sigma_\\mathrm{crit}},$\n", + "\n", + "where\n", + "\n", + "$\\Sigma_\\mathrm{crit}=\\frac{{\\rm c}^2}{4{\\rm \\pi} {\\rm G}}\\frac{D_{\\rm s}}{D_{\\rm l} D_{\\rm ls}},$\n", + "\n", + "and\n", + "\n", + " - `c` is the speed of light.\n", + " - $D_{\\rm l}$, $D_{\\rm s}$, and $D_{\\rm ls}$ are respectively the angular diameter distances to the lens, to the \n", + " source, and from the lens to the source.\n", + "\n", + "For readers less familiar with lensing, we can think of $\\kappa(\\eta)$ as a convenient and\n", + "dimensionless way to describe how light is gravitationally lensed after assuming a cosmology.\n", + "\n", + "For the for the isothermal profile:\n", + "\n", + "$\\kappa(\\eta) = \\frac{1.0}{1 + q} \\bigg( \\frac{\\theta_{\\rm E}}{\\eta} \\bigg)$\n", + "\n", + "Where:\n", + "\n", + " - $\\theta_{\\rm E}$ is the `einstein_radius` (which is rescaled compared to other einstein radius\n", + " definitions)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "shear = al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05)\n", + "\n", + "aplt.plot_array(array=mass.convergence_2d_from(grid=dataset.grid), title=\"Convergence\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From each mass profile we can compute its deflection angles, which describe how due to gravitational lensing\n", + "image-pixels are ray-traced to the source plane.\n", + "\n", + "The deflection angles are computed by integrating $\\kappa$: \n", + "\n", + "$\\vec{{\\alpha}}_{\\rm x,y} (\\vec{x}) = \\frac{1}{\\pi} \\int \\frac{\\vec{x} - \\vec{x'}}{\\left | \\vec{x} - \\vec{x'} \\right |^2} \\kappa(\\vec{x'}) d\\vec{x'} \\, ,$" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deflections_yx_2d = mass.deflections_yx_2d_from(grid=dataset.grid)\n", + "\n", + "deflections = mass.deflections_yx_2d_from(grid=dataset.grid)\n", + "deflections_y = aa.Array2D(values=deflections.slim[:, 0], mask=dataset.grid.mask)\n", + "aplt.plot_array(array=deflections_y, title=\"Deflections Y\")\n", + "deflections = mass.deflections_yx_2d_from(grid=dataset.grid)\n", + "deflections_x = aa.Array2D(values=deflections.slim[:, 1], mask=dataset.grid.mask)\n", + "aplt.plot_array(array=deflections_x, title=\"Deflections X\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "We now combine the light and mass profiles into a single `Galaxy` object for the lens galaxy.\n", + "\n", + "When computing quantities for the light and mass profiles from this object, it computes each individual quantity and \n", + "adds them together. \n", + "\n", + "For example, for the `bulge`, when it computes their 2D images it computes each individually and then adds\n", + "them together." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(redshift=0.5, mass=mass, shear=shear)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Galaxy Light Profile__\n", + "\n", + "The source galaxy is fitted using another analytic light profile, in this example another elliptical Sersic." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Light__\n", + "\n", + "Compute a 2D image of the lens galaxy's light as the sum of its individual light profiles (the an MGE \n", + "bulge). \n", + "\n", + "This computes the `lens_image_2d` of each `LightProfile` and adds them together. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_image_2d = lens_galaxy.image_2d_from(grid=dataset.grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "To perform lensing calculations we ray-trace every 2d (y,x) coordinate $\\theta$ from the image-plane to its (y,x) \n", + "source-plane coordinate $\\beta$ using the summed deflection angles $\\alpha$ of the mass profiles:\n", + "\n", + " $\\beta = \\theta - \\alpha(\\theta)$\n", + "\n", + "The likelihood function of a source light profile ray-traces two grids from the image-plane to the source-plane:\n", + "\n", + " 1) A 2D grid of (y,x) coordinates aligned with the imaging data's image-pixels.\n", + " \n", + " 2) The 2D blurring grid (used for the lens light above) which accounts for pixels at the edge of the mask whose\n", + " light blurs into the mask.\n", + " \n", + "The function below computes the 2D deflection angles of the tracer's lens galaxies and subtracts them from the \n", + "image-plane 2D (y,x) coordinates $\\theta$ of each grid, thus ray-tracing their coordinates to the source plane to \n", + "compute their $\\beta$ values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "# A list of every grid (e.g. image-plane, source-plane) however we only need the source plane grid with index -1.\n", + "traced_grid = tracer.traced_grid_2d_list_from(grid=dataset.grid)[-1]\n", + "\n", + "\n", + "aplt.plot_grid(grid=traced_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Image__\n", + "\n", + "We pass the traced grid of coordinates to the source galaxy to evaluate its 2D image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_image_2d = source_galaxy.image_2d_from(grid=traced_grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens + Source Light Addition__\n", + "\n", + "We add the lens and source galaxy images together, to create an overall image of the strong lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = lens_image_2d + source_image_2d\n", + "\n", + "aplt.plot_array(array=image, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If you are familiar with imaging data, you may have seen that a `blurring_image` of pixels surrounding the mask,\n", + "whose light is convolved into the masked, is also computed at this point.\n", + "\n", + "For interferometer data, this is not necessary as the Fourier transform of the real-space image to the uv-plane \n", + "does not require that the emission from outside the mask is accounted for.\n", + "\n", + "__Fourier Transform__\n", + "\n", + "Fourier Transform the 2D image of the galaxy above using the Non Uniform Fast Fourier Transform (NUFFT)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "visibilities = dataset.transformer.visibilities_from(\n", + " image=image,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The Fourier Transform converts the galaxy image from real-space, which is the observed 2D image of the galaxy we \n", + "see with our eyes, to the uv-plane, where the visibilities are measured.\n", + "\n", + "The visibilities are a grid of 2D values representing the real and imaginary components of the visibilities at each\n", + "uv-plane coordinate.\n", + "\n", + "If you are not familiar with interferometer data and the uv-plane, you will need to read up on interferometry to\n", + "fully understand how this likelihood function works." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_grid(grid=visibilities.in_grid, title=\"\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Likelihood Function__\n", + "\n", + "We now quantify the goodness-of-fit of our galaxy model.\n", + "\n", + "We compute the `log_likelihood` of the fit, which is the value returned by the `log_likelihood_function`.\n", + "\n", + "The likelihood function for parametric galaxy modeling consists of two terms:\n", + "\n", + " $-2 \\mathrm{ln} \\, \\epsilon = \\chi^2 + \\sum_{\\rm j=1}^{J} { \\mathrm{ln}} \\left [2 \\pi (\\sigma_j)^2 \\right] \\, .$\n", + "\n", + "We now explain what each of these terms mean.\n", + "\n", + "__Chi Squared__\n", + "\n", + "The first term is a $\\chi^2$ statistic, which is defined above in our merit function as and is computed as follows:\n", + "\n", + " - `model_data` = `visibilities`\n", + " - `residual_map` = (`data` - `model_data`)\n", + " - `normalized_residual_map` = (`data` - `model_data`) / `noise_map`\n", + " - `chi_squared_map` = (`normalized_residuals`) ** 2.0 = ((`data` - `model_data`)**2.0)/(`variances`)\n", + " - `chi_squared` = sum(`chi_squared_map`)\n", + "\n", + "The chi-squared therefore quantifies if our fit to the data is accurate or not. \n", + "\n", + "High values of chi-squared indicate that there are many image pixels our model did not produce a good fit to the image \n", + "for, corresponding to a fit with a lower likelihood." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_data = visibilities\n", + "\n", + "residual_map = dataset.data - model_data\n", + "normalized_residual_map = residual_map / dataset.noise_map\n", + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = np.sum(chi_squared_map)\n", + "\n", + "print(chi_squared)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `chi_squared_map` indicates which regions of the image we did and did not fit accurately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_map = al.Visibilities(visibilities=chi_squared_map)\n", + "\n", + "aplt.plot_grid(grid=chi_squared_map.in_grid, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "Our likelihood function assumes the imaging data consists of independent Gaussian noise in every image pixel.\n", + "\n", + "The final term in the likelihood function is therefore a `noise_normalization` term, which consists of the sum\n", + "of the log of every noise-map value squared. \n", + "\n", + "Given the `noise_map` is fixed, this term does not change during the galaxy modeling process and has no impact on the \n", + "model we infer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(np.sum(np.log(2 * np.pi * dataset.noise_map**2.0)))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "We can now, finally, compute the `log_likelihood` of the galaxy model, by combining the two terms computed above using\n", + "the likelihood function defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "figure_of_merit = float(-0.5 * (chi_squared + noise_normalization))\n", + "\n", + "print(figure_of_merit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "This process to perform a likelihood function evaluation performed via the `FitInterferometer` object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", + "fit_figure_of_merit = fit.figure_of_merit\n", + "print(fit_figure_of_merit)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=fit)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Modeling__\n", + "\n", + "To fit a lens model to data, the likelihood function illustrated in this tutorial is sampled using a\n", + "non-linear search algorithm.\n", + "\n", + "The default sampler is the nested sampling algorithm `nautilus` (https://github.com/johannesulf/nautilus)\n", + "multiple MCMC and optimization algorithms are supported.\n", + "\n", + "__Wrap Up__\n", + "\n", + "We have presented a visual step-by-step guide to the parametric likelihood function, which uses \n", + "analytic light profiles to fit the galaxy light.\n", + "\n", + "There are a number of other inputs features which slightly change the behaviour of this likelihood function, which\n", + "are described in additional notebooks found in the `guides` package:\n", + "\n", + " - `over_sampling`: Oversampling the image grid into a finer grid of sub-pixels, which are all individually\n", + " ray-traced to the source-plane and used to evaluate the light profile more accurately.\n", + "\n", + "__JAX__\n", + "\n", + "The step-by-step interferometer likelihood you've just walked through\n", + "can be JAX-accelerated by wrapping construction in `@jax.jit`. Pattern\n", + "mirrors the `imaging/likelihood_function.py` walkthrough:\n", + "\n", + "```python\n", + "import jax\n", + "import jax.numpy as jnp\n", + "from autolens.jax import register_tracer_classes\n", + "\n", + "register_tracer_classes(tracer) # one-time pytree setup\n", + "\n", + "@jax.jit\n", + "def my_log_likelihood(instance):\n", + " tracer = al.Tracer(galaxies=instance.galaxies)\n", + " fit = al.FitInterferometer(dataset=dataset, tracer=tracer)\n", + " return fit.log_likelihood\n", + "```\n", + "\n", + "Use `TransformerDFT` (the default) under JAX \u2014 `TransformerNUFFT` is not\n", + "JAX-traceable. To validate the JAX log-likelihood matches the NumPy\n", + "chi-squared you derived above, use `Fitness._vmap(jnp.array([parameters]))`\n", + "(production validation pattern \u2014 single `jax.jit(fn)(concrete)` would\n", + "hide un-threaded `xp` sites).\n", + "\n", + "For the canonical Analysis-driven path (zero JAX code on your side),\n", + "see `start_here.py` / `modeling.py`. For JIT-ing library methods\n", + "directly, see `scripts/guides/lens_calc.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/modeling.ipynb b/notebooks/interferometer/modeling.ipynb index 6854dc31b..a69cf2ab3 100644 --- a/notebooks/interferometer/modeling.ipynb +++ b/notebooks/interferometer/modeling.ipynb @@ -1,772 +1,809 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Start Here\n", - "====================\n", - "\n", - "This script is the starting point for lens modeling of interferometer datasets (e.g. SMA, ALMA) and it\n", - "provides an overview of the lens modeling API. The same workflow scales from a few hundred visibilities\n", - "to many millions, thanks to the JAX-native `TransformerNUFFT` (backed by `nufftax`).\n", - "\n", - "__Contents__\n", - "\n", - "- **Number of Visibilities:** This example fits a low-resolution dataset, but the same workflow scales to many millions of visibilities.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Coordinates:** Coordinate system assumptions for the model-fit.\n", - "- **Improved Lens Model:** The previous model used S\u00e9rsic light profiles for the source galaxy.\n", - "- **Linear Light Profiles:** The MGE model below uses a **linear light profile** for the bulge via the ``lp_linear`` API.\n", - "- **Concise API:** The MGE model composition API is quite long and technical, so we simply load the MGE models for the.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Unique Identifier:** In the path above, the `unique_identifier` appears as a collection of characters, where this.\n", - "- **Iterations Per Update:** Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and.\n", - "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", - "- **VRAM Use:** When running AutoLens with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM.\n", - "- **Run Times:** Profiling the expected run time of the model-fit.\n", - "- **Output Folder Layout:** Description of the structure of the `output` folder where results are written.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Features:** This script gives a concise overview of the basic modeling API, fitting one the simplest lens.\n", - "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", - "- **HowToLens:** This `start_here.py` script, and the features examples above, do not explain many details of how.\n", - "- **Modeling Customization:** The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative.\n", - "\n", - "__Number of Visibilities__\n", - "\n", - "This example fits a **low-resolution interferometric dataset** with a small number of visibilities (273). The\n", - "dataset is intentionally minimal so the example runs quickly and you can become familiar with the API and\n", - "modeling workflow.\n", - "\n", - "The same workflow \u2014 light profiles + `TransformerNUFFT` (backed by `nufftax`, https://github.com/GragasLab/nufftax) \u2014\n", - "scales to high-resolution datasets with **millions to hundreds of millions of visibilities** (e.g. ALMA), with no\n", - "change beyond the transformer choice. The NUFFT runs inside JAX's jit/vmap pipeline, so both run time and VRAM\n", - "stay manageable on a GPU at any visibility count.\n", - "\n", - "Pixelized source reconstructions (see `features/pixelization`) remain the right tool when the source has\n", - "complex, irregular morphology that simple light profiles cannot capture. They are no longer required purely\n", - "because the dataset is large.\n", - "\n", - "__Model__\n", - "\n", - "This script fits `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a Multi Gaussian Expansion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "import numpy as np" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit\n", - "with the lens model.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and\n", - "compare directly to the visibilities. We use `TransformerNUFFT`, a JAX-native Non-Uniform Fast Fourier Transform\n", - "backed by `nufftax`. This is the recommended choice at any visibility count and scales efficiently to ALMA-class\n", - "datasets with tens of millions of visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Over Sampling__\n", - "\n", - "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling is used, \n", - "which evaluates light profiles on a higher resolution grid than the image data to ensure the calculation is accurate.\n", - "\n", - "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", - "calculations are performed without over sampling.\n", - "\n", - "__Model__\n", - "\n", - "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In this \n", - "example our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", - " - An `Sersic` `LightProfile` for the source galaxy's light [7 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14.\n", - "\n", - "__Coordinates__\n", - "\n", - "The model fitting default settings assume that the lens galaxy centre is near the coordinates (0.0\", 0.0\"). \n", - "\n", - "If for your dataset the lens is not centred at (0.0\", 0.0\"), we recommend that you either: \n", - "\n", - " - Reduce your data so that the centre is (`autolens_workspace/*/data_preparation`). \n", - " - Manually override the lens model priors (`autolens_workspace/*/guides/modeling/customize`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = af.Model(al.lp.SersicCore)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "[The `info` below may not display optimally on your computer screen, for example the whitespace between parameter\n", - "names on the left and parameter priors on the right may lead them to appear across multiple lines. This is a\n", - "common issue in Jupyter notebooks.\n", - "\n", - "The`info_whitespace_length` parameter in the file `config/general.yaml` in the [output] section can be changed to \n", - "increase or decrease the amount of whitespace (The Jupyter notebook kernel will need to be reset for this change to \n", - "appear in a notebook).]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Improved Lens Model__\n", - "\n", - "The previous model used S\u00e9rsic light profiles for the source galaxy. This makes the model API concise, readable, and \n", - "easy to follow.\n", - "\n", - "However, single S\u00e9rsic profiles perform poorly for most strong lenses. Symmetric profiles (e.g. elliptical S\u00e9rsics) \n", - "typically leave significant residuals because they cannot capture the irregular and asymmetric morphology of real \n", - "galaxies (e.g. isophotal twists, radially varying ellipticity).\n", - "\n", - "This example therefore uses a lens model that combines two features, described in detail elsewhere (but a brief \n", - "overview is provided below):\n", - "\n", - "- **Linear light profiles** (see ``autolens_workspace/*/imaging/features/linear_light_profiles``)\n", - "- **Multi-Gaussian Expansion (MGE) light profiles** (see ``autolens_workspace/*/imaging/features/multi_gaussian_expansion``)\n", - "\n", - "NOTE: These descriptions are in the `imaging` package as most interferometer users will quickly move on to\n", - "pixelized source reconstructions, which do not use these features. Their use here is therefore mostly to\n", - "given an introduction to lens modeling with interferometer data.\n", - "\n", - "These features avoid wasted effort trying to fit S\u00e9rsic profiles to complex data, which is likely to fail unless the \n", - "lens is extremely simple. This does mean the model composition is more complex and as a user its a steeper learning\n", - "curve to understand the API, but its worth it for the improved accuracy and speed of lens modeling.\n", - "\n", - "__Multi-Gaussian Expansion (MGE)__\n", - "\n", - "A Multi-Gaussian Expansion (MGE) decomposes the source light into ~50\u2013100 Gaussians with varying ellipticities \n", - "and sizes. An MGE captures irregular features far more effectively than S\u00e9rsic profiles, leading to more accurate lens m\n", - "odels.\n", - "\n", - "Remarkably, modeling with MGEs is also significantly faster than using S\u00e9rsics: they remain efficient in JAX (on CPU \n", - "or GPU), require fewer non-linear parameters despite their flexibility, and yield simpler parameter spaces that\n", - "sample in far fewer iterations. \n", - "\n", - "__Linear Light Profiles__\n", - "\n", - "The MGE model below uses a **linear light profile** for the bulge via the ``lp_linear`` API, instead of the \n", - "standard ``lp`` light profiles used above.\n", - "\n", - "A linear light profile solves for the *intensity* of each component via a linear inversion, rather than treating it as \n", - "a free parameter. This reduces the dimensionality of the non-linear parameter space: a model with ~80 Gaussians\n", - "does not introduce ~80 additional free parameters.\n", - "\n", - "Linear light profiles therefore improve speed and accuracy, and they are used by default in all modeling example.\n", - "\n", - "__Concise API__\n", - "\n", - "The MGE model composition API is quite long and technical, so we simply load the MGE models for the lens and source \n", - "below via a utility function `mge_model_from` which hides the API to make the code in this introduction example ready \n", - "to read. We then use the PyAutoLens Model API to compose the over lens model.\n", - "\n", - "The full MGE composition API is given in the `features/multi_gaussian_expansion` package." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing the model info confirms the model has Gaussians for both the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The lens model is fitted to the data using a non-linear search. \n", - "\n", - "All examples in the autolens workspace use the nested sampling algorithm \n", - "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has revealed gives the most \n", - "accurate and efficient modeling results.\n", - "\n", - "Nautilus has one main setting that trades-off accuracy and computational run-time, the number of `live_points`. \n", - "A higher number of live points gives a more accurate result, but increases the run-time. A lower value give \n", - "less reliable lens modeling (e.g. the fit may infer a local maxima), but is faster. \n", - "\n", - "The suitable value depends on the model complexity whereby models with more parameters require more live points. \n", - "The default value of 200 is sufficient for the vast majority of common lens models. Lower values often given reliable\n", - "results though, and speed up the run-times. In this example, given the model is quite simple (N=11 parameters), we \n", - "reduce the number of live points to 75 to speed up the run-time.\n", - "\n", - "__Unique Identifier__\n", - "\n", - "In the path above, the `unique_identifier` appears as a collection of characters, where this identifier is generated \n", - "based on the model, search and dataset that are used in the fit.\n", - " \n", - "An identical combination of model and search generates the same identifier, meaning that rerunning the script will use \n", - "the existing results to resume the model-fit. In contrast, if you change the model or search, a new unique identifier \n", - "will be generated, ensuring that the model-fit results are output into a separate folder.\n", - "\n", - "We additionally want the unique identifier to be specific to the dataset fitted, so that if we fit different datasets\n", - "with the same model and search results are output to a different folder. We achieve this below by passing \n", - "the `dataset_name` to the search's `unique_tag`.\n", - "\n", - "__Iterations Per Update__\n", - "\n", - "Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and its best fit\n", - "image to the Jupyter Notebook display and to hard-disk.\n", - "\n", - "This process takes around ~10 seconds, so we don't want it to happen too often so as to slow down the overall\n", - "fit, but we also want it to happen frequently enough that we can track the progress.\n", - "\n", - "The value of 10000 below means this output happens every few minutes on GPU and every ~10 minutes on CPU, a good balance.\n", - "\n", - "__Live Visual Update__\n", - "\n", - "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", - "live display surface:\n", - "\n", - "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", - " watch the fit converge without leaving your terminal.\n", - "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", - " refreshes in place every `iterations_per_quick_update`.\n", - "\n", - "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", - "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"interferometer\"), # The path where results and output are stored.\n", - " name=\"modeling\", # The name of the fit and folder results are output to.\n", - " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", - " n_live=75, # The number of Nautilus \"live\" points, increase for more complex models.\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=10000, # Every N iterations the max likelihood model is visualized in the Jupter Notebook and output to hard-disk.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We next create an `AnalysisInterferometer` object, which can be given many inputs customizing how the lens model is \n", - "fitted to the data (in this example they are omitted for simplicity).\n", - "\n", - "Internally, this object defines the `log_likelihood_function` used by the non-linear search to fit the model to \n", - "the `Interferometer` dataset. \n", - "\n", - "It is not vital that you as a user understand the details of how the `log_likelihood_function` fits a lens model to \n", - "data, but interested readers can find a step-by-step guide of the likelihood \n", - "function at ``autolens_workspace/*/interferometer/log_likelihood_function`\n", - "\n", - "__JAX__\n", - "\n", - "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU\n", - "support, your fits will run much faster (around 10 minutes instead of an hour). If only a CPU is available,\n", - "JAX will still provide a speed up via multithreading, with fits taking around 20-30 minutes.\n", - "\n", - "If you don\u2019t have a GPU locally, consider Google Colab which provides free GPUs, so your modeling runs are much faster." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisInterferometer(\n", - " dataset=dataset,\n", - " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM Use__\n", - "\n", - "When running AutoLens with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM. If insufficient \n", - "VRAM is available, the analysis will fail with an out-of-memory error, typically during JIT compilation or the \n", - "first likelihood call.\n", - "\n", - "Two factors dictate the VRAM usage of an analysis:\n", - "\n", - "- The number of arrays and other data structures JAX must store in VRAM to fit the model\n", - " to the data in the likelihood function. This is dictated by the model complexity and dataset size.\n", - "\n", - "- The `batch_size` sets how many likelihood evaluations are performed simultaneously.\n", - " Increasing the batch size increases VRAM usage but can reduce overall run time,\n", - " while decreasing it lowers VRAM usage at the cost of slower execution.\n", - "\n", - "Before running an analysis, users should check that the estimated VRAM usage for the\n", - "chosen batch size is comfortably below their GPU\u2019s total VRAM.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits.\n", - "\n", - "With `TransformerNUFFT` (backed by `nufftax`), the dominant contributor to VRAM is usually the real-space image\n", - "and its transforms inside the likelihood function, rather than the visibility count itself. VRAM does not scale\n", - "with batch size for the persistent buffers, so if the analysis fits within VRAM for `batch_size=1` you should be\n", - "able to push the batch size up (e.g. to 50) to maximise GPU throughput without running out of memory.\n", - "\n", - "For an MGE model with the small dataset fitted in this example, VRAM use is modest (~0.3 GB). Larger real-space\n", - "masks (finer pixel scales) and higher visibility counts increase VRAM gradually rather than catastrophically, and\n", - "a single GPU comfortably handles millions of visibilities with this approach.\n", - "\n", - "Pixelized source reconstructions (see `features/pixelization`) take a different VRAM trade-off: they keep VRAM\n", - "use low by exploiting sparsity in the linear inversion, which makes them attractive when the real-space mask is\n", - "very large or the source morphology requires it. They are no longer required purely because the dataset has many\n", - "visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "Lens modeling can be a computationally expensive process. When fitting complex models to high resolution datasets \n", - "run times can be of order hours, days, weeks or even months.\n", - "\n", - "Run times are dictated by two factors:\n", - "\n", - " - The log likelihood evaluation time: the time it takes for a single `instance` of the lens model to be fitted to \n", - " the dataset such that a log likelihood is returned.\n", - " \n", - " - The number of iterations (e.g. log likelihood evaluations) performed by the non-linear search: more complex lens\n", - " models require more iterations to converge to a solution.\n", - " \n", - "For this analysis, the log likelihood evaluation time is < 0.001 seconds on GPU, < 0.01 seconds on CPU, which is \n", - "extremely fast for lens modeling. \n", - "\n", - "To estimate the expected overall run time of the model-fit we multiply the log likelihood evaluation time by an \n", - "estimate of the number of iterations the non-linear search will perform, which is around 10000 to 30000 for this model.\n", - "\n", - "GPU run times are around 10 minutes, CPU run times are around 30 minutes.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We can now begin the model-fit by passing the model and analysis object to the search, which performs the \n", - "Nautilus non-linear search in order to find which models fit the data with the highest likelihood.\n", - "\n", - "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce \n", - "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output Folder Layout__\n", - "\n", - "Now the fit is running you should checkout the `autolens_workspace/output` folder. This is where results are\n", - "written to hard-disk in human-readable formats \u2014 `.json`, `.csv`, `.fits`, `.png` and plain text.\n", - "\n", - "As the fit progresses, results are written on the fly using the highest likelihood model found by the\n", - "non-linear search so far. This means you can inspect the model-fit as it runs, without waiting for the\n", - "non-linear search to terminate.\n", - "\n", - "Each completed fit lives at a path like::\n", - "\n", - " output/interferometer//modeling//\n", - " files/ <- JSON + CSV: loadable Python objects\n", - " tracer.json <- max log likelihood Tracer\n", - " model.json <- fitted af.Collection model\n", - " samples.csv <- full Nautilus samples\n", - " samples_summary.json <- max log likelihood parameter values + errors\n", - " samples_info.json <- metadata about the samples\n", - " search.json <- non-linear search configuration\n", - " settings.json <- search settings\n", - " cosmology.json <- cosmology used for the fit\n", - " covariance.csv <- parameter covariance matrix\n", - " image/ <- FITS + PNG: visibility + image-plane products\n", - " dataset.fits <- visibilities, noise-map and uv-coverage\n", - " fit.fits <- model visibilities, residuals, chi-squared\n", - " dirty_images.fits <- dirty images of data, model and residuals\n", - " tracer.fits <- tracer image-plane images per galaxy\n", - " source_plane_images.fits <- source plane reconstructions\n", - " model_galaxy_images.fits <- per-galaxy model images\n", - " galaxy_images.fits <- per-galaxy images\n", - " dataset.png, fit.png, tracer.png <- visualisations\n", - " model.info <- human-readable model summary\n", - " model.results <- human-readable fit summary\n", - " search.summary <- search run summary\n", - " search_internal/ <- internal files used to resume / visualise the search\n", - " metadata <- run metadata\n", - "\n", - "The `` is a 32-character identifier derived from the model, search and dataset, so re-running the\n", - "same configuration resumes from the existing fit automatically.\n", - "\n", - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", - "\n", - "[Above, we discussed that the `info_whitespace_length` parameter in the config files could b changed to make \n", - "the `model.info` attribute display optimally on your computer. This attribute also controls the whitespace of the\n", - "`result.info` attribute.]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=result.max_log_likelihood_tracer, grid=real_space_mask.derive_grid.unmasked\n", - ")\n", - "\n", - "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", - "aplt.subplot_fit_dirty_images(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The result contains the full posterior information of our non-linear search, including all parameter samples, \n", - "log likelihood values and tools to compute the errors on the lens model. \n", - "\n", - "There are built in visualization tools for plotting this.\n", - "\n", - "The plot is labeled with short hand parameter names (e.g. `sersic_index` is mapped to the short hand \n", - "parameter `n`). These mappings ate specified in the `config/notation.yaml` file and can be customized by users.\n", - "\n", - "The superscripts of labels correspond to the name each component was given in the model (e.g. for the `Isothermal`\n", - "mass its name `mass` defined when making the `Model` above is used)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Science (Magnification, Flux and More)__\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", - "size of the source.\n", - "\n", - "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", - "including examples using a pixelized source reconstruction. \n", - "\n", - "If you want to study the source galaxy after modeling has reconstructed its unlensed, then check out this example.\n", - "\n", - "__Features__\n", - "\n", - "This script gives a concise overview of the basic modeling API, fitting one the simplest lens models possible.\n", - "\n", - "Lets now consider what features you should read about to improve your lens modeling, especially if you are aiming\n", - "to fit more complex models to your data.\n", - "\n", - "The examples in the `autolens_workspace/*/interferometer/features` package illustrate other lens modeling\n", - "features.\n", - "\n", - "We recommend you checkout the `pixelization` feature next, which lets you reconstruct sources with complex,\n", - "irregular morphology that simple light profiles cannot capture:\n", - "\n", - "- ``pixelization``: The source is reconstructed using an adaptive Rectangular mesh or Delaunay mesh.\n", - "\n", - "The files `autolens_workspace/*/guides/modeling/searches` and `autolens_workspace/*/guides/modeling/customize`\n", - "provide guides on how to customize many other aspects of the model-fit. Check them out to see if anything\n", - "sounds useful, but for most users you can get by without using these forms of customization!\n", - " \n", - "__Data Preparation__\n", - "\n", - "If you are looking to fit your own interferometer data of a strong lens, checkout \n", - "the `autolens_workspace/*/interferometer/data_preparation/start_here.ipynb` script for an overview of how data should be \n", - "prepared before being modeled.\n", - "\n", - "__HowToLens__\n", - "\n", - "This `start_here.py` script, and the features examples above, do not explain many details of how lens modeling is \n", - "performed, for example:\n", - "\n", - " - How does PyAutoLens perform ray-tracing and lensing calculations in order to fit a lens model?\n", - " - How is a lens model fitted to data? What quantifies the goodness of fit (e.g. how is a log likelihood computed?).\n", - " - How does Nautilus find the highest likelihood lens models? What exactly is a \"non-linear search\"?\n", - "\n", - "You do not need to be able to answer these questions in order to fit lens models with PyAutoLens and do science.\n", - "However, having a deeper understanding of how it all works is both interesting and will benefit you as a scientist\n", - "\n", - "This deeper insight is offered by the **HowToLens** Jupyter notebook lectures, which live in their own repository:\n", - "https://github.com/PyAutoLabs/HowToLens.\n", - "\n", - "I recommend that you check them out if you are interested in more details!\n", - "\n", - "__Modeling Customization__\n", - "\n", - "The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative non-linear searches,\n", - "other than Nautilus, that can be used to fit lens models. \n", - "\n", - "They also provide details on how to customize the model-fit, for example the priors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Start Here\n", + "====================\n", + "\n", + "This script is the starting point for lens modeling of interferometer datasets (e.g. SMA, ALMA) and it\n", + "provides an overview of the lens modeling API. The same workflow scales from a few hundred visibilities\n", + "to many millions, thanks to the JAX-native `TransformerNUFFT` (backed by `nufftax`).\n", + "\n", + "__Contents__\n", + "\n", + "- **Number of Visibilities:** This example fits a low-resolution dataset, but the same workflow scales to many millions of visibilities.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Coordinates:** Coordinate system assumptions for the model-fit.\n", + "- **Improved Lens Model:** The previous model used S\u00e9rsic light profiles for the source galaxy.\n", + "- **Linear Light Profiles:** The MGE model below uses a **linear light profile** for the bulge via the ``lp_linear`` API.\n", + "- **Concise API:** The MGE model composition API is quite long and technical, so we simply load the MGE models for the.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Unique Identifier:** In the path above, the `unique_identifier` appears as a collection of characters, where this.\n", + "- **Iterations Per Update:** Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and.\n", + "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", + "- **VRAM Use:** When running AutoLens with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM.\n", + "- **Run Times:** Profiling the expected run time of the model-fit.\n", + "- **Output Folder Layout:** Description of the structure of the `output` folder where results are written.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Features:** This script gives a concise overview of the basic modeling API, fitting one the simplest lens.\n", + "- **Data Preparation:** Data standards required for fitting with PyAutoLens.\n", + "- **HowToLens:** This `start_here.py` script, and the features examples above, do not explain many details of how.\n", + "- **Modeling Customization:** The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative.\n", + "\n", + "__Number of Visibilities__\n", + "\n", + "This example fits a **low-resolution interferometric dataset** with a small number of visibilities (273). The\n", + "dataset is intentionally minimal so the example runs quickly and you can become familiar with the API and\n", + "modeling workflow.\n", + "\n", + "The same workflow \u2014 light profiles + `TransformerNUFFT` (backed by `nufftax`, https://github.com/GragasLab/nufftax) \u2014\n", + "scales to high-resolution datasets with **millions to hundreds of millions of visibilities** (e.g. ALMA), with no\n", + "change beyond the transformer choice. The NUFFT runs inside JAX's jit/vmap pipeline, so both run time and VRAM\n", + "stay manageable on a GPU at any visibility count.\n", + "\n", + "Pixelized source reconstructions (see `features/pixelization`) remain the right tool when the source has\n", + "complex, irregular morphology that simple light profiles cannot capture. They are no longer required purely\n", + "because the dataset is large.\n", + "\n", + "__Model__\n", + "\n", + "This script fits `Interferometer` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a Multi Gaussian Expansion." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "import numpy as np" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple` from .fits files, which we will fit\n", + "with the lens model.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and\n", + "compare directly to the visibilities. We use `TransformerNUFFT`, a JAX-native Non-Uniform Fast Fourier Transform\n", + "backed by `nufftax`. This is the recommended choice at any visibility count and scales efficiently to ALMA-class\n", + "datasets with tens of millions of visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Over Sampling__\n", + "\n", + "If you are familiar with using imaging data, you may have seen that a numerical technique called over sampling is used, \n", + "which evaluates light profiles on a higher resolution grid than the image data to ensure the calculation is accurate.\n", + "\n", + "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", + "calculations are performed without over sampling.\n", + "\n", + "__Model__\n", + "\n", + "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In this \n", + "example our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", + " - An `Sersic` `LightProfile` for the source galaxy's light [7 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14.\n", + "\n", + "__Coordinates__\n", + "\n", + "The model fitting default settings assume that the lens galaxy centre is near the coordinates (0.0\", 0.0\"). \n", + "\n", + "If for your dataset the lens is not centred at (0.0\", 0.0\"), we recommend that you either: \n", + "\n", + " - Reduce your data so that the centre is (`autolens_workspace/*/data_preparation`). \n", + " - Manually override the lens model priors (`autolens_workspace/*/guides/modeling/customize`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "[The `info` below may not display optimally on your computer screen, for example the whitespace between parameter\n", + "names on the left and parameter priors on the right may lead them to appear across multiple lines. This is a\n", + "common issue in Jupyter notebooks.\n", + "\n", + "The`info_whitespace_length` parameter in the file `config/general.yaml` in the [output] section can be changed to \n", + "increase or decrease the amount of whitespace (The Jupyter notebook kernel will need to be reset for this change to \n", + "appear in a notebook).]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Improved Lens Model__\n", + "\n", + "The previous model used S\u00e9rsic light profiles for the source galaxy. This makes the model API concise, readable, and \n", + "easy to follow.\n", + "\n", + "However, single S\u00e9rsic profiles perform poorly for most strong lenses. Symmetric profiles (e.g. elliptical S\u00e9rsics) \n", + "typically leave significant residuals because they cannot capture the irregular and asymmetric morphology of real \n", + "galaxies (e.g. isophotal twists, radially varying ellipticity).\n", + "\n", + "This example therefore uses a lens model that combines two features, described in detail elsewhere (but a brief \n", + "overview is provided below):\n", + "\n", + "- **Linear light profiles** (see ``autolens_workspace/*/imaging/features/linear_light_profiles``)\n", + "- **Multi-Gaussian Expansion (MGE) light profiles** (see ``autolens_workspace/*/imaging/features/multi_gaussian_expansion``)\n", + "\n", + "NOTE: These descriptions are in the `imaging` package as most interferometer users will quickly move on to\n", + "pixelized source reconstructions, which do not use these features. Their use here is therefore mostly to\n", + "given an introduction to lens modeling with interferometer data.\n", + "\n", + "These features avoid wasted effort trying to fit S\u00e9rsic profiles to complex data, which is likely to fail unless the \n", + "lens is extremely simple. This does mean the model composition is more complex and as a user its a steeper learning\n", + "curve to understand the API, but its worth it for the improved accuracy and speed of lens modeling.\n", + "\n", + "__Multi-Gaussian Expansion (MGE)__\n", + "\n", + "A Multi-Gaussian Expansion (MGE) decomposes the source light into ~50\u2013100 Gaussians with varying ellipticities \n", + "and sizes. An MGE captures irregular features far more effectively than S\u00e9rsic profiles, leading to more accurate lens m\n", + "odels.\n", + "\n", + "Remarkably, modeling with MGEs is also significantly faster than using S\u00e9rsics: they remain efficient in JAX (on CPU \n", + "or GPU), require fewer non-linear parameters despite their flexibility, and yield simpler parameter spaces that\n", + "sample in far fewer iterations. \n", + "\n", + "__Linear Light Profiles__\n", + "\n", + "The MGE model below uses a **linear light profile** for the bulge via the ``lp_linear`` API, instead of the \n", + "standard ``lp`` light profiles used above.\n", + "\n", + "A linear light profile solves for the *intensity* of each component via a linear inversion, rather than treating it as \n", + "a free parameter. This reduces the dimensionality of the non-linear parameter space: a model with ~80 Gaussians\n", + "does not introduce ~80 additional free parameters.\n", + "\n", + "Linear light profiles therefore improve speed and accuracy, and they are used by default in all modeling example.\n", + "\n", + "__Concise API__\n", + "\n", + "The MGE model composition API is quite long and technical, so we simply load the MGE models for the lens and source \n", + "below via a utility function `mge_model_from` which hides the API to make the code in this introduction example ready \n", + "to read. We then use the PyAutoLens Model API to compose the over lens model.\n", + "\n", + "The full MGE composition API is given in the `features/multi_gaussian_expansion` package." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=5, centre_prior_is_uniform=False\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Printing the model info confirms the model has Gaussians for both the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The lens model is fitted to the data using a non-linear search. \n", + "\n", + "All examples in the autolens workspace use the nested sampling algorithm \n", + "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has revealed gives the most \n", + "accurate and efficient modeling results.\n", + "\n", + "Nautilus has one main setting that trades-off accuracy and computational run-time, the number of `live_points`. \n", + "A higher number of live points gives a more accurate result, but increases the run-time. A lower value give \n", + "less reliable lens modeling (e.g. the fit may infer a local maxima), but is faster. \n", + "\n", + "The suitable value depends on the model complexity whereby models with more parameters require more live points. \n", + "The default value of 200 is sufficient for the vast majority of common lens models. Lower values often given reliable\n", + "results though, and speed up the run-times. In this example, given the model is quite simple (N=11 parameters), we \n", + "reduce the number of live points to 75 to speed up the run-time.\n", + "\n", + "__Unique Identifier__\n", + "\n", + "In the path above, the `unique_identifier` appears as a collection of characters, where this identifier is generated \n", + "based on the model, search and dataset that are used in the fit.\n", + " \n", + "An identical combination of model and search generates the same identifier, meaning that rerunning the script will use \n", + "the existing results to resume the model-fit. In contrast, if you change the model or search, a new unique identifier \n", + "will be generated, ensuring that the model-fit results are output into a separate folder.\n", + "\n", + "We additionally want the unique identifier to be specific to the dataset fitted, so that if we fit different datasets\n", + "with the same model and search results are output to a different folder. We achieve this below by passing \n", + "the `dataset_name` to the search's `unique_tag`.\n", + "\n", + "__Iterations Per Update__\n", + "\n", + "Every `iterations_per_quick_update`, the non-linear search outputs the maximum likelihood model and its best fit\n", + "image to the Jupyter Notebook display and to hard-disk.\n", + "\n", + "This process takes around ~10 seconds, so we don't want it to happen too often so as to slow down the overall\n", + "fit, but we also want it to happen frequently enough that we can track the progress.\n", + "\n", + "The value of 10000 below means this output happens every few minutes on GPU and every ~10 minutes on CPU, a good balance.\n", + "\n", + "__Live Visual Update__\n", + "\n", + "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", + "live display surface:\n", + "\n", + "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", + " watch the fit converge without leaving your terminal.\n", + "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", + " refreshes in place every `iterations_per_quick_update`.\n", + "\n", + "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", + "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"interferometer\"), # The path where results and output are stored.\n", + " name=\"modeling\", # The name of the fit and folder results are output to.\n", + " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", + " n_live=75, # The number of Nautilus \"live\" points, increase for more complex models.\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=10000, # Every N iterations the max likelihood model is visualized in the Jupter Notebook and output to hard-disk.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We next create an `AnalysisInterferometer` object, which can be given many inputs customizing how the lens model is \n", + "fitted to the data (in this example they are omitted for simplicity).\n", + "\n", + "Internally, this object defines the `log_likelihood_function` used by the non-linear search to fit the model to \n", + "the `Interferometer` dataset. \n", + "\n", + "It is not vital that you as a user understand the details of how the `log_likelihood_function` fits a lens model to \n", + "data, but interested readers can find a step-by-step guide of the likelihood \n", + "function at ``autolens_workspace/*/interferometer/log_likelihood_function`\n", + "\n", + "__JAX__\n", + "\n", + "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU\n", + "support, your fits will run much faster (around 10 minutes instead of an hour). If only a CPU is available,\n", + "JAX will still provide a speed up via multithreading, with fits taking around 20-30 minutes.\n", + "\n", + "If you don\u2019t have a GPU locally, consider Google Colab which provides free GPUs, so your modeling runs are much faster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisInterferometer(\n", + " dataset=dataset,\n", + " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM Use__\n", + "\n", + "When running AutoLens with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM. If insufficient \n", + "VRAM is available, the analysis will fail with an out-of-memory error, typically during JIT compilation or the \n", + "first likelihood call.\n", + "\n", + "Two factors dictate the VRAM usage of an analysis:\n", + "\n", + "- The number of arrays and other data structures JAX must store in VRAM to fit the model\n", + " to the data in the likelihood function. This is dictated by the model complexity and dataset size.\n", + "\n", + "- The `batch_size` sets how many likelihood evaluations are performed simultaneously.\n", + " Increasing the batch size increases VRAM usage but can reduce overall run time,\n", + " while decreasing it lowers VRAM usage at the cost of slower execution.\n", + "\n", + "Before running an analysis, users should check that the estimated VRAM usage for the\n", + "chosen batch size is comfortably below their GPU\u2019s total VRAM.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits.\n", + "\n", + "With `TransformerNUFFT` (backed by `nufftax`), the dominant contributor to VRAM is usually the real-space image\n", + "and its transforms inside the likelihood function, rather than the visibility count itself. VRAM does not scale\n", + "with batch size for the persistent buffers, so if the analysis fits within VRAM for `batch_size=1` you should be\n", + "able to push the batch size up (e.g. to 50) to maximise GPU throughput without running out of memory.\n", + "\n", + "For an MGE model with the small dataset fitted in this example, VRAM use is modest (~0.3 GB). Larger real-space\n", + "masks (finer pixel scales) and higher visibility counts increase VRAM gradually rather than catastrophically, and\n", + "a single GPU comfortably handles millions of visibilities with this approach.\n", + "\n", + "Pixelized source reconstructions (see `features/pixelization`) take a different VRAM trade-off: they keep VRAM\n", + "use low by exploiting sparsity in the linear inversion, which makes them attractive when the real-space mask is\n", + "very large or the source morphology requires it. They are no longer required purely because the dataset has many\n", + "visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "Lens modeling can be a computationally expensive process. When fitting complex models to high resolution datasets \n", + "run times can be of order hours, days, weeks or even months.\n", + "\n", + "Run times are dictated by two factors:\n", + "\n", + " - The log likelihood evaluation time: the time it takes for a single `instance` of the lens model to be fitted to \n", + " the dataset such that a log likelihood is returned.\n", + " \n", + " - The number of iterations (e.g. log likelihood evaluations) performed by the non-linear search: more complex lens\n", + " models require more iterations to converge to a solution.\n", + " \n", + "For this analysis, the log likelihood evaluation time is < 0.001 seconds on GPU, < 0.01 seconds on CPU, which is \n", + "extremely fast for lens modeling. \n", + "\n", + "To estimate the expected overall run time of the model-fit we multiply the log likelihood evaluation time by an \n", + "estimate of the number of iterations the non-linear search will perform, which is around 10000 to 30000 for this model.\n", + "\n", + "GPU run times are around 10 minutes, CPU run times are around 30 minutes.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We can now begin the model-fit by passing the model and analysis object to the search, which performs the \n", + "Nautilus non-linear search in order to find which models fit the data with the highest likelihood.\n", + "\n", + "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce \n", + "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output Folder Layout__\n", + "\n", + "Now the fit is running you should checkout the `autolens_workspace/output` folder. This is where results are\n", + "written to hard-disk in human-readable formats \u2014 `.json`, `.csv`, `.fits`, `.png` and plain text.\n", + "\n", + "As the fit progresses, results are written on the fly using the highest likelihood model found by the\n", + "non-linear search so far. This means you can inspect the model-fit as it runs, without waiting for the\n", + "non-linear search to terminate.\n", + "\n", + "Each completed fit lives at a path like::\n", + "\n", + " output/interferometer//modeling//\n", + " files/ <- JSON + CSV: loadable Python objects\n", + " tracer.json <- max log likelihood Tracer\n", + " model.json <- fitted af.Collection model\n", + " samples.csv <- full Nautilus samples\n", + " samples_summary.json <- max log likelihood parameter values + errors\n", + " samples_info.json <- metadata about the samples\n", + " search.json <- non-linear search configuration\n", + " settings.json <- search settings\n", + " cosmology.json <- cosmology used for the fit\n", + " covariance.csv <- parameter covariance matrix\n", + " image/ <- FITS + PNG: visibility + image-plane products\n", + " dataset.fits <- visibilities, noise-map and uv-coverage\n", + " fit.fits <- model visibilities, residuals, chi-squared\n", + " dirty_images.fits <- dirty images of data, model and residuals\n", + " tracer.fits <- tracer image-plane images per galaxy\n", + " source_plane_images.fits <- source plane reconstructions\n", + " model_galaxy_images.fits <- per-galaxy model images\n", + " galaxy_images.fits <- per-galaxy images\n", + " dataset.png, fit.png, tracer.png <- visualisations\n", + " model.info <- human-readable model summary\n", + " model.results <- human-readable fit summary\n", + " search.summary <- search run summary\n", + " search_internal/ <- internal files used to resume / visualise the search\n", + " metadata <- run metadata\n", + "\n", + "The `` is a 32-character identifier derived from the model, search and dataset, so re-running the\n", + "same configuration resumes from the existing fit automatically.\n", + "\n", + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", + "\n", + "[Above, we discussed that the `info_whitespace_length` parameter in the config files could b changed to make \n", + "the `model.info` attribute display optimally on your computer. This attribute also controls the whitespace of the\n", + "`result.info` attribute.]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=result.max_log_likelihood_tracer, grid=real_space_mask.derive_grid.unmasked\n", + ")\n", + "\n", + "aplt.subplot_fit_interferometer(fit=result.max_log_likelihood_fit)\n", + "aplt.subplot_fit_dirty_images(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The result contains the full posterior information of our non-linear search, including all parameter samples, \n", + "log likelihood values and tools to compute the errors on the lens model. \n", + "\n", + "There are built in visualization tools for plotting this.\n", + "\n", + "The plot is labeled with short hand parameter names (e.g. `sersic_index` is mapped to the short hand \n", + "parameter `n`). These mappings ate specified in the `config/notation.yaml` file and can be customized by users.\n", + "\n", + "The superscripts of labels correspond to the name each component was given in the model (e.g. for the `Isothermal`\n", + "mass its name `mass` defined when making the `Model` above is used)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Science (Magnification, Flux and More)__\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using the reconstructed source model, we can compute key quantities such as the magnification, total flux, and intrinsic \n", + "size of the source.\n", + "\n", + "The example `autolens_workspace/*/guides/source_science` gives a complete overview of how to calculate these quantities,\n", + "including examples using a pixelized source reconstruction. \n", + "\n", + "If you want to study the source galaxy after modeling has reconstructed its unlensed, then check out this example.\n", + "\n", + "__Features__\n", + "\n", + "This script gives a concise overview of the basic modeling API, fitting one the simplest lens models possible.\n", + "\n", + "Lets now consider what features you should read about to improve your lens modeling, especially if you are aiming\n", + "to fit more complex models to your data.\n", + "\n", + "The examples in the `autolens_workspace/*/interferometer/features` package illustrate other lens modeling\n", + "features.\n", + "\n", + "We recommend you checkout the `pixelization` feature next, which lets you reconstruct sources with complex,\n", + "irregular morphology that simple light profiles cannot capture:\n", + "\n", + "- ``pixelization``: The source is reconstructed using an adaptive Rectangular mesh or Delaunay mesh.\n", + "\n", + "The files `autolens_workspace/*/guides/modeling/searches` and `autolens_workspace/*/guides/modeling/customize`\n", + "provide guides on how to customize many other aspects of the model-fit. Check them out to see if anything\n", + "sounds useful, but for most users you can get by without using these forms of customization!\n", + " \n", + "__Data Preparation__\n", + "\n", + "If you are looking to fit your own interferometer data of a strong lens, checkout \n", + "the `autolens_workspace/*/interferometer/data_preparation/start_here.ipynb` script for an overview of how data should be \n", + "prepared before being modeled.\n", + "\n", + "__HowToLens__\n", + "\n", + "This `start_here.py` script, and the features examples above, do not explain many details of how lens modeling is \n", + "performed, for example:\n", + "\n", + " - How does PyAutoLens perform ray-tracing and lensing calculations in order to fit a lens model?\n", + " - How is a lens model fitted to data? What quantifies the goodness of fit (e.g. how is a log likelihood computed?).\n", + " - How does Nautilus find the highest likelihood lens models? What exactly is a \"non-linear search\"?\n", + "\n", + "You do not need to be able to answer these questions in order to fit lens models with PyAutoLens and do science.\n", + "However, having a deeper understanding of how it all works is both interesting and will benefit you as a scientist\n", + "\n", + "This deeper insight is offered by the **HowToLens** Jupyter notebook lectures, which live in their own repository:\n", + "https://github.com/PyAutoLabs/HowToLens.\n", + "\n", + "I recommend that you check them out if you are interested in more details!\n", + "\n", + "__Modeling Customization__\n", + "\n", + "The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative non-linear searches,\n", + "other than Nautilus, that can be used to fit lens models. \n", + "\n", + "They also provide details on how to customize the model-fit, for example the priors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/simulator.ipynb b/notebooks/interferometer/simulator.ipynb index 64c2bc13f..377202b3f 100644 --- a/notebooks/interferometer/simulator.ipynb +++ b/notebooks/interferometer/simulator.ipynb @@ -1,596 +1,633 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: SIE\n", - "==============\n", - "\n", - "This script simulates `Interferometer` data of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an `Sersic`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Grid:** Define the 2d grid of (y,x) coordinates that the galaxy images are evaluated and therefore.\n", - "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", - "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) and source galaxy light (elliptical Sersic) for this.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "- **Multiple Images:** Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x).\n", - "- **Many Visibilities:** Simulating interferometer datasets with many visibilities using the JAX-native `TransformerNUFFT`.\n", - "- **High Resolution Dataset:** A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `dataset_type` describes the type of data being simulated (in this case, `Interferometer` data) and `dataset_name` \n", - "gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", - "\n", - " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/image.fits`.\n", - " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.fits`.\n", - " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/psf.fits`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"interferometer\"\n", - "dataset_name = \"simple\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output.\n", - "\n", - "In this example, this is: `/autolens_workspace/dataset/interferometer/simple`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Grid__\n", - "\n", - "Define the 2d grid of (y,x) coordinates that the galaxy images are evaluated and therefore simulated on, via\n", - "the inputs:\n", - "\n", - " - `shape_native`: The (y_pixels, x_pixels) 2D shape of the grid defining the shape of the data that is simulated.\n", - " - `pixel_scales`: The arc-second to pixel conversion factor of the grid and data.\n", - "\n", - "For interferometer data, this image is evaluate in real-space and then transformed to Fourier space.\n", - "\n", - "__Over Sampling__\n", - "\n", - "If you are familiar with using imaging data, you may have seen that a numerical technique called\n", - "over sampling is used, which evaluates light profiles on a higher resolution grid than the image data to ensure the\n", - "calculation is accurate.\n", - "\n", - "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", - "calculations are performed without over sampling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(256, 256), pixel_scales=0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To perform the Fourier transform we need the wavelengths of the baselines, which we'll load from the fits file below.\n", - "\n", - "By default we use baselines from the Square Mile Array (SMA), which produces low resolution interferometer data that\n", - "can be fitted extremely efficiently. The `autolens_workspace` includes ALMA uv_wavelengths files for simulating\n", - "much high resolution datasets (which can be performed by replacing \"sma.fits\" below with \"alma.fits\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "uv_wavelengths_path = Path(\"dataset\", dataset_type, \"uv_wavelengths\")\n", - "uv_wavelengths = al.ndarray_via_fits_from(\n", - " file_path=Path(uv_wavelengths_path, \"sma.fits\"), hdu=0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the interferometer dataset we first create a simulator, which defines the exposure time, noise levels\n", - "and Fourier transform method used in the simulation.\n", - "\n", - "We use `TransformerNUFFT` (backed by `nufftax`, https://github.com/GragasLab/nufftax), a JAX-native Non-Uniform\n", - "Fast Fourier Transform. This is the recommended transformer at any visibility count and is fast enough to\n", - "simulate ALMA-class datasets with millions of visibilities end-to-end on a GPU." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorInterferometer(\n", - " uv_wavelengths=uv_wavelengths,\n", - " exposure_time=300.0,\n", - " noise_sigma=1000.0,\n", - " transformer_class=al.TransformerDFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's mass (SIE+Shear) and source galaxy light (elliptical Sersic) for this simulated lens.\n", - "\n", - "The following should be noted about the parameters below:\n", - "\n", - " - The native units of light and mass profiles distance parameters (e.g. centres, effective_radius) are arc-seconds. \n", - " - The intensity of the light profiles is in units of electrons per second per arc-second squared.\n", - " - The ellipticity of light and mass profiles are defined using the `ell_comps` parameter, however we below use\n", - " the convert module to input the `axis-ratio` (semi-major axis / semi-minor axis = b/a) and positive \n", - " angle (degrees defined counter clockwise from the positive x-axis).\n", - " - The external shear is defined using the (gamma_1, gamma_2) convention.\n", - " - The input redshifts are used to determine which galaxy is the lens (e.g. lower redshift) and which is the \n", - " source (e.g. higher redshift).\n", - " - The source uses a cored Sersic with a radius half the pixel-scale, ensuring that over-sampling is not necessary." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=10.0,\n", - " effective_radius=1.0,\n", - " sersic_index=2.5,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated interferometer dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now pass this simulator a tracer, which creates the ray-traced image plotted above and simulates it as an\n", - "interferometer dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets plot the simulated interferometer dataset before we output it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_interferometer(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_interferometer_dirty_images(\n", - " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multiple Images__\n", - "\n", - "Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x) \n", - "coordinates of multiple images of a source galaxy to positions which are far apart from one another \n", - "in the source plane are penalized in the lens model's overall likelihood.\n", - "\n", - "This speeds up lens modeling, helps the non-linear search avoid local maxima and is vital for inferred \n", - "accurate solutions when using pixelized source reconstructions.\n", - "\n", - "For real data, the multiple image positions are determined by eye from the data, for example\n", - "using a Graphical User Interface (GUI) to mark them with mouse clicks. For simulated data, we can save\n", - "ourselves time by using the `PointSolver` to determine the multiple image positions automatically and\n", - "output to a .json file.\n", - "\n", - "If you have not looked in the `point_source` package, the point solver is the core tool used to find\n", - "multiple image positions for point source lens modeling (e.g. lensed quasars)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " file_path=dataset_path / \"positions.json\",\n", - " obj=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/simple`.\n", - "\n", - "__Many Visibilities__\n", - "\n", - "Simulating interferometer datasets with many visibilities is fast end-to-end when `TransformerNUFFT` is used,\n", - "which is backed by the JAX-native `nufftax` library (https://github.com/GragasLab/nufftax). The NUFFT scales\n", - "efficiently from a few hundred visibilities up to ALMA-class datasets with tens of millions of visibilities,\n", - "all inside the JAX jit/vmap pipeline.\n", - "\n", - "Higher resolution datasets also require a higher resolution real space grid, which is the dominant remaining\n", - "cost \u2014 visibility count alone is no longer the bottleneck.\n", - "\n", - "The code below loads a `uv_wavelengths` file with over 1 million visibilities and simulates the dataset using\n", - "`TransformerNUFFT`.\n", - "\n", - "__High Resolution Dataset__\n", - "\n", - "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large files which\n", - "are too big to include in the main `autolens_workspace` repository:\n", - "\n", - "https://github.com/PyAutoLabs/autolens_workspace_large_files\n", - "\n", - "After downloading the file, place it in the directory:\n", - "\n", - "`autolens_workspace/dataset/interferometer/alma`\n", - "\n", - "You can then simulate and fit this high-resolution ALMA dataset by uncommenting the \n", - "line `dataset_name = \"alma\"` below.\n", - "\n", - "This dataset is particularly useful for testing performance, memory usage, and accuracy when modeling realistic\n", - "ALMA uv-coverage with a very large number of visibilities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"interferometer\"\n", - "# dataset_name = \"alma\"\n", - "\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)\n", - "\n", - "grid = al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.01)\n", - "\n", - "uv_wavelengths_path = Path(\"dataset\", dataset_type, dataset_name)\n", - "uv_wavelengths = al.ndarray_via_fits_from(\n", - " file_path=Path(uv_wavelengths_path, \"uv_wavelengths.fits\"), hdu=0\n", - ")\n", - "\n", - "simulator = al.SimulatorInterferometer(\n", - " uv_wavelengths=uv_wavelengths,\n", - " exposure_time=300.0,\n", - " noise_sigma=1000.0,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The code below is identical to above, outputting images, data, tracer and multiple image\n", - "positions to the dataset folder." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", - "\n", - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=dataset)\n", - "\n", - "aplt.fits_interferometer(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " overwrite=True,\n", - ")\n", - "\n", - "\n", - "aplt.subplot_interferometer_dirty_images(\n", - " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", - ")\n", - "\n", - "al.output_to_json(\n", - " file_path=dataset_path / \"positions.json\",\n", - " obj=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish.\n", - "\n", - "__JAX Variant__\n", - "\n", - "For fast repeated interferometer simulations, instantiate the simulator\n", - "with `use_jax=True` and wrap the call in `@jax.jit`. The simulator handles\n", - "pytree registration internally.\n", - "\n", - "```python\n", - "import jax\n", - "import jax.numpy as jnp\n", - "\n", - "simulator_jax = al.SimulatorInterferometer(\n", - " uv_wavelengths=uv_wavelengths,\n", - " exposure_time=300.0,\n", - " noise_sigma=0.1,\n", - " transformer_class=al.TransformerDFT, # NUFFT (pynufft) is not JAX-traceable\n", - " use_jax=True,\n", - ")\n", - "\n", - "@jax.jit\n", - "def simulate(tracer):\n", - " image = tracer.image_2d_from(grid=real_space_grid, xp=jnp)\n", - " return simulator_jax.via_image_from(image=image)\n", - "\n", - "dataset_jax = simulate(tracer) # Interferometer with jax.Array visibilities\n", - "```\n", - "\n", - "Two notes specific to interferometer:\n", - "\n", - "- Use `TransformerDFT` (the default) under JAX. `TransformerNUFFT` (pynufft)\n", - " is faster on large UV sets but is not JAX-traceable. The `nufftax`\n", - " research path is tracking a JAX-native NUFFT replacement; see\n", - " `autolens_workspace_test/scripts/interferometer/nufft.py` for the\n", - " parity work.\n", - "- Eager `simulator_jax.via_image_from(image)` already runs on JAX without\n", - " the `@jax.jit` wrap; the JIT only matters for repeated calls.\n", - "\n", - "See `scripts/guides/lens_calc.py` for the \"JIT-it-yourself\" pattern\n", - "applied to individual library methods." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: SIE\n", + "==============\n", + "\n", + "This script simulates `Interferometer` data of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an `Sersic`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Grid:** Define the 2d grid of (y,x) coordinates that the galaxy images are evaluated and therefore.\n", + "- **Over Sampling:** Set up the adaptive over-sampling grid for accurate light profile evaluation.\n", + "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) and source galaxy light (elliptical Sersic) for this.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "- **Multiple Images:** Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x).\n", + "- **Many Visibilities:** Simulating interferometer datasets with many visibilities using the JAX-native `TransformerNUFFT`.\n", + "- **High Resolution Dataset:** A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `dataset_type` describes the type of data being simulated (in this case, `Interferometer` data) and `dataset_name` \n", + "gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", + "\n", + " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/image.fits`.\n", + " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.fits`.\n", + " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/psf.fits`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"interferometer\"\n", + "dataset_name = \"simple\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output.\n", + "\n", + "In this example, this is: `/autolens_workspace/dataset/interferometer/simple`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Grid__\n", + "\n", + "Define the 2d grid of (y,x) coordinates that the galaxy images are evaluated and therefore simulated on, via\n", + "the inputs:\n", + "\n", + " - `shape_native`: The (y_pixels, x_pixels) 2D shape of the grid defining the shape of the data that is simulated.\n", + " - `pixel_scales`: The arc-second to pixel conversion factor of the grid and data.\n", + "\n", + "For interferometer data, this image is evaluate in real-space and then transformed to Fourier space.\n", + "\n", + "__Over Sampling__\n", + "\n", + "If you are familiar with using imaging data, you may have seen that a numerical technique called\n", + "over sampling is used, which evaluates light profiles on a higher resolution grid than the image data to ensure the\n", + "calculation is accurate.\n", + "\n", + "Interferometer does not observe galaxies in a way where over sampling is necessary, therefore all interferometer\n", + "calculations are performed without over sampling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(256, 256), pixel_scales=0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To perform the Fourier transform we need the wavelengths of the baselines, which we'll load from the fits file below.\n", + "\n", + "By default we use baselines from the Square Mile Array (SMA), which produces low resolution interferometer data that\n", + "can be fitted extremely efficiently. The `autolens_workspace` includes ALMA uv_wavelengths files for simulating\n", + "much high resolution datasets (which can be performed by replacing \"sma.fits\" below with \"alma.fits\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "uv_wavelengths_path = Path(\"dataset\", dataset_type, \"uv_wavelengths\")\n", + "uv_wavelengths = al.ndarray_via_fits_from(\n", + " file_path=Path(uv_wavelengths_path, \"sma.fits\"), hdu=0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the interferometer dataset we first create a simulator, which defines the exposure time, noise levels\n", + "and Fourier transform method used in the simulation.\n", + "\n", + "We use `TransformerNUFFT` (backed by `nufftax`, https://github.com/GragasLab/nufftax), a JAX-native Non-Uniform\n", + "Fast Fourier Transform. This is the recommended transformer at any visibility count and is fast enough to\n", + "simulate ALMA-class datasets with millions of visibilities end-to-end on a GPU." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorInterferometer(\n", + " uv_wavelengths=uv_wavelengths,\n", + " exposure_time=300.0,\n", + " noise_sigma=1000.0,\n", + " transformer_class=al.TransformerDFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's mass (SIE+Shear) and source galaxy light (elliptical Sersic) for this simulated lens.\n", + "\n", + "The following should be noted about the parameters below:\n", + "\n", + " - The native units of light and mass profiles distance parameters (e.g. centres, effective_radius) are arc-seconds. \n", + " - The intensity of the light profiles is in units of electrons per second per arc-second squared.\n", + " - The ellipticity of light and mass profiles are defined using the `ell_comps` parameter, however we below use\n", + " the convert module to input the `axis-ratio` (semi-major axis / semi-minor axis = b/a) and positive \n", + " angle (degrees defined counter clockwise from the positive x-axis).\n", + " - The external shear is defined using the (gamma_1, gamma_2) convention.\n", + " - The input redshifts are used to determine which galaxy is the lens (e.g. lower redshift) and which is the \n", + " source (e.g. higher redshift).\n", + " - The source uses a cored Sersic with a radius half the pixel-scale, ensuring that over-sampling is not necessary." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=10.0,\n", + " effective_radius=1.0,\n", + " sersic_index=2.5,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated interferometer dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now pass this simulator a tracer, which creates the ray-traced image plotted above and simulates it as an\n", + "interferometer dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets plot the simulated interferometer dataset before we output it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_interferometer(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_interferometer_dirty_images(\n", + " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multiple Images__\n", + "\n", + "Lens modeling can use a \"positions likelihood penalty\", whereby mass models which traces the (y,x) \n", + "coordinates of multiple images of a source galaxy to positions which are far apart from one another \n", + "in the source plane are penalized in the lens model's overall likelihood.\n", + "\n", + "This speeds up lens modeling, helps the non-linear search avoid local maxima and is vital for inferred \n", + "accurate solutions when using pixelized source reconstructions.\n", + "\n", + "For real data, the multiple image positions are determined by eye from the data, for example\n", + "using a Graphical User Interface (GUI) to mark them with mouse clicks. For simulated data, we can save\n", + "ourselves time by using the `PointSolver` to determine the multiple image positions automatically and\n", + "output to a .json file.\n", + "\n", + "If you have not looked in the `point_source` package, the point solver is the core tool used to find\n", + "multiple image positions for point source lens modeling (e.g. lensed quasars)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " file_path=dataset_path / \"positions.json\",\n", + " obj=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/simple`.\n", + "\n", + "__Many Visibilities__\n", + "\n", + "Simulating interferometer datasets with many visibilities is fast end-to-end when `TransformerNUFFT` is used,\n", + "which is backed by the JAX-native `nufftax` library (https://github.com/GragasLab/nufftax). The NUFFT scales\n", + "efficiently from a few hundred visibilities up to ALMA-class datasets with tens of millions of visibilities,\n", + "all inside the JAX jit/vmap pipeline.\n", + "\n", + "Higher resolution datasets also require a higher resolution real space grid, which is the dominant remaining\n", + "cost \u2014 visibility count alone is no longer the bottleneck.\n", + "\n", + "The code below loads a `uv_wavelengths` file with over 1 million visibilities and simulates the dataset using\n", + "`TransformerNUFFT`.\n", + "\n", + "__High Resolution Dataset__\n", + "\n", + "A high-resolution `uv_wavelengths` file for ALMA is available in a separate repository that hosts large files which\n", + "are too big to include in the main `autolens_workspace` repository:\n", + "\n", + "https://github.com/PyAutoLabs/autolens_workspace_large_files\n", + "\n", + "After downloading the file, place it in the directory:\n", + "\n", + "`autolens_workspace/dataset/interferometer/alma`\n", + "\n", + "You can then simulate and fit this high-resolution ALMA dataset by uncommenting the \n", + "line `dataset_name = \"alma\"` below.\n", + "\n", + "This dataset is particularly useful for testing performance, memory usage, and accuracy when modeling realistic\n", + "ALMA uv-coverage with a very large number of visibilities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"interferometer\"\n", + "# dataset_name = \"alma\"\n", + "\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)\n", + "\n", + "grid = al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.01)\n", + "\n", + "uv_wavelengths_path = Path(\"dataset\", dataset_type, dataset_name)\n", + "uv_wavelengths = al.ndarray_via_fits_from(\n", + " file_path=Path(uv_wavelengths_path, \"uv_wavelengths.fits\"), hdu=0\n", + ")\n", + "\n", + "simulator = al.SimulatorInterferometer(\n", + " uv_wavelengths=uv_wavelengths,\n", + " exposure_time=300.0,\n", + " noise_sigma=1000.0,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The code below is identical to above, outputting images, data, tracer and multiple image\n", + "positions to the dataset folder." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", + "\n", + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=dataset)\n", + "\n", + "aplt.fits_interferometer(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " overwrite=True,\n", + ")\n", + "\n", + "\n", + "aplt.subplot_interferometer_dirty_images(\n", + " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.bulge.centre\n", + ")\n", + "\n", + "al.output_to_json(\n", + " file_path=dataset_path / \"positions.json\",\n", + " obj=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish.\n", + "\n", + "__JAX Variant__\n", + "\n", + "For fast repeated interferometer simulations, instantiate the simulator\n", + "with `use_jax=True` and wrap the call in `@jax.jit`. The simulator handles\n", + "pytree registration internally.\n", + "\n", + "```python\n", + "import jax\n", + "import jax.numpy as jnp\n", + "\n", + "simulator_jax = al.SimulatorInterferometer(\n", + " uv_wavelengths=uv_wavelengths,\n", + " exposure_time=300.0,\n", + " noise_sigma=0.1,\n", + " transformer_class=al.TransformerDFT, # NUFFT (pynufft) is not JAX-traceable\n", + " use_jax=True,\n", + ")\n", + "\n", + "@jax.jit\n", + "def simulate(tracer):\n", + " image = tracer.image_2d_from(grid=real_space_grid, xp=jnp)\n", + " return simulator_jax.via_image_from(image=image)\n", + "\n", + "dataset_jax = simulate(tracer) # Interferometer with jax.Array visibilities\n", + "```\n", + "\n", + "Two notes specific to interferometer:\n", + "\n", + "- Use `TransformerDFT` (the default) under JAX. `TransformerNUFFT` (pynufft)\n", + " is faster on large UV sets but is not JAX-traceable. The `nufftax`\n", + " research path is tracking a JAX-native NUFFT replacement; see\n", + " `autolens_workspace_test/scripts/interferometer/nufft.py` for the\n", + " parity work.\n", + "- Eager `simulator_jax.via_image_from(image)` already runs on JAX without\n", + " the `@jax.jit` wrap; the JIT only matters for repeated calls.\n", + "\n", + "See `scripts/guides/lens_calc.py` for the \"JIT-it-yourself\" pattern\n", + "applied to individual library methods." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/interferometer/source_science.ipynb b/notebooks/interferometer/source_science.ipynb index ef09a4450..c88e55c77 100644 --- a/notebooks/interferometer/source_science.ipynb +++ b/notebooks/interferometer/source_science.ipynb @@ -1,470 +1,507 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Source Science\n", - "==============\n", - "\n", - "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", - "\n", - "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", - "size of the source.\n", - "\n", - "This example shows how to perform these calculations using Sersic parametric sources on imaging data, which\n", - "is conceptually the simplest case for source science calculations and a good introduction to the topic.\n", - "\n", - "__Contents__\n", - "\n", - "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", - "- **Loading Data:** We we begin by loading the strong lens dataset `simple` from .fits files, which is the dataset we.\n", - "- **Source Values:** Source science calculations for real lenses are performed using the best-fitting model inferred.\n", - "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", - "- **Source Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", - "- **Tracer:** Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to.\n", - "- **Parametric Source Models:** If your lens modeling uses a parametric source model (e.g." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.5\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(256, 256),\n", - " pixel_scales=0.1,\n", - " radius=mask_radius,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Loading Data__\n", - "\n", - "We we begin by loading the strong lens dataset `simple` from .fits files, which is the dataset \n", - "we will use to demonstrate fitting.\n", - "\n", - "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and\n", - "compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native Non-Uniform Fast Fourier Transform\n", - "backed by `nufftax`, which scales efficiently from a few hundred visibilities to tens of millions.\n", - "\n", - "This dataset was simulated using the `interferometer/simulator` example, read through that to understand how\n", - "the data this example fits was generated." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerNUFFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `aplt.subplot_interferometer_dirty_images` contains a subplot which plots all the key properties of the dataset simultaneously.\n", - "\n", - "This includes the observed visibility data, RMS noise map and other information." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Visibility data is in uv space, making it hard to interpret by eye.\n", - "\n", - "The dirty images of the interferometer dataset can plotted, which use the transformer of the interferometer \n", - "to map the visibilities, noise-map or other quantity to a real-space image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "__Source Values__\n", - "\n", - "Source science calculations for real lenses are performed using the best-fitting model inferred from a dataset, \n", - "and this example demonstrates how to use this below.\n", - "\n", - "However, we for simplicity, we demonstrate these calculations using the Sersic source model used to simulate the dataset, \n", - "which we refer to as the \"true\" source model. When analysing real strong lenses, a true underlying model is not known, \n", - "but for simulated datasets it is.\n", - "\n", - "This allows us to illustrate the calculations in a way that does not depend on the specific details of the data or \n", - "on assumptions about how the lens model is inferred.\n", - "\n", - "The `tracer` below corresponds to the same tracer used to simulate the `simple` dataset, and therefore \n", - "represents the true source model. We also include the 2D grid of (y,x) coordinates which simulate the dataset.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=0.3,\n", - " effective_radius=1.0,\n", - " sersic_index=2.5,\n", - " ),\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By plotting the image of the tracer, we confirm it looks identical to the simulated dataset but does not have\n", - "CCD imaging features such as noise or blurring from a PSF." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Source Flux__\n", - "\n", - "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", - "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", - "\n", - "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", - "we compute the image of the light profile, and then sum the image. \n", - "\n", - "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", - "we will assume everything is in miliJansky per beam (`mJy beam^-1`), which is typical for ALMA data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(f\"Source Galaxy's Intensity {source_galaxy.bulge.intensity} mJy beam^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The total flux, in units of `mJy beam^-1` , is computed by summing the image of the light profile over all pixels.\n", - "\n", - "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", - "and high enough resolution that it captures all the light from the light profile.\n", - "\n", - "Note that we are using the source galaxy's true light profile, which corresponds to its emission in the source-plane.\n", - "For real datasets, we have to infer this via lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", - "\n", - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_flux = np.sum(image) # in units mJy beam^-1 as summed over pixels\n", - "\n", - "print(f\"Total Source Flux: {total_flux} mJy beam^-1\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Below, we will compare how this true source flux compares to the inferred source fluxes we compute using different\n", - "source modeling techniques (e.g. parametric and pixelized source models). Converting the flux to magnitudes or\n", - "other quantities used for tasks like SED fitting is described in the `autolens_workspace/*/guides/units/flux` example.\n", - "\n", - "__Source Magnification__\n", - "\n", - "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", - "total surface brightness in the source-plane.\n", - "\n", - "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", - "We therefore explicitly mention how area folds into the calculation below.\n", - "\n", - "To ensure the magnification is stable and that we resolve all source emission in both the image-plane and source-plane \n", - "we use a very high resolution grid, higher than we used to compute the total flux above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We repeat our calculation of the source's total flux in the source-plane using this higher resolution grid, note\n", - "that we do not take the area into account, the reason for this is explained below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", - "\n", - "total_source_plane_flux = np.sum(image) # in units mJy beam^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now need the total flux of the lensed source in the image-plane, that is how much flux we measure after\n", - "gravitational lensing.\n", - "\n", - "To calculation this, we first ray-trace the grid above from the image-plane to the source-plane using the tracer\n", - "and then pass it to the source galaxy's light profile to compute the lensed image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "source_plane_grid = traced_grid_list[1]\n", - "\n", - "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(\n", - " lensed_source_image\n", - ") # in units mJy beam^-1 as summed over pixels" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now take the ratio of the total image-plane flux to source-plane flux to estimate the magnification.\n", - "\n", - "Because both fluxes were computed on grids with the same total area and area per pixel, we do not need to\n", - "explicitly account for area in this calculation. This is because the area terms cancel out when taking the ratio.\n", - "Were the grid areas different, we would need to include area terms in the calculation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Magnification: {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer__\n", - "\n", - "Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to compute\n", - "source science calculations for real datasets.\n", - "\n", - "The code below shows how using a tracer, composed of any combination of lens and source galaxies, we can\n", - "compute the source flux and magnification. It reproduces the calculations above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", - "\n", - "image_plane_grid = traced_grid_list[0]\n", - "source_plane_grid = traced_grid_list[1]\n", - "\n", - "lensed_source_image = tracer.planes[1].image_2d_from(grid=source_plane_grid)\n", - "source_plane_image = tracer.planes[1].image_2d_from(grid=image_plane_grid)\n", - "\n", - "total_image_plane_flux = np.sum(lensed_source_image)\n", - "total_source_plane_flux = np.sum(source_plane_image)\n", - "\n", - "source_magnification = total_image_plane_flux / total_source_plane_flux\n", - "\n", - "print(f\"Source Plane Total Flux via Tracer: {total_source_plane_flux} mJy beam^-1\")\n", - "print(f\"Source Magnification via Tracer: {source_magnification}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Parametric Source Models__\n", - "\n", - "If your lens modeling uses a parametric source model (e.g. Sersic, Multi Gaussian Expansion), the only object\n", - "you need to perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", - "\n", - "Alternatively, as done above, you can manually set up a tracer using the lens and source galaxies inferred\n", - "by lens modeling.\n", - "\n", - "Therefore, you may now wish to go to your results, extract the `max_log_likelihood_tracer`, and use it to compute\n", - "the source flux and magnification as shown above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Source Science\n", + "==============\n", + "\n", + "Source science focuses on studying the highly magnified properties of the background lensed source galaxy (or galaxies).\n", + "\n", + "Using a source galaxy model, we can compute key quantities such as the magnification, total flux, and intrinsic\n", + "size of the source.\n", + "\n", + "This example shows how to perform these calculations using Sersic parametric sources on imaging data, which\n", + "is conceptually the simplest case for source science calculations and a good introduction to the topic.\n", + "\n", + "__Contents__\n", + "\n", + "- **Mask:** Define the 2D mask applied to the dataset for the model-fit.\n", + "- **Loading Data:** We we begin by loading the strong lens dataset `simple` from .fits files, which is the dataset we.\n", + "- **Source Values:** Source science calculations for real lenses are performed using the best-fitting model inferred.\n", + "- **Source Flux:** A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see.\n", + "- **Source Magnification:** The overall magnification of the source is estimated as the ratio of total surface brightness in.\n", + "- **Tracer:** Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to.\n", + "- **Parametric Source Models:** If your lens modeling uses a parametric source model (e.g." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.5\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(256, 256),\n", + " pixel_scales=0.1,\n", + " radius=mask_radius,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Loading Data__\n", + "\n", + "We we begin by loading the strong lens dataset `simple` from .fits files, which is the dataset \n", + "we will use to demonstrate fitting.\n", + "\n", + "This includes the method used to Fourier transform the real-space image of the strong lens to the uv-plane and\n", + "compare directly to the visibilities. We use `TransformerNUFFT`, the JAX-native Non-Uniform Fast Fourier Transform\n", + "backed by `nufftax`, which scales efficiently from a few hundred visibilities to tens of millions.\n", + "\n", + "This dataset was simulated using the `interferometer/simulator` example, read through that to understand how\n", + "the data this example fits was generated." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"interferometer\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/interferometer/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerNUFFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `aplt.subplot_interferometer_dirty_images` contains a subplot which plots all the key properties of the dataset simultaneously.\n", + "\n", + "This includes the observed visibility data, RMS noise map and other information." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Visibility data is in uv space, making it hard to interpret by eye.\n", + "\n", + "The dirty images of the interferometer dataset can plotted, which use the transformer of the interferometer \n", + "to map the visibilities, noise-map or other quantity to a real-space image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "__Source Values__\n", + "\n", + "Source science calculations for real lenses are performed using the best-fitting model inferred from a dataset, \n", + "and this example demonstrates how to use this below.\n", + "\n", + "However, we for simplicity, we demonstrate these calculations using the Sersic source model used to simulate the dataset, \n", + "which we refer to as the \"true\" source model. When analysing real strong lenses, a true underlying model is not known, \n", + "but for simulated datasets it is.\n", + "\n", + "This allows us to illustrate the calculations in a way that does not depend on the specific details of the data or \n", + "on assumptions about how the lens model is inferred.\n", + "\n", + "The `tracer` below corresponds to the same tracer used to simulate the `simple` dataset, and therefore \n", + "represents the true source model. We also include the 2D grid of (y,x) coordinates which simulate the dataset.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=0.3,\n", + " effective_radius=1.0,\n", + " sersic_index=2.5,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By plotting the image of the tracer, we confirm it looks identical to the simulated dataset but does not have\n", + "CCD imaging features such as noise or blurring from a PSF." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Source Flux__\n", + "\n", + "A key quantity for a source galaxy is its total flux, which can be used to compute magnitudes (see \n", + "`autolens_workspace/*/guides/units/flux`) example for more details on this).\n", + "\n", + "The most simple way to compute the total flux of a light profile is to create a grid of (y,x) coordinates over which\n", + "we compute the image of the light profile, and then sum the image. \n", + "\n", + "The units of the light profile `intensity` are the units of the data the light profile was fitted to. In this example\n", + "we will assume everything is in miliJansky per beam (`mJy beam^-1`), which is typical for ALMA data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"Source Galaxy's Intensity {source_galaxy.bulge.intensity} mJy beam^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The total flux, in units of `mJy beam^-1` , is computed by summing the image of the light profile over all pixels.\n", + "\n", + "Note that we can use a `grid` of any shape and pixel scale here, the important thing is that it is so large\n", + "and high enough resolution that it captures all the light from the light profile.\n", + "\n", + "Note that we are using the source galaxy's true light profile, which corresponds to its emission in the source-plane.\n", + "For real datasets, we have to infer this via lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=0.02)\n", + "\n", + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_flux = np.sum(image) # in units mJy beam^-1 as summed over pixels\n", + "\n", + "print(f\"Total Source Flux: {total_flux} mJy beam^-1\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below, we will compare how this true source flux compares to the inferred source fluxes we compute using different\n", + "source modeling techniques (e.g. parametric and pixelized source models). Converting the flux to magnitudes or\n", + "other quantities used for tasks like SED fitting is described in the `autolens_workspace/*/guides/units/flux` example.\n", + "\n", + "__Source Magnification__\n", + "\n", + "The overall magnification of the source is estimated as the ratio of total surface brightness in the image-plane and \n", + "total surface brightness in the source-plane.\n", + "\n", + "Note that the surface brightness is different to the total flux above, as surface brightness is flux per unit area. \n", + "We therefore explicitly mention how area folds into the calculation below.\n", + "\n", + "To ensure the magnification is stable and that we resolve all source emission in both the image-plane and source-plane \n", + "we use a very high resolution grid, higher than we used to compute the total flux above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(1000, 1000), pixel_scales=0.03)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We repeat our calculation of the source's total flux in the source-plane using this higher resolution grid, note\n", + "that we do not take the area into account, the reason for this is explained below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "image = source_galaxy.bulge.image_2d_from(grid=grid)\n", + "\n", + "total_source_plane_flux = np.sum(image) # in units mJy beam^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now need the total flux of the lensed source in the image-plane, that is how much flux we measure after\n", + "gravitational lensing.\n", + "\n", + "To calculation this, we first ray-trace the grid above from the image-plane to the source-plane using the tracer\n", + "and then pass it to the source galaxy's light profile to compute the lensed image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "source_plane_grid = traced_grid_list[1]\n", + "\n", + "lensed_source_image = source_galaxy.bulge.image_2d_from(grid=source_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(\n", + " lensed_source_image\n", + ") # in units mJy beam^-1 as summed over pixels" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now take the ratio of the total image-plane flux to source-plane flux to estimate the magnification.\n", + "\n", + "Because both fluxes were computed on grids with the same total area and area per pixel, we do not need to\n", + "explicitly account for area in this calculation. This is because the area terms cancel out when taking the ratio.\n", + "Were the grid areas different, we would need to include area terms in the calculation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Magnification: {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "Lens modeling returns a `max_log_likelihood_tracer`, which is likely the object you have at hand to compute\n", + "source science calculations for real datasets.\n", + "\n", + "The code below shows how using a tracer, composed of any combination of lens and source galaxies, we can\n", + "compute the source flux and magnification. It reproduces the calculations above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "traced_grid_list = tracer.traced_grid_2d_list_from(grid=grid)\n", + "\n", + "image_plane_grid = traced_grid_list[0]\n", + "source_plane_grid = traced_grid_list[1]\n", + "\n", + "lensed_source_image = tracer.planes[1].image_2d_from(grid=source_plane_grid)\n", + "source_plane_image = tracer.planes[1].image_2d_from(grid=image_plane_grid)\n", + "\n", + "total_image_plane_flux = np.sum(lensed_source_image)\n", + "total_source_plane_flux = np.sum(source_plane_image)\n", + "\n", + "source_magnification = total_image_plane_flux / total_source_plane_flux\n", + "\n", + "print(f\"Source Plane Total Flux via Tracer: {total_source_plane_flux} mJy beam^-1\")\n", + "print(f\"Source Magnification via Tracer: {source_magnification}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Parametric Source Models__\n", + "\n", + "If your lens modeling uses a parametric source model (e.g. Sersic, Multi Gaussian Expansion), the only object\n", + "you need to perform source science calculations is the `max_log_likelihood_tracer` returned by lens modeling.\n", + "\n", + "Alternatively, as done above, you can manually set up a tracer using the lens and source galaxies inferred\n", + "by lens modeling.\n", + "\n", + "Therefore, you may now wish to go to your results, extract the `max_log_likelihood_tracer`, and use it to compute\n", + "the source flux and magnification as shown above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/dataset_offsets/modeling.ipynb b/notebooks/multi/features/dataset_offsets/modeling.ipynb index 829e13120..45e83226c 100644 --- a/notebooks/multi/features/dataset_offsets/modeling.ipynb +++ b/notebooks/multi/features/dataset_offsets/modeling.ipynb @@ -1,486 +1,523 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Dataset Offsets\n", - "==================================\n", - "\n", - "Multi-wavelength datasets often have offsets between their images, which are due to the different telescope pointings\n", - "during the observations.\n", - "\n", - "These offsets are often accounted for during the data reduction process, which aligns the images, however:\n", - "\n", - " - Certain data reduction pipelines may not perfectly align the images, and the scientist may be unsure what the\n", - " true offset between the images are.\n", - "\n", - " - Even if the reduction process does align the images, there is still a small uncertainty in the offset due to the\n", - " precision of the telescope pointing which for detailed lens models must be accounted for.\n", - "\n", - "This script shows how to include both an offset (two free parameters: y and x shifts) and a rotation (one free\n", - "parameter: roll angle in degrees) in the model for every additional dataset after the first. Together these\n", - "describe the rigid-body misalignment of each dataset relative to the reference (first) dataset, which is the\n", - "common pattern for multi-band space-telescope data (e.g. JWST NIRCam frames at slightly different roll angles).\n", - "\n", - "To apply the misalignment, the code subtracts the offset from the grids aligned to the dataset pixels and then\n", - "rotates them about the offset point before performing lensing calculations. The light and mass model centres\n", - "do not change; only the coordinates of the image pixels input into these profiles are transformed.\n", - "\n", - "__Contents__\n", - "\n", - "- **Advantages & Disadvantages:** If one fits a lens model to one dataset and applies it to other datasets, it is common to see the.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Advantages__\n", - "\n", - "If one fits a lens model to one dataset and applies it to other datasets, it is common to see the lens model fit\n", - "and source reconsturction degrade due to small offsets between the datasets. The same issue persists for simultaneous\n", - "fits to multiple datasets, even when care has been taken to align the datasets.\n", - "\n", - "The advantage is therefore simple, for most multi-wavelength lens modeling, accounting for offsets in this way\n", - "is the only way to ensure the lens model is accurate and the source reconstruction is reliable.\n", - "\n", - "__Disadvantages__\n", - "\n", - "Each offset introduces two additional free parameters into the model for each dataset after the first dataset. For\n", - "4 datasets, this is 6 additional free parameters. This increases the dimensionality of the non-linear parameter space\n", - "and therefore the computational run-time of the model-fit.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a an MGE bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a an MGE.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "\n", - "The strings are used for load each dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixel Scales__\n", - "\n", - "Every multi-wavelength dataset can have its own unique pixel-scale." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors.\n", - "\n", - "The plotted images show that the datasets have a small offset between them, half a pixel based on the resolution of\n", - "the second image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"dataset_offsets\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/multi/features/dataset_offsets/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset_list = [\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " pixel_scales=pixel_scales,\n", - " )\n", - " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", - "\n", - "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", - "absolutely necessary, but provides a more reliable analysis.\n", - "\n", - "The small offset between datasets means that the mask may not contain the exact same area of the image for every\n", - "dataset. \n", - "\n", - "For this dataset's offset of half a pixel (and anything of order a few pixels) this is fine and wont impact the analysis. \n", - "However, for larger offsets the mask may need to be adjusted to ensure the same image area is masked out." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask_list = [\n", - " al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - " for dataset in dataset_list\n", - "]\n", - "\n", - "dataset_list = [\n", - " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We create an `Analysis` object for every dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [al.AnalysisImaging(dataset=dataset) for dataset in dataset_list]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - Parameters which shift (y_offset_0, x_offset_0) and rotate (grid_rotation_angle) the second dataset's image\n", - " relative to the first dataset's image are included via the `DatasetModel` object [3 parameters].\n", - "\n", - " - The lens galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=24." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "dataset_model = af.Model(al.DatasetModel)\n", - "\n", - "model = af.Collection(\n", - " dataset_model=dataset_model, galaxies=af.Collection(lens=lens, source=source)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `DatasetModel` can apply a shift to each dataset, however we do not want this to be applied to both\n", - "datasets as they will then both have free parameters for the same offset, duplicating free parameters.\n", - "\n", - "The default prior on a `DatasetModel`'s offsets are actually not a prior, but fixed values of (0.0, 0.0),\n", - "meaning that if we do not update the model the shift will not be applied to the datasets.\n", - "\n", - "We therefore update the `DatasetModel` below, to only apply a shift to the second dataset, which is the r-band image.\n", - "\n", - "If we add more datasets, the code will apply the shift to each one after the first dataset, which are all shifted\n", - "relative to the first dataset, making it the reference point." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = []\n", - "\n", - "for i, analysis in enumerate(analysis_list):\n", - " model_analysis = model.copy()\n", - "\n", - " if i > 0:\n", - " model_analysis.dataset_model.grid_offset.grid_offset_0 = af.UniformPrior(\n", - " lower_limit=-1.0, upper_limit=1.0\n", - " )\n", - " model_analysis.dataset_model.grid_offset.grid_offset_1 = af.UniformPrior(\n", - " lower_limit=-1.0, upper_limit=1.0\n", - " )\n", - " # Roll-angle misalignment between exposures (degrees, CCW about the offset point).\n", - " model_analysis.dataset_model.grid_rotation_angle = af.UniformPrior(\n", - " lower_limit=-5.0, upper_limit=5.0\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - "\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"multi\", \"modeling\"),\n", - " name=\"dataset_offsets\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", - "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_list[0].max_log_likelihood_instance)\n", - "print(result_list[1].max_log_likelihood_instance)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting each result's tracer shows that the source appears different, owning to its different intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", - "\n", - "Therefore, the samples is identical in every result object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Dataset Offsets\n", + "==================================\n", + "\n", + "Multi-wavelength datasets often have offsets between their images, which are due to the different telescope pointings\n", + "during the observations.\n", + "\n", + "These offsets are often accounted for during the data reduction process, which aligns the images, however:\n", + "\n", + " - Certain data reduction pipelines may not perfectly align the images, and the scientist may be unsure what the\n", + " true offset between the images are.\n", + "\n", + " - Even if the reduction process does align the images, there is still a small uncertainty in the offset due to the\n", + " precision of the telescope pointing which for detailed lens models must be accounted for.\n", + "\n", + "This script shows how to include both an offset (two free parameters: y and x shifts) and a rotation (one free\n", + "parameter: roll angle in degrees) in the model for every additional dataset after the first. Together these\n", + "describe the rigid-body misalignment of each dataset relative to the reference (first) dataset, which is the\n", + "common pattern for multi-band space-telescope data (e.g. JWST NIRCam frames at slightly different roll angles).\n", + "\n", + "To apply the misalignment, the code subtracts the offset from the grids aligned to the dataset pixels and then\n", + "rotates them about the offset point before performing lensing calculations. The light and mass model centres\n", + "do not change; only the coordinates of the image pixels input into these profiles are transformed.\n", + "\n", + "__Contents__\n", + "\n", + "- **Advantages & Disadvantages:** If one fits a lens model to one dataset and applies it to other datasets, it is common to see the.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Advantages__\n", + "\n", + "If one fits a lens model to one dataset and applies it to other datasets, it is common to see the lens model fit\n", + "and source reconsturction degrade due to small offsets between the datasets. The same issue persists for simultaneous\n", + "fits to multiple datasets, even when care has been taken to align the datasets.\n", + "\n", + "The advantage is therefore simple, for most multi-wavelength lens modeling, accounting for offsets in this way\n", + "is the only way to ensure the lens model is accurate and the source reconstruction is reliable.\n", + "\n", + "__Disadvantages__\n", + "\n", + "Each offset introduces two additional free parameters into the model for each dataset after the first dataset. For\n", + "4 datasets, this is 6 additional free parameters. This increases the dimensionality of the non-linear parameter space\n", + "and therefore the computational run-time of the model-fit.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a an MGE bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a an MGE.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "\n", + "The strings are used for load each dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixel Scales__\n", + "\n", + "Every multi-wavelength dataset can have its own unique pixel-scale." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors.\n", + "\n", + "The plotted images show that the datasets have a small offset between them, half a pixel based on the resolution of\n", + "the second image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"dataset_offsets\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/multi/features/dataset_offsets/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_list = [\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " pixel_scales=pixel_scales,\n", + " )\n", + " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", + "\n", + "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", + "absolutely necessary, but provides a more reliable analysis.\n", + "\n", + "The small offset between datasets means that the mask may not contain the exact same area of the image for every\n", + "dataset. \n", + "\n", + "For this dataset's offset of half a pixel (and anything of order a few pixels) this is fine and wont impact the analysis. \n", + "However, for larger offsets the mask may need to be adjusted to ensure the same image area is masked out." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask_list = [\n", + " al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + " for dataset in dataset_list\n", + "]\n", + "\n", + "dataset_list = [\n", + " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We create an `Analysis` object for every dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [al.AnalysisImaging(dataset=dataset) for dataset in dataset_list]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - Parameters which shift (y_offset_0, x_offset_0) and rotate (grid_rotation_angle) the second dataset's image\n", + " relative to the first dataset's image are included via the `DatasetModel` object [3 parameters].\n", + "\n", + " - The lens galaxy's light is an MGE with 1 x 20 Gaussians [6 parameters].\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=24." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "dataset_model = af.Model(al.DatasetModel)\n", + "\n", + "model = af.Collection(\n", + " dataset_model=dataset_model, galaxies=af.Collection(lens=lens, source=source)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `DatasetModel` can apply a shift to each dataset, however we do not want this to be applied to both\n", + "datasets as they will then both have free parameters for the same offset, duplicating free parameters.\n", + "\n", + "The default prior on a `DatasetModel`'s offsets are actually not a prior, but fixed values of (0.0, 0.0),\n", + "meaning that if we do not update the model the shift will not be applied to the datasets.\n", + "\n", + "We therefore update the `DatasetModel` below, to only apply a shift to the second dataset, which is the r-band image.\n", + "\n", + "If we add more datasets, the code will apply the shift to each one after the first dataset, which are all shifted\n", + "relative to the first dataset, making it the reference point." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = []\n", + "\n", + "for i, analysis in enumerate(analysis_list):\n", + " model_analysis = model.copy()\n", + "\n", + " if i > 0:\n", + " model_analysis.dataset_model.grid_offset.grid_offset_0 = af.UniformPrior(\n", + " lower_limit=-1.0, upper_limit=1.0\n", + " )\n", + " model_analysis.dataset_model.grid_offset.grid_offset_1 = af.UniformPrior(\n", + " lower_limit=-1.0, upper_limit=1.0\n", + " )\n", + " # Roll-angle misalignment between exposures (degrees, CCW about the offset point).\n", + " model_analysis.dataset_model.grid_rotation_angle = af.UniformPrior(\n", + " lower_limit=-5.0, upper_limit=5.0\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + "\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"multi\", \"modeling\"),\n", + " name=\"dataset_offsets\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", + "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_list[0].max_log_likelihood_instance)\n", + "print(result_list[1].max_log_likelihood_instance)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting each result's tracer shows that the source appears different, owning to its different intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", + "\n", + "Therefore, the samples is identical in every result object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/dataset_offsets/simulator.ipynb b/notebooks/multi/features/dataset_offsets/simulator.ipynb index ace22b21b..4e89e4d63 100644 --- a/notebooks/multi/features/dataset_offsets/simulator.ipynb +++ b/notebooks/multi/features/dataset_offsets/simulator.ipynb @@ -1,499 +1,536 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Dataset Offsets\n", - "==========================\n", - "\n", - "Multi-wavelength datasets often have offsets between their images, which are due to the different telescope pointings\n", - "during the observations.\n", - "\n", - "These offsets are often accounted for during the data reduction process, which aligns the images, however:\n", - "\n", - " - Certain data reduction pipelines may not perfectly align the images, and the scientist may be unsure what the\n", - " true offset between the images are.\n", - "\n", - " - Even if the reduction process does align the images, there is still a small uncertainty in the offset due to the\n", - " precision of the telescope pointing which for detailed lens models must be accounted for.\n", - "\n", - "This script simulates a multi-wavelength `Imaging` dataset where the two images have a small offset and a small\n", - "relative rotation between them. Both impact the lensing calculation and modeling if not accounted for. The\n", - "`modeling` examples show how to include both as free parameters of a `DatasetModel`.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "- **Dataset Paths:** Overview of dataset paths for this example.\n", - "- **Simulate:** The pixel-scale of each color image is different meaning we make a list of grids for the simulation.\n", - "- **Offset:** Offset the second grid from the first grid by the pixel scale in both the y and x directions.\n", - "- **Ray Tracing:** The lens galaxy light at each wavelength has a different intensity, thus we create two lens.\n", - "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates multi-wavelength `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The two datasets have a small offset of half the pixel scale between them.\n", - " - The lens galaxy's light profile is an `Sersic`, which has a different `intensity` at each wavelength.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", - "\n", - "Two images are simulated, corresponding to a greener ('g' band) redder image (`r` band).\n", - "\n", - "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for simulating images. Thus,\n", - "certain parts of code are not documented to ensure the script is concise." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "\n", - "The strings are used for naming the datasets on output." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"dataset_offsets\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "The pixel-scale of each color image is different meaning we make a list of grids for the simulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12]\n", - "\n", - "grid_list = []\n", - "\n", - "for pixel_scales in pixel_scales_list:\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=(150, 150),\n", - " pixel_scales=pixel_scales,\n", - " )\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " grid_list.append(grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Offset__\n", - "\n", - "Offset the second grid from the first grid by the pixel scale in both the y and x directions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_list[1] -= pixel_scales_list[1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Rotation__\n", - "\n", - "In addition to the small offset, rotate the second grid by 1.5 degrees counter-clockwise about its\n", - "offset point. This emulates a small roll-angle difference between two multi-band exposures (e.g. JWST\n", - "NIRCam frames acquired at slightly different telescope orientations). The modeling script recovers this\n", - "via the `DatasetModel.grid_rotation_angle` parameter alongside `grid_offset`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid_list[1] = grid_list[1].subtracted_and_rotated_from(offset=(0.0, 0.0), angle=1.5)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate simple Gaussian PSFs for the images in the r and g bands." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sigma_list = [0.1, 0.2]\n", - "\n", - "psf_list = [\n", - " al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", - " )\n", - " for grid, sigma in zip(grid_list, sigma_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create separate simulators for the g and r bands." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "background_sky_level_list = [0.1, 0.15]\n", - "\n", - "simulator_list = [\n", - " al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=background_sky_level,\n", - " add_poisson_noise_to_data=True,\n", - " )\n", - " for psf, background_sky_level in zip(psf_list, background_sky_level_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The lens galaxy light at each wavelength has a different intensity, thus we create two lens galaxies for each waveband. \n", - "\n", - "The lens galaxy's mass (SIE+Shear) is identical for each waveband and included in both lens galaxies in the list.." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "intensity_list = [0.05, 1.5]\n", - "\n", - "bulge_list = [\n", - " al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=intensity,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " )\n", - " for intensity in intensity_list\n", - "]\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "lens_galaxy_list = [\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - " )\n", - " for bulge in bulge_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The source galaxy at each wavelength has a different intensity, thus we create two source galaxies for each waveband." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "intensity_list = [0.5, 0.7]\n", - "\n", - "source_galaxy_list = [\n", - " al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=intensity,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - " )\n", - " for intensity in intensity_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging` \n", - "dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_list = [\n", - " al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - " for lens_galaxy, source_galaxy in zip(lens_galaxy_list, source_galaxy_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating.\n", - "\n", - "Close inspection of the images shows they are slightly offset by half a pixel from one another." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for tracer, grid in zip(tracer_list, grid_list):\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - " for grid, simulator, tracer in zip(grid_list, simulator_list, tracer_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " overwrite=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - " aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "for waveband, grid, tracer in zip(waveband_list, grid_list, tracer_list):\n", - " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", - " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)\n", - "\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "[\n", - " al.output_to_json(obj=tracer, file_path=Path(dataset_path, f\"{color}_tracer.json\"))\n", - " for color, tracer in zip(waveband_list, tracer_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/dataset_offsets`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Dataset Offsets\n", + "==========================\n", + "\n", + "Multi-wavelength datasets often have offsets between their images, which are due to the different telescope pointings\n", + "during the observations.\n", + "\n", + "These offsets are often accounted for during the data reduction process, which aligns the images, however:\n", + "\n", + " - Certain data reduction pipelines may not perfectly align the images, and the scientist may be unsure what the\n", + " true offset between the images are.\n", + "\n", + " - Even if the reduction process does align the images, there is still a small uncertainty in the offset due to the\n", + " precision of the telescope pointing which for detailed lens models must be accounted for.\n", + "\n", + "This script simulates a multi-wavelength `Imaging` dataset where the two images have a small offset and a small\n", + "relative rotation between them. Both impact the lensing calculation and modeling if not accounted for. The\n", + "`modeling` examples show how to include both as free parameters of a `DatasetModel`.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "- **Dataset Paths:** Overview of dataset paths for this example.\n", + "- **Simulate:** The pixel-scale of each color image is different meaning we make a list of grids for the simulation.\n", + "- **Offset:** Offset the second grid from the first grid by the pixel scale in both the y and x directions.\n", + "- **Ray Tracing:** The lens galaxy light at each wavelength has a different intensity, thus we create two lens.\n", + "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates multi-wavelength `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The two datasets have a small offset of half the pixel scale between them.\n", + " - The lens galaxy's light profile is an `Sersic`, which has a different `intensity` at each wavelength.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", + "\n", + "Two images are simulated, corresponding to a greener ('g' band) redder image (`r` band).\n", + "\n", + "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for simulating images. Thus,\n", + "certain parts of code are not documented to ensure the script is concise." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "\n", + "The strings are used for naming the datasets on output." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"dataset_offsets\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "The pixel-scale of each color image is different meaning we make a list of grids for the simulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12]\n", + "\n", + "grid_list = []\n", + "\n", + "for pixel_scales in pixel_scales_list:\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=(150, 150),\n", + " pixel_scales=pixel_scales,\n", + " )\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " grid_list.append(grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Offset__\n", + "\n", + "Offset the second grid from the first grid by the pixel scale in both the y and x directions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_list[1] -= pixel_scales_list[1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Rotation__\n", + "\n", + "In addition to the small offset, rotate the second grid by 1.5 degrees counter-clockwise about its\n", + "offset point. This emulates a small roll-angle difference between two multi-band exposures (e.g. JWST\n", + "NIRCam frames acquired at slightly different telescope orientations). The modeling script recovers this\n", + "via the `DatasetModel.grid_rotation_angle` parameter alongside `grid_offset`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid_list[1] = grid_list[1].subtracted_and_rotated_from(offset=(0.0, 0.0), angle=1.5)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate simple Gaussian PSFs for the images in the r and g bands." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sigma_list = [0.1, 0.2]\n", + "\n", + "psf_list = [\n", + " al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", + " )\n", + " for grid, sigma in zip(grid_list, sigma_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create separate simulators for the g and r bands." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "background_sky_level_list = [0.1, 0.15]\n", + "\n", + "simulator_list = [\n", + " al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=background_sky_level,\n", + " add_poisson_noise_to_data=True,\n", + " )\n", + " for psf, background_sky_level in zip(psf_list, background_sky_level_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The lens galaxy light at each wavelength has a different intensity, thus we create two lens galaxies for each waveband. \n", + "\n", + "The lens galaxy's mass (SIE+Shear) is identical for each waveband and included in both lens galaxies in the list.." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "intensity_list = [0.05, 1.5]\n", + "\n", + "bulge_list = [\n", + " al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=intensity,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " )\n", + " for intensity in intensity_list\n", + "]\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "lens_galaxy_list = [\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + " )\n", + " for bulge in bulge_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The source galaxy at each wavelength has a different intensity, thus we create two source galaxies for each waveband." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "intensity_list = [0.5, 0.7]\n", + "\n", + "source_galaxy_list = [\n", + " al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=intensity,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + " )\n", + " for intensity in intensity_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging` \n", + "dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_list = [\n", + " al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + " for lens_galaxy, source_galaxy in zip(lens_galaxy_list, source_galaxy_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating.\n", + "\n", + "Close inspection of the images shows they are slightly offset by half a pixel from one another." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for tracer, grid in zip(tracer_list, grid_list):\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + " for grid, simulator, tracer in zip(grid_list, simulator_list, tracer_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " overwrite=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + " aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "for waveband, grid, tracer in zip(waveband_list, grid_list, tracer_list):\n", + " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", + " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)\n", + "\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "[\n", + " al.output_to_json(obj=tracer, file_path=Path(dataset_path, f\"{color}_tracer.json\"))\n", + " for color, tracer in zip(waveband_list, tracer_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/dataset_offsets`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/imaging_and_interferometer/modeling.ipynb b/notebooks/multi/features/imaging_and_interferometer/modeling.ipynb index 17564c1b5..3e5e88949 100644 --- a/notebooks/multi/features/imaging_and_interferometer/modeling.ipynb +++ b/notebooks/multi/features/imaging_and_interferometer/modeling.ipynb @@ -1,444 +1,481 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Mass Total + Source Parametric\n", - "========================================\n", - "\n", - "This script fits an `Interferometer` and `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is an MGE (but is invisible in the interferometer data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an MGE.\n", - "\n", - "__Contents__\n", - "\n", - "- **Benefits:** A number of benefits are apparently if we combine the analysis of both datasets at both wavelengths.\n", - "- **Interferometer Masking:** We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using.\n", - "- **Interferometer Dataset:** Load and plot the strong lens `Interferometer` dataset `simple__no_lens_light` from .fits files.\n", - "- **Imaging Dataset:** Load and plot the strong lens dataset `simple__no_lens_light` via .fits files, which we will fit.\n", - "- **Imaging Masking:** Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Benefits__\n", - "\n", - " A number of benefits are apparently if we combine the analysis of both datasets at both wavelengths:\n", - "\n", - " - The lens galaxy is invisible at sub-mm wavelengths, making it straight-forward to infer a lens mass model by\n", - " fitting the source at submm wavelengths.\n", - "\n", - " - The source galaxy appears completely different in the g-band and at sub-millimeter wavelengths, providing a lot\n", - " more information with which to constrain the lens galaxy mass model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt\n", - "import numpy as np" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interferometer Masking__\n", - "\n", - "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 4.0\n", - "\n", - "real_space_mask = al.Mask2D.circular(\n", - " shape_native=(800, 800), pixel_scales=0.05, radius=mask_radius\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Interferometer Dataset__\n", - "\n", - "Load and plot the strong lens `Interferometer` dataset `simple__no_lens_light` from .fits files, which we will fit \n", - "with the lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"interferometer\"\n", - "dataset_name = \"simple__no_lens_light\"\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/multi/features/imaging_and_interferometer/simulator.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "interferometer = al.Interferometer.from_fits(\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " real_space_mask=real_space_mask,\n", - " transformer_class=al.TransformerDFT,\n", - ")\n", - "\n", - "aplt.subplot_interferometer_dirty_images(dataset=interferometer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Imaging Dataset__\n", - "\n", - "Load and plot the strong lens dataset `simple__no_lens_light` via .fits files, which we will fit with the lens model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"lens_sersic\"\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name\n", - "\n", - "imaging = al.Imaging.from_fits(\n", - " data_path=Path(dataset_path, \"g_data.fits\"),\n", - " psf_path=Path(dataset_path, \"g_psf.fits\"),\n", - " noise_map_path=Path(dataset_path, \"g_noise_map.fits\"),\n", - " pixel_scales=0.08,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=imaging)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Imaging Masking__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=imaging.shape_native,\n", - " pixel_scales=imaging.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "imaging = imaging.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=imaging)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We create analysis objects for both datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_imaging = al.AnalysisImaging(dataset=imaging)\n", - "analysis_interferometer = al.AnalysisInterferometer(dataset=interferometer)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In this \n", - "example our lens model is:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", - " - with 1 x 20 Gaussians for the source galaxy's light, which is complete different for each waveband. [8 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=21." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now combine them using the factor analysis class, which allows us to fit the two datasets simultaneously.\n", - "\n", - "Imaging and interferometer datasets observe completely different properties of the, such that the galaxy appears \n", - "completely different in the imaging data (e.g. optical emission) and sub-millimeter wavelengths, meaning a completely \n", - "different source model should be used for each dataset.\n", - "\n", - "For this reason, we move lens light and source model composition to the `AnalysisFactor` class, which allows us to fit \n", - "the two datasets simultaneously but with different models.\n", - "\n", - "The benefit of fitting them simultaneously is that the mass model is inferred from both datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = []\n", - "\n", - "for analysis in [analysis_imaging, analysis_interferometer]:\n", - "\n", - " model_analysis = model.copy()\n", - "\n", - " model_analysis.galaxies.lens.bulge = af.Model(al.lp_linear.Sersic)\n", - "\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - " )\n", - "\n", - " model_analysis.galaxies.source.bulge = bulge\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", - "\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` of the model shows us there are two models, one for the imaging dataset and one for the interferometer\n", - "dataset. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"multi\", \"modeling\"),\n", - " name=\"imaging_and_interferometer\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The search returns a result object, which includes: \n", - "\n", - " - The lens model corresponding to the maximum log likelihood solution in parameter space.\n", - " - The corresponding maximum log likelihood `Tracer` and `FitInterferometer` objects.\n", - " - Information on the posterior as estimated by the `Nautilus` non-linear search." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_list[0].max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=result_list[0].max_log_likelihood_tracer,\n", - " grid=real_space_mask.derive_grid.unmasked,\n", - ")\n", - "\n", - "aplt.subplot_fit_imaging(fit=result_list[0].max_log_likelihood_fit)\n", - "\n", - "aplt.subplot_fit_interferometer(fit=result_list[1].max_log_likelihood_fit)\n", - "aplt.subplot_fit_dirty_images(fit=result_list[1].max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result_list.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Mass Total + Source Parametric\n", + "========================================\n", + "\n", + "This script fits an `Interferometer` and `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is an MGE (but is invisible in the interferometer data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an MGE.\n", + "\n", + "__Contents__\n", + "\n", + "- **Benefits:** A number of benefits are apparently if we combine the analysis of both datasets at both wavelengths.\n", + "- **Interferometer Masking:** We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using.\n", + "- **Interferometer Dataset:** Load and plot the strong lens `Interferometer` dataset `simple__no_lens_light` from .fits files.\n", + "- **Imaging Dataset:** Load and plot the strong lens dataset `simple__no_lens_light` via .fits files, which we will fit.\n", + "- **Imaging Masking:** Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Benefits__\n", + "\n", + " A number of benefits are apparently if we combine the analysis of both datasets at both wavelengths:\n", + "\n", + " - The lens galaxy is invisible at sub-mm wavelengths, making it straight-forward to infer a lens mass model by\n", + " fitting the source at submm wavelengths.\n", + "\n", + " - The source galaxy appears completely different in the g-band and at sub-millimeter wavelengths, providing a lot\n", + " more information with which to constrain the lens galaxy mass model." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt\n", + "import numpy as np" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interferometer Masking__\n", + "\n", + "We define the \u2018real_space_mask\u2019 which defines the grid the image the strong lens is evaluated using." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 4.0\n", + "\n", + "real_space_mask = al.Mask2D.circular(\n", + " shape_native=(800, 800), pixel_scales=0.05, radius=mask_radius\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Interferometer Dataset__\n", + "\n", + "Load and plot the strong lens `Interferometer` dataset `simple__no_lens_light` from .fits files, which we will fit \n", + "with the lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"interferometer\"\n", + "dataset_name = \"simple__no_lens_light\"\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/multi/features/imaging_and_interferometer/simulator.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "interferometer = al.Interferometer.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " real_space_mask=real_space_mask,\n", + " transformer_class=al.TransformerDFT,\n", + ")\n", + "\n", + "aplt.subplot_interferometer_dirty_images(dataset=interferometer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imaging Dataset__\n", + "\n", + "Load and plot the strong lens dataset `simple__no_lens_light` via .fits files, which we will fit with the lens model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"lens_sersic\"\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name\n", + "\n", + "imaging = al.Imaging.from_fits(\n", + " data_path=Path(dataset_path, \"g_data.fits\"),\n", + " psf_path=Path(dataset_path, \"g_psf.fits\"),\n", + " noise_map_path=Path(dataset_path, \"g_noise_map.fits\"),\n", + " pixel_scales=0.08,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=imaging)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imaging Masking__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=imaging.shape_native,\n", + " pixel_scales=imaging.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "imaging = imaging.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=imaging)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We create analysis objects for both datasets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_imaging = al.AnalysisImaging(dataset=imaging)\n", + "analysis_interferometer = al.AnalysisInterferometer(dataset=interferometer)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose our lens model using `Model` objects, which represent the galaxies we fit to our data. In this \n", + "example our lens model is:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` with `ExternalShear` [7 parameters].\n", + " - with 1 x 20 Gaussians for the source galaxy's light, which is complete different for each waveband. [8 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=21." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now combine them using the factor analysis class, which allows us to fit the two datasets simultaneously.\n", + "\n", + "Imaging and interferometer datasets observe completely different properties of the, such that the galaxy appears \n", + "completely different in the imaging data (e.g. optical emission) and sub-millimeter wavelengths, meaning a completely \n", + "different source model should be used for each dataset.\n", + "\n", + "For this reason, we move lens light and source model composition to the `AnalysisFactor` class, which allows us to fit \n", + "the two datasets simultaneously but with different models.\n", + "\n", + "The benefit of fitting them simultaneously is that the mass model is inferred from both datasets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = []\n", + "\n", + "for analysis in [analysis_imaging, analysis_interferometer]:\n", + "\n", + " model_analysis = model.copy()\n", + "\n", + " model_analysis.galaxies.lens.bulge = af.Model(al.lp_linear.Sersic)\n", + "\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + " )\n", + "\n", + " model_analysis.galaxies.source.bulge = bulge\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", + "\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` of the model shows us there are two models, one for the imaging dataset and one for the interferometer\n", + "dataset. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"multi\", \"modeling\"),\n", + " name=\"imaging_and_interferometer\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns a result object, which includes: \n", + "\n", + " - The lens model corresponding to the maximum log likelihood solution in parameter space.\n", + " - The corresponding maximum log likelihood `Tracer` and `FitInterferometer` objects.\n", + " - Information on the posterior as estimated by the `Nautilus` non-linear search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_list[0].max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=result_list[0].max_log_likelihood_tracer,\n", + " grid=real_space_mask.derive_grid.unmasked,\n", + ")\n", + "\n", + "aplt.subplot_fit_imaging(fit=result_list[0].max_log_likelihood_fit)\n", + "\n", + "aplt.subplot_fit_interferometer(fit=result_list[1].max_log_likelihood_fit)\n", + "aplt.subplot_fit_dirty_images(fit=result_list[1].max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result_list.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/imaging_and_interferometer/simulator.ipynb b/notebooks/multi/features/imaging_and_interferometer/simulator.ipynb index 10c15f79d..67cdda495 100644 --- a/notebooks/multi/features/imaging_and_interferometer/simulator.ipynb +++ b/notebooks/multi/features/imaging_and_interferometer/simulator.ipynb @@ -1,377 +1,414 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Mutli Interferometer\n", - "===============================\n", - "\n", - "This script simulates `Interferometer` data of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a Sersic.\n", - "\n", - "This dataset is paired with the script `multi/simulators/lens_sersic.py` and therefore\n", - "provides interferometer observations of the same strong lens.\n", - "\n", - "It is used to demonstrate the combination of imaging and interferometer datasets.\n", - "\n", - "__Contents__\n", - "\n", - "- **Simulate:** For simulating interferometer data of a strong lens, we recommend using a Grid2D object with a.\n", - "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) and source galaxy light (elliptical Sersic) for this.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `dataset_type` describes the type of data being simulated (in this case, `Interferometer` data) and `dataset_name` \n", - "gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", - "\n", - " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_label/dataset_name/image.fits`.\n", - " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_label/dataset_name/noise_map.fits`.\n", - " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_label/dataset_name/psf.fits`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"interferometer\"\n", - "dataset_name = \"simple__no_lens_light\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "For simulating interferometer data of a strong lens, we recommend using a Grid2D object with a `sub_size` of 1. This\n", - "simplifies the generation of the strong lens image in real space before it is transformed to Fourier space." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.05)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To perform the Fourier transform we need the wavelengths of the baselines, which we'll load from the fits file below.\n", - "\n", - "By default we use baselines from the Square Mile Array (SMA), which produces low resolution interferometer data that\n", - "can be fitted extremely efficiently. The `autolens_workspace` includes ALMA uv_wavelengths files for simulating\n", - "much high resolution datasets (which can be performed by replacing \"sma.fits\" below with \"alma.fits\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "uv_wavelengths_path = Path(\"dataset\", \"interferometer\", \"uv_wavelengths\")\n", - "uv_wavelengths = al.ndarray_via_fits_from(\n", - " file_path=Path(uv_wavelengths_path, \"sma.fits\"), hdu=0\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the interferometer dataset we first create a simulator, which defines the exposure time, noise levels \n", - "and Fourier transform method used in the simulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorInterferometer(\n", - " uv_wavelengths=uv_wavelengths,\n", - " exposure_time=300.0,\n", - " noise_sigma=1000.0,\n", - " transformer_class=al.TransformerDFT,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's mass (SIE+Shear) and source galaxy light (elliptical Sersic) for this simulated lens.\n", - "\n", - "The lens galaxy's mass model is identical to that of the script `multi/simulators/simple__no_lens_light.py`, because\n", - "we will use this dataset to demonstrate the combination of imaging and interferometer datasets.\n", - "\n", - "the source galaxy morphology is very different to that script, because the sub-milimeter observations observe \n", - "the galaxy's warm dust which is therefore clumpy and irregular." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.2, 0.1),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=90.0),\n", - " intensity=0.2,\n", - " effective_radius=0.5,\n", - " sersic_index=2.0,\n", - " ),\n", - " extra_galaxy_0=al.lp.SersicSph(\n", - " centre=(0.3, 0.4), intensity=0.1, effective_radius=0.3, sersic_index=2.5\n", - " ),\n", - " extra_galaxy_1=al.lp.SersicSph(\n", - " centre=(0.0, 0.05), intensity=0.15, effective_radius=0.2, sersic_index=3.0\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will generate the image for the simulated interferometer dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now pass this simulator a tracer, which creates the ray-traced image plotted above and simulates it as an\n", - "interferometer dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets plot the simulated interferometer dataset before we output it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_interferometer_dirty_images(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_interferometer(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_interferometer_dirty_images(\n", - " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Mutli Interferometer\n", + "===============================\n", + "\n", + "This script simulates `Interferometer` data of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a Sersic.\n", + "\n", + "This dataset is paired with the script `multi/simulators/lens_sersic.py` and therefore\n", + "provides interferometer observations of the same strong lens.\n", + "\n", + "It is used to demonstrate the combination of imaging and interferometer datasets.\n", + "\n", + "__Contents__\n", + "\n", + "- **Simulate:** For simulating interferometer data of a strong lens, we recommend using a Grid2D object with a.\n", + "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) and source galaxy light (elliptical Sersic) for this.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `dataset_type` describes the type of data being simulated (in this case, `Interferometer` data) and `dataset_name` \n", + "gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", + "\n", + " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_label/dataset_name/image.fits`.\n", + " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_label/dataset_name/noise_map.fits`.\n", + " - The psf will be output to `/autolens_workspace/dataset/dataset_type/dataset_label/dataset_name/psf.fits`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"interferometer\"\n", + "dataset_name = \"simple__no_lens_light\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "For simulating interferometer data of a strong lens, we recommend using a Grid2D object with a `sub_size` of 1. This\n", + "simplifies the generation of the strong lens image in real space before it is transformed to Fourier space." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(shape_native=(800, 800), pixel_scales=0.05)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To perform the Fourier transform we need the wavelengths of the baselines, which we'll load from the fits file below.\n", + "\n", + "By default we use baselines from the Square Mile Array (SMA), which produces low resolution interferometer data that\n", + "can be fitted extremely efficiently. The `autolens_workspace` includes ALMA uv_wavelengths files for simulating\n", + "much high resolution datasets (which can be performed by replacing \"sma.fits\" below with \"alma.fits\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "uv_wavelengths_path = Path(\"dataset\", \"interferometer\", \"uv_wavelengths\")\n", + "uv_wavelengths = al.ndarray_via_fits_from(\n", + " file_path=Path(uv_wavelengths_path, \"sma.fits\"), hdu=0\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the interferometer dataset we first create a simulator, which defines the exposure time, noise levels \n", + "and Fourier transform method used in the simulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorInterferometer(\n", + " uv_wavelengths=uv_wavelengths,\n", + " exposure_time=300.0,\n", + " noise_sigma=1000.0,\n", + " transformer_class=al.TransformerDFT,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's mass (SIE+Shear) and source galaxy light (elliptical Sersic) for this simulated lens.\n", + "\n", + "The lens galaxy's mass model is identical to that of the script `multi/simulators/simple__no_lens_light.py`, because\n", + "we will use this dataset to demonstrate the combination of imaging and interferometer datasets.\n", + "\n", + "the source galaxy morphology is very different to that script, because the sub-milimeter observations observe \n", + "the galaxy's warm dust which is therefore clumpy and irregular." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.2, 0.1),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.7, angle=90.0),\n", + " intensity=0.2,\n", + " effective_radius=0.5,\n", + " sersic_index=2.0,\n", + " ),\n", + " extra_galaxy_0=al.lp.SersicSph(\n", + " centre=(0.3, 0.4), intensity=0.1, effective_radius=0.3, sersic_index=2.5\n", + " ),\n", + " extra_galaxy_1=al.lp.SersicSph(\n", + " centre=(0.0, 0.05), intensity=0.15, effective_radius=0.2, sersic_index=3.0\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will generate the image for the simulated interferometer dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now pass this simulator a tracer, which creates the ray-traced image plotted above and simulates it as an\n", + "interferometer dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets plot the simulated interferometer dataset before we output it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_interferometer_dirty_images(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_interferometer(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " uv_wavelengths_path=dataset_path / \"uv_wavelengths.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_interferometer_dirty_images(\n", + " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/one_by_one/modeling.ipynb b/notebooks/multi/features/one_by_one/modeling.ipynb index 82a89dc8c..cec3a9030 100644 --- a/notebooks/multi/features/one_by_one/modeling.ipynb +++ b/notebooks/multi/features/one_by_one/modeling.ipynb @@ -1,513 +1,550 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: One By One\n", - "=============================\n", - "\n", - "Multi-wavelength analysis does not necessarily require us to fit all datasets simultaneously. Instead, we can fit one\n", - "dataset first in order to infer a robust lens and source model, and then fit the next dataset, using the inferred\n", - "model as the starting point.\n", - "\n", - "There are many occasions where this approach is beneficial, for example:\n", - "\n", - "- When certain datasets are worse quality (e.g. lower resolution) than others. Fitting them simultaneously may mean this\n", - " dataset's lower quality makes the model fit less robust. By fitting them one by one, using the inferred model of the\n", - " best dataset first, we can ensure the model-fit is as robust as possible and interpret the results of the lower\n", - " quality datasets more clearly.\n", - "\n", - "- It can often produce faster run times, as although more non-linear searches are performed, each search is faster\n", - " than a search which fits all datasets simultaneously.\n", - "\n", - "- To investigate whether lens modeling results inferred when we model all datasets simultanoeusly are robust. If the\n", - " result disappears for fits to individual datasets, this may suggest the result is not robust.\n", - "\n", - "To perform modeling one-by-one, we have to make decision about how simple or complex we make the model after\n", - "fitting the highest quality dataset. For example, we may:\n", - "\n", - "- Fix the lens mass model and only allow the lens light and source light to vary.\n", - "\n", - "- Fix the lens mass model and the majority of lens light and source light parameters, allowing only the `intensity`\n", - " values to vary.\n", - "\n", - "- Allow all parameters to vary, but use the highest quality dataset's inferred model as the starting point.\n", - "\n", - "- Whether to account for offsets between the datasets, or to assume the datasets are aligned.\n", - "\n", - "We illustrate different examples in this script, with the appropriate choice depending on your specific science case.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Second Dataset Mass Model Fixed:** We now fit the second dataset using the inferred model of the first dataset as the starting point.\n", - "- **Second Dataset Offset:** Multi-wavelength datasets often have offsets between their images, which are due to the different.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a an MGE bulge.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a an MGE.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "\n", - "The strings are used for load each dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixel Scales__\n", - "\n", - "Every multi-wavelength dataset can have its own unique pixel-scale." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors.\n", - "\n", - "The plotted images show that the datasets have a small offset between them, half a pixel based on the resolution of\n", - "the second image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"dataset_offsets\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/multi/features/dataset_offsets/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset_list = [\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " pixel_scales=pixel_scales,\n", - " )\n", - " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", - "\n", - "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", - "absolutely necessary, but provides a more reliable analysis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask_list = [\n", - " al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - " for dataset in dataset_list\n", - "]\n", - "\n", - "dataset_list = [\n", - " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We create an `Analysis` object for every dataset.\n", - "\n", - "We do not sum the analyses, like we do in most other example scripts, as we are going to fit each dataset one-by-one." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [\n", - " al.AnalysisImaging(dataset=dataset, use_jax=True) for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's light is an MGE with 2 x 30 Gaussians, where the `intensity` parameter of the lens galaxy\n", - " is solved for linearly [6 parameters].\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is an MGE with 1 x 20 Gaussians, where the `intensity` parameter of the lens galaxy\n", - " is solved for linearly [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=19." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"multi\", \"modeling\"),\n", - " name=\"one_by_one__main_dataset\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis_list[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result object returned by this model-fit is a `Result` object. It is not a list like other examples, because we \n", - "did not use a combined analysis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting the result's tracer shows the source," - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Second Dataset Mass Model Fixed__\n", - "\n", - "We now fit the second dataset using the inferred model of the first dataset as the starting point.\n", - "\n", - "We compose a simple lens model where the mass model is fixed to the result of the first dataset fit, and the lens\n", - "and source galaxy's light are varied. \n", - "\n", - "This model therefore assumes that the mass does not change over wavelength, but the lens and source light do, which\n", - "is what we expect for a strong lens system.\n", - "\n", - "The code below uses the search chaining API to link the priors between model parameters, if you are not\n", - "familiar with this feature, checkout the `guides/modeling/chaining` package." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# model = af.Collection(\n", - "# galaxies=af.Collection(\n", - "# lens=af.Model(\n", - "# al.Galaxy,\n", - "# redshift=result.instance.galaxies.lens.redshift,\n", - "# bulge=result.model.galaxies.lens.bulge,\n", - "# mass=result.instance.galaxies.lens.mass,\n", - "# shear=result.instance.galaxies.lens.shear,\n", - "# ),\n", - "# source=result.model.galaxies.source,\n", - "# ),\n", - "# )\n", - "#\n", - "# print(model.info)\n", - "#\n", - "# search = af.Nautilus(\n", - "# path_prefix=Path(\"multi\", \"modeling\"),\n", - "# name=\"one_by_one__second_mass_model_fixed\",\n", - "# unique_tag=dataset_name,\n", - "# n_live=100,\n", - "# )\n", - "#\n", - "# result_mass_model_fixed = search.fit(model=model, analysis=analysis_list[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Second Dataset Offset__\n", - "\n", - "Multi-wavelength datasets often have offsets between their images, which are due to the different telescope pointings\n", - "during the observations.\n", - "\n", - "These offsets are often accounted for during the data reduction process, but may not be perfectly corrected and\n", - "have uncertainties associated with them.\n", - "\n", - "Fitting datasets one-by-one offers a straightforward method to account for these offsets, by allowing the offset\n", - "between the datasets to vary during the model-fit as two free parameters (y and x).\n", - "\n", - "We now fit for the offset between datasets, keeping all lens model parameters fixed to the result of the first dataset\n", - "fit. \n", - "\n", - "In this example, the two datasets are not offset, so the model-fit will infer an offset consistent with (0.0\", 0.0\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_model = af.Model(al.DatasetModel)\n", - "\n", - "dataset_model.grid_offset.grid_offset_0 = af.UniformPrior(\n", - " lower_limit=-0.1, upper_limit=0.1\n", - ")\n", - "dataset_model.grid_offset.grid_offset_1 = af.UniformPrior(\n", - " lower_limit=-0.1, upper_limit=0.1\n", - ")\n", - "\n", - "model = af.Collection(\n", - " dataset_model=dataset_model,\n", - " galaxies=result.instance.galaxies,\n", - ")\n", - "\n", - "print(model.info)\n", - "\n", - "search = af.Nautilus(\n", - " path_prefix=Path(\"multi\", \"modeling\"),\n", - " name=\"one_by_one__dataset_offset\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis_list[0])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: One By One\n", + "=============================\n", + "\n", + "Multi-wavelength analysis does not necessarily require us to fit all datasets simultaneously. Instead, we can fit one\n", + "dataset first in order to infer a robust lens and source model, and then fit the next dataset, using the inferred\n", + "model as the starting point.\n", + "\n", + "There are many occasions where this approach is beneficial, for example:\n", + "\n", + "- When certain datasets are worse quality (e.g. lower resolution) than others. Fitting them simultaneously may mean this\n", + " dataset's lower quality makes the model fit less robust. By fitting them one by one, using the inferred model of the\n", + " best dataset first, we can ensure the model-fit is as robust as possible and interpret the results of the lower\n", + " quality datasets more clearly.\n", + "\n", + "- It can often produce faster run times, as although more non-linear searches are performed, each search is faster\n", + " than a search which fits all datasets simultaneously.\n", + "\n", + "- To investigate whether lens modeling results inferred when we model all datasets simultanoeusly are robust. If the\n", + " result disappears for fits to individual datasets, this may suggest the result is not robust.\n", + "\n", + "To perform modeling one-by-one, we have to make decision about how simple or complex we make the model after\n", + "fitting the highest quality dataset. For example, we may:\n", + "\n", + "- Fix the lens mass model and only allow the lens light and source light to vary.\n", + "\n", + "- Fix the lens mass model and the majority of lens light and source light parameters, allowing only the `intensity`\n", + " values to vary.\n", + "\n", + "- Allow all parameters to vary, but use the highest quality dataset's inferred model as the starting point.\n", + "\n", + "- Whether to account for offsets between the datasets, or to assume the datasets are aligned.\n", + "\n", + "We illustrate different examples in this script, with the appropriate choice depending on your specific science case.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Second Dataset Mass Model Fixed:** We now fit the second dataset using the inferred model of the first dataset as the starting point.\n", + "- **Second Dataset Offset:** Multi-wavelength datasets often have offsets between their images, which are due to the different.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a an MGE bulge.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a an MGE.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "\n", + "The strings are used for load each dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixel Scales__\n", + "\n", + "Every multi-wavelength dataset can have its own unique pixel-scale." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors.\n", + "\n", + "The plotted images show that the datasets have a small offset between them, half a pixel based on the resolution of\n", + "the second image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"dataset_offsets\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/multi/features/dataset_offsets/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_list = [\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " pixel_scales=pixel_scales,\n", + " )\n", + " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", + "\n", + "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", + "absolutely necessary, but provides a more reliable analysis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask_list = [\n", + " al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + " for dataset in dataset_list\n", + "]\n", + "\n", + "dataset_list = [\n", + " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We create an `Analysis` object for every dataset.\n", + "\n", + "We do not sum the analyses, like we do in most other example scripts, as we are going to fit each dataset one-by-one." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [\n", + " al.AnalysisImaging(dataset=dataset, use_jax=True) for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's light is an MGE with 2 x 30 Gaussians, where the `intensity` parameter of the lens galaxy\n", + " is solved for linearly [6 parameters].\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is an MGE with 1 x 20 Gaussians, where the `intensity` parameter of the lens galaxy\n", + " is solved for linearly [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=19." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"multi\", \"modeling\"),\n", + " name=\"one_by_one__main_dataset\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis_list[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result object returned by this model-fit is a `Result` object. It is not a list like other examples, because we \n", + "did not use a combined analysis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting the result's tracer shows the source," + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Second Dataset Mass Model Fixed__\n", + "\n", + "We now fit the second dataset using the inferred model of the first dataset as the starting point.\n", + "\n", + "We compose a simple lens model where the mass model is fixed to the result of the first dataset fit, and the lens\n", + "and source galaxy's light are varied. \n", + "\n", + "This model therefore assumes that the mass does not change over wavelength, but the lens and source light do, which\n", + "is what we expect for a strong lens system.\n", + "\n", + "The code below uses the search chaining API to link the priors between model parameters, if you are not\n", + "familiar with this feature, checkout the `guides/modeling/chaining` package." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# model = af.Collection(\n", + "# galaxies=af.Collection(\n", + "# lens=af.Model(\n", + "# al.Galaxy,\n", + "# redshift=result.instance.galaxies.lens.redshift,\n", + "# bulge=result.model.galaxies.lens.bulge,\n", + "# mass=result.instance.galaxies.lens.mass,\n", + "# shear=result.instance.galaxies.lens.shear,\n", + "# ),\n", + "# source=result.model.galaxies.source,\n", + "# ),\n", + "# )\n", + "#\n", + "# print(model.info)\n", + "#\n", + "# search = af.Nautilus(\n", + "# path_prefix=Path(\"multi\", \"modeling\"),\n", + "# name=\"one_by_one__second_mass_model_fixed\",\n", + "# unique_tag=dataset_name,\n", + "# n_live=100,\n", + "# )\n", + "#\n", + "# result_mass_model_fixed = search.fit(model=model, analysis=analysis_list[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Second Dataset Offset__\n", + "\n", + "Multi-wavelength datasets often have offsets between their images, which are due to the different telescope pointings\n", + "during the observations.\n", + "\n", + "These offsets are often accounted for during the data reduction process, but may not be perfectly corrected and\n", + "have uncertainties associated with them.\n", + "\n", + "Fitting datasets one-by-one offers a straightforward method to account for these offsets, by allowing the offset\n", + "between the datasets to vary during the model-fit as two free parameters (y and x).\n", + "\n", + "We now fit for the offset between datasets, keeping all lens model parameters fixed to the result of the first dataset\n", + "fit. \n", + "\n", + "In this example, the two datasets are not offset, so the model-fit will infer an offset consistent with (0.0\", 0.0\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_model = af.Model(al.DatasetModel)\n", + "\n", + "dataset_model.grid_offset.grid_offset_0 = af.UniformPrior(\n", + " lower_limit=-0.1, upper_limit=0.1\n", + ")\n", + "dataset_model.grid_offset.grid_offset_1 = af.UniformPrior(\n", + " lower_limit=-0.1, upper_limit=0.1\n", + ")\n", + "\n", + "model = af.Collection(\n", + " dataset_model=dataset_model,\n", + " galaxies=result.instance.galaxies,\n", + ")\n", + "\n", + "print(model.info)\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"multi\", \"modeling\"),\n", + " name=\"one_by_one__dataset_offset\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis_list[0])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/pixelization/modeling.ipynb b/notebooks/multi/features/pixelization/modeling.ipynb index e554bdd4b..f1c11dc20 100644 --- a/notebooks/multi/features/pixelization/modeling.ipynb +++ b/notebooks/multi/features/pixelization/modeling.ipynb @@ -1,484 +1,521 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Pixelized\n", - "===================\n", - "\n", - "This script fits a multi-wavelength `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's surface-brightness is an `Inversion`.\n", - "\n", - "Two images are fitted, corresponding to a greener ('g' band) redder image (`r` band).\n", - "\n", - "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for lens modeling. Thus,\n", - "certain parts of code are not documented to ensure the script is concise.\n", - "\n", - "__Contents__\n", - "\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Positions:** This fit also uses the arc-second positions of the multiply imaged lensed source galaxy, which were.\n", - "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Result:** Overview of the results of the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "\n", - "The strings are used for load each dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixel Scales__\n", - "\n", - "Every multi-wavelength dataset can have its own unique pixel-scale." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"simple__no_lens_light\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/multi/features/pixelization/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "imaging_dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", - "if not imaging_dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "if not (dataset_path / \"positions.json\").exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [\n", - " sys.executable,\n", - " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", - " ],\n", - " check=True,\n", - " )\n", - "\n", - "dataset_list = [\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " pixel_scales=pixel_scales,\n", - " )\n", - " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", - "\n", - "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", - "absolutely necessary, but provides a more reliable analysis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_list = [\n", - " al.Mask2D.circular(\n", - " shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=3.0\n", - " )\n", - " for dataset in dataset_list\n", - "]\n", - "\n", - "dataset_list = [\n", - " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "This fit also uses the arc-second positions of the multiply imaged lensed source galaxy, which were drawn onto the\n", - "image via the GUI described in the file `autolens_workspace/*/imaging/data_preparation/gui/positions.py`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = al.Grid2DIrregular(\n", - " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mesh Shape__\n", - "\n", - "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We create an `Analysis` object for every dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [al.AnalysisImaging(dataset=dataset) for dataset in dataset_list]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - This pixelization is regularized using a `Constant` scheme which smooths every source pixel \n", - " equally, where its `regularization_coefficient` varies across the datasets [2 parameter]. \n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=9." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", - ")\n", - "\n", - "pixelization = af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization=al.reg.Constant,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now combine them using the factor analysis class, which allows us to fit the two datasets simultaneously.\n", - "\n", - "We make the regularization coefficient a free parameter across every analysis object, ensuring different levels\n", - "of regularization are applied to each wavelength of data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = []\n", - "\n", - "for analysis in analysis_list:\n", - " model_analysis = model.copy()\n", - " model_analysis.galaxies.source.pixelization.regularization.coefficient = (\n", - " af.LogUniformPrior(lower_limit=1e-4, upper_limit=1e4)\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", - "\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` of the model shows us there are two models each with their own regularization coefficient as a free parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"multi\") / \"features\",\n", - " name=\"pixelized\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", - "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_list[0].max_log_likelihood_instance)\n", - "print(result_list[1].max_log_likelihood_instance)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting each result's tracer shows that the source appears different, owning to its different intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", - "\n", - "Therefore, the samples is identical in every result object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Pixelized\n", + "===================\n", + "\n", + "This script fits a multi-wavelength `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's surface-brightness is an `Inversion`.\n", + "\n", + "Two images are fitted, corresponding to a greener ('g' band) redder image (`r` band).\n", + "\n", + "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for lens modeling. Thus,\n", + "certain parts of code are not documented to ensure the script is concise.\n", + "\n", + "__Contents__\n", + "\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Positions:** This fit also uses the arc-second positions of the multiply imaged lensed source galaxy, which were.\n", + "- **Mesh Shape:** As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Result:** Overview of the results of the model-fit." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "\n", + "The strings are used for load each dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixel Scales__\n", + "\n", + "Every multi-wavelength dataset can have its own unique pixel-scale." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"simple__no_lens_light\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/multi/features/pixelization/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "imaging_dataset_path = Path(\"dataset\") / \"imaging\" / dataset_name\n", + "if not imaging_dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/imaging/features/no_lens_light/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "if not (dataset_path / \"positions.json\").exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [\n", + " sys.executable,\n", + " \"scripts/imaging/data_preparation/examples/optional/positions.py\",\n", + " ],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_list = [\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " pixel_scales=pixel_scales,\n", + " )\n", + " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", + "\n", + "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", + "absolutely necessary, but provides a more reliable analysis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_list = [\n", + " al.Mask2D.circular(\n", + " shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=3.0\n", + " )\n", + " for dataset in dataset_list\n", + "]\n", + "\n", + "dataset_list = [\n", + " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "This fit also uses the arc-second positions of the multiply imaged lensed source galaxy, which were drawn onto the\n", + "image via the GUI described in the file `autolens_workspace/*/imaging/data_preparation/gui/positions.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = al.Grid2DIrregular(\n", + " al.from_json(file_path=Path(dataset_path, \"positions.json\"))\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mesh Shape__\n", + "\n", + "As discussed in the `features/pixelization/modeling` example, the mesh shape is fixed before modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We create an `Analysis` object for every dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [al.AnalysisImaging(dataset=dataset) for dataset in dataset_list]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - This pixelization is regularized using a `Constant` scheme which smooths every source pixel \n", + " equally, where its `regularization_coefficient` varies across the datasets [2 parameter]. \n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=9." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, shear=al.mp.ExternalShear\n", + ")\n", + "\n", + "pixelization = af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization=al.reg.Constant,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now combine them using the factor analysis class, which allows us to fit the two datasets simultaneously.\n", + "\n", + "We make the regularization coefficient a free parameter across every analysis object, ensuring different levels\n", + "of regularization are applied to each wavelength of data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = []\n", + "\n", + "for analysis in analysis_list:\n", + " model_analysis = model.copy()\n", + " model_analysis.galaxies.source.pixelization.regularization.coefficient = (\n", + " af.LogUniformPrior(lower_limit=1e-4, upper_limit=1e4)\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", + "\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` of the model shows us there are two models each with their own regularization coefficient as a free parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"multi\") / \"features\",\n", + " name=\"pixelized\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", + "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_list[0].max_log_likelihood_instance)\n", + "print(result_list[1].max_log_likelihood_instance)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting each result's tracer shows that the source appears different, owning to its different intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", + "\n", + "Therefore, the samples is identical in every result object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/pixelization/simulator.ipynb b/notebooks/multi/features/pixelization/simulator.ipynb index 159b687b8..af914569c 100644 --- a/notebooks/multi/features/pixelization/simulator.ipynb +++ b/notebooks/multi/features/pixelization/simulator.ipynb @@ -1,451 +1,488 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: SIE\n", - "==============\n", - "\n", - "This script simulates multi-wavelength `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", - "\n", - "Two images are simulated, corresponding to a greener ('g' band) redder image (`r` band).\n", - "\n", - "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for simulating images. Thus,\n", - "certain parts of code are not documented to ensure the script is concise.\n", - "\n", - "__Contents__\n", - "\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "- **Dataset Paths:** Overview of dataset paths for this example.\n", - "- **Simulate:** The pixel-scale of each color image is different meaning we make a list of grids for the simulation.\n", - "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) for this simulated lens.\n", - "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "\n", - "The strings are used for naming the datasets on output." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"simple__no_lens_light\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "The pixel-scale of each color image is different meaning we make a list of grids for the simulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12]\n", - "\n", - "grid_list = []\n", - "\n", - "for pixel_scales in pixel_scales_list:\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=(150, 150),\n", - " pixel_scales=pixel_scales,\n", - " )\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " grid_list.append(grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate simple Gaussian PSFs for the images in the r and g bands." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sigma_list = [0.1, 0.2]\n", - "\n", - "psf_list = [\n", - " al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", - " )\n", - " for grid, sigma in zip(grid_list, sigma_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create separate simulators for the g and r bands." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "background_sky_level_list = [0.1, 0.15]\n", - "\n", - "simulator_list = [\n", - " al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=background_sky_level,\n", - " add_poisson_noise_to_data=True,\n", - " )\n", - " for psf, background_sky_level in zip(psf_list, background_sky_level_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's mass (SIE+Shear) for this simulated lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The source galaxy at each wavelength has a different intensity, thus we create two source galaxies for each waveband." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "intensity_list = [0.3, 0.2]\n", - "\n", - "source_galaxy_list = [\n", - " al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=intensity,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - " )\n", - " for intensity in intensity_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging` \n", - "dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_list = [\n", - " al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - " for source_galaxy in source_galaxy_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for tracer, grid in zip(tracer_list, grid_list):\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - " for grid, simulator, tracer in zip(grid_list, simulator_list, tracer_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " overwrite=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - " aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "for waveband, grid, tracer in zip(waveband_list, grid_list, tracer_list):\n", - " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", - " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "[\n", - " al.output_to_json(\n", - " obj=tracer, file_path=Path(dataset_path, f\"{waveband}_tracer.json\")\n", - " )\n", - " for color, tracer in zip(waveband_list, tracer_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Positions__\n", - "\n", - "Use a `PointSolver` to compute the multiple image positions of the source galaxy and output them\n", - "to `positions.json`. Pixelization-based lens modeling uses these positions to penalise solutions\n", - "where the source plane reconstruction does not produce the observed multiple images.\n", - "\n", - "The positions depend only on the lens mass and source-plane coordinate, so a single waveband's\n", - "grid and tracer is sufficient." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.PointSolver.for_grid(\n", - " grid=grid_list[0], pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer_list[0],\n", - " source_plane_coordinate=source_galaxy_list[0].bulge.centre,\n", - ")\n", - "\n", - "al.output_to_json(\n", - " file_path=dataset_path / \"positions.json\",\n", - " obj=positions,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: SIE\n", + "==============\n", + "\n", + "This script simulates multi-wavelength `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", + "\n", + "Two images are simulated, corresponding to a greener ('g' band) redder image (`r` band).\n", + "\n", + "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for simulating images. Thus,\n", + "certain parts of code are not documented to ensure the script is concise.\n", + "\n", + "__Contents__\n", + "\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "- **Dataset Paths:** Overview of dataset paths for this example.\n", + "- **Simulate:** The pixel-scale of each color image is different meaning we make a list of grids for the simulation.\n", + "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) for this simulated lens.\n", + "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "\n", + "The strings are used for naming the datasets on output." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"simple__no_lens_light\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "The pixel-scale of each color image is different meaning we make a list of grids for the simulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12]\n", + "\n", + "grid_list = []\n", + "\n", + "for pixel_scales in pixel_scales_list:\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=(150, 150),\n", + " pixel_scales=pixel_scales,\n", + " )\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " grid_list.append(grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate simple Gaussian PSFs for the images in the r and g bands." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sigma_list = [0.1, 0.2]\n", + "\n", + "psf_list = [\n", + " al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", + " )\n", + " for grid, sigma in zip(grid_list, sigma_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create separate simulators for the g and r bands." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "background_sky_level_list = [0.1, 0.15]\n", + "\n", + "simulator_list = [\n", + " al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=background_sky_level,\n", + " add_poisson_noise_to_data=True,\n", + " )\n", + " for psf, background_sky_level in zip(psf_list, background_sky_level_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's mass (SIE+Shear) for this simulated lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The source galaxy at each wavelength has a different intensity, thus we create two source galaxies for each waveband." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "intensity_list = [0.3, 0.2]\n", + "\n", + "source_galaxy_list = [\n", + " al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=intensity,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + " )\n", + " for intensity in intensity_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging` \n", + "dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_list = [\n", + " al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + " for source_galaxy in source_galaxy_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for tracer, grid in zip(tracer_list, grid_list):\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + " for grid, simulator, tracer in zip(grid_list, simulator_list, tracer_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " overwrite=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + " aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "for waveband, grid, tracer in zip(waveband_list, grid_list, tracer_list):\n", + " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", + " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "[\n", + " al.output_to_json(\n", + " obj=tracer, file_path=Path(dataset_path, f\"{waveband}_tracer.json\")\n", + " )\n", + " for color, tracer in zip(waveband_list, tracer_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Positions__\n", + "\n", + "Use a `PointSolver` to compute the multiple image positions of the source galaxy and output them\n", + "to `positions.json`. Pixelization-based lens modeling uses these positions to penalise solutions\n", + "where the source plane reconstruction does not produce the observed multiple images.\n", + "\n", + "The positions depend only on the lens mass and source-plane coordinate, so a single waveband's\n", + "grid and tracer is sufficient." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.PointSolver.for_grid(\n", + " grid=grid_list[0], pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer_list[0],\n", + " source_plane_coordinate=source_galaxy_list[0].bulge.centre,\n", + ")\n", + "\n", + "al.output_to_json(\n", + " file_path=dataset_path / \"positions.json\",\n", + " obj=positions,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/same_wavelength/modeling.ipynb b/notebooks/multi/features/same_wavelength/modeling.ipynb index 794009f01..c16b99ed7 100644 --- a/notebooks/multi/features/same_wavelength/modeling.ipynb +++ b/notebooks/multi/features/same_wavelength/modeling.ipynb @@ -1,423 +1,460 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Same Wavelength\n", - "=========================\n", - "\n", - "This script fits a multiple `Imaging` datasets observed at the same wavelength of a 'galaxy-scale' strong lens with a\n", - "model where:\n", - "\n", - " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an MGE.\n", - "\n", - "This script demonstrates how PyAutoLens's multi-dataset modeling tools can also simultaneously analyse datasets\n", - "observed at the same wavelength.\n", - "\n", - "An example use case might be analysing undithered HST images before they are combined via the multidrizzing process,\n", - "to remove correlated noise in the data.\n", - "\n", - "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for lens modeling. Thus,\n", - "certain parts of code are not documented to ensure the script is concise.\n", - "\n", - "__Contents__\n", - "\n", - "- **Pixel Scales:** If observed at the same wavelength, it is likely the datasets have the same pixel-scale.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Result:** Overview of the results of the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixel Scales__\n", - "\n", - "If observed at the same wavelength, it is likely the datasets have the same pixel-scale.\n", - "\n", - "Nevertheless, we specify this as a list as there could be an exception." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.1, 0.1]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"same_wavelength\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/multi/features/same_wavelength/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset_list = [\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_path, f\"image_{i}.fits\"),\n", - " psf_path=Path(dataset_path, f\"psf_{i}.fits\"),\n", - " noise_map_path=Path(dataset_path, f\"noise_map_{i}.fits\"),\n", - " pixel_scales=pixel_scales,\n", - " )\n", - " for i, pixel_scales in enumerate(pixel_scales_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", - "\n", - "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", - "absolutely necessary, but provides a more reliable analysis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask_list = [\n", - " al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - " for dataset in dataset_list\n", - "]\n", - "\n", - "dataset_list = [\n", - " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We create an `Analysis` object for every dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [al.AnalysisImaging(dataset=dataset) for dataset in dataset_list]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "shear = af.Model(al.mp.ExternalShear)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", - "\n", - "# Source:\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now combine them using the factor analysis class, which allows us to fit the two datasets simultaneously.\n", - "\n", - "Unlike other examples, no customization to the model is applied that, for example, adds more free parameters,\n", - "given that the datasets do not vary over wavelength." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = []\n", - "\n", - "for analysis in analysis_list:\n", - " bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - " )\n", - " disk = af.Model(al.lp_linear.Exponential)\n", - "\n", - " galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, disk=disk)\n", - "\n", - " model_analysis = af.Collection(galaxies=af.Collection(galaxy=galaxy))\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", - "\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` of the model shows us there are two models each with linear light profiles." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"multi\", \"modeling\"),\n", - " name=\"same_wavelength\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", - "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_list[0].max_log_likelihood_instance)\n", - "print(result_list[1].max_log_likelihood_instance)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting each result's tracer shows that the source appears different, owning to its different intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", - "\n", - "Therefore, the samples is identical in every result object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result_list.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Same Wavelength\n", + "=========================\n", + "\n", + "This script fits a multiple `Imaging` datasets observed at the same wavelength of a 'galaxy-scale' strong lens with a\n", + "model where:\n", + "\n", + " - The lens galaxy's light is omitted (and is not present in the simulated data).\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an MGE.\n", + "\n", + "This script demonstrates how PyAutoLens's multi-dataset modeling tools can also simultaneously analyse datasets\n", + "observed at the same wavelength.\n", + "\n", + "An example use case might be analysing undithered HST images before they are combined via the multidrizzing process,\n", + "to remove correlated noise in the data.\n", + "\n", + "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for lens modeling. Thus,\n", + "certain parts of code are not documented to ensure the script is concise.\n", + "\n", + "__Contents__\n", + "\n", + "- **Pixel Scales:** If observed at the same wavelength, it is likely the datasets have the same pixel-scale.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Result:** Overview of the results of the model-fit." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixel Scales__\n", + "\n", + "If observed at the same wavelength, it is likely the datasets have the same pixel-scale.\n", + "\n", + "Nevertheless, we specify this as a list as there could be an exception." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.1, 0.1]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"same_wavelength\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/multi/features/same_wavelength/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_list = [\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_path, f\"image_{i}.fits\"),\n", + " psf_path=Path(dataset_path, f\"psf_{i}.fits\"),\n", + " noise_map_path=Path(dataset_path, f\"noise_map_{i}.fits\"),\n", + " pixel_scales=pixel_scales,\n", + " )\n", + " for i, pixel_scales in enumerate(pixel_scales_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", + "\n", + "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", + "absolutely necessary, but provides a more reliable analysis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask_list = [\n", + " al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + " for dataset in dataset_list\n", + "]\n", + "\n", + "dataset_list = [\n", + " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We create an `Analysis` object for every dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [al.AnalysisImaging(dataset=dataset) for dataset in dataset_list]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=14." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "shear = af.Model(al.mp.ExternalShear)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass, shear=shear)\n", + "\n", + "# Source:\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now combine them using the factor analysis class, which allows us to fit the two datasets simultaneously.\n", + "\n", + "Unlike other examples, no customization to the model is applied that, for example, adds more free parameters,\n", + "given that the datasets do not vary over wavelength." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = []\n", + "\n", + "for analysis in analysis_list:\n", + " bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + " )\n", + " disk = af.Model(al.lp_linear.Exponential)\n", + "\n", + " galaxy = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, disk=disk)\n", + "\n", + " model_analysis = af.Collection(galaxies=af.Collection(galaxy=galaxy))\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", + "\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` of the model shows us there are two models each with linear light profiles." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"multi\", \"modeling\"),\n", + " name=\"same_wavelength\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", + "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_list[0].max_log_likelihood_instance)\n", + "print(result_list[1].max_log_likelihood_instance)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting each result's tracer shows that the source appears different, owning to its different intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", + "\n", + "Therefore, the samples is identical in every result object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result_list.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/same_wavelength/simulator.ipynb b/notebooks/multi/features/same_wavelength/simulator.ipynb index b27511993..f87041329 100644 --- a/notebooks/multi/features/same_wavelength/simulator.ipynb +++ b/notebooks/multi/features/same_wavelength/simulator.ipynb @@ -1,398 +1,435 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Wavelength Dependent\n", - "===============================\n", - "\n", - "This script simulates multiple `Imaging` datasets of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", - "\n", - "Unlike other `multi` simulators, all datasets are at the same wavelength and therefore the source does not change\n", - "its appearance in each dataset.\n", - "\n", - "This dataset demonstrates how PyAutoLens's multi-dataset modeling tools can also simultaneously analyse datasets\n", - "observed at the same wavelength.\n", - "\n", - "An example use case might be analysing undithered HST images before they are combined via the multidrizzing process,\n", - "to remove correlated noise in the data.\n", - "\n", - "TODO: NEED TO INCLUDE DIFFERENT POINTING / CENTERINGS.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** Overview of dataset paths for this example.\n", - "- **Simulate:** If observed at the same wavelength, it is likely the datasets have the same pixel-scale.\n", - "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) for this simulated lens.\n", - "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"same_wavelength\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "If observed at the same wavelength, it is likely the datasets have the same pixel-scale.\n", - "\n", - "Nevertheless, we specify this as a list as there could be an exception." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.1, 0.1]\n", - "\n", - "grid_list = []\n", - "\n", - "for pixel_scales in pixel_scales_list:\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=(150, 150),\n", - " pixel_scales=pixel_scales,\n", - " )\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " grid_list.append(grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate simple Gaussian PSFs for the images, which we assume slightly vary (e.g. due to different bserving conditions\n", - "for each image)" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sigma_list = [0.09, 0.11]\n", - "\n", - "psf_list = [\n", - " al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", - " )\n", - " for grid, sigma in zip(grid_list, sigma_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create separate simulators for the images, which we will assume have slightly different exposure times and background\n", - "sky levels." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "exposure_time_list = [300.0, 350.0]\n", - "background_sky_level_list = [0.1, 0.12]\n", - "\n", - "simulator_list = [\n", - " al.SimulatorImaging(\n", - " exposure_time=exposure_time,\n", - " psf=psf,\n", - " background_sky_level=background_sky_level,\n", - " add_poisson_noise_to_data=True,\n", - " )\n", - " for psf, exposure_time, background_sky_level in zip(\n", - " psf_list, exposure_time_list, background_sky_level_list\n", - " )\n", - "]\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's mass (SIE+Shear) for this simulated lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The source galaxy is observed att he same wavelength in each image thus its intensity does not vary across the datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=4.0,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging` \n", - "dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for grid in grid_list:\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - " for grid, simulator in zip(grid_list, simulator_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for i, dataset in enumerate(dataset_list):\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=Path(dataset_path, f\"image_{i}.fits\"),\n", - " psf_path=Path(dataset_path, f\"psf_{i}.fits\"),\n", - " noise_map_path=Path(dataset_path, f\"noise_map_{i}.fits\"),\n", - " overwrite=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for i, dataset in enumerate(dataset_list):\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - " aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "for i, grid in enumerate(grid_list):\n", - " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", - " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/same_wavelength/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Wavelength Dependent\n", + "===============================\n", + "\n", + "This script simulates multiple `Imaging` datasets of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", + "\n", + "Unlike other `multi` simulators, all datasets are at the same wavelength and therefore the source does not change\n", + "its appearance in each dataset.\n", + "\n", + "This dataset demonstrates how PyAutoLens's multi-dataset modeling tools can also simultaneously analyse datasets\n", + "observed at the same wavelength.\n", + "\n", + "An example use case might be analysing undithered HST images before they are combined via the multidrizzing process,\n", + "to remove correlated noise in the data.\n", + "\n", + "TODO: NEED TO INCLUDE DIFFERENT POINTING / CENTERINGS.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** Overview of dataset paths for this example.\n", + "- **Simulate:** If observed at the same wavelength, it is likely the datasets have the same pixel-scale.\n", + "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) for this simulated lens.\n", + "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"same_wavelength\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "If observed at the same wavelength, it is likely the datasets have the same pixel-scale.\n", + "\n", + "Nevertheless, we specify this as a list as there could be an exception." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.1, 0.1]\n", + "\n", + "grid_list = []\n", + "\n", + "for pixel_scales in pixel_scales_list:\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=(150, 150),\n", + " pixel_scales=pixel_scales,\n", + " )\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " grid_list.append(grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate simple Gaussian PSFs for the images, which we assume slightly vary (e.g. due to different bserving conditions\n", + "for each image)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sigma_list = [0.09, 0.11]\n", + "\n", + "psf_list = [\n", + " al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", + " )\n", + " for grid, sigma in zip(grid_list, sigma_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create separate simulators for the images, which we will assume have slightly different exposure times and background\n", + "sky levels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "exposure_time_list = [300.0, 350.0]\n", + "background_sky_level_list = [0.1, 0.12]\n", + "\n", + "simulator_list = [\n", + " al.SimulatorImaging(\n", + " exposure_time=exposure_time,\n", + " psf=psf,\n", + " background_sky_level=background_sky_level,\n", + " add_poisson_noise_to_data=True,\n", + " )\n", + " for psf, exposure_time, background_sky_level in zip(\n", + " psf_list, exposure_time_list, background_sky_level_list\n", + " )\n", + "]\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's mass (SIE+Shear) for this simulated lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The source galaxy is observed att he same wavelength in each image thus its intensity does not vary across the datasets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging` \n", + "dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for grid in grid_list:\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + " for grid, simulator in zip(grid_list, simulator_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for i, dataset in enumerate(dataset_list):\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=Path(dataset_path, f\"image_{i}.fits\"),\n", + " psf_path=Path(dataset_path, f\"psf_{i}.fits\"),\n", + " noise_map_path=Path(dataset_path, f\"noise_map_{i}.fits\"),\n", + " overwrite=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for i, dataset in enumerate(dataset_list):\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + " aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "for i, grid in enumerate(grid_list):\n", + " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", + " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/same_wavelength/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/slam/independent.ipynb b/notebooks/multi/features/slam/independent.ipynb index d9cb422a9..a08a53dd5 100644 --- a/notebooks/multi/features/slam/independent.ipynb +++ b/notebooks/multi/features/slam/independent.ipynb @@ -1,967 +1,1004 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM: Multi Wavelength Independent\n", - "==================================\n", - "\n", - "This example shows how to use the SLaM pipeline to fit a lens dataset at one wavelength, and then fit other data of\n", - "the same lens at different wavelengths with the same mass model fitted to the original dataset.\n", - "\n", - "The first dataset fitted is regarded as the \"main\" dataset, meaning it should have the highest resolution and\n", - "signal-to-noise. This will ensure the mass model is the most accurate.\n", - "\n", - "The remaining datasets are then fitted, which may be similar quality to the main dataset or lower resolution and\n", - "signal-to-noise. These datasets are fitted with the following approach:\n", - "\n", - "- The mass model (e.g. SIE +Shear) is fixed to the result of the VIS fit.\n", - "\n", - "- The lens light (Multi Gaussian Expansion) has the `intensity` values of the Gaussians updated using linear algebra.\n", - " to capture changes in the lens light over wavelength, but it does not update the Gaussian parameters (e.g. `centre`,\n", - " `elliptical_comps`, `sigma`) themselves due to the lower resolution of the data.\n", - "\n", - "- The source reconstruction (RectangularAdaptDensity adaptive mesh) is updated using linear algebra to reconstruct\n", - " the source, but again fixes the source pixelization parameters themselves.\n", - "\n", - "- Sub-pixel offsets between the datasets are fully modeled as free parameters, because the precision of a lens model\n", - "can often be less than the requirements on astrometry.\n", - "\n", - "The restrictive nature of the lens mass, light and source models mean that much lower quality multi-wavelength data\n", - "can be fitted provided the first dataset is of high quality. This is key for upcoming surveys such as Euclid, where\n", - "the VIS instrument will be high resolution but many other wavebands will be lower resolution.\n", - "\n", - "The subsequent fits to the lower resolution data use a reduced and simplified SLaM pipeline with the mass model\n", - "fixed to the result of the VIS fit.\n", - "\n", - "__Contents__\n", - "\n", - "- **Preqrequisites:** Before using this SLaM pipeline, you should be familiar with.\n", - "- **This Script:** Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **SLaM Pipeline Functions:** Overview of slam pipeline functions for this example.\n", - "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", - "- **Second Dataset Fits:** We now fit the secondary multi-wavelength datasets, which are lower resolution than the main.\n", - "- **Dataset Wavebands:** The following list gives the names of the wavebands we are going to fit.\n", - "- **Result Dict:** The results of each fit are stored in a dictionary, which is used to pass the results of each fit.\n", - "\n", - "__Preqrequisites__\n", - "\n", - "Before using this SLaM pipeline, you should be familiar with:\n", - "\n", - "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", - " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", - " and how they integrate into strong-lens modeling.\n", - "\n", - "- **Multi**: The `autolens_workspace/*/advanced/multi` package describes many different ways that multiple datasets\n", - " can be modeled in a single analysis, including the example script `one_by_one.ipynb` which fits a primary dataset\n", - " and then follows it up with fits to lower resolution datasets.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this SLaM modeling\n", - "script fits `Imaging` dataset of a strong lens system where in the final model:\n", - "\n", - " - The lens galaxy's light is a bulge with Multiple Gaussian Expansion (MGE) light profile.\n", - " - The lens galaxy's total mass distribution is an `PowerLaw` plus an `ExternalShear`.\n", - " - The source galaxy's light is a `Pixelization`.\n", - " - Two extra galaxies are included in the model, each with their light represented as a bulge with MGE light profile\n", - " and their mass as a `IsothermalSph` profile.\n", - "\n", - "This modeling script uses the SLaM pipelines:\n", - "\n", - " `source_lp`\n", - " `source_pix`\n", - " `light_lp`\n", - " `mass_total`\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "Everything below is identical to `start_here.py` and thus not commented, as it is the same code.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__ \n", - "\n", - "Load, plot and mask the `Imaging` data.\n", - "\n", - "We load a dataset with the waveband \"g\", which is the highest resolution data in this multi-wavelength example. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"lens_sersic\"\n", - "dataset_main_path = Path(\"dataset\", \"multi\", \"imaging\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_main_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/multi/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset_waveband = \"g\"\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=Path(dataset_main_path, f\"{dataset_waveband}_data.fits\"),\n", - " noise_map_path=Path(dataset_main_path, f\"{dataset_waveband}_noise_map.fits\"),\n", - " psf_path=Path(dataset_main_path, f\"{dataset_waveband}_psf.fits\"),\n", - " pixel_scales=0.08,\n", - ")\n", - "\n", - "mask_radius = 3.0\n", - "\n", - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"slam\", \"multi\", \"independent\"),\n", - " unique_tag=f\"{dataset_name}_data_{dataset_waveband}\",\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline Functions__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search,\n", - " analysis,\n", - " lens_bulge,\n", - " source_bulge,\n", - " redshift_lens,\n", - " redshift_source,\n", - " mass_centre=(0.0, 0.0),\n", - " n_batch=50,\n", - "):\n", - " \"\"\"\n", - " SOURCE LP PIPELINE: fits an initial lens model using a parametric source to establish a robust\n", - " lens light, mass and source model for the main high-resolution dataset.\n", - " \"\"\"\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = mass_centre\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def source_pix_1(\n", - " settings_search,\n", - " analysis,\n", - " source_lp_result,\n", - " mesh_shape,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SOURCE PIX PIPELINE 1: initializes a pixelized source model with mass priors from SOURCE LP PIPELINE.\n", - " \"\"\"\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result.model.galaxies.lens.mass,\n", - " mass_result=source_lp_result.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_lp_result.model.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def source_pix_2(\n", - " settings_search,\n", - " analysis,\n", - " source_lp_result,\n", - " source_pix_result_1,\n", - " mesh_shape,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SOURCE PIX PIPELINE 2: fits an improved pixelized source using adapt images from SOURCE PIX PIPELINE 1,\n", - " with fixed lens mass.\n", - " \"\"\"\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def light_lp(\n", - " settings_search,\n", - " analysis,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " lens_bulge,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " LIGHT LP PIPELINE: fits the lens galaxy light with mass and source fixed from SOURCE PIX PIPELINE.\n", - " \"\"\"\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source, source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def mass_total(\n", - " settings_search,\n", - " analysis,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " light_result,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " MASS TOTAL PIPELINE: fits a PowerLaw total mass model with priors from SOURCE PIX PIPELINE and\n", - " lens light fixed from LIGHT LP PIPELINE.\n", - " \"\"\"\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " mass = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(result=source_result_for_source)\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", - " bulge=light_result.instance.galaxies.lens.bulge,\n", - " disk=light_result.instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=source_result_for_lens.model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def source_lp_secondary(\n", - " settings_search,\n", - " analysis,\n", - " light_result,\n", - " mass_result,\n", - " source_bulge,\n", - " dataset_model,\n", - " redshift_lens=0.5,\n", - " redshift_source=1.0,\n", - " n_batch=50,\n", - "):\n", - " \"\"\"\n", - " SOURCE LP PIPELINE (secondary dataset): fits the source for a secondary waveband dataset with the lens\n", - " light and mass fixed to the results of the main dataset pipeline.\n", - " \"\"\"\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=light_result.instance.galaxies.lens.bulge,\n", - " disk=None,\n", - " mass=mass_result.instance.galaxies.lens.mass,\n", - " shear=mass_result.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def source_pix_1_secondary(\n", - " settings_search,\n", - " analysis,\n", - " source_lp_result,\n", - " mesh_shape,\n", - " dataset_model,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SOURCE PIX PIPELINE 1 (secondary dataset): initializes a pixelized source with fixed mass from\n", - " the main dataset result, updating source reconstruction and dataset offsets.\n", - " \"\"\"\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_lp_result.instance.galaxies.lens.mass,\n", - " shear=source_lp_result.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", - "\n", - "\n", - "def source_pix_2_secondary(\n", - " settings_search,\n", - " analysis,\n", - " source_lp_result,\n", - " source_pix_result_1,\n", - " mesh_shape,\n", - " dataset_model,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SOURCE PIX PIPELINE 2 (secondary dataset): fits an improved pixelized source using adapt images from\n", - " SOURCE PIX PIPELINE 1, with all lens parameters fixed.\n", - " \"\"\"\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", - " disk=source_lp_result.instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result.instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", - "\n", - "# Lens Light\n", - "\n", - "lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "# Source Light\n", - "\n", - "source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", - "\n", - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " lens_bulge=lens_bulge,\n", - " source_bulge=source_bulge,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_lp_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - " use_jax=True,\n", - ")\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_lp_result=source_lp_result,\n", - " mesh_shape=mesh_shape,\n", - ")\n", - "\n", - "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - ")\n", - "\n", - "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - ")\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh_shape=mesh_shape,\n", - ")\n", - "\n", - "lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - ")\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " lens_bulge=lens_bulge,\n", - ")\n", - "\n", - "analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[\n", - " source_pix_result_2.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", - " ],\n", - ")\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Second Dataset Fits__\n", - "\n", - "We now fit the secondary multi-wavelength datasets, which are lower resolution than the main dataset.\n", - "\n", - "This uses a for loop to iterate over every waveband of every dataset, load and mask the data and fit it.\n", - "\n", - "Each fit uses a fixed mass model, the lens and source light models update via linear algebra and offsets are\n", - "included (see full description above)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"lens_sersic\"\n", - "dataset_main_path = Path(\"dataset\", \"multi\", \"imaging\", dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Wavebands__\n", - "\n", - "The following list gives the names of the wavebands we are going to fit.\n", - "\n", - "The data for each waveband is loaded from a folder in the dataset folder with that name." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_waveband_list = [\"r\"]\n", - "pixel_scale_list = [0.12]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result Dict__\n", - "\n", - "The results of each fit are stored in a dictionary, which is used to pass the results of each fit to the\n", - "visualization functions." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "multi_result_dict = {\"g\": mass_result}\n", - "\n", - "for dataset_waveband, pixel_scale in zip(dataset_waveband_list, pixel_scale_list):\n", - "\n", - " dataset = al.Imaging.from_fits(\n", - " data_path=Path(dataset_main_path, f\"{dataset_waveband}_data.fits\"),\n", - " noise_map_path=Path(dataset_main_path, f\"{dataset_waveband}_noise_map.fits\"),\n", - " psf_path=Path(dataset_main_path, f\"{dataset_waveband}_psf.fits\"),\n", - " pixel_scales=pixel_scale,\n", - " )\n", - "\n", - " mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=mask)\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.1, 0.3],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"slam\", \"multi\", \"independent\"),\n", - " unique_tag=f\"{dataset_name}_data_{dataset_waveband}\",\n", - " info=None,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Dataset Model__\n", - "\n", - " For each secondary dataset, the (y,x) offset relative to the primary data is a free parameter.\n", - " \"\"\"\n", - " dataset_model = af.Model(al.DatasetModel)\n", - "\n", - " dataset_model.grid_offset.grid_offset_0 = af.UniformPrior(\n", - " lower_limit=-0.2, upper_limit=0.2\n", - " )\n", - " dataset_model.grid_offset.grid_offset_1 = af.UniformPrior(\n", - " lower_limit=-0.2, upper_limit=0.2\n", - " )\n", - "\n", - " centre_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - " centre_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - "\n", - " total_gaussians = 20\n", - " gaussian_per_basis = 1\n", - "\n", - " log10_sigma_list = np.linspace(-3, np.log10(1.0), total_gaussians)\n", - "\n", - " bulge_gaussian_list = []\n", - "\n", - " for j in range(gaussian_per_basis):\n", - " gaussian_list = af.Collection(\n", - " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", - " )\n", - "\n", - " for i, gaussian in enumerate(gaussian_list):\n", - " gaussian.centre.centre_0 = centre_0\n", - " gaussian.centre.centre_1 = centre_1\n", - " gaussian.ell_comps = gaussian_list[0].ell_comps\n", - " gaussian.sigma = 10 ** log10_sigma_list[i]\n", - "\n", - " bulge_gaussian_list += gaussian_list\n", - "\n", - " source_bulge = af.Model(\n", - " al.lp_basis.Basis,\n", - " profile_list=bulge_gaussian_list,\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(dataset=dataset)\n", - "\n", - " source_lp_result = source_lp_secondary(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " light_result=light_result,\n", - " mass_result=mass_result,\n", - " source_bulge=source_bulge,\n", - " dataset_model=dataset_model,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - " )\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_lp_result\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " dataset_model.grid_offset.grid_offset_0 = (\n", - " source_lp_result.instance.dataset_model.grid_offset[0]\n", - " )\n", - " dataset_model.grid_offset.grid_offset_1 = (\n", - " source_lp_result.instance.dataset_model.grid_offset[1]\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " raise_inversion_positions_likelihood_exception=False,\n", - " )\n", - "\n", - " source_pix_result_1 = source_pix_1_secondary(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_lp_result=source_lp_result,\n", - " mesh_shape=mesh_shape,\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " source_pix_result_1.max_log_likelihood_fit.inversion.cls_list_from(cls=al.Mapper)[\n", - " 0\n", - " ].extent_from()\n", - "\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", - " result=source_pix_result_1\n", - " )\n", - "\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - "\n", - " dataset_model.grid_offset.grid_offset_0 = (\n", - " source_lp_result.instance.dataset_model.grid_offset[0]\n", - " )\n", - " dataset_model.grid_offset.grid_offset_1 = (\n", - " source_lp_result.instance.dataset_model.grid_offset[1]\n", - " )\n", - "\n", - " analysis = al.AnalysisImaging(\n", - " dataset=dataset,\n", - " adapt_images=adapt_images,\n", - " )\n", - "\n", - " multi_result = source_pix_2_secondary(\n", - " settings_search=settings_search,\n", - " analysis=analysis,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh_shape=mesh_shape,\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " multi_result_dict[dataset_waveband] = multi_result\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM: Multi Wavelength Independent\n", + "==================================\n", + "\n", + "This example shows how to use the SLaM pipeline to fit a lens dataset at one wavelength, and then fit other data of\n", + "the same lens at different wavelengths with the same mass model fitted to the original dataset.\n", + "\n", + "The first dataset fitted is regarded as the \"main\" dataset, meaning it should have the highest resolution and\n", + "signal-to-noise. This will ensure the mass model is the most accurate.\n", + "\n", + "The remaining datasets are then fitted, which may be similar quality to the main dataset or lower resolution and\n", + "signal-to-noise. These datasets are fitted with the following approach:\n", + "\n", + "- The mass model (e.g. SIE +Shear) is fixed to the result of the VIS fit.\n", + "\n", + "- The lens light (Multi Gaussian Expansion) has the `intensity` values of the Gaussians updated using linear algebra.\n", + " to capture changes in the lens light over wavelength, but it does not update the Gaussian parameters (e.g. `centre`,\n", + " `elliptical_comps`, `sigma`) themselves due to the lower resolution of the data.\n", + "\n", + "- The source reconstruction (RectangularAdaptDensity adaptive mesh) is updated using linear algebra to reconstruct\n", + " the source, but again fixes the source pixelization parameters themselves.\n", + "\n", + "- Sub-pixel offsets between the datasets are fully modeled as free parameters, because the precision of a lens model\n", + "can often be less than the requirements on astrometry.\n", + "\n", + "The restrictive nature of the lens mass, light and source models mean that much lower quality multi-wavelength data\n", + "can be fitted provided the first dataset is of high quality. This is key for upcoming surveys such as Euclid, where\n", + "the VIS instrument will be high resolution but many other wavebands will be lower resolution.\n", + "\n", + "The subsequent fits to the lower resolution data use a reduced and simplified SLaM pipeline with the mass model\n", + "fixed to the result of the VIS fit.\n", + "\n", + "__Contents__\n", + "\n", + "- **Preqrequisites:** Before using this SLaM pipeline, you should be familiar with.\n", + "- **This Script:** Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **SLaM Pipeline Functions:** Overview of slam pipeline functions for this example.\n", + "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", + "- **Second Dataset Fits:** We now fit the secondary multi-wavelength datasets, which are lower resolution than the main.\n", + "- **Dataset Wavebands:** The following list gives the names of the wavebands we are going to fit.\n", + "- **Result Dict:** The results of each fit are stored in a dictionary, which is used to pass the results of each fit.\n", + "\n", + "__Preqrequisites__\n", + "\n", + "Before using this SLaM pipeline, you should be familiar with:\n", + "\n", + "- **SLaM Start Here** (`guides/modeling/slam_start_here`)\n", + " An introduction to the goals, structure, and design philosophy behind SLaM pipelines\n", + " and how they integrate into strong-lens modeling.\n", + "\n", + "- **Multi**: The `autolens_workspace/*/advanced/multi` package describes many different ways that multiple datasets\n", + " can be modeled in a single analysis, including the example script `one_by_one.ipynb` which fits a primary dataset\n", + " and then follows it up with fits to lower resolution datasets.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this SLaM modeling\n", + "script fits `Imaging` dataset of a strong lens system where in the final model:\n", + "\n", + " - The lens galaxy's light is a bulge with Multiple Gaussian Expansion (MGE) light profile.\n", + " - The lens galaxy's total mass distribution is an `PowerLaw` plus an `ExternalShear`.\n", + " - The source galaxy's light is a `Pixelization`.\n", + " - Two extra galaxies are included in the model, each with their light represented as a bulge with MGE light profile\n", + " and their mass as a `IsothermalSph` profile.\n", + "\n", + "This modeling script uses the SLaM pipelines:\n", + "\n", + " `source_lp`\n", + " `source_pix`\n", + " `light_lp`\n", + " `mass_total`\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "Everything below is identical to `start_here.py` and thus not commented, as it is the same code.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__ \n", + "\n", + "Load, plot and mask the `Imaging` data.\n", + "\n", + "We load a dataset with the waveband \"g\", which is the highest resolution data in this multi-wavelength example. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"lens_sersic\"\n", + "dataset_main_path = Path(\"dataset\", \"multi\", \"imaging\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_main_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/multi/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_waveband = \"g\"\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=Path(dataset_main_path, f\"{dataset_waveband}_data.fits\"),\n", + " noise_map_path=Path(dataset_main_path, f\"{dataset_waveband}_noise_map.fits\"),\n", + " psf_path=Path(dataset_main_path, f\"{dataset_waveband}_psf.fits\"),\n", + " pixel_scales=0.08,\n", + ")\n", + "\n", + "mask_radius = 3.0\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"slam\", \"multi\", \"independent\"),\n", + " unique_tag=f\"{dataset_name}_data_{dataset_waveband}\",\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline Functions__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search,\n", + " analysis,\n", + " lens_bulge,\n", + " source_bulge,\n", + " redshift_lens,\n", + " redshift_source,\n", + " mass_centre=(0.0, 0.0),\n", + " n_batch=50,\n", + "):\n", + " \"\"\"\n", + " SOURCE LP PIPELINE: fits an initial lens model using a parametric source to establish a robust\n", + " lens light, mass and source model for the main high-resolution dataset.\n", + " \"\"\"\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = mass_centre\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def source_pix_1(\n", + " settings_search,\n", + " analysis,\n", + " source_lp_result,\n", + " mesh_shape,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SOURCE PIX PIPELINE 1: initializes a pixelized source model with mass priors from SOURCE LP PIPELINE.\n", + " \"\"\"\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result.model.galaxies.lens.mass,\n", + " mass_result=source_lp_result.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_lp_result.model.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def source_pix_2(\n", + " settings_search,\n", + " analysis,\n", + " source_lp_result,\n", + " source_pix_result_1,\n", + " mesh_shape,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SOURCE PIX PIPELINE 2: fits an improved pixelized source using adapt images from SOURCE PIX PIPELINE 1,\n", + " with fixed lens mass.\n", + " \"\"\"\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def light_lp(\n", + " settings_search,\n", + " analysis,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " lens_bulge,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " LIGHT LP PIPELINE: fits the lens galaxy light with mass and source fixed from SOURCE PIX PIPELINE.\n", + " \"\"\"\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source, source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens.instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens.instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def mass_total(\n", + " settings_search,\n", + " analysis,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " light_result,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " MASS TOTAL PIPELINE: fits a PowerLaw total mass model with priors from SOURCE PIX PIPELINE and\n", + " lens light fixed from LIGHT LP PIPELINE.\n", + " \"\"\"\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " mass = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens.model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(result=source_result_for_source)\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens.instance.galaxies.lens.redshift,\n", + " bulge=light_result.instance.galaxies.lens.bulge,\n", + " disk=light_result.instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=source_result_for_lens.model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def source_lp_secondary(\n", + " settings_search,\n", + " analysis,\n", + " light_result,\n", + " mass_result,\n", + " source_bulge,\n", + " dataset_model,\n", + " redshift_lens=0.5,\n", + " redshift_source=1.0,\n", + " n_batch=50,\n", + "):\n", + " \"\"\"\n", + " SOURCE LP PIPELINE (secondary dataset): fits the source for a secondary waveband dataset with the lens\n", + " light and mass fixed to the results of the main dataset pipeline.\n", + " \"\"\"\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=light_result.instance.galaxies.lens.bulge,\n", + " disk=None,\n", + " mass=mass_result.instance.galaxies.lens.mass,\n", + " shear=mass_result.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def source_pix_1_secondary(\n", + " settings_search,\n", + " analysis,\n", + " source_lp_result,\n", + " mesh_shape,\n", + " dataset_model,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SOURCE PIX PIPELINE 1 (secondary dataset): initializes a pixelized source with fixed mass from\n", + " the main dataset result, updating source reconstruction and dataset offsets.\n", + " \"\"\"\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_lp_result.instance.galaxies.lens.mass,\n", + " shear=source_lp_result.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptDensity, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n", + "\n", + "\n", + "def source_pix_2_secondary(\n", + " settings_search,\n", + " analysis,\n", + " source_lp_result,\n", + " source_pix_result_1,\n", + " mesh_shape,\n", + " dataset_model,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SOURCE PIX PIPELINE 2 (secondary dataset): fits an improved pixelized source using adapt images from\n", + " SOURCE PIX PIPELINE 1, with all lens parameters fixed.\n", + " \"\"\"\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result.instance.galaxies.lens.bulge,\n", + " disk=source_lp_result.instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1.instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1.instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result.instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=model, analysis=analysis, **settings_search.fit_dict)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", + "\n", + "# Lens Light\n", + "\n", + "lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "# Source Light\n", + "\n", + "source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(dataset=dataset, use_jax=True)\n", + "\n", + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " lens_bulge=lens_bulge,\n", + " source_bulge=source_bulge,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_lp_result.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + " use_jax=True,\n", + ")\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_lp_result=source_lp_result,\n", + " mesh_shape=mesh_shape,\n", + ")\n", + "\n", + "galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + ")\n", + "\n", + "adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + ")\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh_shape=mesh_shape,\n", + ")\n", + "\n", + "lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + ")\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " lens_bulge=lens_bulge,\n", + ")\n", + "\n", + "analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[\n", + " source_pix_result_2.positions_likelihood_from(factor=3.0, minimum_threshold=0.2)\n", + " ],\n", + ")\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Second Dataset Fits__\n", + "\n", + "We now fit the secondary multi-wavelength datasets, which are lower resolution than the main dataset.\n", + "\n", + "This uses a for loop to iterate over every waveband of every dataset, load and mask the data and fit it.\n", + "\n", + "Each fit uses a fixed mass model, the lens and source light models update via linear algebra and offsets are\n", + "included (see full description above)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"lens_sersic\"\n", + "dataset_main_path = Path(\"dataset\", \"multi\", \"imaging\", dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Wavebands__\n", + "\n", + "The following list gives the names of the wavebands we are going to fit.\n", + "\n", + "The data for each waveband is loaded from a folder in the dataset folder with that name." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_waveband_list = [\"r\"]\n", + "pixel_scale_list = [0.12]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result Dict__\n", + "\n", + "The results of each fit are stored in a dictionary, which is used to pass the results of each fit to the\n", + "visualization functions." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "multi_result_dict = {\"g\": mass_result}\n", + "\n", + "for dataset_waveband, pixel_scale in zip(dataset_waveband_list, pixel_scale_list):\n", + "\n", + " dataset = al.Imaging.from_fits(\n", + " data_path=Path(dataset_main_path, f\"{dataset_waveband}_data.fits\"),\n", + " noise_map_path=Path(dataset_main_path, f\"{dataset_waveband}_noise_map.fits\"),\n", + " psf_path=Path(dataset_main_path, f\"{dataset_waveband}_psf.fits\"),\n", + " pixel_scales=pixel_scale,\n", + " )\n", + "\n", + " mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=mask)\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.1, 0.3],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"slam\", \"multi\", \"independent\"),\n", + " unique_tag=f\"{dataset_name}_data_{dataset_waveband}\",\n", + " info=None,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Dataset Model__\n", + "\n", + " For each secondary dataset, the (y,x) offset relative to the primary data is a free parameter.\n", + " \"\"\"\n", + " dataset_model = af.Model(al.DatasetModel)\n", + "\n", + " dataset_model.grid_offset.grid_offset_0 = af.UniformPrior(\n", + " lower_limit=-0.2, upper_limit=0.2\n", + " )\n", + " dataset_model.grid_offset.grid_offset_1 = af.UniformPrior(\n", + " lower_limit=-0.2, upper_limit=0.2\n", + " )\n", + "\n", + " centre_0 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + " centre_1 = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + "\n", + " total_gaussians = 20\n", + " gaussian_per_basis = 1\n", + "\n", + " log10_sigma_list = np.linspace(-3, np.log10(1.0), total_gaussians)\n", + "\n", + " bulge_gaussian_list = []\n", + "\n", + " for j in range(gaussian_per_basis):\n", + " gaussian_list = af.Collection(\n", + " af.Model(al.lp_linear.Gaussian) for _ in range(total_gaussians)\n", + " )\n", + "\n", + " for i, gaussian in enumerate(gaussian_list):\n", + " gaussian.centre.centre_0 = centre_0\n", + " gaussian.centre.centre_1 = centre_1\n", + " gaussian.ell_comps = gaussian_list[0].ell_comps\n", + " gaussian.sigma = 10 ** log10_sigma_list[i]\n", + "\n", + " bulge_gaussian_list += gaussian_list\n", + "\n", + " source_bulge = af.Model(\n", + " al.lp_basis.Basis,\n", + " profile_list=bulge_gaussian_list,\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(dataset=dataset)\n", + "\n", + " source_lp_result = source_lp_secondary(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " light_result=light_result,\n", + " mass_result=mass_result,\n", + " source_bulge=source_bulge,\n", + " dataset_model=dataset_model,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + " )\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_lp_result\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " dataset_model.grid_offset.grid_offset_0 = (\n", + " source_lp_result.instance.dataset_model.grid_offset[0]\n", + " )\n", + " dataset_model.grid_offset.grid_offset_1 = (\n", + " source_lp_result.instance.dataset_model.grid_offset[1]\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " raise_inversion_positions_likelihood_exception=False,\n", + " )\n", + "\n", + " source_pix_result_1 = source_pix_1_secondary(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_lp_result=source_lp_result,\n", + " mesh_shape=mesh_shape,\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " source_pix_result_1.max_log_likelihood_fit.inversion.cls_list_from(cls=al.Mapper)[\n", + " 0\n", + " ].extent_from()\n", + "\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(\n", + " result=source_pix_result_1\n", + " )\n", + "\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + "\n", + " dataset_model.grid_offset.grid_offset_0 = (\n", + " source_lp_result.instance.dataset_model.grid_offset[0]\n", + " )\n", + " dataset_model.grid_offset.grid_offset_1 = (\n", + " source_lp_result.instance.dataset_model.grid_offset[1]\n", + " )\n", + "\n", + " analysis = al.AnalysisImaging(\n", + " dataset=dataset,\n", + " adapt_images=adapt_images,\n", + " )\n", + "\n", + " multi_result = source_pix_2_secondary(\n", + " settings_search=settings_search,\n", + " analysis=analysis,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh_shape=mesh_shape,\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " multi_result_dict[dataset_waveband] = multi_result\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/slam/simultaneous.ipynb b/notebooks/multi/features/slam/simultaneous.ipynb index 00c7bee2a..bc2f05ed2 100644 --- a/notebooks/multi/features/slam/simultaneous.ipynb +++ b/notebooks/multi/features/slam/simultaneous.ipynb @@ -1,880 +1,917 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SLaM: Multi Wavelength Simultaneous\n", - "===================================\n", - "\n", - "This example shows how to use the SLaM pipeline to fit a lens dataset at multiple wavelengths simultaneously.\n", - "\n", - "Simultaneous multi-dataset fits are currently built into the SLaM pipeline without user input or customization.\n", - "Therefore, as long as lists of `Analysis` objects are created, summed and passed to the SLaM pipelines, the analysis\n", - "will fit every dataset simultaneously and it will adapt the model as follows:\n", - "\n", - "- Sub-pixel offsets between the datasets are fully modeled as free parameters in each stage of the pipeline, assuming\n", - " broad uniform priors for every step. This is because the precision of a lens model can often be less than the\n", - " requirements on astrometry.\n", - "\n", - "- The regularization parameters are free for every dataset in the `source_pix[1]` and `source_pix[2]` stages. This is because\n", - " the source morphology can be different between datasets, and the regularization scheme adapts to this.\n", - "\n", - "- From the `light_lp` stage onwards, the regularization scheme for each dataset is different fixed to that inferred\n", - " for the `source_pix[2]` stage.\n", - "\n", - "Simultaneous fitting SLaM pipelines are not designed for customization, for example changing the model from the\n", - "set up above. This is because we are still figuring out the best way to perform multi-wavelength modeling, but have\n", - "so far figured the above settings are important.\n", - "\n", - "If you need customization of the model or pipeline, you should pick apart the SLaM pipeline and customize\n", - "them as you see fit.\n", - "\n", - "__Contents__\n", - "\n", - "- **Preqrequisites:** Before reading this script, you should have familiarity with the following key concepts.\n", - "- **This Script:** Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", - "- **Redshifts:** The redshifts of the lens and source galaxies.\n", - "- **SLaM Pipeline Functions:** Overview of slam pipeline functions for this example.\n", - "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", - "\n", - "__Preqrequisites__\n", - "\n", - "Before reading this script, you should have familiarity with the following key concepts:\n", - "\n", - "- **Multi**: The `autolens_workspace/*/advanced/multi` package describes many different ways that multiple datasets\n", - " can be modeled in a single analysis.\n", - "\n", - "__This Script__\n", - "\n", - "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this SLaM modeling\n", - "script fits `Imaging` dataset of a strong lens system where in the final model:\n", - "\n", - " - The lens galaxy's light is a bulge with Multiple Gaussian Expansion (MGE) light profile.\n", - " - The lens galaxy's total mass distribution is an `PowerLaw` plus an `ExternalShear`.\n", - " - The source galaxy's light is a `Pixelization`.\n", - "\n", - "This modeling script uses the SLaM pipelines:\n", - "\n", - " `source_lp`\n", - " `source_pix`\n", - " `light_lp`\n", - " `mass_total`\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__ \n", - "\n", - "Load, plot and mask the `Imaging` data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_waveband_list = [\"g\", \"r\"]\n", - "pixel_scale_list = [0.12, 0.08]\n", - "\n", - "dataset_name = \"lens_sersic\"\n", - "dataset_main_path = Path(\"dataset\", \"multi\", \"imaging\", dataset_name)\n", - "dataset_path = Path(dataset_main_path, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/multi/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "\n", - "dataset_list = []\n", - "\n", - "for dataset_waveband, pixel_scale in zip(dataset_waveband_list, pixel_scale_list):\n", - " dataset = al.Imaging.from_fits(\n", - " data_path=Path(dataset_main_path, f\"{dataset_waveband}_data.fits\"),\n", - " noise_map_path=Path(dataset_main_path, f\"{dataset_waveband}_noise_map.fits\"),\n", - " psf_path=Path(dataset_main_path, f\"{dataset_waveband}_psf.fits\"),\n", - " pixel_scales=pixel_scale,\n", - " )\n", - "\n", - " mask_radius = 3.0\n", - "\n", - " mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - "\n", - " dataset = dataset.apply_mask(mask=mask)\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=dataset.grid,\n", - " sub_size_list=[4, 2, 1],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", - "\n", - " dataset_list.append(dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Settings AutoFit__\n", - "\n", - "The settings of autofit, which controls the output paths, parallelization, database use, etc." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "settings_search = af.SettingsSearch(\n", - " path_prefix=Path(\"slam\", \"multi\", \"simultaneous\"),\n", - " unique_tag=dataset_name,\n", - " info=None,\n", - " session=None,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Redshifts__\n", - "\n", - "The redshifts of the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "redshift_lens = 0.5\n", - "redshift_source = 1.0" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline Functions__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def source_lp(\n", - " settings_search,\n", - " analysis_list,\n", - " lens_bulge,\n", - " source_bulge,\n", - " redshift_lens,\n", - " redshift_source,\n", - " dataset_model,\n", - " mass_centre=(0.0, 0.0),\n", - " n_batch=50,\n", - "):\n", - " \"\"\"\n", - " SOURCE LP PIPELINE: fits an initial lens model using a parametric source, shared across all datasets\n", - " via a factor graph so that the model parameters are the same for every waveband.\n", - " \"\"\"\n", - " mass = af.Model(al.mp.Isothermal)\n", - " mass.centre = mass_centre\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_lens,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=mass,\n", - " shear=af.Model(al.mp.ExternalShear),\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=redshift_source,\n", - " bulge=source_bulge,\n", - " ),\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " analysis_factor_list = []\n", - "\n", - " for analysis in analysis_list:\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_lp[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "\n", - "def source_pix_1(\n", - " settings_search,\n", - " analysis_list,\n", - " source_lp_result,\n", - " mesh_shape,\n", - " dataset_model,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SOURCE PIX PIPELINE 1: initializes a pixelized source model, with per-dataset models that share\n", - " the mass centre across wavebands to ensure consistent lens mass geometry.\n", - " \"\"\"\n", - " analysis_factor_list = []\n", - "\n", - " for i, analysis in enumerate(analysis_list):\n", - " mass = al.util.chaining.mass_from(\n", - " mass=source_lp_result[i].model.galaxies.lens.mass,\n", - " mass_result=source_lp_result[i].model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " if i > 0:\n", - " mass.centre = model.galaxies.lens.mass.centre\n", - "\n", - " shear = source_lp_result[i].model.galaxies.lens.shear\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result[i].instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result[i].instance.galaxies.lens.bulge,\n", - " disk=source_lp_result[i].instance.galaxies.lens.disk,\n", - " mass=mass,\n", - " shear=shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result[i].instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(\n", - " al.mesh.RectangularAdaptDensity, shape=mesh_shape\n", - " ),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "\n", - "def source_pix_2(\n", - " settings_search,\n", - " analysis_list,\n", - " source_lp_result,\n", - " source_pix_result_1,\n", - " mesh_shape,\n", - " dataset_model,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SOURCE PIX PIPELINE 2: fits a pixelized source using adapt images from SOURCE PIX PIPELINE 1,\n", - " with fixed lens mass and improved mesh and regularization.\n", - " \"\"\"\n", - " analysis_factor_list = []\n", - "\n", - " for i, analysis in enumerate(analysis_list):\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result[i].instance.galaxies.lens.redshift,\n", - " bulge=source_lp_result[i].instance.galaxies.lens.bulge,\n", - " disk=source_lp_result[i].instance.galaxies.lens.disk,\n", - " mass=source_pix_result_1[i].instance.galaxies.lens.mass,\n", - " shear=source_pix_result_1[i].instance.galaxies.lens.shear,\n", - " ),\n", - " source=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_lp_result[i].instance.galaxies.source.redshift,\n", - " pixelization=af.Model(\n", - " al.Pixelization,\n", - " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", - " regularization=al.reg.Adapt,\n", - " ),\n", - " ),\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"source_pix[2]\",\n", - " **settings_search.search_dict,\n", - " n_live=75,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "\n", - "def light_lp(\n", - " settings_search,\n", - " analysis_list,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " lens_bulge,\n", - " dataset_model,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " LIGHT LP PIPELINE: fits the lens galaxy light with mass and source fixed from SOURCE PIX PIPELINE,\n", - " using per-dataset models sharing the same lens bulge model across wavebands.\n", - " \"\"\"\n", - " analysis_factor_list = []\n", - "\n", - " for i, analysis in enumerate(analysis_list):\n", - " source = al.util.chaining.source_custom_model_from(\n", - " result=source_result_for_source[i], source_is_model=False\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens[i].instance.galaxies.lens.redshift,\n", - " bulge=lens_bulge,\n", - " disk=None,\n", - " mass=source_result_for_lens[i].instance.galaxies.lens.mass,\n", - " shear=source_result_for_lens[i].instance.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"light[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=250,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "\n", - "def mass_total(\n", - " settings_search,\n", - " analysis_list,\n", - " source_result_for_lens,\n", - " source_result_for_source,\n", - " light_result,\n", - " dataset_model,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " MASS TOTAL PIPELINE: fits a PowerLaw total mass model using priors from SOURCE PIX PIPELINE,\n", - " with lens light fixed from LIGHT LP PIPELINE.\n", - " \"\"\"\n", - " # Total mass model for the lens galaxy.\n", - " mass = af.Model(al.mp.PowerLaw)\n", - "\n", - " analysis_factor_list = []\n", - "\n", - " for i, analysis in enumerate(analysis_list):\n", - " mass_i = al.util.chaining.mass_from(\n", - " mass=mass,\n", - " mass_result=source_result_for_lens[i].model.galaxies.lens.mass,\n", - " unfix_mass_centre=True,\n", - " )\n", - "\n", - " source = al.util.chaining.source_from(\n", - " result=source_result_for_source[i],\n", - " )\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=af.Model(\n", - " al.Galaxy,\n", - " redshift=source_result_for_lens[i].instance.galaxies.lens.redshift,\n", - " bulge=light_result[i].instance.galaxies.lens.bulge,\n", - " disk=light_result[i].instance.galaxies.lens.disk,\n", - " mass=mass_i,\n", - " shear=source_result_for_lens[i].model.galaxies.lens.shear,\n", - " ),\n", - " source=source,\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"mass_total[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=150,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "\n", - "def subhalo_no_subhalo(\n", - " settings_search,\n", - " analysis_list,\n", - " mass_result,\n", - " dataset_model,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SUBHALO PIPELINE 1: fits the lens model without a subhalo to provide Bayesian evidence for\n", - " model comparison with the subhalo detection searches.\n", - " \"\"\"\n", - " analysis_factor_list = []\n", - "\n", - " for i, analysis in enumerate(analysis_list):\n", - " source = al.util.chaining.source_from(result=mass_result[i])\n", - " lens = mass_result[i].model.galaxies.lens\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source=source),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[1]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "\n", - "def subhalo_grid_search(\n", - " settings_search,\n", - " analysis_list,\n", - " mass_result,\n", - " subhalo_result_1,\n", - " subhalo_mass,\n", - " grid_dimension_arcsec=3.0,\n", - " number_of_steps=2,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SUBHALO PIPELINE 2: performs a grid search over subhalo positions to detect a dark matter subhalo,\n", - " using per-dataset models with shared subhalo parameters.\n", - " \"\"\"\n", - " subhalo = af.Model(al.Galaxy, mass=subhalo_mass)\n", - "\n", - " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", - " subhalo.mass.centre_0 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - " )\n", - " subhalo.mass.centre_1 = af.UniformPrior(\n", - " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", - " )\n", - "\n", - " subhalo.redshift = subhalo_result_1[0].instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_object = subhalo_result_1[0].instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_source = subhalo_result_1[0].instance.galaxies.source.redshift\n", - "\n", - " analysis_factor_list = []\n", - "\n", - " for i, analysis in enumerate(analysis_list):\n", - " lens = mass_result[i].model.galaxies.lens\n", - " source = al.util.chaining.source_from(result=mass_result[i])\n", - "\n", - " model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, subhalo=subhalo, source=source),\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[2]_[search_lens_plane]\",\n", - " **settings_search.search_dict,\n", - " n_live=200,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " subhalo_grid_search = af.SearchGridSearch(\n", - " search=search,\n", - " number_of_steps=number_of_steps,\n", - " )\n", - "\n", - " return subhalo_grid_search.fit(\n", - " model=model,\n", - " analysis=analysis,\n", - " grid_priors=[\n", - " model.galaxies.subhalo.mass.centre_1,\n", - " model.galaxies.subhalo.mass.centre_0,\n", - " ],\n", - " info=settings_search.info,\n", - " )\n", - "\n", - "\n", - "def subhalo_refine(\n", - " settings_search,\n", - " analysis_list,\n", - " subhalo_result_1,\n", - " subhalo_grid_search_result_2,\n", - " subhalo_mass,\n", - " dataset_model,\n", - " n_batch=20,\n", - "):\n", - " \"\"\"\n", - " SUBHALO PIPELINE 3: refines the subhalo model parameters using priors initialized from the\n", - " highest-evidence cell of the grid search.\n", - " \"\"\"\n", - " subhalo = af.Model(\n", - " al.Galaxy,\n", - " redshift=subhalo_result_1[0].instance.galaxies.lens.redshift,\n", - " mass=subhalo_mass,\n", - " )\n", - "\n", - " subhalo.redshift = subhalo_result_1[0].instance.galaxies.lens.redshift\n", - " subhalo.mass.redshift_object = subhalo_result_1[0].instance.galaxies.lens.redshift\n", - " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", - " subhalo.mass.centre = subhalo_grid_search_result_2.model_centred_absolute(\n", - " a=1.0\n", - " ).galaxies.subhalo.mass.centre\n", - " subhalo.redshift = subhalo_grid_search_result_2.model.galaxies.subhalo.redshift\n", - " subhalo.mass.redshift_object = subhalo.redshift\n", - "\n", - " analysis_factor_list = []\n", - "\n", - " for i, analysis in enumerate(analysis_list):\n", - " model = af.Collection(\n", - " galaxies=af.Collection(\n", - " lens=subhalo_grid_search_result_2.model.galaxies.lens,\n", - " subhalo=subhalo,\n", - " source=subhalo_grid_search_result_2.model.galaxies.source,\n", - " ),\n", - " dataset_model=dataset_model,\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " analysis_factor_list.append(analysis_factor)\n", - "\n", - " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - " search = af.Nautilus(\n", - " name=\"subhalo[3]_[single_plane_refine]\",\n", - " **settings_search.search_dict,\n", - " n_live=600,\n", - " n_batch=n_batch,\n", - " )\n", - "\n", - " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__SLaM Pipeline__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "mesh_pixels_yx = 28\n", - "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", - "\n", - "dataset_model = af.Model(al.DatasetModel)\n", - "\n", - "# Lens Light\n", - "\n", - "lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "# Source Light\n", - "\n", - "source_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", - ")\n", - "\n", - "analysis_list = [\n", - " al.AnalysisImaging(dataset=dataset, use_jax=True) for dataset in dataset_list\n", - "]\n", - "\n", - "source_lp_result = source_lp(\n", - " settings_search=settings_search,\n", - " analysis_list=analysis_list,\n", - " lens_bulge=lens_bulge,\n", - " source_bulge=source_bulge,\n", - " redshift_lens=redshift_lens,\n", - " redshift_source=redshift_source,\n", - " dataset_model=dataset_model,\n", - ")\n", - "\n", - "positions_likelihood = source_lp_result.positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - ")\n", - "\n", - "adapt_images_list = []\n", - "\n", - "for result in source_lp_result:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result)\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - " adapt_images_list.append(adapt_images)\n", - "\n", - "analysis_list = [\n", - " al.AnalysisImaging(\n", - " dataset=result.max_log_likelihood_fit.dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " use_jax=True,\n", - " )\n", - " for result, adapt_images in zip(source_lp_result, adapt_images_list)\n", - "]\n", - "\n", - "source_pix_result_1 = source_pix_1(\n", - " settings_search=settings_search,\n", - " analysis_list=analysis_list,\n", - " source_lp_result=source_lp_result,\n", - " mesh_shape=mesh_shape,\n", - " dataset_model=dataset_model,\n", - ")\n", - "\n", - "adapt_images_list = []\n", - "\n", - "for result in source_pix_result_1:\n", - " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result)\n", - " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", - " adapt_images_list.append(adapt_images)\n", - "\n", - "analysis_list = [\n", - " al.AnalysisImaging(\n", - " dataset=result.max_log_likelihood_fit.dataset,\n", - " adapt_images=adapt_images,\n", - " use_jax=True,\n", - " )\n", - " for result, adapt_images in zip(source_pix_result_1, adapt_images_list)\n", - "]\n", - "\n", - "source_pix_result_2 = source_pix_2(\n", - " settings_search=settings_search,\n", - " analysis_list=analysis_list,\n", - " source_lp_result=source_lp_result,\n", - " source_pix_result_1=source_pix_result_1,\n", - " mesh_shape=mesh_shape,\n", - " dataset_model=dataset_model,\n", - ")\n", - "\n", - "lens_bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=2,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "analysis_list = [\n", - " al.AnalysisImaging(\n", - " dataset=result.max_log_likelihood_fit.dataset,\n", - " adapt_images=adapt_images,\n", - " raise_inversion_positions_likelihood_exception=False,\n", - " )\n", - " for result, adapt_images in zip(source_pix_result_1, adapt_images_list)\n", - "]\n", - "\n", - "light_result = light_lp(\n", - " settings_search=settings_search,\n", - " analysis_list=analysis_list,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " lens_bulge=lens_bulge,\n", - " dataset_model=dataset_model,\n", - ")\n", - "\n", - "positions_likelihood = source_pix_result_1[0].positions_likelihood_from(\n", - " factor=3.0, minimum_threshold=0.2\n", - ")\n", - "\n", - "analysis_list = [\n", - " al.AnalysisImaging(\n", - " dataset=result.max_log_likelihood_fit.dataset,\n", - " adapt_images=adapt_images,\n", - " positions_likelihood_list=[positions_likelihood],\n", - " )\n", - " for result, adapt_images in zip(source_pix_result_1, adapt_images_list)\n", - "]\n", - "\n", - "mass_result = mass_total(\n", - " settings_search=settings_search,\n", - " analysis_list=analysis_list,\n", - " source_result_for_lens=source_pix_result_1,\n", - " source_result_for_source=source_pix_result_2,\n", - " light_result=light_result,\n", - " dataset_model=dataset_model,\n", - ")\n", - "\n", - "subhalo_result_1 = subhalo_no_subhalo(\n", - " settings_search=settings_search,\n", - " analysis_list=analysis_list,\n", - " mass_result=mass_result,\n", - " dataset_model=dataset_model,\n", - ")\n", - "\n", - "subhalo_grid_search_result_2 = subhalo_grid_search(\n", - " settings_search=settings_search,\n", - " analysis_list=analysis_list,\n", - " mass_result=mass_result,\n", - " subhalo_result_1=subhalo_result_1,\n", - " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", - " grid_dimension_arcsec=3.0,\n", - " number_of_steps=2,\n", - ")\n", - "\n", - "subhalo_result_3 = subhalo_refine(\n", - " settings_search=settings_search,\n", - " analysis_list=analysis_list,\n", - " subhalo_result_1=subhalo_result_1,\n", - " subhalo_grid_search_result_2=subhalo_grid_search_result_2,\n", - " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", - " dataset_model=dataset_model,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finish." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SLaM: Multi Wavelength Simultaneous\n", + "===================================\n", + "\n", + "This example shows how to use the SLaM pipeline to fit a lens dataset at multiple wavelengths simultaneously.\n", + "\n", + "Simultaneous multi-dataset fits are currently built into the SLaM pipeline without user input or customization.\n", + "Therefore, as long as lists of `Analysis` objects are created, summed and passed to the SLaM pipelines, the analysis\n", + "will fit every dataset simultaneously and it will adapt the model as follows:\n", + "\n", + "- Sub-pixel offsets between the datasets are fully modeled as free parameters in each stage of the pipeline, assuming\n", + " broad uniform priors for every step. This is because the precision of a lens model can often be less than the\n", + " requirements on astrometry.\n", + "\n", + "- The regularization parameters are free for every dataset in the `source_pix[1]` and `source_pix[2]` stages. This is because\n", + " the source morphology can be different between datasets, and the regularization scheme adapts to this.\n", + "\n", + "- From the `light_lp` stage onwards, the regularization scheme for each dataset is different fixed to that inferred\n", + " for the `source_pix[2]` stage.\n", + "\n", + "Simultaneous fitting SLaM pipelines are not designed for customization, for example changing the model from the\n", + "set up above. This is because we are still figuring out the best way to perform multi-wavelength modeling, but have\n", + "so far figured the above settings are important.\n", + "\n", + "If you need customization of the model or pipeline, you should pick apart the SLaM pipeline and customize\n", + "them as you see fit.\n", + "\n", + "__Contents__\n", + "\n", + "- **Preqrequisites:** Before reading this script, you should have familiarity with the following key concepts.\n", + "- **This Script:** Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Settings AutoFit:** The settings of autofit, which controls the output paths, parallelization, database use, etc.\n", + "- **Redshifts:** The redshifts of the lens and source galaxies.\n", + "- **SLaM Pipeline Functions:** Overview of slam pipeline functions for this example.\n", + "- **SLaM Pipeline:** Overview of slam pipeline for this example.\n", + "\n", + "__Preqrequisites__\n", + "\n", + "Before reading this script, you should have familiarity with the following key concepts:\n", + "\n", + "- **Multi**: The `autolens_workspace/*/advanced/multi` package describes many different ways that multiple datasets\n", + " can be modeled in a single analysis.\n", + "\n", + "__This Script__\n", + "\n", + "Using a SOURCE LP PIPELINE, SOURCE PIX PIPELINE, LIGHT LP PIPELINE and TOTAL MASS PIPELINE this SLaM modeling\n", + "script fits `Imaging` dataset of a strong lens system where in the final model:\n", + "\n", + " - The lens galaxy's light is a bulge with Multiple Gaussian Expansion (MGE) light profile.\n", + " - The lens galaxy's total mass distribution is an `PowerLaw` plus an `ExternalShear`.\n", + " - The source galaxy's light is a `Pixelization`.\n", + "\n", + "This modeling script uses the SLaM pipelines:\n", + "\n", + " `source_lp`\n", + " `source_pix`\n", + " `light_lp`\n", + " `mass_total`\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `guides/modeling/slam_start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__ \n", + "\n", + "Load, plot and mask the `Imaging` data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_waveband_list = [\"g\", \"r\"]\n", + "pixel_scale_list = [0.12, 0.08]\n", + "\n", + "dataset_name = \"lens_sersic\"\n", + "dataset_main_path = Path(\"dataset\", \"multi\", \"imaging\", dataset_name)\n", + "dataset_path = Path(dataset_main_path, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/multi/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "\n", + "dataset_list = []\n", + "\n", + "for dataset_waveband, pixel_scale in zip(dataset_waveband_list, pixel_scale_list):\n", + " dataset = al.Imaging.from_fits(\n", + " data_path=Path(dataset_main_path, f\"{dataset_waveband}_data.fits\"),\n", + " noise_map_path=Path(dataset_main_path, f\"{dataset_waveband}_noise_map.fits\"),\n", + " psf_path=Path(dataset_main_path, f\"{dataset_waveband}_psf.fits\"),\n", + " pixel_scales=pixel_scale,\n", + " )\n", + "\n", + " mask_radius = 3.0\n", + "\n", + " mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + "\n", + " dataset = dataset.apply_mask(mask=mask)\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=dataset.grid,\n", + " sub_size_list=[4, 2, 1],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)\n", + "\n", + " dataset_list.append(dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Settings AutoFit__\n", + "\n", + "The settings of autofit, which controls the output paths, parallelization, database use, etc." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "settings_search = af.SettingsSearch(\n", + " path_prefix=Path(\"slam\", \"multi\", \"simultaneous\"),\n", + " unique_tag=dataset_name,\n", + " info=None,\n", + " session=None,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Redshifts__\n", + "\n", + "The redshifts of the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "redshift_lens = 0.5\n", + "redshift_source = 1.0" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline Functions__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def source_lp(\n", + " settings_search,\n", + " analysis_list,\n", + " lens_bulge,\n", + " source_bulge,\n", + " redshift_lens,\n", + " redshift_source,\n", + " dataset_model,\n", + " mass_centre=(0.0, 0.0),\n", + " n_batch=50,\n", + "):\n", + " \"\"\"\n", + " SOURCE LP PIPELINE: fits an initial lens model using a parametric source, shared across all datasets\n", + " via a factor graph so that the model parameters are the same for every waveband.\n", + " \"\"\"\n", + " mass = af.Model(al.mp.Isothermal)\n", + " mass.centre = mass_centre\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_lens,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=mass,\n", + " shear=af.Model(al.mp.ExternalShear),\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=redshift_source,\n", + " bulge=source_bulge,\n", + " ),\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " analysis_factor_list = []\n", + "\n", + " for analysis in analysis_list:\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_lp[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "\n", + "def source_pix_1(\n", + " settings_search,\n", + " analysis_list,\n", + " source_lp_result,\n", + " mesh_shape,\n", + " dataset_model,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SOURCE PIX PIPELINE 1: initializes a pixelized source model, with per-dataset models that share\n", + " the mass centre across wavebands to ensure consistent lens mass geometry.\n", + " \"\"\"\n", + " analysis_factor_list = []\n", + "\n", + " for i, analysis in enumerate(analysis_list):\n", + " mass = al.util.chaining.mass_from(\n", + " mass=source_lp_result[i].model.galaxies.lens.mass,\n", + " mass_result=source_lp_result[i].model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " if i > 0:\n", + " mass.centre = model.galaxies.lens.mass.centre\n", + "\n", + " shear = source_lp_result[i].model.galaxies.lens.shear\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result[i].instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result[i].instance.galaxies.lens.bulge,\n", + " disk=source_lp_result[i].instance.galaxies.lens.disk,\n", + " mass=mass,\n", + " shear=shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result[i].instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(\n", + " al.mesh.RectangularAdaptDensity, shape=mesh_shape\n", + " ),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "\n", + "def source_pix_2(\n", + " settings_search,\n", + " analysis_list,\n", + " source_lp_result,\n", + " source_pix_result_1,\n", + " mesh_shape,\n", + " dataset_model,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SOURCE PIX PIPELINE 2: fits a pixelized source using adapt images from SOURCE PIX PIPELINE 1,\n", + " with fixed lens mass and improved mesh and regularization.\n", + " \"\"\"\n", + " analysis_factor_list = []\n", + "\n", + " for i, analysis in enumerate(analysis_list):\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result[i].instance.galaxies.lens.redshift,\n", + " bulge=source_lp_result[i].instance.galaxies.lens.bulge,\n", + " disk=source_lp_result[i].instance.galaxies.lens.disk,\n", + " mass=source_pix_result_1[i].instance.galaxies.lens.mass,\n", + " shear=source_pix_result_1[i].instance.galaxies.lens.shear,\n", + " ),\n", + " source=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_lp_result[i].instance.galaxies.source.redshift,\n", + " pixelization=af.Model(\n", + " al.Pixelization,\n", + " mesh=af.Model(al.mesh.RectangularAdaptImage, shape=mesh_shape),\n", + " regularization=al.reg.Adapt,\n", + " ),\n", + " ),\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"source_pix[2]\",\n", + " **settings_search.search_dict,\n", + " n_live=75,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "\n", + "def light_lp(\n", + " settings_search,\n", + " analysis_list,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " lens_bulge,\n", + " dataset_model,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " LIGHT LP PIPELINE: fits the lens galaxy light with mass and source fixed from SOURCE PIX PIPELINE,\n", + " using per-dataset models sharing the same lens bulge model across wavebands.\n", + " \"\"\"\n", + " analysis_factor_list = []\n", + "\n", + " for i, analysis in enumerate(analysis_list):\n", + " source = al.util.chaining.source_custom_model_from(\n", + " result=source_result_for_source[i], source_is_model=False\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens[i].instance.galaxies.lens.redshift,\n", + " bulge=lens_bulge,\n", + " disk=None,\n", + " mass=source_result_for_lens[i].instance.galaxies.lens.mass,\n", + " shear=source_result_for_lens[i].instance.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"light[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=250,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "\n", + "def mass_total(\n", + " settings_search,\n", + " analysis_list,\n", + " source_result_for_lens,\n", + " source_result_for_source,\n", + " light_result,\n", + " dataset_model,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " MASS TOTAL PIPELINE: fits a PowerLaw total mass model using priors from SOURCE PIX PIPELINE,\n", + " with lens light fixed from LIGHT LP PIPELINE.\n", + " \"\"\"\n", + " # Total mass model for the lens galaxy.\n", + " mass = af.Model(al.mp.PowerLaw)\n", + "\n", + " analysis_factor_list = []\n", + "\n", + " for i, analysis in enumerate(analysis_list):\n", + " mass_i = al.util.chaining.mass_from(\n", + " mass=mass,\n", + " mass_result=source_result_for_lens[i].model.galaxies.lens.mass,\n", + " unfix_mass_centre=True,\n", + " )\n", + "\n", + " source = al.util.chaining.source_from(\n", + " result=source_result_for_source[i],\n", + " )\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=af.Model(\n", + " al.Galaxy,\n", + " redshift=source_result_for_lens[i].instance.galaxies.lens.redshift,\n", + " bulge=light_result[i].instance.galaxies.lens.bulge,\n", + " disk=light_result[i].instance.galaxies.lens.disk,\n", + " mass=mass_i,\n", + " shear=source_result_for_lens[i].model.galaxies.lens.shear,\n", + " ),\n", + " source=source,\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"mass_total[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=150,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "\n", + "def subhalo_no_subhalo(\n", + " settings_search,\n", + " analysis_list,\n", + " mass_result,\n", + " dataset_model,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SUBHALO PIPELINE 1: fits the lens model without a subhalo to provide Bayesian evidence for\n", + " model comparison with the subhalo detection searches.\n", + " \"\"\"\n", + " analysis_factor_list = []\n", + "\n", + " for i, analysis in enumerate(analysis_list):\n", + " source = al.util.chaining.source_from(result=mass_result[i])\n", + " lens = mass_result[i].model.galaxies.lens\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source=source),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[1]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "\n", + "def subhalo_grid_search(\n", + " settings_search,\n", + " analysis_list,\n", + " mass_result,\n", + " subhalo_result_1,\n", + " subhalo_mass,\n", + " grid_dimension_arcsec=3.0,\n", + " number_of_steps=2,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SUBHALO PIPELINE 2: performs a grid search over subhalo positions to detect a dark matter subhalo,\n", + " using per-dataset models with shared subhalo parameters.\n", + " \"\"\"\n", + " subhalo = af.Model(al.Galaxy, mass=subhalo_mass)\n", + "\n", + " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", + " subhalo.mass.centre_0 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + " )\n", + " subhalo.mass.centre_1 = af.UniformPrior(\n", + " lower_limit=-grid_dimension_arcsec, upper_limit=grid_dimension_arcsec\n", + " )\n", + "\n", + " subhalo.redshift = subhalo_result_1[0].instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_object = subhalo_result_1[0].instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_source = subhalo_result_1[0].instance.galaxies.source.redshift\n", + "\n", + " analysis_factor_list = []\n", + "\n", + " for i, analysis in enumerate(analysis_list):\n", + " lens = mass_result[i].model.galaxies.lens\n", + " source = al.util.chaining.source_from(result=mass_result[i])\n", + "\n", + " model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, subhalo=subhalo, source=source),\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[2]_[search_lens_plane]\",\n", + " **settings_search.search_dict,\n", + " n_live=200,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " subhalo_grid_search = af.SearchGridSearch(\n", + " search=search,\n", + " number_of_steps=number_of_steps,\n", + " )\n", + "\n", + " return subhalo_grid_search.fit(\n", + " model=model,\n", + " analysis=analysis,\n", + " grid_priors=[\n", + " model.galaxies.subhalo.mass.centre_1,\n", + " model.galaxies.subhalo.mass.centre_0,\n", + " ],\n", + " info=settings_search.info,\n", + " )\n", + "\n", + "\n", + "def subhalo_refine(\n", + " settings_search,\n", + " analysis_list,\n", + " subhalo_result_1,\n", + " subhalo_grid_search_result_2,\n", + " subhalo_mass,\n", + " dataset_model,\n", + " n_batch=20,\n", + "):\n", + " \"\"\"\n", + " SUBHALO PIPELINE 3: refines the subhalo model parameters using priors initialized from the\n", + " highest-evidence cell of the grid search.\n", + " \"\"\"\n", + " subhalo = af.Model(\n", + " al.Galaxy,\n", + " redshift=subhalo_result_1[0].instance.galaxies.lens.redshift,\n", + " mass=subhalo_mass,\n", + " )\n", + "\n", + " subhalo.redshift = subhalo_result_1[0].instance.galaxies.lens.redshift\n", + " subhalo.mass.redshift_object = subhalo_result_1[0].instance.galaxies.lens.redshift\n", + " subhalo.mass.mass_at_200 = af.LogUniformPrior(lower_limit=1.0e6, upper_limit=1.0e11)\n", + " subhalo.mass.centre = subhalo_grid_search_result_2.model_centred_absolute(\n", + " a=1.0\n", + " ).galaxies.subhalo.mass.centre\n", + " subhalo.redshift = subhalo_grid_search_result_2.model.galaxies.subhalo.redshift\n", + " subhalo.mass.redshift_object = subhalo.redshift\n", + "\n", + " analysis_factor_list = []\n", + "\n", + " for i, analysis in enumerate(analysis_list):\n", + " model = af.Collection(\n", + " galaxies=af.Collection(\n", + " lens=subhalo_grid_search_result_2.model.galaxies.lens,\n", + " subhalo=subhalo,\n", + " source=subhalo_grid_search_result_2.model.galaxies.source,\n", + " ),\n", + " dataset_model=dataset_model,\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " analysis_factor_list.append(analysis_factor)\n", + "\n", + " factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + " search = af.Nautilus(\n", + " name=\"subhalo[3]_[single_plane_refine]\",\n", + " **settings_search.search_dict,\n", + " n_live=600,\n", + " n_batch=n_batch,\n", + " )\n", + "\n", + " return search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__SLaM Pipeline__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "mesh_pixels_yx = 28\n", + "mesh_shape = (mesh_pixels_yx, mesh_pixels_yx)\n", + "\n", + "dataset_model = af.Model(al.DatasetModel)\n", + "\n", + "# Lens Light\n", + "\n", + "lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "# Source Light\n", + "\n", + "source_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False\n", + ")\n", + "\n", + "analysis_list = [\n", + " al.AnalysisImaging(dataset=dataset, use_jax=True) for dataset in dataset_list\n", + "]\n", + "\n", + "source_lp_result = source_lp(\n", + " settings_search=settings_search,\n", + " analysis_list=analysis_list,\n", + " lens_bulge=lens_bulge,\n", + " source_bulge=source_bulge,\n", + " redshift_lens=redshift_lens,\n", + " redshift_source=redshift_source,\n", + " dataset_model=dataset_model,\n", + ")\n", + "\n", + "positions_likelihood = source_lp_result.positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + ")\n", + "\n", + "adapt_images_list = []\n", + "\n", + "for result in source_lp_result:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result)\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + " adapt_images_list.append(adapt_images)\n", + "\n", + "analysis_list = [\n", + " al.AnalysisImaging(\n", + " dataset=result.max_log_likelihood_fit.dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " use_jax=True,\n", + " )\n", + " for result, adapt_images in zip(source_lp_result, adapt_images_list)\n", + "]\n", + "\n", + "source_pix_result_1 = source_pix_1(\n", + " settings_search=settings_search,\n", + " analysis_list=analysis_list,\n", + " source_lp_result=source_lp_result,\n", + " mesh_shape=mesh_shape,\n", + " dataset_model=dataset_model,\n", + ")\n", + "\n", + "adapt_images_list = []\n", + "\n", + "for result in source_pix_result_1:\n", + " galaxy_image_name_dict = al.galaxy_name_image_dict_via_result_from(result=result)\n", + " adapt_images = al.AdaptImages(galaxy_name_image_dict=galaxy_image_name_dict)\n", + " adapt_images_list.append(adapt_images)\n", + "\n", + "analysis_list = [\n", + " al.AnalysisImaging(\n", + " dataset=result.max_log_likelihood_fit.dataset,\n", + " adapt_images=adapt_images,\n", + " use_jax=True,\n", + " )\n", + " for result, adapt_images in zip(source_pix_result_1, adapt_images_list)\n", + "]\n", + "\n", + "source_pix_result_2 = source_pix_2(\n", + " settings_search=settings_search,\n", + " analysis_list=analysis_list,\n", + " source_lp_result=source_lp_result,\n", + " source_pix_result_1=source_pix_result_1,\n", + " mesh_shape=mesh_shape,\n", + " dataset_model=dataset_model,\n", + ")\n", + "\n", + "lens_bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=2,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "analysis_list = [\n", + " al.AnalysisImaging(\n", + " dataset=result.max_log_likelihood_fit.dataset,\n", + " adapt_images=adapt_images,\n", + " raise_inversion_positions_likelihood_exception=False,\n", + " )\n", + " for result, adapt_images in zip(source_pix_result_1, adapt_images_list)\n", + "]\n", + "\n", + "light_result = light_lp(\n", + " settings_search=settings_search,\n", + " analysis_list=analysis_list,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " lens_bulge=lens_bulge,\n", + " dataset_model=dataset_model,\n", + ")\n", + "\n", + "positions_likelihood = source_pix_result_1[0].positions_likelihood_from(\n", + " factor=3.0, minimum_threshold=0.2\n", + ")\n", + "\n", + "analysis_list = [\n", + " al.AnalysisImaging(\n", + " dataset=result.max_log_likelihood_fit.dataset,\n", + " adapt_images=adapt_images,\n", + " positions_likelihood_list=[positions_likelihood],\n", + " )\n", + " for result, adapt_images in zip(source_pix_result_1, adapt_images_list)\n", + "]\n", + "\n", + "mass_result = mass_total(\n", + " settings_search=settings_search,\n", + " analysis_list=analysis_list,\n", + " source_result_for_lens=source_pix_result_1,\n", + " source_result_for_source=source_pix_result_2,\n", + " light_result=light_result,\n", + " dataset_model=dataset_model,\n", + ")\n", + "\n", + "subhalo_result_1 = subhalo_no_subhalo(\n", + " settings_search=settings_search,\n", + " analysis_list=analysis_list,\n", + " mass_result=mass_result,\n", + " dataset_model=dataset_model,\n", + ")\n", + "\n", + "subhalo_grid_search_result_2 = subhalo_grid_search(\n", + " settings_search=settings_search,\n", + " analysis_list=analysis_list,\n", + " mass_result=mass_result,\n", + " subhalo_result_1=subhalo_result_1,\n", + " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", + " grid_dimension_arcsec=3.0,\n", + " number_of_steps=2,\n", + ")\n", + "\n", + "subhalo_result_3 = subhalo_refine(\n", + " settings_search=settings_search,\n", + " analysis_list=analysis_list,\n", + " subhalo_result_1=subhalo_result_1,\n", + " subhalo_grid_search_result_2=subhalo_grid_search_result_2,\n", + " subhalo_mass=af.Model(al.mp.NFWMCRLudlowSph),\n", + " dataset_model=dataset_model,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finish." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/wavelength_dependence/modeling.ipynb b/notebooks/multi/features/wavelength_dependence/modeling.ipynb index 410361144..3337d6538 100644 --- a/notebooks/multi/features/wavelength_dependence/modeling.ipynb +++ b/notebooks/multi/features/wavelength_dependence/modeling.ipynb @@ -1,492 +1,529 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Mass Total + Source Parametric\n", - "========================================\n", - "\n", - "This script fits a multi-wavelength `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is an MGE bulge where the `ell_comps` varies across wavelength.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an MGE.\n", - "\n", - "Three images are fitted, corresponding to a green ('g' band), red (`r` band) and near infrared ('I' band) images.\n", - "\n", - "This script assumes previous knowledge of the `multi` modeling API found in other scripts in the `multi/modeling`\n", - "package. If anything is unclear check those scripts out.\n", - "\n", - "__Contents__\n", - "\n", - "- **Effective Radius vs Wavelength:** Unlike other `multi` modeling scripts, the effective radius of the lens and source galaxies as a.\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "- **Wavelengths:** The effective_radius of each source galaxy is parameterized as a function of wavelength.\n", - "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Effective Radius vs Wavelength__\n", - "\n", - "Unlike other `multi` modeling scripts, the effective radius of the lens and source galaxies as a user defined function of\n", - "wavelength, for example following a relation `y = (m * x) + c` -> `effective_radius = (m * wavelength) + c`.\n", - "\n", - "By using a linear relation `y = mx + c` the free parameters are `m` and `c`, which does not scale with the number\n", - "of datasets. For datasets with multi-wavelength images (e.g. 5 or more) this allows us to parameterize the variation\n", - "of parameters across the datasets in a way that does not lead to a very complex parameter space.\n", - "\n", - "For example, in other scripts, a free `effective_radius` is created for every datasets, which would add 5+ free parameters\n", - "to the model for 5+ datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "\n", - "The strings are used for load each dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"] # , \"I\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wavelengths__\n", - "\n", - "The effective_radius of each source galaxy is parameterized as a function of wavelength.\n", - "\n", - "Therefore we define a list of wavelengths of each color above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "wavelength_list = [464, 658, 806]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixel Scales__\n", - "\n", - "Every multi-wavelength dataset can have its own unique pixel-scale." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12, 0.12]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"wavelength_dependence\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name\n", - "\n", - "dataset_list = [\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " pixel_scales=pixel_scales,\n", - " )\n", - " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", - "\n", - "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", - "absolutely necessary, but provides a more reliable analysis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_list = [\n", - " al.Mask2D.circular(\n", - " shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=3.0\n", - " )\n", - " for dataset in dataset_list\n", - "]\n", - "\n", - "dataset_list = [\n", - " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "We create an `Analysis` object for every dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [al.AnalysisImaging(dataset=dataset) for dataset in dataset_list]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=15." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=al.lp_linear.Sersic,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model + Analysis__\n", - "\n", - "We now make the lens and source `effective_radius` a free parameter across every analysis object.\n", - "\n", - "Unlike other scripts, where the `effective_radius` for every dataset is created as a free parameter, we will assume that \n", - "the `effective_radius` of the lens and source galaxies linearly varies as a function of wavelength, and therefore compute \n", - "the `effective_radius` value for each color image using a linear relation `y = mx + c`.\n", - "\n", - "The function below is not used to compose the model, but illustrates how the `effective_radius` values were computed\n", - "in the corresponding `wavelength_dependence` simulator script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def lens_effective_radius_from(wavelength):\n", - " m = 1.0 / 100.0 # lens appears brighter with wavelength\n", - " c = 3\n", - "\n", - " return m * wavelength + c\n", - "\n", - "\n", - "def source_effective_radius_from(wavelength):\n", - " m = -(1.2 / 100.0) # source appears fainter with wavelength\n", - " c = 10\n", - "\n", - " return m * wavelength + c\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To parameterize the above relation as a model, we compose `m` and `c` as priors and use PyAutoFit's prior arithmatic\n", - "to compose a model as a linear relation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_m = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "lens_c = af.UniformPrior(lower_limit=-10.0, upper_limit=10.0)\n", - "\n", - "source_m = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "source_c = af.UniformPrior(lower_limit=-10.0, upper_limit=10.0)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The free parameters of our model there are no longer `effective_radius` values, but the parameters `m` and `c` in the relation\n", - "above. \n", - "\n", - "The model complexity therefore does not increase as we add more parameters to the model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "analysis_factor_list = []\n", - "\n", - "for wavelength, analysis in zip(wavelength_list, analysis_list):\n", - " lens_effective_radius = (wavelength * lens_m) + lens_c\n", - " source_effective_radius = (wavelength * source_m) + source_c\n", - "\n", - " model_analysis = model.copy()\n", - "\n", - " model_analysis.galaxies.lens.bulge.effective_radius = lens_effective_radius\n", - " model_analysis.galaxies.source.bulge.effective_radius = source_effective_radius\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", - "\n", - " analysis_factor_list.append(analysis_factor)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The factor graph is created and its info can be printed after the relational model has been defined." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", - "\n", - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"multi\", \"modeling\"),\n", - " name=\"wavelength_dependence\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", - "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength.\n", - "\n", - "For example, close inspection of the `max_log_likelihood_instance` of the two results shows that all parameters,\n", - "except the `effective_radius` of the source galaxy's `bulge`, are identical." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_list[0].max_log_likelihood_instance)\n", - "print(result_list[1].max_log_likelihood_instance)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting each result's tracer shows that the source appears different, owning to its different intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", - "\n", - "Therefore, the samples is identical in every result object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Mass Total + Source Parametric\n", + "========================================\n", + "\n", + "This script fits a multi-wavelength `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is an MGE bulge where the `ell_comps` varies across wavelength.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an MGE.\n", + "\n", + "Three images are fitted, corresponding to a green ('g' band), red (`r` band) and near infrared ('I' band) images.\n", + "\n", + "This script assumes previous knowledge of the `multi` modeling API found in other scripts in the `multi/modeling`\n", + "package. If anything is unclear check those scripts out.\n", + "\n", + "__Contents__\n", + "\n", + "- **Effective Radius vs Wavelength:** Unlike other `multi` modeling scripts, the effective radius of the lens and source galaxies as a.\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "- **Wavelengths:** The effective_radius of each source galaxy is parameterized as a function of wavelength.\n", + "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Effective Radius vs Wavelength__\n", + "\n", + "Unlike other `multi` modeling scripts, the effective radius of the lens and source galaxies as a user defined function of\n", + "wavelength, for example following a relation `y = (m * x) + c` -> `effective_radius = (m * wavelength) + c`.\n", + "\n", + "By using a linear relation `y = mx + c` the free parameters are `m` and `c`, which does not scale with the number\n", + "of datasets. For datasets with multi-wavelength images (e.g. 5 or more) this allows us to parameterize the variation\n", + "of parameters across the datasets in a way that does not lead to a very complex parameter space.\n", + "\n", + "For example, in other scripts, a free `effective_radius` is created for every datasets, which would add 5+ free parameters\n", + "to the model for 5+ datasets." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "\n", + "The strings are used for load each dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"] # , \"I\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wavelengths__\n", + "\n", + "The effective_radius of each source galaxy is parameterized as a function of wavelength.\n", + "\n", + "Therefore we define a list of wavelengths of each color above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "wavelength_list = [464, 658, 806]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixel Scales__\n", + "\n", + "Every multi-wavelength dataset can have its own unique pixel-scale." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12, 0.12]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"wavelength_dependence\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name\n", + "\n", + "dataset_list = [\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " pixel_scales=pixel_scales,\n", + " )\n", + " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", + "\n", + "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", + "absolutely necessary, but provides a more reliable analysis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_list = [\n", + " al.Mask2D.circular(\n", + " shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=3.0\n", + " )\n", + " for dataset in dataset_list\n", + "]\n", + "\n", + "dataset_list = [\n", + " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We create an `Analysis` object for every dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [al.AnalysisImaging(dataset=dataset) for dataset in dataset_list]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is an MGE with 1 x 20 Gaussians [4 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=15." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=al.lp_linear.Sersic,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp_linear.SersicCore)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model + Analysis__\n", + "\n", + "We now make the lens and source `effective_radius` a free parameter across every analysis object.\n", + "\n", + "Unlike other scripts, where the `effective_radius` for every dataset is created as a free parameter, we will assume that \n", + "the `effective_radius` of the lens and source galaxies linearly varies as a function of wavelength, and therefore compute \n", + "the `effective_radius` value for each color image using a linear relation `y = mx + c`.\n", + "\n", + "The function below is not used to compose the model, but illustrates how the `effective_radius` values were computed\n", + "in the corresponding `wavelength_dependence` simulator script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def lens_effective_radius_from(wavelength):\n", + " m = 1.0 / 100.0 # lens appears brighter with wavelength\n", + " c = 3\n", + "\n", + " return m * wavelength + c\n", + "\n", + "\n", + "def source_effective_radius_from(wavelength):\n", + " m = -(1.2 / 100.0) # source appears fainter with wavelength\n", + " c = 10\n", + "\n", + " return m * wavelength + c\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To parameterize the above relation as a model, we compose `m` and `c` as priors and use PyAutoFit's prior arithmatic\n", + "to compose a model as a linear relation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_m = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "lens_c = af.UniformPrior(lower_limit=-10.0, upper_limit=10.0)\n", + "\n", + "source_m = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "source_c = af.UniformPrior(lower_limit=-10.0, upper_limit=10.0)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The free parameters of our model there are no longer `effective_radius` values, but the parameters `m` and `c` in the relation\n", + "above. \n", + "\n", + "The model complexity therefore does not increase as we add more parameters to the model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "analysis_factor_list = []\n", + "\n", + "for wavelength, analysis in zip(wavelength_list, analysis_list):\n", + " lens_effective_radius = (wavelength * lens_m) + lens_c\n", + " source_effective_radius = (wavelength * source_m) + source_c\n", + "\n", + " model_analysis = model.copy()\n", + "\n", + " model_analysis.galaxies.lens.bulge.effective_radius = lens_effective_radius\n", + " model_analysis.galaxies.source.bulge.effective_radius = source_effective_radius\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", + "\n", + " analysis_factor_list.append(analysis_factor)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The factor graph is created and its info can be printed after the relational model has been defined." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)\n", + "\n", + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"multi\", \"modeling\"),\n", + " name=\"wavelength_dependence\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", + "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength.\n", + "\n", + "For example, close inspection of the `max_log_likelihood_instance` of the two results shows that all parameters,\n", + "except the `effective_radius` of the source galaxy's `bulge`, are identical." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_list[0].max_log_likelihood_instance)\n", + "print(result_list[1].max_log_likelihood_instance)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting each result's tracer shows that the source appears different, owning to its different intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", + "\n", + "Therefore, the samples is identical in every result object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/features/wavelength_dependence/simulator.ipynb b/notebooks/multi/features/wavelength_dependence/simulator.ipynb index 127b4e563..e9b02a7a5 100644 --- a/notebooks/multi/features/wavelength_dependence/simulator.ipynb +++ b/notebooks/multi/features/wavelength_dependence/simulator.ipynb @@ -1,518 +1,555 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Wavelength Dependent\n", - "===============================\n", - "\n", - "This script simulates multi-wavelength `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's light is a Sersic bulge where the `intensity` varies across wavelength.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", - "\n", - "Unlike other `multi` simulators, the intensity of the source galaxy is a linear function of wavelength following\n", - "a relation `y = mx + c`.\n", - "\n", - "This image is used to demonstrate multi-wavelength fitting where a user specified function (e.g. `y = mx+c`) can be\n", - "used to parameterize the wavelength variation, as opposed to simply making every `intensity` a free parameter.\n", - "\n", - "Three images are simulated, corresponding green g band (wavelength=464nm), red r-band (wavelength=658nm) and\n", - "infrared I-band (wavelength=806nm) observations.\n", - "\n", - "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for simulating images. Thus,\n", - "certain parts of code are not documented to ensure the script is concise.\n", - "\n", - "__Contents__\n", - "\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band), red (r-band) and.\n", - "- **Wavelengths:** The intensity of each source galaxy is parameterized as a function of wavelength.\n", - "- **Dataset Paths:** Overview of dataset paths for this example.\n", - "- **Simulate:** The pixel-scale of each color image is different meaning we make a list of grids for the simulation.\n", - "- **Intensity vs Wavelength:** We will assume that the `intensity` of the lens and source galaxies linearly varies as a function.\n", - "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) for this simulated lens.\n", - "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band), red (r-band) and infrared (I-band).\n", - "\n", - "The strings are used for naming the datasets on output." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\", \"I\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wavelengths__\n", - "\n", - "The intensity of each source galaxy is parameterized as a function of wavelength.\n", - "\n", - "Therefore we define a list of wavelengths of each color above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "wavelength_list = [464, 658, 806]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"wavelength_dependence\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "The pixel-scale of each color image is different meaning we make a list of grids for the simulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12, 0.12]\n", - "\n", - "grid_list = []\n", - "\n", - "for pixel_scales in pixel_scales_list:\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=(150, 150),\n", - " pixel_scales=pixel_scales,\n", - " )\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " grid_list.append(grid)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate simple Gaussian PSFs for the images in the r and g bands." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sigma_list = [0.1, 0.2, 0.25]\n", - "\n", - "psf_list = [\n", - " al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", - " )\n", - " for grid, sigma in zip(grid_list, sigma_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create separate simulators for the g and r bands." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "background_sky_level_list = [0.1, 0.15, 0.1]\n", - "\n", - "simulator_list = [\n", - " al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=background_sky_level,\n", - " add_poisson_noise_to_data=True,\n", - " )\n", - " for psf, background_sky_level in zip(psf_list, background_sky_level_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensity vs Wavelength__\n", - "\n", - "We will assume that the `intensity` of the lens and source galaxies linearly varies as a function of wavelength, and \n", - "therefore compute the `intensity` value for each color image using a linear relation.\n", - "\n", - "The relation below is not realistic and has been chosen to make it straight forward to illustrate this functionality." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "\n", - "def lens_intensity_from(wavelength):\n", - " m = 1.0 / 100.0\n", - " c = 3\n", - "\n", - " return m * wavelength + c\n", - "\n", - "\n", - "def source_intensity_from(wavelength):\n", - " m = -(1.2 / 100.0)\n", - " c = 10\n", - "\n", - " return m * wavelength + c\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's mass (SIE+Shear) for this simulated lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_intensity_list = [\n", - " lens_intensity_from(wavelength=wavelength) for wavelength in wavelength_list\n", - "]\n", - "\n", - "bulge_list = [\n", - " al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=intensity,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " )\n", - " for intensity in lens_intensity_list\n", - "]\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "lens_galaxy_list = [\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - " )\n", - " for bulge in bulge_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Intensity vs Wavelength__\n", - "\n", - "We will assume that the `intensity` of the source galaxy linearly varies as a function of wavelength, and therefore\n", - "compute the `intensity` value for each color image using a linear relation.\n", - "\n", - "The relation below is not realistic and has been chosen to make it straight forward to illustrate this functionality." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_intensity_list = [\n", - " source_intensity_from(wavelength=wavelength) for wavelength in wavelength_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The source galaxy at each wavelength has a different intensity, thus we create three source galaxies for each waveband." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy_list = [\n", - " al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=intensity,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - " )\n", - " for intensity in source_intensity_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging` \n", - "dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_list = [\n", - " al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - " for lens_galaxy, source_galaxy in zip(lens_galaxy_list, source_galaxy_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for tracer, grid in zip(tracer_list, grid_list):\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - " for grid, simulator, tracer in zip(grid_list, simulator_list, tracer_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " overwrite=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - " aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "for waveband, grid, tracer in zip(waveband_list, grid_list, tracer_list):\n", - " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", - " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "[\n", - " al.output_to_json(obj=tracer, file_path=Path(dataset_path, f\"{color}_tracer.json\"))\n", - " for color, tracer in zip(waveband_list, tracer_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/simple__no_lens_light`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Wavelength Dependent\n", + "===============================\n", + "\n", + "This script simulates multi-wavelength `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's light is a Sersic bulge where the `intensity` varies across wavelength.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", + "\n", + "Unlike other `multi` simulators, the intensity of the source galaxy is a linear function of wavelength following\n", + "a relation `y = mx + c`.\n", + "\n", + "This image is used to demonstrate multi-wavelength fitting where a user specified function (e.g. `y = mx+c`) can be\n", + "used to parameterize the wavelength variation, as opposed to simply making every `intensity` a free parameter.\n", + "\n", + "Three images are simulated, corresponding green g band (wavelength=464nm), red r-band (wavelength=658nm) and\n", + "infrared I-band (wavelength=806nm) observations.\n", + "\n", + "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for simulating images. Thus,\n", + "certain parts of code are not documented to ensure the script is concise.\n", + "\n", + "__Contents__\n", + "\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band), red (r-band) and.\n", + "- **Wavelengths:** The intensity of each source galaxy is parameterized as a function of wavelength.\n", + "- **Dataset Paths:** Overview of dataset paths for this example.\n", + "- **Simulate:** The pixel-scale of each color image is different meaning we make a list of grids for the simulation.\n", + "- **Intensity vs Wavelength:** We will assume that the `intensity` of the lens and source galaxies linearly varies as a function.\n", + "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) for this simulated lens.\n", + "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band), red (r-band) and infrared (I-band).\n", + "\n", + "The strings are used for naming the datasets on output." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\", \"I\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wavelengths__\n", + "\n", + "The intensity of each source galaxy is parameterized as a function of wavelength.\n", + "\n", + "Therefore we define a list of wavelengths of each color above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "wavelength_list = [464, 658, 806]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"wavelength_dependence\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "The pixel-scale of each color image is different meaning we make a list of grids for the simulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12, 0.12]\n", + "\n", + "grid_list = []\n", + "\n", + "for pixel_scales in pixel_scales_list:\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=(150, 150),\n", + " pixel_scales=pixel_scales,\n", + " )\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " grid_list.append(grid)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate simple Gaussian PSFs for the images in the r and g bands." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sigma_list = [0.1, 0.2, 0.25]\n", + "\n", + "psf_list = [\n", + " al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", + " )\n", + " for grid, sigma in zip(grid_list, sigma_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create separate simulators for the g and r bands." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "background_sky_level_list = [0.1, 0.15, 0.1]\n", + "\n", + "simulator_list = [\n", + " al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=background_sky_level,\n", + " add_poisson_noise_to_data=True,\n", + " )\n", + " for psf, background_sky_level in zip(psf_list, background_sky_level_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensity vs Wavelength__\n", + "\n", + "We will assume that the `intensity` of the lens and source galaxies linearly varies as a function of wavelength, and \n", + "therefore compute the `intensity` value for each color image using a linear relation.\n", + "\n", + "The relation below is not realistic and has been chosen to make it straight forward to illustrate this functionality." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "\n", + "def lens_intensity_from(wavelength):\n", + " m = 1.0 / 100.0\n", + " c = 3\n", + "\n", + " return m * wavelength + c\n", + "\n", + "\n", + "def source_intensity_from(wavelength):\n", + " m = -(1.2 / 100.0)\n", + " c = 10\n", + "\n", + " return m * wavelength + c\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's mass (SIE+Shear) for this simulated lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_intensity_list = [\n", + " lens_intensity_from(wavelength=wavelength) for wavelength in wavelength_list\n", + "]\n", + "\n", + "bulge_list = [\n", + " al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=intensity,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " )\n", + " for intensity in lens_intensity_list\n", + "]\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "lens_galaxy_list = [\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + " )\n", + " for bulge in bulge_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Intensity vs Wavelength__\n", + "\n", + "We will assume that the `intensity` of the source galaxy linearly varies as a function of wavelength, and therefore\n", + "compute the `intensity` value for each color image using a linear relation.\n", + "\n", + "The relation below is not realistic and has been chosen to make it straight forward to illustrate this functionality." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_intensity_list = [\n", + " source_intensity_from(wavelength=wavelength) for wavelength in wavelength_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The source galaxy at each wavelength has a different intensity, thus we create three source galaxies for each waveband." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy_list = [\n", + " al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=intensity,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + " )\n", + " for intensity in source_intensity_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging` \n", + "dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_list = [\n", + " al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + " for lens_galaxy, source_galaxy in zip(lens_galaxy_list, source_galaxy_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for tracer, grid in zip(tracer_list, grid_list):\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + " for grid, simulator, tracer in zip(grid_list, simulator_list, tracer_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " overwrite=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + " aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "for waveband, grid, tracer in zip(waveband_list, grid_list, tracer_list):\n", + " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", + " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "[\n", + " al.output_to_json(obj=tracer, file_path=Path(dataset_path, f\"{color}_tracer.json\"))\n", + " for color, tracer in zip(waveband_list, tracer_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/simple__no_lens_light`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/modeling.ipynb b/notebooks/multi/modeling.ipynb index 485840e42..7b592eeda 100644 --- a/notebooks/multi/modeling.ipynb +++ b/notebooks/multi/modeling.ipynb @@ -1,660 +1,697 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Multi Modeling\n", - "========================\n", - "\n", - "This script fits a multi-wavelength `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a MGE bulge where the `ell_comps` varies across wavelength.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is a an MGE where the `ell_comps` varies across wavelength.\n", - "\n", - "Two images are fitted, corresponding to a greener ('g' band) redder image (`r` band).\n", - "\n", - "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for lens modeling. Thus,\n", - "certain parts of code are not documented to ensure the script is concise.\n", - "\n", - "__Contents__\n", - "\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Model Extension:** Galaxies change appearance across wavelength, for example their ellipticities.\n", - "- **Linear Light Profiles:** As an advanced user you should be familiar wiht linear light profiles, see elsewhere in the.\n", - "- **Analysis List:** Set up two instances of the `Analysis` class object, one for each dataset.\n", - "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", - "- **Analysis Factor:** Each analysis object is wrapped in an `AnalysisFactor`, which pairs it with the model and prepares.\n", - "- **Factor Graph:** All `AnalysisFactor` objects are combined into a `FactorGraphModel`, which represents a global.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", - "- **VRAM Use:** The `modeling` examples of individual dataset types explain how VRAM is used during GPU-based.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Wrap Up:** Summary of the script and next steps." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "\n", - "The strings are used for load each dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Pixel Scales__\n", - "\n", - "Every multi-wavelength dataset can have its own unique pixel-scale." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors.\n", - "\n", - "Note how the lens and source appear different brightnesses in each wavelength. Multi-wavelength image can therefore \n", - "better separate the lens and source galaxies." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"lens_sersic\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if al.util.dataset.should_simulate(str(dataset_path)):\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/multi/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset_list = [\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " pixel_scales=pixel_scales,\n", - " )\n", - " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxies Noise Scaling__\n", - "\n", - "Before masking, we must deal with any extra galaxies in the data: nearby galaxies (or foreground stars, or\n", - "data-reduction artefacts) whose emission is not associated with the strong lens but blends into the field. If\n", - "their light is left in the data it will contaminate the model-fit and bias the inferred lens model. It is too\n", - "easy to skip straight to modeling without checking for these, so we make this step an explicit part of the\n", - "workflow.\n", - "\n", - "To prevent extra galaxies from impacting the fit, we do not mask them entirely from the fit. Instead, the pixels\n", - "are kept in the fit but their data values are scaled to zero and their noise-map values increased to very large\n", - "values, so they contribute negligibly to the likelihood. This is preferable to removing the pixels entirely\n", - "(e.g. for a pixelized source reconstruction, removing pixels can produce discontinuities in the pixelization).\n", - "\n", - "The `lens_sersic` dataset includes a faint extra galaxy, and a per-waveband `{waveband}_mask_extra_galaxies.fits`\n", - "covering it is shipped with the dataset (created by the simulator). If you are modeling your own data with an\n", - "extra galaxy, you must either create such a mask using the data-preparation tools, or shrink the circular mask\n", - "below so the extra galaxy lies outside it and is removed from the fit entirely.\n", - "\n", - "**Multi-wavelength Specific:** the noise scaling is applied to every waveband one-by-one, loading the mask whose\n", - "pixel scale and shape match that waveband's dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_scaled_list = []\n", - "\n", - "for dataset, waveband in zip(dataset_list, waveband_list):\n", - "\n", - " mask_extra_galaxies = al.Mask2D.from_fits(\n", - " file_path=Path(dataset_path) / f\"{waveband}_mask_extra_galaxies.fits\",\n", - " pixel_scales=dataset.pixel_scales,\n", - " invert=True, # `True` means a pixel is scaled.\n", - " )\n", - "\n", - " dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " dataset_scaled_list.append(dataset)\n", - "\n", - "dataset_list = dataset_scaled_list" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", - "\n", - "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", - "absolutely necessary, but provides a more reliable analysis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask_radius = 3.0\n", - "\n", - "mask_list = [\n", - " al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " radius=mask_radius,\n", - " )\n", - " for dataset in dataset_list\n", - "]\n", - "\n", - "\n", - "dataset_list = [\n", - " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's light is an MGE with 2 x 30 Gaussians, where the `intensity` parameter of the lens galaxy\n", - " for each individual waveband of imaging is a different free parameter [6 parameters].\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", - "\n", - " - The source galaxy's light is a an MGE, where the `intensity` parameter of the source galaxy\n", - " for each individual waveband of imaging is a different free parameter [8 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=23.\n", - "\n", - "__Model Extension__\n", - "\n", - "Galaxies change appearance across wavelength, for example their ellipticities.\n", - "\n", - "Models applied to combined analyses can be extended to include free parameters specific to each dataset. In this example,\n", - "we will make the galaxy's ellipticity vary across the g and r-band datasets, which will be illustrated below.\n", - "\n", - "__Linear Light Profiles__\n", - "\n", - "As an advanced user you should be familiar wiht linear light profiles, see elsewhere in the workspace for informaiton\n", - "if not.\n", - "\n", - "For multi wavelength dataset modeling, the `lp_linear` API is extremely powerful as the `ell_comps` varies across\n", - "the datasets, meaning that making it linear reduces the dimensionality of parameter space significantly." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_gaussians = 20\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=total_gaussians,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=True,\n", - ")\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - ")\n", - "\n", - "bulge = al.model_util.mge_model_from(\n", - " mask_radius=mask_radius,\n", - " total_gaussians=20,\n", - " gaussian_per_basis=1,\n", - " centre_prior_is_uniform=False,\n", - ")\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis List__\n", - "\n", - "Set up two instances of the `Analysis` class object, one for each dataset.\n", - "\n", - "__JAX__\n", - "\n", - "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU\n", - "support, your fits will run much faster (around 10 minutes instead of an hour). If only a CPU is available,\n", - "JAX will still provide a speed up via multithreading, with fits taking around 20-30 minutes.\n", - "\n", - "If you don\u2019t have a GPU locally, consider Google Colab which provides free GPUs, so your modeling runs are much faster." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [\n", - " al.AnalysisImaging(\n", - " dataset=dataset,\n", - " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", - " )\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis Factor__\n", - "\n", - "Each analysis object is wrapped in an `AnalysisFactor`, which pairs it with the model and prepares it for use in a \n", - "factor graph. This step allows us to flexibly define how each dataset relates to the model.\n", - "\n", - "The term \"Factor\" comes from factor graphs, a type of probabilistic graphical model. In this context, each factor \n", - "represents the connection between one dataset and the shared model.\n", - "\n", - "The API for extending the model across datasets is shown below, by overwriting the `ell_comps`\n", - "variables of the model passed to each `AnalysisFactor` object with new priors, making each dataset have its own\n", - "`ell_comps` free parameter.\n", - "\n", - "NOTE: Other aspects of galaxies may vary across wavelength, none of which are included in this example. The API below \n", - "can easily be extended to include these additional parameters, and the `features` package explains other tools for \n", - "extending the model across datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = []\n", - "\n", - "for analysis in analysis_list:\n", - "\n", - " model_analysis = model.copy()\n", - "\n", - " ell_comps_0_prior = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - " ell_comps_1_prior = af.GaussianPrior(mean=0.0, sigma=0.3)\n", - "\n", - " for i in range(len(model_analysis.galaxies.lens.bulge.profile_list)):\n", - "\n", - " model_analysis.galaxies.lens.bulge.profile_list[i].ell_comps.ell_comps_0 = (\n", - " ell_comps_0_prior\n", - " )\n", - " model_analysis.galaxies.lens.bulge.profile_list[i].ell_comps.ell_comps_1 = (\n", - " ell_comps_1_prior\n", - " )\n", - "\n", - " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", - "\n", - " analysis_factor_list.append(analysis_factor)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Factor Graph__\n", - "\n", - "All `AnalysisFactor` objects are combined into a `FactorGraphModel`, which represents a global model fit to \n", - "multiple datasets using a graphical model structure.\n", - "\n", - "The key outcomes of this setup are:\n", - "\n", - " - The individual log likelihoods from each `Analysis` object are summed to form the total log likelihood \n", - " evaluated during the model-fitting process.\n", - "\n", - " - Results from all datasets are output to a unified directory, with subdirectories for visualizations \n", - " from each analysis object, as defined by their `visualize` methods.\n", - "\n", - "This is a basic use of **PyAutoFit**'s graphical modeling capabilities, which support advanced hierarchical \n", - "and probabilistic modeling for large, multi-dataset analyses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To inspect this new model, with extra parameters for each dataset created, we \n", - "print `factor_graph.global_prior_model.info`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "__Live Visual Update__\n", - "\n", - "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", - "live display surface:\n", - "\n", - "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", - " watch the fit converge without leaving your terminal.\n", - "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", - " refreshes in place every `iterations_per_quick_update`.\n", - "\n", - "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", - "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\n", - " \"multi_wavelength\"\n", - " ), # The path where results and output are stored.\n", - " name=\"modeling\", # The name of the fit and folder results are output to.\n", - " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", - " n_live=150, # The number of Nautilus \"live\" points, increase for more complex models.\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=10000, # Every N iterations the max likelihood model is visualized in the Jupter Notebook and output to hard-disk.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM Use__\n", - "\n", - "The `modeling` examples of individual dataset types explain how VRAM is used during GPU-based fitting and how to \n", - "print the estimated VRAM required by a model.\n", - "\n", - "When multiple datasets are fitted simultaneously, as in this example, VRAM usage increases with each\n", - "dataset, as their data structures must all be stored in VRAM.\n", - "\n", - "Given VRAM use is an important consideration, we print out the estimated VRAM required for this\n", - "model-fit and advise you do this for your own pixelization model-fits.\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "factor_graph.print_vram_use(\n", - " model=factor_graph.global_prior_model, batch_size=search.batch_size\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__\n", - "\n", - "To fit multiple datasets, we pass the `FactorGraphModel` to a non-linear search.\n", - "\n", - "Unlike single-dataset fitting, we now pass the `factor_graph.global_prior_model` as the model and \n", - "the `factor_graph` itself as the analysis object.\n", - "\n", - "This structure enables simultaneous fitting of multiple datasets in a consistent and scalable way.\n", - "\n", - "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce \n", - "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", - "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength.\n", - "\n", - "For example, close inspection of the `max_log_likelihood_instance` of the two results shows that all parameters,\n", - "except the `effective_radius` of the source galaxy's `bulge`, are identical." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result_list[0].max_log_likelihood_instance)\n", - "print(result_list[1].max_log_likelihood_instance)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plotting each result's tracer shows that the source appears different, owning to its different intensities." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", - "\n", - "Therefore, the samples is identical in every result object." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This simple example introduces the API for fitting multiple datasets with a shared model.\n", - "\n", - "It should already be quite intuitive how this API can be adapted to fit more complex models, or fit different\n", - "datasets with different models. For example, an `AnalysisImaging` and `AnalysisInterferometer` can be combined, into\n", - "a single factor graph model, to simultaneously fit a imaging and interferometric data.\n", - "\n", - "The `advanced/multi/modeling` package has more examples of how to fit multiple datasets with different models,\n", - "including relational models that vary parameters across datasets as a function of wavelength." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Multi Modeling\n", + "========================\n", + "\n", + "This script fits a multi-wavelength `Imaging` dataset of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a MGE bulge where the `ell_comps` varies across wavelength.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is a an MGE where the `ell_comps` varies across wavelength.\n", + "\n", + "Two images are fitted, corresponding to a greener ('g' band) redder image (`r` band).\n", + "\n", + "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for lens modeling. Thus,\n", + "certain parts of code are not documented to ensure the script is concise.\n", + "\n", + "__Contents__\n", + "\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "- **Pixel Scales:** Every multi-wavelength dataset can have its own unique pixel-scale.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Model Extension:** Galaxies change appearance across wavelength, for example their ellipticities.\n", + "- **Linear Light Profiles:** As an advanced user you should be familiar wiht linear light profiles, see elsewhere in the.\n", + "- **Analysis List:** Set up two instances of the `Analysis` class object, one for each dataset.\n", + "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", + "- **Analysis Factor:** Each analysis object is wrapped in an `AnalysisFactor`, which pairs it with the model and prepares.\n", + "- **Factor Graph:** All `AnalysisFactor` objects are combined into a `FactorGraphModel`, which represents a global.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", + "- **VRAM Use:** The `modeling` examples of individual dataset types explain how VRAM is used during GPU-based.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Wrap Up:** Summary of the script and next steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "\n", + "The strings are used for load each dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Pixel Scales__\n", + "\n", + "Every multi-wavelength dataset can have its own unique pixel-scale." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot each multi-wavelength strong lens dataset, using a list of their waveband colors.\n", + "\n", + "Note how the lens and source appear different brightnesses in each wavelength. Multi-wavelength image can therefore \n", + "better separate the lens and source galaxies." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"lens_sersic\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/multi/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_list = [\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " pixel_scales=pixel_scales,\n", + " )\n", + " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxies Noise Scaling__\n", + "\n", + "Before masking, we must deal with any extra galaxies in the data: nearby galaxies (or foreground stars, or\n", + "data-reduction artefacts) whose emission is not associated with the strong lens but blends into the field. If\n", + "their light is left in the data it will contaminate the model-fit and bias the inferred lens model. It is too\n", + "easy to skip straight to modeling without checking for these, so we make this step an explicit part of the\n", + "workflow.\n", + "\n", + "To prevent extra galaxies from impacting the fit, we do not mask them entirely from the fit. Instead, the pixels\n", + "are kept in the fit but their data values are scaled to zero and their noise-map values increased to very large\n", + "values, so they contribute negligibly to the likelihood. This is preferable to removing the pixels entirely\n", + "(e.g. for a pixelized source reconstruction, removing pixels can produce discontinuities in the pixelization).\n", + "\n", + "The `lens_sersic` dataset includes a faint extra galaxy, and a per-waveband `{waveband}_mask_extra_galaxies.fits`\n", + "covering it is shipped with the dataset (created by the simulator). If you are modeling your own data with an\n", + "extra galaxy, you must either create such a mask using the data-preparation tools, or shrink the circular mask\n", + "below so the extra galaxy lies outside it and is removed from the fit entirely.\n", + "\n", + "**Multi-wavelength Specific:** the noise scaling is applied to every waveband one-by-one, loading the mask whose\n", + "pixel scale and shape match that waveband's dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_scaled_list = []\n", + "\n", + "for dataset, waveband in zip(dataset_list, waveband_list):\n", + "\n", + " mask_extra_galaxies = al.Mask2D.from_fits(\n", + " file_path=Path(dataset_path) / f\"{waveband}_mask_extra_galaxies.fits\",\n", + " pixel_scales=dataset.pixel_scales,\n", + " invert=True, # `True` means a pixel is scaled.\n", + " )\n", + "\n", + " dataset = dataset.apply_noise_scaling(mask=mask_extra_galaxies)\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " dataset_scaled_list.append(dataset)\n", + "\n", + "dataset_list = dataset_scaled_list" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and source galaxies.\n", + "\n", + "For multi-wavelength lens modeling, we use the same mask for every dataset whenever possible. This is not\n", + "absolutely necessary, but provides a more reliable analysis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask_radius = 3.0\n", + "\n", + "mask_list = [\n", + " al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " radius=mask_radius,\n", + " )\n", + " for dataset in dataset_list\n", + "]\n", + "\n", + "\n", + "dataset_list = [\n", + " dataset.apply_mask(mask=mask) for imaging, mask in zip(dataset_list, mask_list)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's light is an MGE with 2 x 30 Gaussians, where the `intensity` parameter of the lens galaxy\n", + " for each individual waveband of imaging is a different free parameter [6 parameters].\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear` [7 parameters].\n", + "\n", + " - The source galaxy's light is a an MGE, where the `intensity` parameter of the source galaxy\n", + " for each individual waveband of imaging is a different free parameter [8 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=23.\n", + "\n", + "__Model Extension__\n", + "\n", + "Galaxies change appearance across wavelength, for example their ellipticities.\n", + "\n", + "Models applied to combined analyses can be extended to include free parameters specific to each dataset. In this example,\n", + "we will make the galaxy's ellipticity vary across the g and r-band datasets, which will be illustrated below.\n", + "\n", + "__Linear Light Profiles__\n", + "\n", + "As an advanced user you should be familiar wiht linear light profiles, see elsewhere in the workspace for informaiton\n", + "if not.\n", + "\n", + "For multi wavelength dataset modeling, the `lp_linear` API is extremely powerful as the `ell_comps` varies across\n", + "the datasets, meaning that making it linear reduces the dimensionality of parameter space significantly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_gaussians = 20\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=total_gaussians,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=True,\n", + ")\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + ")\n", + "\n", + "bulge = al.model_util.mge_model_from(\n", + " mask_radius=mask_radius,\n", + " total_gaussians=20,\n", + " gaussian_per_basis=1,\n", + " centre_prior_is_uniform=False,\n", + ")\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis List__\n", + "\n", + "Set up two instances of the `Analysis` class object, one for each dataset.\n", + "\n", + "__JAX__\n", + "\n", + "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU\n", + "support, your fits will run much faster (around 10 minutes instead of an hour). If only a CPU is available,\n", + "JAX will still provide a speed up via multithreading, with fits taking around 20-30 minutes.\n", + "\n", + "If you don\u2019t have a GPU locally, consider Google Colab which provides free GPUs, so your modeling runs are much faster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [\n", + " al.AnalysisImaging(\n", + " dataset=dataset,\n", + " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", + " )\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis Factor__\n", + "\n", + "Each analysis object is wrapped in an `AnalysisFactor`, which pairs it with the model and prepares it for use in a \n", + "factor graph. This step allows us to flexibly define how each dataset relates to the model.\n", + "\n", + "The term \"Factor\" comes from factor graphs, a type of probabilistic graphical model. In this context, each factor \n", + "represents the connection between one dataset and the shared model.\n", + "\n", + "The API for extending the model across datasets is shown below, by overwriting the `ell_comps`\n", + "variables of the model passed to each `AnalysisFactor` object with new priors, making each dataset have its own\n", + "`ell_comps` free parameter.\n", + "\n", + "NOTE: Other aspects of galaxies may vary across wavelength, none of which are included in this example. The API below \n", + "can easily be extended to include these additional parameters, and the `features` package explains other tools for \n", + "extending the model across datasets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = []\n", + "\n", + "for analysis in analysis_list:\n", + "\n", + " model_analysis = model.copy()\n", + "\n", + " ell_comps_0_prior = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + " ell_comps_1_prior = af.GaussianPrior(mean=0.0, sigma=0.3)\n", + "\n", + " for i in range(len(model_analysis.galaxies.lens.bulge.profile_list)):\n", + "\n", + " model_analysis.galaxies.lens.bulge.profile_list[i].ell_comps.ell_comps_0 = (\n", + " ell_comps_0_prior\n", + " )\n", + " model_analysis.galaxies.lens.bulge.profile_list[i].ell_comps.ell_comps_1 = (\n", + " ell_comps_1_prior\n", + " )\n", + "\n", + " analysis_factor = af.AnalysisFactor(prior_model=model_analysis, analysis=analysis)\n", + "\n", + " analysis_factor_list.append(analysis_factor)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Factor Graph__\n", + "\n", + "All `AnalysisFactor` objects are combined into a `FactorGraphModel`, which represents a global model fit to \n", + "multiple datasets using a graphical model structure.\n", + "\n", + "The key outcomes of this setup are:\n", + "\n", + " - The individual log likelihoods from each `Analysis` object are summed to form the total log likelihood \n", + " evaluated during the model-fitting process.\n", + "\n", + " - Results from all datasets are output to a unified directory, with subdirectories for visualizations \n", + " from each analysis object, as defined by their `visualize` methods.\n", + "\n", + "This is a basic use of **PyAutoFit**'s graphical modeling capabilities, which support advanced hierarchical \n", + "and probabilistic modeling for large, multi-dataset analyses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To inspect this new model, with extra parameters for each dataset created, we \n", + "print `factor_graph.global_prior_model.info`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "__Live Visual Update__\n", + "\n", + "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", + "live display surface:\n", + "\n", + "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", + " watch the fit converge without leaving your terminal.\n", + "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", + " refreshes in place every `iterations_per_quick_update`.\n", + "\n", + "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", + "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\n", + " \"multi_wavelength\"\n", + " ), # The path where results and output are stored.\n", + " name=\"modeling\", # The name of the fit and folder results are output to.\n", + " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", + " n_live=150, # The number of Nautilus \"live\" points, increase for more complex models.\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=10000, # Every N iterations the max likelihood model is visualized in the Jupter Notebook and output to hard-disk.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM Use__\n", + "\n", + "The `modeling` examples of individual dataset types explain how VRAM is used during GPU-based fitting and how to \n", + "print the estimated VRAM required by a model.\n", + "\n", + "When multiple datasets are fitted simultaneously, as in this example, VRAM usage increases with each\n", + "dataset, as their data structures must all be stored in VRAM.\n", + "\n", + "Given VRAM use is an important consideration, we print out the estimated VRAM required for this\n", + "model-fit and advise you do this for your own pixelization model-fits.\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph.print_vram_use(\n", + " model=factor_graph.global_prior_model, batch_size=search.batch_size\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__\n", + "\n", + "To fit multiple datasets, we pass the `FactorGraphModel` to a non-linear search.\n", + "\n", + "Unlike single-dataset fitting, we now pass the `factor_graph.global_prior_model` as the model and \n", + "the `factor_graph` itself as the analysis object.\n", + "\n", + "This structure enables simultaneous fitting of multiple datasets in a consistent and scalable way.\n", + "\n", + "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce \n", + "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result object returned by this model-fit is a list of `Result` objects, because we used a factor graph.\n", + "Each result corresponds to each analysis, and therefore corresponds to the model-fit at that wavelength.\n", + "\n", + "For example, close inspection of the `max_log_likelihood_instance` of the two results shows that all parameters,\n", + "except the `effective_radius` of the source galaxy's `bulge`, are identical." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result_list[0].max_log_likelihood_instance)\n", + "print(result_list[1].max_log_likelihood_instance)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plotting each result's tracer shows that the source appears different, owning to its different intensities." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + " aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Samples` object still has the dimensions of the overall non-linear search (in this case N=15). \n", + "\n", + "Therefore, the samples is identical in every result object." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This simple example introduces the API for fitting multiple datasets with a shared model.\n", + "\n", + "It should already be quite intuitive how this API can be adapted to fit more complex models, or fit different\n", + "datasets with different models. For example, an `AnalysisImaging` and `AnalysisInterferometer` can be combined, into\n", + "a single factor graph model, to simultaneously fit a imaging and interferometric data.\n", + "\n", + "The `advanced/multi/modeling` package has more examples of how to fit multiple datasets with different models,\n", + "including relational models that vary parameters across datasets as a function of wavelength." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/plot.ipynb b/notebooks/multi/plot.ipynb index 2baa79f60..3af13187e 100644 --- a/notebooks/multi/plot.ipynb +++ b/notebooks/multi/plot.ipynb @@ -1,253 +1,290 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plots: Multi Figure Plotter\n", - "============================\n", - "\n", - "This example shows how to plot the same figure from multiple datasets on a single subplot.\n", - "\n", - "The specific example loads multi-wavelength imaging datasets and plots the data image from\n", - "each dataset side-by-side.\n", - "\n", - "In the old API, this was done using a `MultiFigurePlotter` object with a list of `Imaging`\n", - "objects. Both `MultiFigurePlotter` and `Imaging` have been removed.\n", - "\n", - "In the new API, we load each dataset and use matplotlib subplots directly.\n", - "\n", - "The dedicated `aplt.subplot_imaging_dataset()` function is also shown for single-dataset plots.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Single Dataset Subplots:** Plot the full subplot overview of each dataset using `aplt.subplot_imaging_dataset()`.\n", - "- **Multi Dataset Plot:** Plot the data image from each dataset side-by-side on the same matplotlib figure.\n", - "- **Multi Dataset Array Plot:** We can also call `aplt.plot_array()` for each dataset separately.\n", - "- **Multi Fits:** We can also output a list of figures to a single `.fits` file, where each image goes in each HDU.\n", - "- **Wrap Up:** Summary of the script and next steps.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to `plot/start_here.ipynb`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import matplotlib.pyplot as plt\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the multi-wavelength `lens_sersic` datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"]\n", - "\n", - "pixel_scales_list = [0.08, 0.12]\n", - "\n", - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"lens_sersic\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/multi/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset_list = [\n", - " al.Imaging.from_fits(\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " pixel_scales=pixel_scales,\n", - " )\n", - " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Single Dataset Subplots__\n", - "\n", - "Plot the full subplot overview of each dataset using `aplt.subplot_imaging_dataset()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multi Dataset Plot__\n", - "\n", - "Plot the data image from each dataset side-by-side on the same matplotlib figure." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fig, axes = plt.subplots(1, len(dataset_list), figsize=(12, 5))\n", - "\n", - "for ax, dataset, waveband in zip(axes, dataset_list, waveband_list):\n", - " im = ax.imshow(dataset.data.native, origin=\"upper\", cmap=\"gray\")\n", - " ax.set_title(f\"Data ({waveband}-band)\", fontsize=12)\n", - " ax.axis(\"off\")\n", - " fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)\n", - "\n", - "plt.suptitle(\"Multi-Wavelength Data\", fontsize=14)\n", - "plt.tight_layout()\n", - "plt.show()\n", - "plt.close()" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multi Dataset Array Plot__\n", - "\n", - "We can also call `aplt.plot_array()` for each dataset separately." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for dataset, waveband in zip(dataset_list, waveband_list):\n", - " aplt.plot_array(array=dataset.data, title=f\"Data ({waveband}-band)\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Multi Fits__\n", - "\n", - "We can also output a list of figures to a single `.fits` file, where each image goes in\n", - "each HDU extension." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "from autoconf.fitsable import hdu_list_for_output_from\n", - "\n", - "dataset = dataset_list[-1]\n", - "\n", - "image_list = [dataset.data, dataset.noise_map]\n", - "\n", - "hdu_list = hdu_list_for_output_from(\n", - " values_list=[image_list[0].mask.astype(\"float\")] + image_list,\n", - " ext_name_list=[\"mask\"] + [\"data\", \"noise_map\"],\n", - " header_dict=dataset.mask.header_dict,\n", - ")\n", - "\n", - "hdu_list.writeto(\"dataset.fits\", overwrite=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "The new API uses direct `aplt.plot_array()` calls and matplotlib subplots for combining\n", - "multiple figures from different datasets or objects." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plots: Multi Figure Plotter\n", + "============================\n", + "\n", + "This example shows how to plot the same figure from multiple datasets on a single subplot.\n", + "\n", + "The specific example loads multi-wavelength imaging datasets and plots the data image from\n", + "each dataset side-by-side.\n", + "\n", + "In the old API, this was done using a `MultiFigurePlotter` object with a list of `Imaging`\n", + "objects. Both `MultiFigurePlotter` and `Imaging` have been removed.\n", + "\n", + "In the new API, we load each dataset and use matplotlib subplots directly.\n", + "\n", + "The dedicated `aplt.subplot_imaging_dataset()` function is also shown for single-dataset plots.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Single Dataset Subplots:** Plot the full subplot overview of each dataset using `aplt.subplot_imaging_dataset()`.\n", + "- **Multi Dataset Plot:** Plot the data image from each dataset side-by-side on the same matplotlib figure.\n", + "- **Multi Dataset Array Plot:** We can also call `aplt.plot_array()` for each dataset separately.\n", + "- **Multi Fits:** We can also output a list of figures to a single `.fits` file, where each image goes in each HDU.\n", + "- **Wrap Up:** Summary of the script and next steps.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to `plot/start_here.ipynb`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the multi-wavelength `lens_sersic` datasets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"]\n", + "\n", + "pixel_scales_list = [0.08, 0.12]\n", + "\n", + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"lens_sersic\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/multi/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_list = [\n", + " al.Imaging.from_fits(\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " pixel_scales=pixel_scales,\n", + " )\n", + " for waveband, pixel_scales in zip(waveband_list, pixel_scales_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Single Dataset Subplots__\n", + "\n", + "Plot the full subplot overview of each dataset using `aplt.subplot_imaging_dataset()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multi Dataset Plot__\n", + "\n", + "Plot the data image from each dataset side-by-side on the same matplotlib figure." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fig, axes = plt.subplots(1, len(dataset_list), figsize=(12, 5))\n", + "\n", + "for ax, dataset, waveband in zip(axes, dataset_list, waveband_list):\n", + " im = ax.imshow(dataset.data.native, origin=\"upper\", cmap=\"gray\")\n", + " ax.set_title(f\"Data ({waveband}-band)\", fontsize=12)\n", + " ax.axis(\"off\")\n", + " fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)\n", + "\n", + "plt.suptitle(\"Multi-Wavelength Data\", fontsize=14)\n", + "plt.tight_layout()\n", + "plt.show()\n", + "plt.close()" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multi Dataset Array Plot__\n", + "\n", + "We can also call `aplt.plot_array()` for each dataset separately." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for dataset, waveband in zip(dataset_list, waveband_list):\n", + " aplt.plot_array(array=dataset.data, title=f\"Data ({waveband}-band)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Multi Fits__\n", + "\n", + "We can also output a list of figures to a single `.fits` file, where each image goes in\n", + "each HDU extension." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autoconf.fitsable import hdu_list_for_output_from\n", + "\n", + "dataset = dataset_list[-1]\n", + "\n", + "image_list = [dataset.data, dataset.noise_map]\n", + "\n", + "hdu_list = hdu_list_for_output_from(\n", + " values_list=[image_list[0].mask.astype(\"float\")] + image_list,\n", + " ext_name_list=[\"mask\"] + [\"data\", \"noise_map\"],\n", + " header_dict=dataset.mask.header_dict,\n", + ")\n", + "\n", + "hdu_list.writeto(\"dataset.fits\", overwrite=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The new API uses direct `aplt.plot_array()` calls and matplotlib subplots for combining\n", + "multiple figures from different datasets or objects." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/multi/simulator.ipynb b/notebooks/multi/simulator.ipynb index 49b7a3bf4..b8f764d67 100644 --- a/notebooks/multi/simulator.ipynb +++ b/notebooks/multi/simulator.ipynb @@ -1,528 +1,565 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: SIE\n", - "==============\n", - "\n", - "This script simulates multi-wavelength `Imaging` of a 'galaxy-scale' strong lens where:\n", - "\n", - " - The lens galaxy's light profile is an `Sersic`, which has a different `intensity` at each wavelength.\n", - " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", - " - A faint extra galaxy is included offset from the lens, whose emission must be removed via noise scaling\n", - " (a per-waveband `{waveband}_mask_extra_galaxies.fits` covering it is written below).\n", - "\n", - "Two images are simulated, corresponding to a greener ('g' band) redder image (`r` band).\n", - "\n", - "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for simulating images. Thus,\n", - "certain parts of code are not documented to ensure the script is concise.\n", - "\n", - "__Contents__\n", - "\n", - "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "- **Dataset Paths:** Overview of dataset paths for this example.\n", - "- **Simulate:** The pixel-scale of each color image is different meaning we make a list of grids for the simulation.\n", - "- **Ray Tracing:** The lens galaxy light at each wavelength has a different intensity, thus we create two lens.\n", - "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", - "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Colors__\n", - "\n", - "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", - "\n", - "The strings are used for naming the datasets on output." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "waveband_list = [\"g\", \"r\"]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"multi\"\n", - "dataset_label = \"imaging\"\n", - "dataset_name = \"lens_sersic\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "The pixel-scale of each color image is different meaning we make a list of grids for the simulation." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "pixel_scales_list = [0.08, 0.12]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The centre of a faint extra galaxy, placed inside the 3.0\" modeling mask but clear of the lensed source arcs\n", - "(Einstein radius ~1.6\"). It is reused for over-sampling, the galaxy itself and the per-waveband\n", - "`mask_extra_galaxies.fits` written further down." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_galaxy_centre = (2.2, 1.6)\n", - "\n", - "grid_list = []\n", - "\n", - "for pixel_scales in pixel_scales_list:\n", - " grid = al.Grid2D.uniform(\n", - " shape_native=(150, 150),\n", - " pixel_scales=pixel_scales,\n", - " )\n", - "\n", - " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0), extra_galaxy_centre],\n", - " )\n", - "\n", - " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", - "\n", - " grid_list.append(grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate simple Gaussian PSFs for the images in the r and g bands." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "sigma_list = [0.1, 0.2]\n", - "\n", - "psf_list = [\n", - " al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", - " )\n", - " for grid, sigma in zip(grid_list, sigma_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create separate simulators for the g and r bands." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "background_sky_level_list = [0.1, 0.15]\n", - "\n", - "simulator_list = [\n", - " al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=background_sky_level,\n", - " add_poisson_noise_to_data=True,\n", - " )\n", - " for psf, background_sky_level in zip(psf_list, background_sky_level_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The lens galaxy light at each wavelength has a different intensity, thus we create two lens galaxies for each waveband. \n", - "\n", - "The lens galaxy's mass (SIE+Shear) is identical for each waveband and included in both lens galaxies in the list.." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "intensity_list = [0.05, 1.5]\n", - "\n", - "bulge_list = [\n", - " al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=intensity,\n", - " effective_radius=0.8,\n", - " sersic_index=4.0,\n", - " )\n", - " for intensity in intensity_list\n", - "]\n", - "\n", - "mass = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "lens_galaxy_list = [\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " mass=mass,\n", - " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", - " )\n", - " for bulge in bulge_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "The source galaxy at each wavelength has a different intensity, thus we create two source galaxies for each waveband." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "intensity_list = [0.5, 0.7]\n", - "\n", - "source_galaxy_list = [\n", - " al.Galaxy(\n", - " redshift=1.0,\n", - " bulge=al.lp.SersicCore(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " intensity=intensity,\n", - " effective_radius=0.1,\n", - " sersic_index=1.0,\n", - " ),\n", - " )\n", - " for intensity in intensity_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Extra Galaxy__\n", - "\n", - "A single faint extra galaxy offset from the lens, representing a nearby contaminating object whose emission is\n", - "removed in the modeling example via the `__Extra Galaxies Noise Scaling__` step. Its intensity differs per\n", - "waveband (like the lens and source) and it has a light profile only (no mass), so the lensed source arcs are\n", - "unchanged." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "extra_intensity_list = [0.4, 1.0]\n", - "\n", - "extra_galaxy_list = [\n", - " al.Galaxy(\n", - " redshift=0.5,\n", - " light=al.lp.ExponentialSph(\n", - " centre=extra_galaxy_centre, intensity=intensity, effective_radius=0.3\n", - " ),\n", - " )\n", - " for intensity in extra_intensity_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging`\n", - "dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer_list = [\n", - " al.Tracer(galaxies=[lens_galaxy, extra_galaxy, source_galaxy])\n", - " for lens_galaxy, extra_galaxy, source_galaxy in zip(\n", - " lens_galaxy_list, extra_galaxy_list, source_galaxy_list\n", - " )\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lets look at the tracer`s image, this is the image we'll be simulating." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for tracer, grid in zip(tracer_list, grid_list):\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - " for grid, simulator, tracer in zip(grid_list, simulator_list, tracer_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for dataset in dataset_list:\n", - " aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", - " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", - " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", - " overwrite=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask Extra Galaxies__\n", - "\n", - "Build and output a per-waveband `{waveband}_mask_extra_galaxies.fits` covering the extra galaxy, so the modeling\n", - "example (`multi/modeling.py`) can load each one and apply noise scaling. The mask is built per dataset because the\n", - "wavebands have different pixel scales and therefore different `shape_native`. The circle is sized to ~3x the\n", - "galaxy's `effective_radius`, derived from the same `extra_galaxy_centre` defined above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " mask_extra_galaxies = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native,\n", - " pixel_scales=dataset.pixel_scales,\n", - " centre=extra_galaxy_centre,\n", - " radius=3.0 * 0.3,\n", - " invert=True, # `True` inside the circle, i.e. the region whose noise is scaled.\n", - " )\n", - "\n", - " aplt.fits_array(\n", - " array=mask_extra_galaxies,\n", - " file_path=Path(dataset_path) / f\"{waveband}_mask_extra_galaxies.fits\",\n", - " overwrite=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for waveband, dataset in zip(waveband_list, dataset_list):\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - " aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "\n", - "for waveband, grid, tracer in zip(waveband_list, grid_list, tracer_list):\n", - " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", - " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)\n", - "\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "[\n", - " al.output_to_json(\n", - " obj=tracer, file_path=Path(dataset_path, f\"{waveband}_tracer.json\")\n", - " )\n", - " for color, tracer in zip(waveband_list, tracer_list)\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/lens_sersic`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: SIE\n", + "==============\n", + "\n", + "This script simulates multi-wavelength `Imaging` of a 'galaxy-scale' strong lens where:\n", + "\n", + " - The lens galaxy's light profile is an `Sersic`, which has a different `intensity` at each wavelength.\n", + " - The lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The source galaxy's light is an `Sersic`, which has a different `intensity` at each wavelength.\n", + " - A faint extra galaxy is included offset from the lens, whose emission must be removed via noise scaling\n", + " (a per-waveband `{waveband}_mask_extra_galaxies.fits` covering it is written below).\n", + "\n", + "Two images are simulated, corresponding to a greener ('g' band) redder image (`r` band).\n", + "\n", + "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for simulating images. Thus,\n", + "certain parts of code are not documented to ensure the script is concise.\n", + "\n", + "__Contents__\n", + "\n", + "- **Colors:** The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "- **Dataset Paths:** Overview of dataset paths for this example.\n", + "- **Simulate:** The pixel-scale of each color image is different meaning we make a list of grids for the simulation.\n", + "- **Ray Tracing:** The lens galaxy light at each wavelength has a different intensity, thus we create two lens.\n", + "- **Output:** Output each simulated dataset to the dataset path as .fits files, with a tag describing its color.\n", + "- **Visualize:** Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Colors__\n", + "\n", + "The colors of the multi-wavelength image, which in this case are green (g-band) and red (r-band).\n", + "\n", + "The strings are used for naming the datasets on output." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "waveband_list = [\"g\", \"r\"]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"multi\"\n", + "dataset_label = \"imaging\"\n", + "dataset_name = \"lens_sersic\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "The pixel-scale of each color image is different meaning we make a list of grids for the simulation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "pixel_scales_list = [0.08, 0.12]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The centre of a faint extra galaxy, placed inside the 3.0\" modeling mask but clear of the lensed source arcs\n", + "(Einstein radius ~1.6\"). It is reused for over-sampling, the galaxy itself and the per-waveband\n", + "`mask_extra_galaxies.fits` written further down." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_galaxy_centre = (2.2, 1.6)\n", + "\n", + "grid_list = []\n", + "\n", + "for pixel_scales in pixel_scales_list:\n", + " grid = al.Grid2D.uniform(\n", + " shape_native=(150, 150),\n", + " pixel_scales=pixel_scales,\n", + " )\n", + "\n", + " over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0), extra_galaxy_centre],\n", + " )\n", + "\n", + " grid = grid.apply_over_sampling(over_sample_size=over_sample_size)\n", + "\n", + " grid_list.append(grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate simple Gaussian PSFs for the images in the r and g bands." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sigma_list = [0.1, 0.2]\n", + "\n", + "psf_list = [\n", + " al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=sigma, pixel_scales=grid.pixel_scales\n", + " )\n", + " for grid, sigma in zip(grid_list, sigma_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create separate simulators for the g and r bands." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "background_sky_level_list = [0.1, 0.15]\n", + "\n", + "simulator_list = [\n", + " al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=background_sky_level,\n", + " add_poisson_noise_to_data=True,\n", + " )\n", + " for psf, background_sky_level in zip(psf_list, background_sky_level_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The lens galaxy light at each wavelength has a different intensity, thus we create two lens galaxies for each waveband. \n", + "\n", + "The lens galaxy's mass (SIE+Shear) is identical for each waveband and included in both lens galaxies in the list.." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "intensity_list = [0.05, 1.5]\n", + "\n", + "bulge_list = [\n", + " al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=intensity,\n", + " effective_radius=0.8,\n", + " sersic_index=4.0,\n", + " )\n", + " for intensity in intensity_list\n", + "]\n", + "\n", + "mass = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "lens_galaxy_list = [\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " mass=mass,\n", + " shear=al.mp.ExternalShear(gamma_1=0.05, gamma_2=0.05),\n", + " )\n", + " for bulge in bulge_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The source galaxy at each wavelength has a different intensity, thus we create two source galaxies for each waveband." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "intensity_list = [0.5, 0.7]\n", + "\n", + "source_galaxy_list = [\n", + " al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " intensity=intensity,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + " )\n", + " for intensity in intensity_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Extra Galaxy__\n", + "\n", + "A single faint extra galaxy offset from the lens, representing a nearby contaminating object whose emission is\n", + "removed in the modeling example via the `__Extra Galaxies Noise Scaling__` step. Its intensity differs per\n", + "waveband (like the lens and source) and it has a light profile only (no mass), so the lensed source arcs are\n", + "unchanged." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extra_intensity_list = [0.4, 1.0]\n", + "\n", + "extra_galaxy_list = [\n", + " al.Galaxy(\n", + " redshift=0.5,\n", + " light=al.lp.ExponentialSph(\n", + " centre=extra_galaxy_centre, intensity=intensity, effective_radius=0.3\n", + " ),\n", + " )\n", + " for intensity in extra_intensity_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup tracers at each waveband, which will generate each image for the simulated `Imaging`\n", + "dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer_list = [\n", + " al.Tracer(galaxies=[lens_galaxy, extra_galaxy, source_galaxy])\n", + " for lens_galaxy, extra_galaxy, source_galaxy in zip(\n", + " lens_galaxy_list, extra_galaxy_list, source_galaxy_list\n", + " )\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Lets look at the tracer`s image, this is the image we'll be simulating." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for tracer, grid in zip(tracer_list, grid_list):\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + " for grid, simulator, tracer in zip(grid_list, simulator_list, tracer_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for dataset in dataset_list:\n", + " aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output each simulated dataset to the dataset path as .fits files, with a tag describing its color." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=Path(dataset_path) / f\"{waveband}_data.fits\",\n", + " psf_path=Path(dataset_path) / f\"{waveband}_psf.fits\",\n", + " noise_map_path=Path(dataset_path) / f\"{waveband}_noise_map.fits\",\n", + " overwrite=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask Extra Galaxies__\n", + "\n", + "Build and output a per-waveband `{waveband}_mask_extra_galaxies.fits` covering the extra galaxy, so the modeling\n", + "example (`multi/modeling.py`) can load each one and apply noise scaling. The mask is built per dataset because the\n", + "wavebands have different pixel scales and therefore different `shape_native`. The circle is sized to ~3x the\n", + "galaxy's `effective_radius`, derived from the same `extra_galaxy_centre` defined above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " mask_extra_galaxies = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native,\n", + " pixel_scales=dataset.pixel_scales,\n", + " centre=extra_galaxy_centre,\n", + " radius=3.0 * 0.3,\n", + " invert=True, # `True` inside the circle, i.e. the region whose noise is scaled.\n", + " )\n", + "\n", + " aplt.fits_array(\n", + " array=mask_extra_galaxies,\n", + " file_path=Path(dataset_path) / f\"{waveband}_mask_extra_galaxies.fits\",\n", + " overwrite=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for waveband, dataset in zip(waveband_list, dataset_list):\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + " aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "\n", + "for waveband, grid, tracer in zip(waveband_list, grid_list, tracer_list):\n", + " aplt.subplot_tracer(tracer=tracer, grid=grid)\n", + " aplt.subplot_galaxies_images(tracer=tracer, grid=grid)\n", + "\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "[\n", + " al.output_to_json(\n", + " obj=tracer, file_path=Path(dataset_path, f\"{waveband}_tracer.json\")\n", + " )\n", + " for color, tracer in zip(waveband_list, tracer_list)\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The dataset can be viewed in the folder `autolens_workspace/imaging/multi/lens_sersic`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/features/deblending/modeling.ipynb b/notebooks/point_source/features/deblending/modeling.ipynb index 85bd64b51..271314f6e 100644 --- a/notebooks/point_source/features/deblending/modeling.ipynb +++ b/notebooks/point_source/features/deblending/modeling.ipynb @@ -1,545 +1,582 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling Features: Deblending\n", - "=============================\n", - "\n", - "The image-plane multiple-image positions of a lensed point source (e.g. a quasar or supernova) are used as the\n", - "dataset in point-source modeling. For example, simulated values were input into the `PointDataset` object in the\n", - "`point_source/modeling/start_here.ipynb` example.\n", - "\n", - "These positions must first be measured from imaging data of the lensed point-source. A simple way to do this is\n", - "to locate the brightest 2 or 4 pixels of the lensed point-source (e.g. via a GUI or ds9) and use these values\n", - "as the positions.\n", - "\n", - "For many users this will be sufficient, however it has downsides:\n", - "\n", - "- It does not provide sub-pixel precision on the positions.\n", - "\n", - "- It does not account for the Point Spread Function.\n", - "\n", - "It also does not measure the following quantities at all:\n", - "\n", - "- The flux of each lensed point source (it provides an estimate via the brightest pixel fluxes, but proper deblending\n", - "of the PSF is key for accurate flux measurements).\n", - "\n", - "- Any properties of the lens galaxy's light, which is blended with the lensed point source images.\n", - "\n", - "In this example, we perform this deblending so that we can accurately measure the point-source positions, fluxes and\n", - "properties of the lens galaxy's light.\n", - "\n", - "__Contents__\n", - "\n", - "- **Image Plane Multiple Images:** When fitting the `Imaging` dataset in order to deblend the lensed point-source images and lens.\n", - "- **Point Source Host Galaxy:** For high quality imaging of a lensed point source, the light from the point source's host galaxy.\n", - "- **Imaging:** This example script fits `Imaging` data, using many of the features described in the.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", - "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Run Time:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Point Source:** After the analysis above is complete, the lens model infers the following information.\n", - "\n", - "__Image Plane Multiple Images__\n", - "\n", - "When fitting the `Imaging` dataset in order to deblend the lensed point-source images and lens galaxies, the four\n", - "multiple images of the lensed point source are modeled in the image-plane, using four independent light profiles.\n", - "\n", - "This means the model does not place a point-source in the source-plane, and does not use ray-tracing to determine its\n", - "image-plane multiple images and fluxes.\n", - "\n", - "The reason for this is due to a phenomenon called 'micro lensing'. In brief, each multiple image of the lensed\n", - "point source will have its magnification boosted or reduced by stars in the lens galaxy. This lensing effect is\n", - "extremely difficult to model accurately in the lens galaxy's mass model. This means that if we modeled the lensed point\n", - "source in the source-plane, we would not be able to accurately measure its fluxes.\n", - "\n", - "By fitting each multiple image in the image-plane, the effects of micro lensing on each multiple image are accounted\n", - "for in the deblending process by the `intensity` of each light profile being free parameters in the model. Micro\n", - "lensing is also why `fluxes` are typically not used to fit point source lens models.\n", - "\n", - "__Point Source Host Galaxy__\n", - "\n", - "For high quality imaging of a lensed point source, the light from the point source's host galaxy may also be visible.\n", - "The deblending procedure illustrated in this script can also therefore be used to extract and model the host galaxy's\n", - "light.\n", - "\n", - "We do not perform any deblending of the lensed point source's host source galaxy, because it requires a more\n", - "sophisticated analysis. An example script for doing this is not currently available, but if it would be useful for\n", - "you please contact me on SLACK and I can write it!\n", - "\n", - "__Imaging__\n", - "\n", - "This example script fits `Imaging` data, using many of the features described in the `imaging/modeling` workspace\n", - "examples.\n", - "\n", - "It also uses the following features described in the `modeling/features` workspace examples:\n", - "\n", - "- `linear_light_profiles.py`: The model includes light profiles which use linear algebra to solve for their\n", - " intensity, reducing model complexity.\n", - "\n", - "- `advanced/operated_light_profiles.py`: There are light profiles which are assumed to already be convolved with\n", - " the instrumental PSF (e.g. point sources), commonly used for modeling bright AGN in the centre of a galaxy.\n", - "\n", - "It is recommended you are familiar with imaging modeling and these features before reading this example.\n", - "\n", - "However, you would probably be able to use and adapt this script to your use-case even if you are not.\n", - "\n", - "__Model__\n", - "\n", - "This script fits an `Imaging` dataset of a 'galaxy-scale' point-source strong lens with a model where:\n", - "\n", - " - The lens galaxy's light is a linear `Sersic` bulge.\n", - " - The multiple images of the lensed source are each fitted with a `Gaussian` operated linear light profile.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load and plot the strong lens dataset `deblending` via .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"deblending\"\n", - "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not Path(dataset_path).exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/point_source/features/deblending/simulator.py\"],\n", - " check=True,\n", - " )\n", - "\n", - "dataset = al.Imaging.from_fits(\n", - " data_path=Path(dataset_path) / \"data.fits\",\n", - " psf_path=Path(dataset_path) / \"psf.fits\",\n", - " noise_map_path=Path(dataset_path) / \"noise_map.fits\",\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Mask__\n", - "\n", - "Define a 3.0\" circular mask, which includes the emission of the lens and lensed point-sources." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mask = al.Mask2D.circular(\n", - " shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=3.0\n", - ")\n", - "\n", - "dataset = dataset.apply_mask(mask=mask)\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's light is a linear `Sersic` bulge [6 parameters].\n", - " \n", - " - The four image-plane multiple images of the lensed source are each fitted with a `Gaussian` operated linear light \n", - " profile [4 x 5 = 20 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=26.\n", - "\n", - "Note how all light profiles are linear light profiles, meaning that the `intensity` parameter of all profiles are\n", - "not free parameters in the fit but instead are solved via linear algebra. This reduces the dimensionality of the\n", - "non-linear parameter space by N=5.\n", - "\n", - "We note that our lens model therefore does not include:\n", - "\n", - " - A lens galaxy with a total mass distribution.\n", - " - A source galaxy's with a light profile or point source.\n", - "\n", - "__Model Cookbook__\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "bulge = af.Model(al.lp_linear.Sersic)\n", - "\n", - "# Lensed Source Multiple Images (Image-Plane):\n", - "\n", - "multiple_image_0 = af.Model(al.lp_linear_operated.Gaussian)\n", - "multiple_image_1 = af.Model(al.lp_linear_operated.Gaussian)\n", - "multiple_image_2 = af.Model(al.lp_linear_operated.Gaussian)\n", - "multiple_image_3 = af.Model(al.lp_linear_operated.Gaussian)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The model has N=26 free parameters, and is quite a complex model to fit. \n", - "\n", - "The different multiple image light profiles can also produce identical solutions, because each `Gaussian` point source \n", - "can change its `centre` to fit different lensed point source images.\n", - "\n", - "To simplify the model and remove identical solutions, we manually set the priors on the `centre` of the lens galaxy \n", - "light profile and each multiple image light profile to narrow uniform priors. The values of these priors are based on\n", - "where the peak fluxes of each image appear to be located in the image plotted above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge.centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "bulge.centre_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", - "\n", - "multiple_image_0_estimate = (1.148, -1.148)\n", - "multiple_image_1_estimate = (1.190, 1.190)\n", - "multiple_image_2_estimate = (-1.190, -1.190)\n", - "multiple_image_3_estimate = (-1.148, 1.148)\n", - "\n", - "multiple_image_width = dataset.pixel_scales[0]\n", - "\n", - "multiple_image_0.centre_0 = af.UniformPrior(\n", - " lower_limit=multiple_image_0_estimate[0] - multiple_image_width,\n", - " upper_limit=multiple_image_0_estimate[0] + multiple_image_width,\n", - ")\n", - "multiple_image_0.centre_1 = af.UniformPrior(\n", - " lower_limit=multiple_image_0_estimate[1] - multiple_image_width,\n", - " upper_limit=multiple_image_0_estimate[1] + multiple_image_width,\n", - ")\n", - "\n", - "multiple_image_1.centre_0 = af.UniformPrior(\n", - " lower_limit=multiple_image_1_estimate[0] - multiple_image_width,\n", - " upper_limit=multiple_image_1_estimate[0] + multiple_image_width,\n", - ")\n", - "multiple_image_1.centre_1 = af.UniformPrior(\n", - " lower_limit=multiple_image_1_estimate[1] - multiple_image_width,\n", - " upper_limit=multiple_image_1_estimate[1] + multiple_image_width,\n", - ")\n", - "\n", - "multiple_image_2.centre_0 = af.UniformPrior(\n", - " lower_limit=multiple_image_2_estimate[0] - multiple_image_width,\n", - " upper_limit=multiple_image_2_estimate[0] + multiple_image_width,\n", - ")\n", - "multiple_image_2.centre_1 = af.UniformPrior(\n", - " lower_limit=multiple_image_2_estimate[1] - multiple_image_width,\n", - " upper_limit=multiple_image_2_estimate[1] + multiple_image_width,\n", - ")\n", - "\n", - "multiple_image_3.centre_0 = af.UniformPrior(\n", - " lower_limit=multiple_image_3_estimate[0] - multiple_image_width,\n", - " upper_limit=multiple_image_3_estimate[0] + multiple_image_width,\n", - ")\n", - "multiple_image_3.centre_1 = af.UniformPrior(\n", - " lower_limit=multiple_image_3_estimate[1] - multiple_image_width,\n", - " upper_limit=multiple_image_3_estimate[1] + multiple_image_width,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now create the model, using the standard model composition API.\n", - "\n", - "Note that we put each multiple image light profile inside the `lens`. This is a bit of a strange syntax, but \n", - "functionally it works.\n", - "\n", - "Future versions of PyAutoLens will have a more intuitive API for this, but for now we have to do it this way!" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " bulge=bulge,\n", - " multiple_image_0=multiple_image_0,\n", - " multiple_image_1=multiple_image_1,\n", - " multiple_image_2=multiple_image_2,\n", - " multiple_image_3=multiple_image_3,\n", - ")\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that the model only has a `lens` and that it has different components for each multiple image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description).\n", - "\n", - "In the `start_here.py` example 150 live points (`n_live=150`) were used to sample parameter space. For this fit\n", - "we have a much more complex parameter space with N=26 free parameters, therefore we use 400 live points to ensure\n", - "we thoroughly sample parameter space." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"point_source\") / \"modeling\",\n", - " name=\"deblending\",\n", - " unique_tag=dataset_name,\n", - " n_live=400,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisImaging(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Time__\n", - "\n", - "For standard light profiles, the log likelihood evaluation time is of order ~0.0001 seconds for this dataset.\n", - "\n", - "For linear light profiles, the log likelihood evaluation increases to around ~0.001 seconds per likelihood evaluation.\n", - "This is still fast, but it does mean that the fit may take around five times longer to run.\n", - "\n", - "The run time to perform deblending modeling re around 30 minutes on CPU, under 10 minutes on GPU.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that `intensity` parameters are not inferred by the model-fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "The lens and source galaxies appear similar to those in the data, confirming that the `intensity` values inferred by\n", - "the inversion process are accurate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", - "\n", - "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", - "\n", - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Source__\n", - "\n", - "After the analysis above is complete, the lens model infers the following information: \n", - "\n", - "- The `centre` parameter of each multiple image `Gaussian` is the (y,x) image-plane coordinate of each lensed point \n", - " source. These values provide sub-pixel precision, because they fully account for the shape and blending of the PSF. \n", - " Using these as a `positions` of point-source mass modeling will also produce a more accurate lens model.\n", - "\n", - "- The `intensity` value of each `Gaussian` estimates the flux of each point source. The `fluxes` are typically not\n", - " used in point-source modeling as they are subject to microlensing, but this analysis nevertheless does measure them.\n", - "\n", - " - The lens galaxy's properties are measured via this analysis.\n", - "\n", - "The lensed source image-plane positions, inferred to sub-pixel precision, are printed below and output to a \n", - "`PointDataset` object and .json file.\n", - "\n", - "They can be used as input positions in a point-source model-fit, using an identical API to \n", - "the `point_source/modeling/start_here.ipynb` example, to perform mass modeling of the point source dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling Features: Deblending\n", + "=============================\n", + "\n", + "The image-plane multiple-image positions of a lensed point source (e.g. a quasar or supernova) are used as the\n", + "dataset in point-source modeling. For example, simulated values were input into the `PointDataset` object in the\n", + "`point_source/modeling/start_here.ipynb` example.\n", + "\n", + "These positions must first be measured from imaging data of the lensed point-source. A simple way to do this is\n", + "to locate the brightest 2 or 4 pixels of the lensed point-source (e.g. via a GUI or ds9) and use these values\n", + "as the positions.\n", + "\n", + "For many users this will be sufficient, however it has downsides:\n", + "\n", + "- It does not provide sub-pixel precision on the positions.\n", + "\n", + "- It does not account for the Point Spread Function.\n", + "\n", + "It also does not measure the following quantities at all:\n", + "\n", + "- The flux of each lensed point source (it provides an estimate via the brightest pixel fluxes, but proper deblending\n", + "of the PSF is key for accurate flux measurements).\n", + "\n", + "- Any properties of the lens galaxy's light, which is blended with the lensed point source images.\n", + "\n", + "In this example, we perform this deblending so that we can accurately measure the point-source positions, fluxes and\n", + "properties of the lens galaxy's light.\n", + "\n", + "__Contents__\n", + "\n", + "- **Image Plane Multiple Images:** When fitting the `Imaging` dataset in order to deblend the lensed point-source images and lens.\n", + "- **Point Source Host Galaxy:** For high quality imaging of a lensed point source, the light from the point source's host galaxy.\n", + "- **Imaging:** This example script fits `Imaging` data, using many of the features described in the.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset & Mask:** Standard set up of the dataset and mask that is fitted.\n", + "- **Model Cookbook:** A full description of model composition is provided by the model cookbook.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Run Time:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Point Source:** After the analysis above is complete, the lens model infers the following information.\n", + "\n", + "__Image Plane Multiple Images__\n", + "\n", + "When fitting the `Imaging` dataset in order to deblend the lensed point-source images and lens galaxies, the four\n", + "multiple images of the lensed point source are modeled in the image-plane, using four independent light profiles.\n", + "\n", + "This means the model does not place a point-source in the source-plane, and does not use ray-tracing to determine its\n", + "image-plane multiple images and fluxes.\n", + "\n", + "The reason for this is due to a phenomenon called 'micro lensing'. In brief, each multiple image of the lensed\n", + "point source will have its magnification boosted or reduced by stars in the lens galaxy. This lensing effect is\n", + "extremely difficult to model accurately in the lens galaxy's mass model. This means that if we modeled the lensed point\n", + "source in the source-plane, we would not be able to accurately measure its fluxes.\n", + "\n", + "By fitting each multiple image in the image-plane, the effects of micro lensing on each multiple image are accounted\n", + "for in the deblending process by the `intensity` of each light profile being free parameters in the model. Micro\n", + "lensing is also why `fluxes` are typically not used to fit point source lens models.\n", + "\n", + "__Point Source Host Galaxy__\n", + "\n", + "For high quality imaging of a lensed point source, the light from the point source's host galaxy may also be visible.\n", + "The deblending procedure illustrated in this script can also therefore be used to extract and model the host galaxy's\n", + "light.\n", + "\n", + "We do not perform any deblending of the lensed point source's host source galaxy, because it requires a more\n", + "sophisticated analysis. An example script for doing this is not currently available, but if it would be useful for\n", + "you please contact me on SLACK and I can write it!\n", + "\n", + "__Imaging__\n", + "\n", + "This example script fits `Imaging` data, using many of the features described in the `imaging/modeling` workspace\n", + "examples.\n", + "\n", + "It also uses the following features described in the `modeling/features` workspace examples:\n", + "\n", + "- `linear_light_profiles.py`: The model includes light profiles which use linear algebra to solve for their\n", + " intensity, reducing model complexity.\n", + "\n", + "- `advanced/operated_light_profiles.py`: There are light profiles which are assumed to already be convolved with\n", + " the instrumental PSF (e.g. point sources), commonly used for modeling bright AGN in the centre of a galaxy.\n", + "\n", + "It is recommended you are familiar with imaging modeling and these features before reading this example.\n", + "\n", + "However, you would probably be able to use and adapt this script to your use-case even if you are not.\n", + "\n", + "__Model__\n", + "\n", + "This script fits an `Imaging` dataset of a 'galaxy-scale' point-source strong lens with a model where:\n", + "\n", + " - The lens galaxy's light is a linear `Sersic` bulge.\n", + " - The multiple images of the lensed source are each fitted with a `Gaussian` operated linear light profile.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load and plot the strong lens dataset `deblending` via .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"deblending\"\n", + "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not Path(dataset_path).exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/point_source/features/deblending/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.Imaging.from_fits(\n", + " data_path=Path(dataset_path) / \"data.fits\",\n", + " psf_path=Path(dataset_path) / \"psf.fits\",\n", + " noise_map_path=Path(dataset_path) / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mask__\n", + "\n", + "Define a 3.0\" circular mask, which includes the emission of the lens and lensed point-sources." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=dataset.shape_native, pixel_scales=dataset.pixel_scales, radius=3.0\n", + ")\n", + "\n", + "dataset = dataset.apply_mask(mask=mask)\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's light is a linear `Sersic` bulge [6 parameters].\n", + " \n", + " - The four image-plane multiple images of the lensed source are each fitted with a `Gaussian` operated linear light \n", + " profile [4 x 5 = 20 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=26.\n", + "\n", + "Note how all light profiles are linear light profiles, meaning that the `intensity` parameter of all profiles are\n", + "not free parameters in the fit but instead are solved via linear algebra. This reduces the dimensionality of the\n", + "non-linear parameter space by N=5.\n", + "\n", + "We note that our lens model therefore does not include:\n", + "\n", + " - A lens galaxy with a total mass distribution.\n", + " - A source galaxy's with a light profile or point source.\n", + "\n", + "__Model Cookbook__\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "bulge = af.Model(al.lp_linear.Sersic)\n", + "\n", + "# Lensed Source Multiple Images (Image-Plane):\n", + "\n", + "multiple_image_0 = af.Model(al.lp_linear_operated.Gaussian)\n", + "multiple_image_1 = af.Model(al.lp_linear_operated.Gaussian)\n", + "multiple_image_2 = af.Model(al.lp_linear_operated.Gaussian)\n", + "multiple_image_3 = af.Model(al.lp_linear_operated.Gaussian)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The model has N=26 free parameters, and is quite a complex model to fit. \n", + "\n", + "The different multiple image light profiles can also produce identical solutions, because each `Gaussian` point source \n", + "can change its `centre` to fit different lensed point source images.\n", + "\n", + "To simplify the model and remove identical solutions, we manually set the priors on the `centre` of the lens galaxy \n", + "light profile and each multiple image light profile to narrow uniform priors. The values of these priors are based on\n", + "where the peak fluxes of each image appear to be located in the image plotted above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge.centre_0 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "bulge.centre_1 = af.UniformPrior(lower_limit=-0.1, upper_limit=0.1)\n", + "\n", + "multiple_image_0_estimate = (1.148, -1.148)\n", + "multiple_image_1_estimate = (1.190, 1.190)\n", + "multiple_image_2_estimate = (-1.190, -1.190)\n", + "multiple_image_3_estimate = (-1.148, 1.148)\n", + "\n", + "multiple_image_width = dataset.pixel_scales[0]\n", + "\n", + "multiple_image_0.centre_0 = af.UniformPrior(\n", + " lower_limit=multiple_image_0_estimate[0] - multiple_image_width,\n", + " upper_limit=multiple_image_0_estimate[0] + multiple_image_width,\n", + ")\n", + "multiple_image_0.centre_1 = af.UniformPrior(\n", + " lower_limit=multiple_image_0_estimate[1] - multiple_image_width,\n", + " upper_limit=multiple_image_0_estimate[1] + multiple_image_width,\n", + ")\n", + "\n", + "multiple_image_1.centre_0 = af.UniformPrior(\n", + " lower_limit=multiple_image_1_estimate[0] - multiple_image_width,\n", + " upper_limit=multiple_image_1_estimate[0] + multiple_image_width,\n", + ")\n", + "multiple_image_1.centre_1 = af.UniformPrior(\n", + " lower_limit=multiple_image_1_estimate[1] - multiple_image_width,\n", + " upper_limit=multiple_image_1_estimate[1] + multiple_image_width,\n", + ")\n", + "\n", + "multiple_image_2.centre_0 = af.UniformPrior(\n", + " lower_limit=multiple_image_2_estimate[0] - multiple_image_width,\n", + " upper_limit=multiple_image_2_estimate[0] + multiple_image_width,\n", + ")\n", + "multiple_image_2.centre_1 = af.UniformPrior(\n", + " lower_limit=multiple_image_2_estimate[1] - multiple_image_width,\n", + " upper_limit=multiple_image_2_estimate[1] + multiple_image_width,\n", + ")\n", + "\n", + "multiple_image_3.centre_0 = af.UniformPrior(\n", + " lower_limit=multiple_image_3_estimate[0] - multiple_image_width,\n", + " upper_limit=multiple_image_3_estimate[0] + multiple_image_width,\n", + ")\n", + "multiple_image_3.centre_1 = af.UniformPrior(\n", + " lower_limit=multiple_image_3_estimate[1] - multiple_image_width,\n", + " upper_limit=multiple_image_3_estimate[1] + multiple_image_width,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now create the model, using the standard model composition API.\n", + "\n", + "Note that we put each multiple image light profile inside the `lens`. This is a bit of a strange syntax, but \n", + "functionally it works.\n", + "\n", + "Future versions of PyAutoLens will have a more intuitive API for this, but for now we have to do it this way!" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " bulge=bulge,\n", + " multiple_image_0=multiple_image_0,\n", + " multiple_image_1=multiple_image_1,\n", + " multiple_image_2=multiple_image_2,\n", + " multiple_image_3=multiple_image_3,\n", + ")\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that the model only has a `lens` and that it has different components for each multiple image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description).\n", + "\n", + "In the `start_here.py` example 150 live points (`n_live=150`) were used to sample parameter space. For this fit\n", + "we have a much more complex parameter space with N=26 free parameters, therefore we use 400 live points to ensure\n", + "we thoroughly sample parameter space." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"point_source\") / \"modeling\",\n", + " name=\"deblending\",\n", + " unique_tag=dataset_name,\n", + " n_live=400,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisImaging` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisImaging(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Time__\n", + "\n", + "For standard light profiles, the log likelihood evaluation time is of order ~0.0001 seconds for this dataset.\n", + "\n", + "For linear light profiles, the log likelihood evaluation increases to around ~0.001 seconds per likelihood evaluation.\n", + "This is still fast, but it does mean that the fit may take around five times longer to run.\n", + "\n", + "The run time to perform deblending modeling re around 30 minutes on CPU, under 10 minutes on GPU.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that `intensity` parameters are not inferred by the model-fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "The lens and source galaxies appear similar to those in the data, confirming that the `intensity` values inferred by\n", + "the inversion process are accurate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grids.lp)\n", + "\n", + "aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Source__\n", + "\n", + "After the analysis above is complete, the lens model infers the following information: \n", + "\n", + "- The `centre` parameter of each multiple image `Gaussian` is the (y,x) image-plane coordinate of each lensed point \n", + " source. These values provide sub-pixel precision, because they fully account for the shape and blending of the PSF. \n", + " Using these as a `positions` of point-source mass modeling will also produce a more accurate lens model.\n", + "\n", + "- The `intensity` value of each `Gaussian` estimates the flux of each point source. The `fluxes` are typically not\n", + " used in point-source modeling as they are subject to microlensing, but this analysis nevertheless does measure them.\n", + "\n", + " - The lens galaxy's properties are measured via this analysis.\n", + "\n", + "The lensed source image-plane positions, inferred to sub-pixel precision, are printed below and output to a \n", + "`PointDataset` object and .json file.\n", + "\n", + "They can be used as input positions in a point-source model-fit, using an identical API to \n", + "the `point_source/modeling/start_here.ipynb` example, to perform mass modeling of the point source dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/features/deblending/simulator.ipynb b/notebooks/point_source/features/deblending/simulator.ipynb index 69b053612..c6b7a88a6 100644 --- a/notebooks/point_source/features/deblending/simulator.ipynb +++ b/notebooks/point_source/features/deblending/simulator.ipynb @@ -1,640 +1,677 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Debelending\n", - "======================\n", - "\n", - "This script simulations a `Point` dataset of a galaxy-scale strong lens which is identical to the dataset\n", - "simulated in the `start_here.ipynb` example, but where an image of the multiply imaged lensed point source (e.g.\n", - "the quasar) and its lens galaxy are included.\n", - "\n", - "It is used in `autolens_workspace/notebooks/point_source/modeling/features/deblending.ipynb` to illustrate how to\n", - "perform deblending of a point source dataset, in order to measure the image-plane multiple image positions, fluxes\n", - "and lens galaxy light.\n", - "\n", - "The simulation procedure in this script simulates the lens in two steps:\n", - "\n", - "1) Simulate the point-source dataset, in an identical fashion to the `start_here.ipynb` example.\n", - "2) Use this result to simulate the imaging dataset of the lensed point source and lens galaxy.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Point Solver:** We use a `PointSolver` to locate the multiple images.\n", - "- **Fluxes:** Use the positions to compute the magnification of the `Tracer` at every position.\n", - "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `Imaging` and `PointDataset` data of a strong lens where:\n", - "\n", - " - The lens galaxy's light profile is a `Sersic`.\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The source `Galaxy` is a `Point`.\n", - " - The multiple images of each lensed point source are `Gaussian` which already represent the PSF convolved images.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import numpy as np\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"point_source\"\n", - "dataset_name = \"deblending\"\n", - "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing (Point Source)__\n", - "\n", - "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " point_0=al.ps.Point(centre=(0.0, 0.0)),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will compute the multiple image positions of the simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "We use a `PointSolver` to locate the multiple images. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(200, 200),\n", - " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now pass the `Tracer` to the solver. This will then find the image-plane coordinates that map directly to the\n", - "source-plane coordinate (0.0\", 0.0\").\n", - "\n", - "Position noise is set to 0.005\" (5 mas), reflecting realistic PSF-centroiding precision on HST imaging\n", - "rather than the imaging pixel scale. Flux noise is set to 5% relative, reflecting that lensed-quasar\n", - "flux uncertainties are dominated by microlensing systematics rather than photon noise. See\n", - "`scripts/point_source/simulator.py` for a full discussion of these values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "position_noise = 0.005\n", - "flux_rel_noise = 0.05\n", - "\n", - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.point_0.centre\n", - ")\n", - "\n", - "positions_with_noise = positions + np.random.normal(\n", - " loc=0.0, scale=position_noise, size=positions.shape\n", - ")\n", - "\n", - "positions_with_noise = al.Grid2DIrregular(\n", - " values=positions_with_noise,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fluxes__\n", - "\n", - "Use the positions to compute the magnification of the `Tracer` at every position." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnifications = al.LensCalc.from_tracer(\n", - " tracer=tracer\n", - ").magnification_2d_via_hessian_from(grid=positions)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now compute the observed fluxes of the `Point`, give we know how much each is magnified." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "flux = 1.0\n", - "fluxes = [flux * np.abs(magnification) for magnification in magnifications]\n", - "fluxes = al.ArrayIrregular(values=fluxes)\n", - "\n", - "fluxes_with_noise = fluxes + np.random.normal(\n", - " loc=0.0, scale=flux_rel_noise * np.asarray(fluxes), size=len(fluxes)\n", - ")\n", - "\n", - "fluxes_with_noise = al.ArrayIrregular(values=fluxes_with_noise)\n", - "\n", - "fluxes_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes))\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Datasets (Point Source)__\n", - "\n", - "Create the `PointDataset` and `PointDataset` objects using identical code to the `start_here.ipynb` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_with_noise,\n", - " positions_noise_map=position_noise,\n", - " fluxes=fluxes_with_noise,\n", - " fluxes_noise_map=fluxes_noise_map,\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=dataset,\n", - " file_path=Path(dataset_path, \"point_dataset_positions_only.json\"),\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize (Point Source)__\n", - "\n", - "Visualize the `PointDataset` using identical code to the `start_here.ipynb` example." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_point_dataset(\n", - " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output subplots of the tracer's images, including the positions of the multiple images on the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json (Point Source)__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer_point.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate (Imaging)__\n", - "\n", - "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", - " grid=grid,\n", - " sub_size_list=[32, 8, 2],\n", - " radial_list=[0.3, 0.6],\n", - " centre_list=[(0.0, 0.0)],\n", - ")\n", - "\n", - "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulate a simple Gaussian PSF for the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf_sigma = 0.1\n", - "\n", - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=psf_sigma, pixel_scales=grid.pixel_scales\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lensed Source Image (Imaging)__\n", - "\n", - "The `positions` and `fluxes` above represent the location and brightnesses of the multiple images in the image-plane.\n", - "\n", - "To include these multiple images in the imaging simulation, we add each multiple image individually in the image-plane. \n", - "These multiple images are assumed to have already been convolved with the PSF, which is why they use the `lp_operated` \n", - "profile (see `autolens_workspace/*/notebooks/modeling/features/advanced/operated_light_profiles.py`).\n", - "\n", - "The `Imaging` simulation procedure therefore does not place a point-source in the source-plane, and use ray-tracing\n", - "to determine its image-plane multiple images. It is effectively doing this, because it uses the `positions` and\n", - "`fluxes` above to add the multiple images in the image-plane, but the `Tracer` below does not explicitly perform\n", - "this ray-tracing calculation.\n", - "\n", - "The reason we choose this approach is because it is closer to how we model the multiple images of actual lensed point \n", - "sources, where each multiple image is modeled in the image-plane as a separate light \n", - "profile (see `point_source/modeling/features/debeleing.ipynb` for a description of why)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "point_image_kwargs = {\n", - " f\"point_image_{i}\": al.lp_operated.Gaussian(\n", - " centre=positions[i], intensity=fluxes[i], sigma=psf_sigma\n", - " )\n", - " for i in range(len(positions))\n", - "}\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " bulge=al.lp.Sersic(\n", - " centre=(0.0, 0.0),\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " intensity=2.0,\n", - " effective_radius=0.6,\n", - " sersic_index=3.0,\n", - " ),\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - " **point_image_kwargs,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The source galaxy now long uses a `Point` component as the multiple images are included in the image-plane instead." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - ")\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Imaging` simulation now uses the normal API for simulating images." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "# %%\n", - "'''\n", - "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset.\n", - "'''" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Plot the simulated `Imaging` dataset before outputting it to fits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_imaging_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Output the simulated dataset to the dataset path as .fits files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "We now output the image of this strong lens to `.fits` which can be used for visualize when performing point-source \n", - "modeling and to `.png` for general inspection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_imaging_dataset(dataset=dataset)\n", - "aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - "aplt.plot_array(\n", - " array=tracer.image_2d_from(grid=grid),\n", - " title=\"Image\",\n", - " output_path=dataset_path,\n", - " output_format=\"png\",\n", - ")\n", - "\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=Path(dataset_path, \"tracer_imaging.json\"),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Debelending\n", + "======================\n", + "\n", + "This script simulations a `Point` dataset of a galaxy-scale strong lens which is identical to the dataset\n", + "simulated in the `start_here.ipynb` example, but where an image of the multiply imaged lensed point source (e.g.\n", + "the quasar) and its lens galaxy are included.\n", + "\n", + "It is used in `autolens_workspace/notebooks/point_source/modeling/features/deblending.ipynb` to illustrate how to\n", + "perform deblending of a point source dataset, in order to measure the image-plane multiple image positions, fluxes\n", + "and lens galaxy light.\n", + "\n", + "The simulation procedure in this script simulates the lens in two steps:\n", + "\n", + "1) Simulate the point-source dataset, in an identical fashion to the `start_here.ipynb` example.\n", + "2) Use this result to simulate the imaging dataset of the lensed point source and lens galaxy.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Point Solver:** We use a `PointSolver` to locate the multiple images.\n", + "- **Fluxes:** Use the positions to compute the magnification of the `Tracer` at every position.\n", + "- **Output:** Output the simulated dataset to the dataset path as .fits files.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `Imaging` and `PointDataset` data of a strong lens where:\n", + "\n", + " - The lens galaxy's light profile is a `Sersic`.\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The source `Galaxy` is a `Point`.\n", + " - The multiple images of each lensed point source are `Gaussian` which already represent the PSF convolved images.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import numpy as np\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"point_source\"\n", + "dataset_name = \"deblending\"\n", + "dataset_path = Path(\"dataset\", dataset_type, dataset_name)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing (Point Source)__\n", + "\n", + "Setup the lens galaxy's light, mass and source galaxy light for this simulated lens." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " point_0=al.ps.Point(centre=(0.0, 0.0)),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will compute the multiple image positions of the simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "We use a `PointSolver` to locate the multiple images. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(200, 200),\n", + " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now pass the `Tracer` to the solver. This will then find the image-plane coordinates that map directly to the\n", + "source-plane coordinate (0.0\", 0.0\").\n", + "\n", + "Position noise is set to 0.005\" (5 mas), reflecting realistic PSF-centroiding precision on HST imaging\n", + "rather than the imaging pixel scale. Flux noise is set to 5% relative, reflecting that lensed-quasar\n", + "flux uncertainties are dominated by microlensing systematics rather than photon noise. See\n", + "`scripts/point_source/simulator.py` for a full discussion of these values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "position_noise = 0.005\n", + "flux_rel_noise = 0.05\n", + "\n", + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.point_0.centre\n", + ")\n", + "\n", + "positions_with_noise = positions + np.random.normal(\n", + " loc=0.0, scale=position_noise, size=positions.shape\n", + ")\n", + "\n", + "positions_with_noise = al.Grid2DIrregular(\n", + " values=positions_with_noise,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fluxes__\n", + "\n", + "Use the positions to compute the magnification of the `Tracer` at every position." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnifications = al.LensCalc.from_tracer(\n", + " tracer=tracer\n", + ").magnification_2d_via_hessian_from(grid=positions)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now compute the observed fluxes of the `Point`, give we know how much each is magnified." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "flux = 1.0\n", + "fluxes = [flux * np.abs(magnification) for magnification in magnifications]\n", + "fluxes = al.ArrayIrregular(values=fluxes)\n", + "\n", + "fluxes_with_noise = fluxes + np.random.normal(\n", + " loc=0.0, scale=flux_rel_noise * np.asarray(fluxes), size=len(fluxes)\n", + ")\n", + "\n", + "fluxes_with_noise = al.ArrayIrregular(values=fluxes_with_noise)\n", + "\n", + "fluxes_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes))\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Datasets (Point Source)__\n", + "\n", + "Create the `PointDataset` and `PointDataset` objects using identical code to the `start_here.ipynb` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_with_noise,\n", + " positions_noise_map=position_noise,\n", + " fluxes=fluxes_with_noise,\n", + " fluxes_noise_map=fluxes_noise_map,\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=dataset,\n", + " file_path=Path(dataset_path, \"point_dataset_positions_only.json\"),\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize (Point Source)__\n", + "\n", + "Visualize the `PointDataset` using identical code to the `start_here.ipynb` example." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_point_dataset(\n", + " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output subplots of the tracer's images, including the positions of the multiple images on the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json (Point Source)__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer_point.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate (Imaging)__\n", + "\n", + "Simulate the image using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "over_sample_size = al.util.over_sample.over_sample_size_via_radial_bins_from(\n", + " grid=grid,\n", + " sub_size_list=[32, 8, 2],\n", + " radial_list=[0.3, 0.6],\n", + " centre_list=[(0.0, 0.0)],\n", + ")\n", + "\n", + "grid = grid.apply_over_sampling(over_sample_size=over_sample_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulate a simple Gaussian PSF for the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf_sigma = 0.1\n", + "\n", + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=psf_sigma, pixel_scales=grid.pixel_scales\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the simulator for the imaging data, which defines the exposure time, background sky, noise levels and psf." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lensed Source Image (Imaging)__\n", + "\n", + "The `positions` and `fluxes` above represent the location and brightnesses of the multiple images in the image-plane.\n", + "\n", + "To include these multiple images in the imaging simulation, we add each multiple image individually in the image-plane. \n", + "These multiple images are assumed to have already been convolved with the PSF, which is why they use the `lp_operated` \n", + "profile (see `autolens_workspace/*/notebooks/modeling/features/advanced/operated_light_profiles.py`).\n", + "\n", + "The `Imaging` simulation procedure therefore does not place a point-source in the source-plane, and use ray-tracing\n", + "to determine its image-plane multiple images. It is effectively doing this, because it uses the `positions` and\n", + "`fluxes` above to add the multiple images in the image-plane, but the `Tracer` below does not explicitly perform\n", + "this ray-tracing calculation.\n", + "\n", + "The reason we choose this approach is because it is closer to how we model the multiple images of actual lensed point \n", + "sources, where each multiple image is modeled in the image-plane as a separate light \n", + "profile (see `point_source/modeling/features/debeleing.ipynb` for a description of why)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "point_image_kwargs = {\n", + " f\"point_image_{i}\": al.lp_operated.Gaussian(\n", + " centre=positions[i], intensity=fluxes[i], sigma=psf_sigma\n", + " )\n", + " for i in range(len(positions))\n", + "}\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " bulge=al.lp.Sersic(\n", + " centre=(0.0, 0.0),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " intensity=2.0,\n", + " effective_radius=0.6,\n", + " sersic_index=3.0,\n", + " ),\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + " **point_image_kwargs,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The source galaxy now long uses a `Point` component as the multiple images are included in the image-plane instead." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Imaging` simulation now uses the normal API for simulating images." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# %%\n", + "'''\n", + "Pass the simulator a tracer, which creates the image which is simulated as an imaging dataset.\n", + "'''" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Plot the simulated `Imaging` dataset before outputting it to fits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Output the simulated dataset to the dataset path as .fits files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "We now output the image of this strong lens to `.fits` which can be used for visualize when performing point-source \n", + "modeling and to `.png` for general inspection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_imaging_dataset(dataset=dataset)\n", + "aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + "aplt.plot_array(\n", + " array=tracer.image_2d_from(grid=grid),\n", + " title=\"Image\",\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")\n", + "\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=Path(dataset_path, \"tracer_imaging.json\"),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/features/fluxes.ipynb b/notebooks/point_source/features/fluxes.ipynb index 4fc287ae4..b13ca3a40 100644 --- a/notebooks/point_source/features/fluxes.ipynb +++ b/notebooks/point_source/features/fluxes.ipynb @@ -1,406 +1,443 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Fluxes\n", - "================\n", - "\n", - "A measurable quantity of a point source is its flux\u2014the total amount of light received from each multiple image\n", - "of the point source (e.g., the quasar images).\n", - "\n", - "In practice, fluxes are often measured but not used directly when analyzing lensed point sources such as quasars or\n", - "supernovae. This is because fluxes can be significantly affected by microlensing, which many lens models do not\n", - "accurately capture. However, in this simulation, microlensing is not included, so the fluxes can be simulated and\n", - "fitted reliably.\n", - "\n", - "Nevertheless, this script describes how to perform point source lens modeling using the fluxes of the point source\n", - "dataset as additional information on top of the positions of the point source, in case you are studying microlensing\n", - "or confident the fluxes are not affected by it.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Point Solver:** We set up the `PointSolver`, which is used to compute the multiple images of the point source in.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Run Times:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "\n", - "__Model__\n", - "\n", - "This script fits a `PointDataset` data of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The source `Galaxy` is a point source with flux, a `PointFlux`.\n", - "\n", - "The `ExternalShear` is also not included in the mass model, where it is for the `imaging` and `interferometer` examples.\n", - "For a quadruply imaged point source (8 data points) there is insufficient information to fully constain a model with\n", - "an `Isothermal` and `ExternalShear` (9 parameters).\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens point-source dataset `simple`, which is the dataset we will use to perform point source \n", - "lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/point_source/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now load the point source dataset we will fit using point source modeling. \n", - "\n", - "We load this data as a `PointDataset`, which contains the positions and fluxes of every point source. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.from_json(\n", - " file_path=dataset_path / \"point_dataset_with_fluxes.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can print this dictionary to see the dataset's `name`, `positions` and `fluxes` and noise-map values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Point Dataset Info:\")\n", - "print(dataset.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also plot the positions and fluxes of the `PointDataset`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_point_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We next load an image of the dataset and plot the point source data over it, because as described in \n", - "the `modeling/start_here.ipynb` notebook, it is useful for visualizing the point source dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.05)\n", - "\n", - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "We set up the `PointSolver`, which is used to compute the multiple images of the point source in the image-plane.\n", - "\n", - "There are no special settings or inputs for the fitting of fluxes, therefore the `PointSolver` is set up in the same way\n", - "as in the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", - " - The source galaxy's light is a point `PointFlux` [3 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=8.\n", - "\n", - "Name pairing is used as before to pair the `PointDataset` to the `Point` in the model, which is discussed below.\n", - "\n", - "To fit fluxes, our model point source also needs a flux parameter, which is done by using the `PointFlux`\n", - "component instead of the `Point` component. This has a free parameter `flux`, which is the flux of the point source\n", - "in the source-plane. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal)\n", - "\n", - "# Source:\n", - "\n", - "point_0 = af.Model(al.ps.PointFlux)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format, which now includes the `flux` parameter of the point source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description).\n", - "\n", - "In the `start_here.py` example 100 live points (`n_live=100`) were used to sample parameter space. We increase this\n", - "to 150, to account for the additional free parameters in the model that is the source flux." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"point_source\") / \"features\",\n", - " name=\"fluxes\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see modeling examples.\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisPoint` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisPoint(\n", - " dataset=dataset,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePairRepeat, # Image-plane chi-squared with repeat image pairs.\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "For the positions-only fit, the run time of the log likelihood function was ~0.01 seconds, which is a modest run-time.\n", - "\n", - "Evaluating the time delays does not increase this much, with a value of around ~0.01 seconds still expected.\n", - "\n", - "Overall modeling run times should therefore be around 20 minutes on CPU, under 5 minutes on GPU.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this).\n", - "\n", - "This confirms that `flux` parameters of the source is inferred by the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Fluxes\n", + "================\n", + "\n", + "A measurable quantity of a point source is its flux\u2014the total amount of light received from each multiple image\n", + "of the point source (e.g., the quasar images).\n", + "\n", + "In practice, fluxes are often measured but not used directly when analyzing lensed point sources such as quasars or\n", + "supernovae. This is because fluxes can be significantly affected by microlensing, which many lens models do not\n", + "accurately capture. However, in this simulation, microlensing is not included, so the fluxes can be simulated and\n", + "fitted reliably.\n", + "\n", + "Nevertheless, this script describes how to perform point source lens modeling using the fluxes of the point source\n", + "dataset as additional information on top of the positions of the point source, in case you are studying microlensing\n", + "or confident the fluxes are not affected by it.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Point Solver:** We set up the `PointSolver`, which is used to compute the multiple images of the point source in.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Run Times:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "\n", + "__Model__\n", + "\n", + "This script fits a `PointDataset` data of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The source `Galaxy` is a point source with flux, a `PointFlux`.\n", + "\n", + "The `ExternalShear` is also not included in the mass model, where it is for the `imaging` and `interferometer` examples.\n", + "For a quadruply imaged point source (8 data points) there is insufficient information to fully constain a model with\n", + "an `Isothermal` and `ExternalShear` (9 parameters).\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens point-source dataset `simple`, which is the dataset we will use to perform point source \n", + "lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/point_source/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now load the point source dataset we will fit using point source modeling. \n", + "\n", + "We load this data as a `PointDataset`, which contains the positions and fluxes of every point source. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.from_json(\n", + " file_path=dataset_path / \"point_dataset_with_fluxes.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can print this dictionary to see the dataset's `name`, `positions` and `fluxes` and noise-map values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Point Dataset Info:\")\n", + "print(dataset.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also plot the positions and fluxes of the `PointDataset`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_point_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We next load an image of the dataset and plot the point source data over it, because as described in \n", + "the `modeling/start_here.ipynb` notebook, it is useful for visualizing the point source dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.05)\n", + "\n", + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "We set up the `PointSolver`, which is used to compute the multiple images of the point source in the image-plane.\n", + "\n", + "There are no special settings or inputs for the fitting of fluxes, therefore the `PointSolver` is set up in the same way\n", + "as in the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + " - The source galaxy's light is a point `PointFlux` [3 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=8.\n", + "\n", + "Name pairing is used as before to pair the `PointDataset` to the `Point` in the model, which is discussed below.\n", + "\n", + "To fit fluxes, our model point source also needs a flux parameter, which is done by using the `PointFlux`\n", + "component instead of the `Point` component. This has a free parameter `flux`, which is the flux of the point source\n", + "in the source-plane. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal)\n", + "\n", + "# Source:\n", + "\n", + "point_0 = af.Model(al.ps.PointFlux)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format, which now includes the `flux` parameter of the point source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description).\n", + "\n", + "In the `start_here.py` example 100 live points (`n_live=100`) were used to sample parameter space. We increase this\n", + "to 150, to account for the additional free parameters in the model that is the source flux." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"point_source\") / \"features\",\n", + " name=\"fluxes\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see modeling examples.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisPoint` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisPoint(\n", + " dataset=dataset,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePairRepeat, # Image-plane chi-squared with repeat image pairs.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "For the positions-only fit, the run time of the log likelihood function was ~0.01 seconds, which is a modest run-time.\n", + "\n", + "Evaluating the time delays does not increase this much, with a value of around ~0.01 seconds still expected.\n", + "\n", + "Overall modeling run times should therefore be around 20 minutes on CPU, under 5 minutes on GPU.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this).\n", + "\n", + "This confirms that `flux` parameters of the source is inferred by the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/features/multiple_sources/modeling.ipynb b/notebooks/point_source/features/multiple_sources/modeling.ipynb index d1ef9abad..1eb9055fe 100644 --- a/notebooks/point_source/features/multiple_sources/modeling.ipynb +++ b/notebooks/point_source/features/multiple_sources/modeling.ipynb @@ -1,523 +1,560 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Multiple Sources\n", - "==========================\n", - "\n", - "This script fits a `PointDataset` of a 'galaxy-scale' strong lens with multiple lensed point sources at different\n", - "redshifts. The lens system is multi-plane: a foreground lens at z=0.5 deflects both background sources, while\n", - "source_0 at z=1.0 is itself a deflector for source_1 at z=2.0 (the \"double Einstein cross\" configuration). Each\n", - "source's multiple images are stored in their own `PointDataset`, and the two datasets are fitted jointly using\n", - "the multi/factor-graph API:\n", - "\n", - " - The foreground lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The first source `Galaxy` (at z=1.0) is itself a `Galaxy` with a mass profile and a `Point`.\n", - " - The second source `Galaxy` (at z=2.0) is a `Point`-only galaxy.\n", - "\n", - "Two `PointDataset`s are fitted simultaneously, one per lensed source. The fit uses one `AnalysisPoint` per\n", - "dataset, each wrapped in an `AnalysisFactor` that pairs it with the shared lens model. The factors are combined\n", - "into a `FactorGraphModel`, which sums the individual log-likelihoods to form the global log-likelihood the\n", - "non-linear search optimises. Multi-plane lensing is handled automatically inside `AnalysisPoint`, which uses each\n", - "`Point`'s plane redshift in the tracer when solving for image-plane positions.\n", - "\n", - "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for point-source modeling\n", - "(see `point_source/modeling.py`) and the multi/factor-graph API (see `multi/modeling.py` and `cluster/modeling.py`).\n", - "Common boilerplate is therefore not re-explained in detail here.\n", - "\n", - "__Currently Blocked By PyAutoLens #480__\n", - "\n", - "This script does not run end-to-end on the current PyAutoLens release. The `PointSolver` magnification filter\n", - "uses the tracer's last-plane magnification instead of the requested `plane_redshift`'s magnification, so every\n", - "likelihood evaluation finds 0 image positions for source_0 (whose plane z=1.0 is intermediate). See\n", - "https://github.com/PyAutoLabs/PyAutoLens/issues/480 \u2014 both this script and `simulator.py` are listed in\n", - "`config/build/no_run.yaml` until that bug is fixed. The script is left here in its intended form so the example\n", - "is correct as soon as #480 lands; no script changes will be needed once it does.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load the list of `PointDataset` objects, one per lensed source.\n", - "- **Point Solver:** We set up the `PointSolver`, which determines the multiple images of each point source.\n", - "- **Model:** Compose the multi-plane lens model fitted to the data.\n", - "- **Name Pairing:** Each `PointDataset` name is paired with a `Point` model component of the same name.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis List:** Set up one `AnalysisPoint` per dataset.\n", - "- **Analysis Factor:** Each analysis is wrapped in an `AnalysisFactor` paired with the shared lens model.\n", - "- **Factor Graph:** All `AnalysisFactor` objects are combined into a `FactorGraphModel`.\n", - "- **Model-Fit:** Pass the factor graph to the non-linear search.\n", - "- **Result:** Iterate the per-analysis results returned by the factor-graph fit.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `point_source/modeling.ipynb` notebook for the single-source\n", - "case, and `multi/modeling.ipynb` for the factor-graph API used to combine multiple datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens point-source dataset `multiple_sources`, which is the dataset we will fit. The simulator\n", - "writes one `PointDataset` per lensed source (`point_dataset_0.json` and `point_dataset_1.json`), one for each\n", - "source-plane redshift." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"multiple_sources\"\n", - "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it is created by running the corresponding simulator script.\n", - "This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/point_source/features/multiple_sources/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We load each point-source dataset as a `PointDataset` and place them into a list. Every entry in the list is\n", - "fitted by its own `AnalysisPoint` further down in the script." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_list = [\n", - " al.from_json(file_path=dataset_path / f\"point_dataset_{i}.json\") for i in range(2)\n", - "]\n", - "\n", - "for dataset in dataset_list:\n", - " print(\"Point Dataset Info:\")\n", - " print(dataset.info)\n", - " aplt.subplot_point_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "For point-source modeling we require a `PointSolver`, which determines the multiple images of the mass model for\n", - "a point source at location (y,x) in the source plane. It does this by ray tracing triangles from the image-plane\n", - "to the source-plane and refining the multiple images to sub-pixel precision.\n", - "\n", - "The solver requires a starting grid of (y,x) image-plane coordinates and a `pixel_scale_precision` controlling\n", - "the precision of the converged multiple images. The grid below matches the simulator so the solver covers the\n", - "same region of sky used to generate the data.\n", - "\n", - "Strong lens mass models have a \"central image\" which is nearly always significantly demagnified and not observed.\n", - "Setting `magnification_threshold=0.1` discards this image so it does not contaminate the fit." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(200, 200),\n", - " pixel_scales=0.05,\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid,\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a multi-plane lens model where:\n", - "\n", - " - The lens galaxy at z=0.5 has an `Isothermal` mass distribution with `ExternalShear` [7 parameters].\n", - "\n", - " - The first source galaxy at z=1.0 has its own `Isothermal` mass distribution and a `Point` source [7 parameters].\n", - " The mass of this galaxy is what makes the system genuinely multi-plane: it lenses the further source behind it\n", - " in addition to the foreground lens, doubling the number of images of source_1.\n", - "\n", - " - The second source galaxy at z=2.0 is a `Point` only [2 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=16.\n", - "\n", - "__Name Pairing__\n", - "\n", - "Every `PointDataset` carries a `name` (here `\"point_0\"` and `\"point_1\"`). This `name` pairs the dataset with the\n", - "`Point` in the model that has the same attribute name. Below, `source_0` has a `point_0` component which pairs\n", - "with `dataset_0`, and `source_1` has a `point_1` component which pairs with `dataset_1`.\n", - "\n", - "If a model contains a `Point` whose name has no matching dataset, or vice versa, **PyAutoLens** raises an error.\n", - "The factor graph below ensures every dataset sees the full model, so the name pairs match across both analyses.\n", - "\n", - "__Coordinates__\n", - "\n", - "The model's prior centres assume the lens galaxy is near (0.0\", 0.0\"). If your dataset's lens is not at the\n", - "origin, recentre the data (`autolens_workspace/*/data_preparation`) or override the priors manually\n", - "(`autolens_workspace/*/guides/modeling/customize`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens (z=0.5):\n", - "\n", - "lens = af.Model(\n", - " al.Galaxy,\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal,\n", - " shear=al.mp.ExternalShear,\n", - ")\n", - "\n", - "# Source 0 (z=1.0) \u2014 itself a lens for source 1 behind it:\n", - "\n", - "source_0 = af.Model(\n", - " al.Galaxy,\n", - " redshift=1.0,\n", - " mass=al.mp.Isothermal,\n", - " point_0=al.ps.Point,\n", - ")\n", - "\n", - "# Source 1 (z=2.0):\n", - "\n", - "source_1 = af.Model(al.Galaxy, redshift=2.0, point_1=al.ps.Point)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1)\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using a non-linear search. All examples in the autolens workspace use the nested\n", - "sampling algorithm Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has\n", - "revealed gives the most accurate and efficient modeling results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"point_source\"),\n", - " name=\"modeling\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50,\n", - " iterations_per_quick_update=10000,\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis List__\n", - "\n", - "Set up one `AnalysisPoint` per dataset. We use an \"image-plane chi-squared\" via `FitPositionsImagePairRepeat`,\n", - "which is the most robust likelihood for point-source modeling. See `point_source/modeling.py` for an in-depth\n", - "discussion of the chi-squared options.\n", - "\n", - "__JAX__\n", - "\n", - "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU support, the fit\n", - "runs much faster (~10 minutes instead of an hour). On CPU-only systems JAX still provides a speed-up via\n", - "multithreading, with fits taking around 20-30 minutes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_list = [\n", - " al.AnalysisPoint(\n", - " dataset=dataset,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePairRepeat,\n", - " use_jax=True,\n", - " )\n", - " for dataset in dataset_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis Factor__\n", - "\n", - "Each analysis object is wrapped in an `AnalysisFactor`, which pairs it with the shared model and prepares it for\n", - "use in a factor graph. For this example every factor uses the same lens model, because all sources are lensed by\n", - "the same lens galaxy.\n", - "\n", - "The term \"Factor\" comes from factor graphs, a type of probabilistic graphical model. In this context, each factor\n", - "represents the connection between one `PointDataset` and the shared model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis_factor_list = [\n", - " af.AnalysisFactor(prior_model=model, analysis=analysis)\n", - " for analysis in analysis_list\n", - "]" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Factor Graph__\n", - "\n", - "All `AnalysisFactor` objects are combined into a `FactorGraphModel`, which represents a global model fit to\n", - "multiple datasets using a graphical model structure.\n", - "\n", - "The key outcomes of this setup are:\n", - "\n", - " - The individual log likelihoods from each `Analysis` object are summed to form the total log likelihood\n", - " evaluated during the model-fitting process.\n", - "\n", - " - Results from all datasets are output to a unified directory, with subdirectories for visualisations from each\n", - " analysis object, as defined by their `visualize` methods.\n", - "\n", - "This is a basic use of **PyAutoFit**'s graphical modeling capabilities, which support advanced hierarchical and\n", - "probabilistic modeling for large, multi-dataset analyses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To inspect the global model the factor graph fits, print `factor_graph.global_prior_model.info`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(factor_graph.global_prior_model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__\n", - "\n", - "To fit multiple datasets, we pass the `FactorGraphModel` to a non-linear search.\n", - "\n", - "Unlike single-dataset fitting, we now pass the `factor_graph.global_prior_model` as the model and the\n", - "`factor_graph` itself as the analysis object. This structure enables simultaneous fitting of multiple datasets\n", - "in a consistent and scalable way.\n", - "\n", - "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may\n", - "produce an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell will progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The result returned by a factor-graph fit is a list of `Result` objects, one per `AnalysisFactor`. Each entry\n", - "corresponds to the model-fit of one dataset. Because every factor sees the same global model, the\n", - "`max_log_likelihood_instance` is identical across results \u2014 the per-result objects differ only in their\n", - "analysis-specific data and visualisations." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "for result in result_list:\n", - " print(result.max_log_likelihood_instance)\n", - "\n", - " aplt.subplot_tracer(\n", - " tracer=result.max_log_likelihood_tracer,\n", - " grid=grid,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `Samples` object has the dimensions of the overall non-linear search and is identical in every result, so it\n", - "is sufficient to plot the corner from only the first result." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result_list[0].samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Wrap Up__\n", - "\n", - "This example introduces the API for fitting multiple lensed point sources at different redshifts with the\n", - "multi/factor-graph API. The same pattern can be extended to many more sources (see `cluster/modeling.py`) or to\n", - "combine point-source data with imaging or interferometer data via heterogeneous analyses in a single factor graph." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Multiple Sources\n", + "==========================\n", + "\n", + "This script fits a `PointDataset` of a 'galaxy-scale' strong lens with multiple lensed point sources at different\n", + "redshifts. The lens system is multi-plane: a foreground lens at z=0.5 deflects both background sources, while\n", + "source_0 at z=1.0 is itself a deflector for source_1 at z=2.0 (the \"double Einstein cross\" configuration). Each\n", + "source's multiple images are stored in their own `PointDataset`, and the two datasets are fitted jointly using\n", + "the multi/factor-graph API:\n", + "\n", + " - The foreground lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The first source `Galaxy` (at z=1.0) is itself a `Galaxy` with a mass profile and a `Point`.\n", + " - The second source `Galaxy` (at z=2.0) is a `Point`-only galaxy.\n", + "\n", + "Two `PointDataset`s are fitted simultaneously, one per lensed source. The fit uses one `AnalysisPoint` per\n", + "dataset, each wrapped in an `AnalysisFactor` that pairs it with the shared lens model. The factors are combined\n", + "into a `FactorGraphModel`, which sums the individual log-likelihoods to form the global log-likelihood the\n", + "non-linear search optimises. Multi-plane lensing is handled automatically inside `AnalysisPoint`, which uses each\n", + "`Point`'s plane redshift in the tracer when solving for image-plane positions.\n", + "\n", + "This is an advanced script and assumes previous knowledge of the core **PyAutoLens** API for point-source modeling\n", + "(see `point_source/modeling.py`) and the multi/factor-graph API (see `multi/modeling.py` and `cluster/modeling.py`).\n", + "Common boilerplate is therefore not re-explained in detail here.\n", + "\n", + "__Currently Blocked By PyAutoLens #480__\n", + "\n", + "This script does not run end-to-end on the current PyAutoLens release. The `PointSolver` magnification filter\n", + "uses the tracer's last-plane magnification instead of the requested `plane_redshift`'s magnification, so every\n", + "likelihood evaluation finds 0 image positions for source_0 (whose plane z=1.0 is intermediate). See\n", + "https://github.com/PyAutoLabs/PyAutoLens/issues/480 \u2014 both this script and `simulator.py` are listed in\n", + "`config/build/no_run.yaml` until that bug is fixed. The script is left here in its intended form so the example\n", + "is correct as soon as #480 lands; no script changes will be needed once it does.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load the list of `PointDataset` objects, one per lensed source.\n", + "- **Point Solver:** We set up the `PointSolver`, which determines the multiple images of each point source.\n", + "- **Model:** Compose the multi-plane lens model fitted to the data.\n", + "- **Name Pairing:** Each `PointDataset` name is paired with a `Point` model component of the same name.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis List:** Set up one `AnalysisPoint` per dataset.\n", + "- **Analysis Factor:** Each analysis is wrapped in an `AnalysisFactor` paired with the shared lens model.\n", + "- **Factor Graph:** All `AnalysisFactor` objects are combined into a `FactorGraphModel`.\n", + "- **Model-Fit:** Pass the factor graph to the non-linear search.\n", + "- **Result:** Iterate the per-analysis results returned by the factor-graph fit.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `point_source/modeling.ipynb` notebook for the single-source\n", + "case, and `multi/modeling.ipynb` for the factor-graph API used to combine multiple datasets." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens point-source dataset `multiple_sources`, which is the dataset we will fit. The simulator\n", + "writes one `PointDataset` per lensed source (`point_dataset_0.json` and `point_dataset_1.json`), one for each\n", + "source-plane redshift." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"multiple_sources\"\n", + "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it is created by running the corresponding simulator script.\n", + "This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/point_source/features/multiple_sources/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We load each point-source dataset as a `PointDataset` and place them into a list. Every entry in the list is\n", + "fitted by its own `AnalysisPoint` further down in the script." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_list = [\n", + " al.from_json(file_path=dataset_path / f\"point_dataset_{i}.json\") for i in range(2)\n", + "]\n", + "\n", + "for dataset in dataset_list:\n", + " print(\"Point Dataset Info:\")\n", + " print(dataset.info)\n", + " aplt.subplot_point_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "For point-source modeling we require a `PointSolver`, which determines the multiple images of the mass model for\n", + "a point source at location (y,x) in the source plane. It does this by ray tracing triangles from the image-plane\n", + "to the source-plane and refining the multiple images to sub-pixel precision.\n", + "\n", + "The solver requires a starting grid of (y,x) image-plane coordinates and a `pixel_scale_precision` controlling\n", + "the precision of the converged multiple images. The grid below matches the simulator so the solver covers the\n", + "same region of sky used to generate the data.\n", + "\n", + "Strong lens mass models have a \"central image\" which is nearly always significantly demagnified and not observed.\n", + "Setting `magnification_threshold=0.1` discards this image so it does not contaminate the fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(200, 200),\n", + " pixel_scales=0.05,\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid,\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a multi-plane lens model where:\n", + "\n", + " - The lens galaxy at z=0.5 has an `Isothermal` mass distribution with `ExternalShear` [7 parameters].\n", + "\n", + " - The first source galaxy at z=1.0 has its own `Isothermal` mass distribution and a `Point` source [7 parameters].\n", + " The mass of this galaxy is what makes the system genuinely multi-plane: it lenses the further source behind it\n", + " in addition to the foreground lens, doubling the number of images of source_1.\n", + "\n", + " - The second source galaxy at z=2.0 is a `Point` only [2 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=16.\n", + "\n", + "__Name Pairing__\n", + "\n", + "Every `PointDataset` carries a `name` (here `\"point_0\"` and `\"point_1\"`). This `name` pairs the dataset with the\n", + "`Point` in the model that has the same attribute name. Below, `source_0` has a `point_0` component which pairs\n", + "with `dataset_0`, and `source_1` has a `point_1` component which pairs with `dataset_1`.\n", + "\n", + "If a model contains a `Point` whose name has no matching dataset, or vice versa, **PyAutoLens** raises an error.\n", + "The factor graph below ensures every dataset sees the full model, so the name pairs match across both analyses.\n", + "\n", + "__Coordinates__\n", + "\n", + "The model's prior centres assume the lens galaxy is near (0.0\", 0.0\"). If your dataset's lens is not at the\n", + "origin, recentre the data (`autolens_workspace/*/data_preparation`) or override the priors manually\n", + "(`autolens_workspace/*/guides/modeling/customize`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens (z=0.5):\n", + "\n", + "lens = af.Model(\n", + " al.Galaxy,\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal,\n", + " shear=al.mp.ExternalShear,\n", + ")\n", + "\n", + "# Source 0 (z=1.0) \u2014 itself a lens for source 1 behind it:\n", + "\n", + "source_0 = af.Model(\n", + " al.Galaxy,\n", + " redshift=1.0,\n", + " mass=al.mp.Isothermal,\n", + " point_0=al.ps.Point,\n", + ")\n", + "\n", + "# Source 1 (z=2.0):\n", + "\n", + "source_1 = af.Model(al.Galaxy, redshift=2.0, point_1=al.ps.Point)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source_0=source_0, source_1=source_1)\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using a non-linear search. All examples in the autolens workspace use the nested\n", + "sampling algorithm Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has\n", + "revealed gives the most accurate and efficient modeling results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"point_source\"),\n", + " name=\"modeling\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50,\n", + " iterations_per_quick_update=10000,\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis List__\n", + "\n", + "Set up one `AnalysisPoint` per dataset. We use an \"image-plane chi-squared\" via `FitPositionsImagePairRepeat`,\n", + "which is the most robust likelihood for point-source modeling. See `point_source/modeling.py` for an in-depth\n", + "discussion of the chi-squared options.\n", + "\n", + "__JAX__\n", + "\n", + "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU support, the fit\n", + "runs much faster (~10 minutes instead of an hour). On CPU-only systems JAX still provides a speed-up via\n", + "multithreading, with fits taking around 20-30 minutes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_list = [\n", + " al.AnalysisPoint(\n", + " dataset=dataset,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePairRepeat,\n", + " use_jax=True,\n", + " )\n", + " for dataset in dataset_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis Factor__\n", + "\n", + "Each analysis object is wrapped in an `AnalysisFactor`, which pairs it with the shared model and prepares it for\n", + "use in a factor graph. For this example every factor uses the same lens model, because all sources are lensed by\n", + "the same lens galaxy.\n", + "\n", + "The term \"Factor\" comes from factor graphs, a type of probabilistic graphical model. In this context, each factor\n", + "represents the connection between one `PointDataset` and the shared model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_list = [\n", + " af.AnalysisFactor(prior_model=model, analysis=analysis)\n", + " for analysis in analysis_list\n", + "]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Factor Graph__\n", + "\n", + "All `AnalysisFactor` objects are combined into a `FactorGraphModel`, which represents a global model fit to\n", + "multiple datasets using a graphical model structure.\n", + "\n", + "The key outcomes of this setup are:\n", + "\n", + " - The individual log likelihoods from each `Analysis` object are summed to form the total log likelihood\n", + " evaluated during the model-fitting process.\n", + "\n", + " - Results from all datasets are output to a unified directory, with subdirectories for visualisations from each\n", + " analysis object, as defined by their `visualize` methods.\n", + "\n", + "This is a basic use of **PyAutoFit**'s graphical modeling capabilities, which support advanced hierarchical and\n", + "probabilistic modeling for large, multi-dataset analyses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph = af.FactorGraphModel(*analysis_factor_list, use_jax=True)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To inspect the global model the factor graph fits, print `factor_graph.global_prior_model.info`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(factor_graph.global_prior_model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__\n", + "\n", + "To fit multiple datasets, we pass the `FactorGraphModel` to a non-linear search.\n", + "\n", + "Unlike single-dataset fitting, we now pass the `factor_graph.global_prior_model` as the model and the\n", + "`factor_graph` itself as the analysis object. This structure enables simultaneous fitting of multiple datasets\n", + "in a consistent and scalable way.\n", + "\n", + "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may\n", + "produce an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell will progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The result returned by a factor-graph fit is a list of `Result` objects, one per `AnalysisFactor`. Each entry\n", + "corresponds to the model-fit of one dataset. Because every factor sees the same global model, the\n", + "`max_log_likelihood_instance` is identical across results \u2014 the per-result objects differ only in their\n", + "analysis-specific data and visualisations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "for result in result_list:\n", + " print(result.max_log_likelihood_instance)\n", + "\n", + " aplt.subplot_tracer(\n", + " tracer=result.max_log_likelihood_tracer,\n", + " grid=grid,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Samples` object has the dimensions of the overall non-linear search and is identical in every result, so it\n", + "is sufficient to plot the corner from only the first result." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result_list[0].samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This example introduces the API for fitting multiple lensed point sources at different redshifts with the\n", + "multi/factor-graph API. The same pattern can be extended to many more sources (see `cluster/modeling.py`) or to\n", + "combine point-source data with imaging or interferometer data via heterogeneous analyses in a single factor graph." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/features/multiple_sources/simulator.ipynb b/notebooks/point_source/features/multiple_sources/simulator.ipynb index 4317eb41e..7c513ccbf 100644 --- a/notebooks/point_source/features/multiple_sources/simulator.ipynb +++ b/notebooks/point_source/features/multiple_sources/simulator.ipynb @@ -1,495 +1,532 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Multiple Sources\n", - "===========================\n", - "\n", - "This script simulates `PointDataset` data of a strong lens which has multiple lensed point sources at different\n", - "redshifts (the example below has two sources, but the same approach extends to any number).\n", - "\n", - "The simulated system is a multi-plane lens where source_0 is itself a deflector for source_1 behind it: a\n", - "foreground galaxy lens (z=0.5) lenses both background sources, while source_0's own mass at z=1.0 additionally\n", - "lenses source_1 at z=2.0. The result is two interleaved sets of multiple images (a \"double Einstein cross\"-like\n", - "configuration). Each source's multiple images are stored as a separate `PointDataset`, and the companion\n", - "`modeling.py` script fits both datasets jointly using the multi/factor-graph API.\n", - "\n", - "The lens model below uses:\n", - "\n", - " - The foreground lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", - " - The first source `Galaxy` (at z=1.0) is itself a `Galaxy` with a mass profile and a `Point`.\n", - " - The second source `Galaxy` (at z=2.0) is a `Point`-only galaxy.\n", - "\n", - "__Currently Blocked By PyAutoLens #480__\n", - "\n", - "This script does not run end-to-end on the current PyAutoLens release. The `PointSolver` magnification filter\n", - "uses the tracer's last-plane magnification instead of the requested `plane_redshift`'s magnification, so every\n", - "candidate image position for source_0 (which is at the intermediate plane z=1.0) is rejected and the solver\n", - "returns 0 positions. See https://github.com/PyAutoLabs/PyAutoLens/issues/480 \u2014 the simulator and modeling scripts\n", - "in this folder are listed in `config/build/no_run.yaml` until that bug is fixed. The script is left here in its\n", - "intended form so the example is correct as soon as #480 lands; no script changes will be needed once it does.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) and source galaxy `Point` for this simulated lens.\n", - "- **Point Solver:** We use a `PointSolver` to locate the multiple images.\n", - "- **Point Datasets:** Create a point-source data object and output this to a `.json` file, which is the format used to.\n", - "- **Visualize:** Output a subplot of the simulated point source dictionary and the tracer's quantities to the.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import numpy as np\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"point_source\"\n", - "dataset_name = \"multiple_sources\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's mass (SIE+Shear) and source galaxy `Point` for this simulated lens. We include a\n", - "faint disk in the source for purely visualisation purposes to show where the multiple images appear.\n", - "\n", - "For lens modeling, defining ellipticity in terms of the `ell_comps` improves the model-fitting procedure. However,\n", - "for simulating a strong lens you may find it more intuitive to define the elliptical geometry using the\n", - "axis-ratio of the profile (axis_ratio = semi-major axis / semi-minor axis = b/a) and position angle, where angle\n", - "is in degrees and defined counter clockwise from the positive x-axis.\n", - "\n", - "We can use the `convert` module to determine the elliptical components from the axis-ratio and angle." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy_0 = al.Galaxy(\n", - " redshift=1.0,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.02, 0.03),\n", - " einstein_radius=0.2,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", - " ),\n", - " light=al.lp.ExponentialCore(\n", - " centre=(0.02, 0.03), intensity=0.1, effective_radius=0.02\n", - " ),\n", - " point_0=al.ps.Point(centre=(0.02, 0.03)),\n", - ")\n", - "\n", - "\n", - "source_galaxy_1 = al.Galaxy(\n", - " redshift=2.0,\n", - " light=al.lp.ExponentialCore(\n", - " centre=(0.0, 0.0), intensity=0.1, effective_radius=0.02\n", - " ),\n", - " point_1=al.ps.Point(centre=(0.0, 0.0)),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will compute the multiple image positions of the simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "We use a `PointSolver` to locate the multiple images. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(200, 200),\n", - " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now pass the `Tracer` to the solver.\n", - "\n", - "This finds the image-plane coordinates that map directly to the source-plane centres (0.02\", 0.03\") and\n", - "(0.0\", 0.0\"). A double Einstein cross is a multi-plane lensing system, therefore for each source we also pass\n", - "their redshift into the solver as `plane_redshift` so that it finds the multiple images while properly accounting\n", - "for the multi-plane lensing.\n", - "\n", - "Position noise is set to 0.005\" (5 mas, realistic PSF-centroiding precision on HST imaging) and flux noise\n", - "to 5% relative (microlensing-dominated regime). See `scripts/point_source/simulator.py` for a full\n", - "discussion of these values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "position_noise = 0.005\n", - "flux_rel_noise = 0.05\n", - "\n", - "positions_0 = solver.solve(\n", - " tracer=tracer,\n", - " source_plane_coordinate=source_galaxy_0.point_0.centre,\n", - " plane_redshift=source_galaxy_0.redshift,\n", - ")\n", - "\n", - "positions_0_with_noise = positions_0 + np.random.normal(\n", - " loc=0.0, scale=position_noise, size=positions_0.shape\n", - ")\n", - "\n", - "positions_0_with_noise = al.Grid2DIrregular(\n", - " values=positions_0_with_noise,\n", - ")\n", - "\n", - "positions_1 = solver.solve(\n", - " tracer=tracer,\n", - " source_plane_coordinate=source_galaxy_1.point_1.centre,\n", - " plane_redshift=source_galaxy_1.redshift,\n", - ")\n", - "\n", - "positions_1_with_noise = positions_1 + np.random.normal(\n", - " loc=0.0, scale=position_noise, size=positions_1.shape\n", - ")\n", - "\n", - "positions_1_with_noise = al.Grid2DIrregular(\n", - " values=positions_1_with_noise,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use the positions to compute the magnification of the `Tracer` at every position." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnifications_0 = al.LensCalc.from_tracer(\n", - " tracer=tracer\n", - ").magnification_2d_via_hessian_from(grid=positions_0)\n", - "magnifications_1 = al.LensCalc.from_tracer(\n", - " tracer=tracer\n", - ").magnification_2d_via_hessian_from(grid=positions_1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can now compute the observed fluxes of the `Point`, give we know how much each is magnified." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "flux = 1.0\n", - "fluxes_0 = [flux * np.abs(magnification) for magnification in magnifications_0]\n", - "fluxes_0 = al.ArrayIrregular(values=fluxes_0)\n", - "fluxes_0_with_noise = fluxes_0 + np.random.normal(\n", - " loc=0.0, scale=flux_rel_noise * np.asarray(fluxes_0), size=len(fluxes_0)\n", - ")\n", - "fluxes_0_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes_0))\n", - "\n", - "fluxes_1 = [flux * np.abs(magnification) for magnification in magnifications_1]\n", - "fluxes_1 = al.ArrayIrregular(values=fluxes_1)\n", - "fluxes_1_with_noise = fluxes_1 + np.random.normal(\n", - " loc=0.0, scale=flux_rel_noise * np.asarray(fluxes_1), size=len(fluxes_1)\n", - ")\n", - "fluxes_1_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes_1))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now output the image of this strong lens to `.fits` which can be used for visualize when performing point-source \n", - "modeling and to `.png` for general inspection." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Datasets__\n", - "\n", - "Create a point-source data object and output this to a `.json` file, which is the format used to load and\n", - "analyse the dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_0 = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_0_with_noise,\n", - " positions_noise_map=position_noise,\n", - " fluxes=fluxes_0_with_noise,\n", - " fluxes_noise_map=fluxes_0_noise_map,\n", - ")\n", - "dataset_1 = al.PointDataset(\n", - " name=\"point_1\",\n", - " positions=positions_1_with_noise,\n", - " positions_noise_map=position_noise,\n", - " fluxes=fluxes_1_with_noise,\n", - " fluxes_noise_map=fluxes_1_noise_map,\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now output the point datasets to the dataset path as a .json file, which is loaded in the point source modeling\n", - "examples." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=dataset_0,\n", - " file_path=dataset_path / \"point_dataset_0.json\",\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=dataset_1,\n", - " file_path=dataset_path / \"point_dataset_1.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__CSV Output__\n", - "\n", - "Both datasets can also be saved to a single CSV \u2014 one row per observed image grouped by\n", - "``name`` \u2014 as the hand-editable spreadsheet format alternative to per-source JSON." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_csv(\n", - " datasets=[dataset_0, dataset_1],\n", - " file_path=dataset_path / \"point_datasets.csv\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated point source dictionary and the tracer's quantities to the dataset path as .png files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_point_dataset(\n", - " dataset=dataset_0, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_point_dataset(\n", - " dataset=dataset_1, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=dataset_path / \"tracer.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Multiple Sources\n", + "===========================\n", + "\n", + "This script simulates `PointDataset` data of a strong lens which has multiple lensed point sources at different\n", + "redshifts (the example below has two sources, but the same approach extends to any number).\n", + "\n", + "The simulated system is a multi-plane lens where source_0 is itself a deflector for source_1 behind it: a\n", + "foreground galaxy lens (z=0.5) lenses both background sources, while source_0's own mass at z=1.0 additionally\n", + "lenses source_1 at z=2.0. The result is two interleaved sets of multiple images (a \"double Einstein cross\"-like\n", + "configuration). Each source's multiple images are stored as a separate `PointDataset`, and the companion\n", + "`modeling.py` script fits both datasets jointly using the multi/factor-graph API.\n", + "\n", + "The lens model below uses:\n", + "\n", + " - The foreground lens galaxy's total mass distribution is an `Isothermal` and `ExternalShear`.\n", + " - The first source `Galaxy` (at z=1.0) is itself a `Galaxy` with a mass profile and a `Point`.\n", + " - The second source `Galaxy` (at z=2.0) is a `Point`-only galaxy.\n", + "\n", + "__Currently Blocked By PyAutoLens #480__\n", + "\n", + "This script does not run end-to-end on the current PyAutoLens release. The `PointSolver` magnification filter\n", + "uses the tracer's last-plane magnification instead of the requested `plane_redshift`'s magnification, so every\n", + "candidate image position for source_0 (which is at the intermediate plane z=1.0) is rejected and the solver\n", + "returns 0 positions. See https://github.com/PyAutoLabs/PyAutoLens/issues/480 \u2014 the simulator and modeling scripts\n", + "in this folder are listed in `config/build/no_run.yaml` until that bug is fixed. The script is left here in its\n", + "intended form so the example is correct as soon as #480 lands; no script changes will be needed once it does.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Ray Tracing:** Setup the lens galaxy's mass (SIE+Shear) and source galaxy `Point` for this simulated lens.\n", + "- **Point Solver:** We use a `PointSolver` to locate the multiple images.\n", + "- **Point Datasets:** Create a point-source data object and output this to a `.json` file, which is the format used to.\n", + "- **Visualize:** Output a subplot of the simulated point source dictionary and the tracer's quantities to the.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import numpy as np\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"point_source\"\n", + "dataset_name = \"multiple_sources\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's mass (SIE+Shear) and source galaxy `Point` for this simulated lens. We include a\n", + "faint disk in the source for purely visualisation purposes to show where the multiple images appear.\n", + "\n", + "For lens modeling, defining ellipticity in terms of the `ell_comps` improves the model-fitting procedure. However,\n", + "for simulating a strong lens you may find it more intuitive to define the elliptical geometry using the\n", + "axis-ratio of the profile (axis_ratio = semi-major axis / semi-minor axis = b/a) and position angle, where angle\n", + "is in degrees and defined counter clockwise from the positive x-axis.\n", + "\n", + "We can use the `convert` module to determine the elliptical components from the axis-ratio and angle." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy_0 = al.Galaxy(\n", + " redshift=1.0,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.02, 0.03),\n", + " einstein_radius=0.2,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=60.0),\n", + " ),\n", + " light=al.lp.ExponentialCore(\n", + " centre=(0.02, 0.03), intensity=0.1, effective_radius=0.02\n", + " ),\n", + " point_0=al.ps.Point(centre=(0.02, 0.03)),\n", + ")\n", + "\n", + "\n", + "source_galaxy_1 = al.Galaxy(\n", + " redshift=2.0,\n", + " light=al.lp.ExponentialCore(\n", + " centre=(0.0, 0.0), intensity=0.1, effective_radius=0.02\n", + " ),\n", + " point_1=al.ps.Point(centre=(0.0, 0.0)),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will compute the multiple image positions of the simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy_0, source_galaxy_1])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "We use a `PointSolver` to locate the multiple images. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(200, 200),\n", + " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now pass the `Tracer` to the solver.\n", + "\n", + "This finds the image-plane coordinates that map directly to the source-plane centres (0.02\", 0.03\") and\n", + "(0.0\", 0.0\"). A double Einstein cross is a multi-plane lensing system, therefore for each source we also pass\n", + "their redshift into the solver as `plane_redshift` so that it finds the multiple images while properly accounting\n", + "for the multi-plane lensing.\n", + "\n", + "Position noise is set to 0.005\" (5 mas, realistic PSF-centroiding precision on HST imaging) and flux noise\n", + "to 5% relative (microlensing-dominated regime). See `scripts/point_source/simulator.py` for a full\n", + "discussion of these values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "position_noise = 0.005\n", + "flux_rel_noise = 0.05\n", + "\n", + "positions_0 = solver.solve(\n", + " tracer=tracer,\n", + " source_plane_coordinate=source_galaxy_0.point_0.centre,\n", + " plane_redshift=source_galaxy_0.redshift,\n", + ")\n", + "\n", + "positions_0_with_noise = positions_0 + np.random.normal(\n", + " loc=0.0, scale=position_noise, size=positions_0.shape\n", + ")\n", + "\n", + "positions_0_with_noise = al.Grid2DIrregular(\n", + " values=positions_0_with_noise,\n", + ")\n", + "\n", + "positions_1 = solver.solve(\n", + " tracer=tracer,\n", + " source_plane_coordinate=source_galaxy_1.point_1.centre,\n", + " plane_redshift=source_galaxy_1.redshift,\n", + ")\n", + "\n", + "positions_1_with_noise = positions_1 + np.random.normal(\n", + " loc=0.0, scale=position_noise, size=positions_1.shape\n", + ")\n", + "\n", + "positions_1_with_noise = al.Grid2DIrregular(\n", + " values=positions_1_with_noise,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use the positions to compute the magnification of the `Tracer` at every position." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnifications_0 = al.LensCalc.from_tracer(\n", + " tracer=tracer\n", + ").magnification_2d_via_hessian_from(grid=positions_0)\n", + "magnifications_1 = al.LensCalc.from_tracer(\n", + " tracer=tracer\n", + ").magnification_2d_via_hessian_from(grid=positions_1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now compute the observed fluxes of the `Point`, give we know how much each is magnified." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "flux = 1.0\n", + "fluxes_0 = [flux * np.abs(magnification) for magnification in magnifications_0]\n", + "fluxes_0 = al.ArrayIrregular(values=fluxes_0)\n", + "fluxes_0_with_noise = fluxes_0 + np.random.normal(\n", + " loc=0.0, scale=flux_rel_noise * np.asarray(fluxes_0), size=len(fluxes_0)\n", + ")\n", + "fluxes_0_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes_0))\n", + "\n", + "fluxes_1 = [flux * np.abs(magnification) for magnification in magnifications_1]\n", + "fluxes_1 = al.ArrayIrregular(values=fluxes_1)\n", + "fluxes_1_with_noise = fluxes_1 + np.random.normal(\n", + " loc=0.0, scale=flux_rel_noise * np.asarray(fluxes_1), size=len(fluxes_1)\n", + ")\n", + "fluxes_1_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes_1))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now output the image of this strong lens to `.fits` which can be used for visualize when performing point-source \n", + "modeling and to `.png` for general inspection." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Datasets__\n", + "\n", + "Create a point-source data object and output this to a `.json` file, which is the format used to load and\n", + "analyse the dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_0 = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_0_with_noise,\n", + " positions_noise_map=position_noise,\n", + " fluxes=fluxes_0_with_noise,\n", + " fluxes_noise_map=fluxes_0_noise_map,\n", + ")\n", + "dataset_1 = al.PointDataset(\n", + " name=\"point_1\",\n", + " positions=positions_1_with_noise,\n", + " positions_noise_map=position_noise,\n", + " fluxes=fluxes_1_with_noise,\n", + " fluxes_noise_map=fluxes_1_noise_map,\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now output the point datasets to the dataset path as a .json file, which is loaded in the point source modeling\n", + "examples." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=dataset_0,\n", + " file_path=dataset_path / \"point_dataset_0.json\",\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=dataset_1,\n", + " file_path=dataset_path / \"point_dataset_1.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__CSV Output__\n", + "\n", + "Both datasets can also be saved to a single CSV \u2014 one row per observed image grouped by\n", + "``name`` \u2014 as the hand-editable spreadsheet format alternative to per-source JSON." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_csv(\n", + " datasets=[dataset_0, dataset_1],\n", + " file_path=dataset_path / \"point_datasets.csv\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated point source dictionary and the tracer's quantities to the dataset path as .png files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_point_dataset(\n", + " dataset=dataset_0, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_point_dataset(\n", + " dataset=dataset_1, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=dataset_path / \"tracer.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/features/time_delays.ipynb b/notebooks/point_source/features/time_delays.ipynb index 3abea8a90..2b21fa7a4 100644 --- a/notebooks/point_source/features/time_delays.ipynb +++ b/notebooks/point_source/features/time_delays.ipynb @@ -1,516 +1,553 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Time Delays\n", - "=====================\n", - "\n", - "A measurable quantity of a point source is its time delay\u2014the time it takes for light to travel from the source to the\n", - "observer for each multiple image of the point source (e.g., the quasar images). This is often expressed as the relative\n", - "time delay between each image and the image with the shortest time delay, which is often referred to as\n", - "the \"reference image.\"\n", - "\n", - "Time delays are commonly used in strong lensing analyses, for example to measure the Hubble constant, since\n", - "they are less affected by microlensing and can provide robust cosmological constraints.\n", - "\n", - "This script describes how to perform point source lens modeling using the time delays of the point source dataset\n", - "as additional information on top of the positions of the point source, in case you are studying the Hubble constant\n", - "or another measureable quantity that uses time delays.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Point Solver:** We set up the `PointSolver`, which is used to compute the multiple images of the point source in.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **Run Times:** Profiling the expected run time of the model-fit.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Cosmology:** Time Delay lenses allow the Hubble constant to be constrained, because the difference between the.\n", - "\n", - "__Model__\n", - "\n", - "This script fits a `PointDataset` data of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The source `Galaxy` is a point source `Point`.\n", - "\n", - "The `ExternalShear` is also not included in the mass model, where it is for the `imaging` and `interferometer` examples.\n", - "For a quadruply imaged point source (8 data points) there is insufficient information to fully constain a model with\n", - "an `Isothermal` and `ExternalShear` (9 parameters).\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens point-source dataset `simple`, which is the dataset we will use to perform point source \n", - "lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/point_source/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now load the point source dataset we will fit using point source modeling. \n", - "\n", - "We load this data as a `PointDataset`, which contains the positions and time_delays of every point source. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.from_json(\n", - " file_path=dataset_path / \"point_dataset_with_time_delays.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can print this dictionary to see the dataset's `name`, `positions` and `time_delays` and noise-map values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Point Dataset Info:\")\n", - "print(dataset.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also plot the positions of the `PointDataset`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_point_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We next load an image of the dataset and plot the point source data over it, because as described in \n", - "the `modeling/start_here.ipynb` notebook, it is useful for visualizing the point source dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.05)\n", - "\n", - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "We set up the `PointSolver`, which is used to compute the multiple images of the point source in the image-plane.\n", - "\n", - "There are no special settings or inputs for the fitting of time_delays, therefore the `PointSolver` is set up in the same way\n", - "as in the `modeling/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid,\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", - " - The source galaxy's light is a point `Point` [2 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=7.\n", - "\n", - "Name pairing is used as before to pair the `PointDataset` to the `Point` in the model, which is discussed below.\n", - "\n", - "If you have fitted fluxes in the `fluxes` example, you will have seen that a `PointFlux` model component was used\n", - "which had the `flux` of the point source as an additional free parameter. For time delays, there is no special\n", - "model component or extra free parameters, because the time delays are a propety of the mass model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal)\n", - "\n", - "# Source:\n", - "\n", - "point_0 = af.Model(al.ps.Point)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"point_source\") / \"features\",\n", - " name=\"time_delays\",\n", - " unique_tag=dataset_name,\n", - " n_live=100,\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see modeling examples\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Analysis__\n", - "\n", - "Create the `AnalysisPoint` object defining how the via Nautilus the model is fitted to the data." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisPoint(\n", - " dataset=dataset,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePairRepeat, # Image-plane chi-squared with repeat image pairs.\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "For the positions-only fit, the run time of the log likelihood function was ~0.01 seconds, which is fast\n", - "\n", - "Evaluating the time delays does not increase this much, with a value of around ~0.01 seconds still expected.\n", - "\n", - "Overall modeling run times should therefore be around 20 minutes on CPU, under 5 minutes on GPU.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Result__\n", - "\n", - "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", - "`start_here.ipynb` for a description of how to fix this)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Cosmology__\n", - "\n", - "Time Delay lenses allow the Hubble constant to be constrained, because the difference between the geometric\n", - "time delay and the physical time delay is proportional to the Hubble constant.\n", - "\n", - "We therefore create a Cosmology as a `Model` object in order to make the cosmological parameter Omega_m a free \n", - "parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "cosmology = af.Model(al.cosmo.FlatLambdaCDM)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By default, all parameters of a cosmology model are initialized as fixed values based on the Planck18 cosmology.\n", - "\n", - "In order to make the Hubble constant, we override the default value of the Hubble constant with uniform prior." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "cosmology.H0 = af.UniformPrior(lower_limit=0.0, upper_limit=150.0)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "lens.mass.centre.centre_0 = 0.0\n", - "lens.mass.centre.centre_1 = 0.0\n", - "lens.mass.einstein_radius = 1.6\n", - "\n", - "model = af.Collection(\n", - " galaxies=af.Collection(lens=lens, source=source), cosmology=cosmology\n", - ")\n" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format\n", - "\n", - "This confirms the model includes the Cosmology, which has the Hubble constant as a free parameter." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", - "full description)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"point_source\") / \"features\",\n", - " name=\"time_delays_hubble_constant2\",\n", - " unique_tag=dataset_name,\n", - " n_live=150,\n", - " n_batch=50,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "result = search.fit(model=model, analysis=analysis)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Time Delays\n", + "=====================\n", + "\n", + "A measurable quantity of a point source is its time delay\u2014the time it takes for light to travel from the source to the\n", + "observer for each multiple image of the point source (e.g., the quasar images). This is often expressed as the relative\n", + "time delay between each image and the image with the shortest time delay, which is often referred to as\n", + "the \"reference image.\"\n", + "\n", + "Time delays are commonly used in strong lensing analyses, for example to measure the Hubble constant, since\n", + "they are less affected by microlensing and can provide robust cosmological constraints.\n", + "\n", + "This script describes how to perform point source lens modeling using the time delays of the point source dataset\n", + "as additional information on top of the positions of the point source, in case you are studying the Hubble constant\n", + "or another measureable quantity that uses time delays.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Point Solver:** We set up the `PointSolver`, which is used to compute the multiple images of the point source in.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **Run Times:** Profiling the expected run time of the model-fit.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Cosmology:** Time Delay lenses allow the Hubble constant to be constrained, because the difference between the.\n", + "\n", + "__Model__\n", + "\n", + "This script fits a `PointDataset` data of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The source `Galaxy` is a point source `Point`.\n", + "\n", + "The `ExternalShear` is also not included in the mass model, where it is for the `imaging` and `interferometer` examples.\n", + "For a quadruply imaged point source (8 data points) there is insufficient information to fully constain a model with\n", + "an `Isothermal` and `ExternalShear` (9 parameters).\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens point-source dataset `simple`, which is the dataset we will use to perform point source \n", + "lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/point_source/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now load the point source dataset we will fit using point source modeling. \n", + "\n", + "We load this data as a `PointDataset`, which contains the positions and time_delays of every point source. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.from_json(\n", + " file_path=dataset_path / \"point_dataset_with_time_delays.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can print this dictionary to see the dataset's `name`, `positions` and `time_delays` and noise-map values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Point Dataset Info:\")\n", + "print(dataset.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also plot the positions of the `PointDataset`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_point_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We next load an image of the dataset and plot the point source data over it, because as described in \n", + "the `modeling/start_here.ipynb` notebook, it is useful for visualizing the point source dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.05)\n", + "\n", + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "We set up the `PointSolver`, which is used to compute the multiple images of the point source in the image-plane.\n", + "\n", + "There are no special settings or inputs for the fitting of time_delays, therefore the `PointSolver` is set up in the same way\n", + "as in the `modeling/start_here.ipynb` notebook." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid,\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + " - The source galaxy's light is a point `Point` [2 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=7.\n", + "\n", + "Name pairing is used as before to pair the `PointDataset` to the `Point` in the model, which is discussed below.\n", + "\n", + "If you have fitted fluxes in the `fluxes` example, you will have seen that a `PointFlux` model component was used\n", + "which had the `flux` of the point source as an additional free parameter. For time delays, there is no special\n", + "model component or extra free parameters, because the time delays are a propety of the mass model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal)\n", + "\n", + "# Source:\n", + "\n", + "point_0 = af.Model(al.ps.Point)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"point_source\") / \"features\",\n", + " name=\"time_delays\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see modeling examples\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Create the `AnalysisPoint` object defining how the via Nautilus the model is fitted to the data." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisPoint(\n", + " dataset=dataset,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePairRepeat, # Image-plane chi-squared with repeat image pairs.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "For the positions-only fit, the run time of the log likelihood function was ~0.01 seconds, which is fast\n", + "\n", + "Evaluating the time delays does not increase this much, with a value of around ~0.01 seconds still expected.\n", + "\n", + "Overall modeling run times should therefore be around 20 minutes on CPU, under 5 minutes on GPU.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The `info` attribute shows the model in a readable format (if this does not display clearly on your screen refer to\n", + "`start_here.ipynb` for a description of how to fix this)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Cosmology__\n", + "\n", + "Time Delay lenses allow the Hubble constant to be constrained, because the difference between the geometric\n", + "time delay and the physical time delay is proportional to the Hubble constant.\n", + "\n", + "We therefore create a Cosmology as a `Model` object in order to make the cosmological parameter Omega_m a free \n", + "parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cosmology = af.Model(al.cosmo.FlatLambdaCDM)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "By default, all parameters of a cosmology model are initialized as fixed values based on the Planck18 cosmology.\n", + "\n", + "In order to make the Hubble constant, we override the default value of the Hubble constant with uniform prior." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "cosmology.H0 = af.UniformPrior(lower_limit=0.0, upper_limit=150.0)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "lens.mass.centre.centre_0 = 0.0\n", + "lens.mass.centre.centre_1 = 0.0\n", + "lens.mass.einstein_radius = 1.6\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens, source=source), cosmology=cosmology\n", + ")\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format\n", + "\n", + "This confirms the model includes the Cosmology, which has the Hubble constant as a free parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm Nautilus (see `start.here.py` for a \n", + "full description)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"point_source\") / \"features\",\n", + " name=\"time_delays_hubble_constant2\",\n", + " unique_tag=dataset_name,\n", + " n_live=150,\n", + " n_batch=50,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = search.fit(model=model, analysis=analysis)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/fit.ipynb b/notebooks/point_source/fit.ipynb index d81949eb4..2f5eb6880 100644 --- a/notebooks/point_source/fit.ipynb +++ b/notebooks/point_source/fit.ipynb @@ -1,941 +1,978 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Guide: Point Sources\n", - "--------------------\n", - "\n", - "The examples covered so far have focused on strongly lensed galaxies, where extended surface brightness is warped into\n", - "stunning giant arcs and Einstein rings visible in high-quality telescope images. For these extended sources, light\n", - "profile objects\u2014such as analytic Sersic profiles\u2014are used to represent their surface brightness.\n", - "\n", - "However, some observed sources are extremely small, spanning just light weeks or days across. In these cases, only the\n", - "source\u2019s central point of light is detected in each multiple image. Such sources are called **point sources**,\n", - "which typically include quasars, supernovae, or stars.\n", - "\n", - "Strictly speaking, a point source does have a finite size\u2014on the order of light weeks or days\u2014but it is considered a\n", - "point source because its size is orders of magnitude smaller than the resolution of the telescope. As a result, it\n", - "appears as a single point of light, with all the flux of each multiple image effectively contained within a single pixel.\n", - "\n", - "Point sources affect lensing calculations differently than extended sources, requiring dedicated methods and\n", - "functionality. This functionality is described here and used throughout the `point_source` simulation and modeling\n", - "examples.\n", - "\n", - "If you are new to analyzing strong lenses with point sources, this guide is the ideal place to start!\n", - "\n", - "__Contents__\n", - "\n", - "- **Lensed Point Source:** To begin, we create a strong lens image using an isothermal mass model and a source with a compact.\n", - "- **Point Source:** The image above visually illustrates where the source\u2019s light is traced on the image plane.\n", - "- **Point Solver:** For a point source, our goal is to find the (y, x) coordinates in the image plane that map directly.\n", - "- **Number of Solutions:** The number of solutions (e.g.\n", - "- **Solving the Lens Equation:** In the literature, the process of finding the multiple images of a source in the image-plane is.\n", - "- **Triangle Tracing:** Computing the multiple image positions of a point source is a non-linear problem.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Name Pairing:** The names of the point-source datasets have an even more important role, the names are used to pair.\n", - "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", - "- **Chi Squared:** For point-source modeling, there are many different ways to define the likelihood function, broadly.\n", - "- **Fluxes:** Another measurable quantity of a point source is its flux\u2014the total amount of light received from.\n", - "- **Flux Point Dataset:** The fluxes are not input a `PointDataset` object, alongside the image-plane coordinates of the.\n", - "- **Flux Fitting:** Above, we used a `FitPointDataset` to fit the positions of the point source in the image-plane.\n", - "- **Time Delays:** Another measurable quantity of a point source is its time delay\u2014the time it takes for light to.\n", - "- **Time Delay Fitting:** We can also use the `FitPointDataset` to fit the time delays of the point source, which is done by.\n", - "- **New User Wrap Up:** The `point_source` package of the `autolens_workspace` contains numerous example scripts for.\n", - "- **Shape Solver:** All calculations above assumed the source was a point source with no size.\n", - "\n", - "__JAX__\n", - "\n", - "This guide walks through point-source fitting at a low level (no\n", - "non-linear search) \u2014 the `FitPointDataset` and `PointSolver` objects work\n", - "on either NumPy or JAX. For the standard search-driven modeling path\n", - "where `AnalysisPoint` auto-enables `use_jax=True` and the JIT happens\n", - "inside the search driver, see `start_here.py` / `modeling.py`. For the\n", - "explicit `@jax.jit + PointSolver(use_jax=True)` pattern (useful for fast\n", - "forward-solving in custom code), see `simulator.py`'s `__JAX Variant__`\n", - "section.\n", - "\n", - "Note: `PointSolver(use_jax=True)` works inside `@jax.jit`. The triangle-\n", - "refinement loop operates on raw arrays and doesn't go through `.native`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Lensed Point Source__\n", - "\n", - "To begin, we create a strong lens image using an isothermal mass model and a source with a compact exponential light profile.\n", - "\n", - "Although our goal is to demonstrate solving for the multiple image positions of a point source, simulating the data \n", - "with a compact extended source makes the visualization of the point solver\u2019s solutions clearer." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "isothermal_mass_profile = al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - ")\n", - "\n", - "exponential_light_profile = al.lp.ExponentialCore(\n", - " centre=(0.07, 0.07), intensity=0.1, effective_radius=0.1\n", - ")\n", - "\n", - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=isothermal_mass_profile,\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " light=exponential_light_profile,\n", - ")\n", - "\n", - "tracer_extended = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the image of our strongly lensed source galaxy.\n", - "\n", - "Clearly visible are four multiple images arranged in a cross configuration. The brightest pixels correspond to the \n", - "four (y, x) multiple image positions that our point source solver aims to identify." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.plot_array(array=tracer_extended.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Source__\n", - "\n", - "The image above visually illustrates where the source\u2019s light is traced on the image plane.\n", - "\n", - "We now treat this source as a point source by defining a source galaxy using the `Point` class.\n", - "\n", - "This point source shares the same center as the compact source above, ensuring that the multiple image positions \n", - "coincide with those previously shown in the image plane." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "point_source = al.ps.Point(centre=(0.07, 0.07))\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, point_0=point_source)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "For a point source, our goal is to find the (y, x) coordinates in the image plane that map directly to the center of \n", - "the point source in the source plane\u2014these are its \"multiple images.\" This is achieved using a `PointSolver`, which \n", - "determines the multiple images of the mass model for a point source located at a given (y, x) position in the \n", - "source plane.\n", - "\n", - "The solver works by ray tracing triangles from the image plane back to the source plane and checking whether the \n", - "source-plane (y, x) center lies inside each triangle. It iteratively refines this process by ray tracing progressively \n", - "smaller triangles, allowing the multiple image positions to be determined with sub-pixel precision.\n", - "\n", - "The `PointSolver` requires an initial grid of (y, x) coordinates in the image plane, which defines the first set of \n", - "triangles to ray trace. It also needs a `pixel_scale_precision` parameter, specifying the resolution at which the \n", - "multiple images are computed. Smaller values increase precision but require longer computation times. The value \n", - "of 0.001 used here balances efficiency and accuracy.\n", - "\n", - "Strong lens mass models often predict a \"central image,\" a multiple image that is usually heavily demagnified and thus \n", - "not observed. Since the `PointSolver` finds all valid multiple images, it will locate this central image regardless of \n", - "its visibility. To avoid including this unobservable image, we set a `magnification_threshold=0.1`, which discards any \n", - "images with magnifications below this value.\n", - "\n", - "If your dataset does include a detectable central image, you should lower this threshold accordingly to include it in \n", - "your analysis.\n", - "\n", - "We now compute the multiple image positions by creating a `PointSolver` object and passing it the tracer of our \n", - "strong lens system." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now pass the tracer to the solver, to determine the image-plane multiple images for the source centre.\n", - "\n", - "The solver will find the image-plane coordinates that map directly to the source-plane coordinate (0.07\", 0.07\"), \n", - "which we plot below.\n", - "\n", - "The plot shows the four solved multiple image positions (with the central image excluded) as a scatter plot. To make \n", - "the positions clearer, we increase the marker size and use asterisks\u2014PyAutoLens\u2019s standard symbol for denoting \n", - "multiple images of strong lenses." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = solver.solve(tracer=tracer, source_plane_coordinate=(0.07, 0.07))\n", - "\n", - "aplt.plot_grid(grid=positions, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The plot above makes it difficult to directly compare the multiple image positions with the image of the strong lens itself.\n", - "\n", - "To improve clarity, we overplot the multiple image positions on the strong lens image. This clearly shows that the \n", - "multiple images coincide with the centers of the brightest pixels of the lensed source galaxy." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=tracer_extended.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Number of Solutions__\n", - "\n", - "The number of solutions (e.g. the number of image-plane multiple images that map to the source centre) depends\n", - "on the mass model of the lens: \n", - "\n", - " - For spherical mass profiles, there are three unique solutions, including a demagnified central image.\n", - "\n", - " - For elliptical mass profiles, there are five unique solutions, again including a demagnified central image.\n", - "\n", - " - For lenses with multiple mass profiles (e.g. two galaxies) and more exotic mass distributions, the number of \n", - " solutions can be even higher. \n", - "\n", - "__Solving the Lens Equation__\n", - "\n", - "In the literature, the process of finding the multiple images of a source in the image-plane is often referred to as\n", - "'solving the lens equation'.\n", - "\n", - "There lens equation is a fundamental equation in lensing, which describes how light rays are deflected from the\n", - "image-plane to the source-plane. It is given by:\n", - "\n", - "$\\beta = \\theta - \\hat{\\alpha}(\\theta)$\n", - "\n", - "Where:\n", - "\n", - "$\\beta$ is the source-plane (y,x) coordinate.\n", - "$\\theta$ is the image-plane (y,x) coordinate.\n", - "$\\hat{\\alpha}(\\theta)$ is the deflection angle at image-plane (y,x) coordinate $\\theta$.\n", - "\n", - "The lens equation is non-linear, as the deflection angle $\\hat{\\alpha}$ depends on the mass model of the lens galaxy.\n", - "\n", - "It is therefore called solving the lens equation because we are trying to find the image-plane (y,x) coordinates $\\theta$\n", - "that satisfies the equation above for a given source-plane (y,x) coordinate $\\beta$.\n", - "\n", - "__Triangle Tracing__\n", - "\n", - "Computing the multiple image positions of a point source is a non-linear problem. Given a source-plane (y,x) coordinate,\n", - "there are multiple image-plane (y,x) coordinates that trace to that source-plane coordinate, and there is no simple\n", - "analytic solution to determine these image-plane coordinates.\n", - "\n", - "The solver therefore uses a triangulation approach to find the multiple image positions. It first overlays a grid of\n", - "triangles over the image-plane, and uses the mass model to trace these triangles to the source-plane. If a triangle\n", - "contains the source-plane (y,x) coordinate, it is retained and its image-plane coordinates are assigned as a multiple\n", - "image of the source.\n", - "\n", - "We require the grid of triangles to be fine enough such that the source-plane (y,x) coordinate is contained within\n", - "one of the triangles to a sufficient precision for our science case. This is controlled by the `pixel_scale_precision`\n", - "input, which sets the target pixel scale of the grid. \n", - "\n", - "Triangles of iteratively finer resolution are created until this precision is met, therefore a lower value of\n", - "`pixel_scale_precision` will lead to a more precise estimate of the multiple image positions at the expense of\n", - "increased computational overhead.\n", - "\n", - "Here is a visualization of the triangulation approach:\n", - "\n", - "[CODE]\n", - "\n", - "__Dataset__\n", - "\n", - "We first create a `PointDataset` object, which is similar to an `Imaging` or `Interferometer` object but contains the\n", - "positions of the multiple images of the point source and their noise-map values.\n", - "\n", - "The noise values are the centroid uncertainty with which each multiple image is localised in the data. This is *not*\n", - "the imaging pixel scale \u2014 that is the detector's sampling, not its centroiding precision. Bright point sources are\n", - "localised by fitting the instrumental PSF to the image, and the resulting centroid uncertainty is typically a small\n", - "fraction of a pixel: ~0.005\" (5 mas) for HST or adaptive-optics imaging of lensed quasars in the strong-lensing\n", - "literature (CASTLES, TDCOSMO/H0LiCOW). See `scripts/point_source/simulator.py` for a full discussion.\n", - "\n", - "We construct a `PointDataset` from the multiple-image positions we just solved for, with a small (5 mas) noise\n", - "value per position. In a real analysis the positions would come from your data reduction pipeline (e.g. a\n", - "centroid fit to each multiple image in the calibrated image), but for this demo we reuse the solver's output\n", - "so the script adapts to whatever number of multiple images the mass model produces.\n", - "\n", - "The demagnified central image is excluded by the solver's `magnification_threshold=0.1` setting, so it does not\n", - "appear in the dataset either \u2014 standard practice in point-source modeling.\n", - "\n", - "The dataset's name `point_0` is an important label, as explained in more detail below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions_data = al.Grid2DIrregular(positions)\n", - "\n", - "positions_noise_map = al.ArrayIrregular([0.005] * len(positions))\n", - "\n", - "dataset = al.PointDataset(\n", - " name=\"point_0\", positions=positions_data, positions_noise_map=positions_noise_map\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can print this dictionary to see the dataset's `name`, `positions` and `fluxes` and noise-map values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Point Dataset Info:\")\n", - "print(dataset.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The positions can be plotted over the observed image, to make sure they overlap with the multiple images we expect." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=tracer_extended.image_2d_from(grid=grid), title=\"Image\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Name Pairing__\n", - "\n", - "The names of the point-source datasets have an even more important role, the names are used to pair each dataset to the\n", - "point sources in the lens model used to fit it.\n", - "\n", - "For example, when creating the tracer at the beginning of this script, we named the point source `point_0`:\n", - "\n", - "point_source = al.ps.Point(centre=(0.07, 0.07))\n", - "source_galaxy = al.Galaxy(redshift=1.0, point_0=point_source)\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "When we fit the point source dataset using this tracer, the name is again used in order to pair the dataset to the\n", - "this point source. This means that point source with a centre of (0.07\", 0.07\") is used to fit the dataset with the\n", - "name `point_0`.\n", - "\n", - "If there is no point-source in the model that has the same name as a `PointDataset`, that data is not used in\n", - "the model-fit. If a point-source is included in the model whose name has no corresponding entry in \n", - "the `PointDataset` an error will be raised.\n", - "\n", - "In this example, where there is just one source, name pairing is redundant. However, point-source datasets may\n", - "have many source galaxies in them, and name pairing allows us to extend the point-source modeling to systems with\n", - "many point sources.\n", - "\n", - "__Fitting__\n", - "\n", - "Just like we used a `Tracer` to fit imaging and interferometer data, we can use it to fit point-source data via the\n", - "`FitPoint` object.\n", - "\n", - "The name pairing described above is used internally into the `FitPointDataset` object to ensure that the correct point\n", - "source is fitted to each dataset. \n", - "\n", - "The fit is returned as a dictionary which mirrors the `PointDataset`, where its keys are again the names of the datasets." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitPointDataset(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePairRepeat, # This input is describe below\n", - ")\n", - "\n", - "print(fit.positions.residual_map)\n", - "print(fit.positions.normalized_residual_map)\n", - "print(fit.positions.chi_squared_map)\n", - "print(fit.positions.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Chi Squared__\n", - "\n", - "For point-source modeling, there are many different ways to define the likelihood function, broadly referred to a\n", - "an `image-plane chi-squared` or `source-plane chi-squared`. This determines whether the multiple images of the point\n", - "source are used to compute the likelihood in the source-plane or image-plane.\n", - "\n", - "The default settings used above use the image-plane chi-squared, which uses the `PointSolver` to determine the \n", - "multiple images of the point source in the image-plane for the given mass model and compares the positions of these \n", - "model images to the observed images to compute the chi-squared and likelihood.\n", - "\n", - "There are still many different ways the image-plane chi-squared can be computed, for example do we allow for \n", - "repeat image-pairs (i.e. the same multiple image being observed multiple times)? Do we pair all possible combinations\n", - "of multiple images to observed images? This default settings use the simplest approach, which pair each multiple image\n", - "with the observed image that is closest to it, allowing for repeat image pairs. \n", - "\n", - "For example, we can repeat the fit above whilst not allowing for repeat image pairs as follows:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitPointDataset(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePair, # Different input to the one used above\n", - ")\n", - "\n", - "print(\n", - " \"Minimum Distance Between Observed Multiple Images and Model Multiple Images Without Repeats:\"\n", - ")\n", - "print(fit.positions.residual_map)\n", - "\n", - "print(\"Log Likelihood Without Repeats:\")\n", - "print(fit.positions.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can allow for repeat image pairs by using the `FitPositionsImagePairRepeat` class, which is the default input." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitPointDataset(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePairRepeat, # Different input to the one used above\n", - ")\n", - "\n", - "print(\n", - " \"Minimum Distance Between Observed Multiple Images and Model Multiple Images With Repeats:\"\n", - ")\n", - "print(fit.positions.residual_map)\n", - "\n", - "print(\"Log Likelihood With Repeats:\")\n", - "print(fit.positions.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For a \"source-plane chi-squared\", the likelihood is computed in the source-plane. The analysis is simpler, it ray-traces\n", - "the multiple images back to the source-plane and defines a chi-squared metric. For example, the default implementation \n", - "sums the Euclidean distance between the image positions and the point source centre in the source-plane.\n", - "\n", - "The source-plane chi-squared is significantly faster to compute than the image-plane chi-squared, however it is \n", - "less robust than the image-plane chi-squared and can lead to biased lens model results. \n", - "\n", - "Here is an example of how to use the source-plane chi-squared:" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitPointDataset(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsSource, # Different input to the one used above\n", - ")\n", - "\n", - "print(\n", - " \"Minimum Distance Between Source Plane Centre and Model Source Plane Images After Ray Tracing:\"\n", - ")\n", - "print(fit.positions.residual_map)\n", - "\n", - "print(\"Log Likelihood in the Source Plane:\")\n", - "print(fit.positions.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fluxes__\n", - "\n", - "Another measurable quantity of a point source is its flux\u2014the total amount of light received from each multiple image \n", - "of the point source (e.g., the quasar images).\n", - "\n", - "In practice, fluxes are often measured but not used directly when analyzing lensed point sources such as quasars or \n", - "supernovae. This is because fluxes can be significantly affected by microlensing, which many lens models do not \n", - "accurately capture. However, in this simulation, microlensing is not included, so the fluxes can be simulated and \n", - "fitted reliably.\n", - "\n", - "We now simulate the fluxes of the multiple images of this point source.\n", - "\n", - "Given a mass model and the (y, x) image-plane coordinates of each image, the magnification at each point can be \n", - "calculated.\n", - "\n", - "Below, we compute the magnification for every multiple image coordinate, which will then be used to simulate their \n", - "fluxes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnifications = al.LensCalc.from_tracer(\n", - " tracer=tracer\n", - ").magnification_2d_via_hessian_from(grid=positions)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the fluxes, we assume the source galaxy point-source has a total flux of 1.0.\n", - "\n", - "Each observed image has a flux that is the source's flux multiplied by the magnification at that image-plane coordinate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "flux = 1.0\n", - "fluxes = [flux * np.abs(magnification) for magnification in magnifications]\n", - "fluxes = al.ArrayIrregular(values=fluxes)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We set the flux noise to 5% of each measured flux. For lensed quasars and supernovae, photometric flux\n", - "uncertainties are dominated by microlensing systematics rather than photon noise, so models that exclude\n", - "microlensing typically assume a few-percent flux uncertainty per image. See `scripts/point_source/simulator.py`\n", - "for a fuller discussion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "flux_rel_noise = 0.05\n", - "\n", - "fluxes_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Flux Point Dataset__\n", - "\n", - "The fluxes are not input a `PointDataset` object, alongside the image-plane coordinates of the multiple images\n", - "and their associated noise-map values. \n", - "\n", - "We again give the dataset the name `point_0`, which is a label given to the dataset to indicate that it is a dataset \n", - "of a single point-source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_data,\n", - " positions_noise_map=positions_noise_map,\n", - " fluxes=fluxes,\n", - " fluxes_noise_map=fluxes_noise_map,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Flux Fitting__\n", - "\n", - "Above, we used a `FitPointDataset` to fit the positions of the point source in the image-plane.\n", - "\n", - "We can also use it to fit the fluxes of the point source, which is done by passing the new dataset also containing\n", - "the `fluxes` and `fluxes_noise_map` to the fit.\n", - "\n", - "To fit fluxes, our model point source also needs a flux parameter, which is done by using the `PointFlux`\n", - "component instead of the `Point` component. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "point = al.ps.PointFlux(centre=(0.07, 0.07), flux=1.0)\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, point_0=point)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitPointDataset(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePairRepeat, # This input is describe below\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit now contains both a `positions` and `fluxes` attribute, which contain the fit of the positions and fluxes\n", - "of the point source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.positions.residual_map)\n", - "print(fit.positions.normalized_residual_map)\n", - "print(fit.positions.chi_squared_map)\n", - "print(fit.positions.log_likelihood)\n", - "\n", - "print(fit.flux.residual_map)\n", - "print(fit.flux.normalized_residual_map)\n", - "print(fit.flux.chi_squared_map)\n", - "print(fit.flux.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Time Delays__\n", - "\n", - "Another measurable quantity of a point source is its time delay\u2014the time it takes for light to travel from the\n", - "source to the observer for each multiple image of the point source (e.g., the quasar images). This is often expressed\n", - "as the relative time delay between each image and the image with the shortest time delay, which is often referred to as\n", - "the \"reference image.\"\n", - "\n", - "Time delays are commonly used in strong lensing analyses, for example to measure the Hubble constant, since\n", - "they are less affected by microlensing and can provide robust cosmological constraints.\n", - "\n", - "We now simulate the same point source dataset, but this time including the time delays of the multiple images.\n", - "\n", - "Given a mass model and (y, x) image-plane coordinates, the time delay at each image-plane position can be\n", - "calculated from the mass model. It includes the contribution of both the geometric time delay (the time it takes\n", - "different light rays to travel from the source to the observer) and the Shapiro time delay (the time it takes\n", - "light to travel through the gravitational potential of the lens galaxy)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "time_delays = tracer.time_delays_from(grid=positions)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In real observations, time delays are measured by photometrically monitoring the multiple images over months\n", - "to years and cross-correlating their light curves. State-of-the-art monitoring campaigns (COSMOGRAIL, TDCOSMO)\n", - "achieve ~1\u20133% relative precision on the longest delays in well-sampled quad systems. For simplicity we adopt\n", - "a 5% relative uncertainty here. Real-world uncertainties are not strictly proportional to the delay magnitude\n", - "(they depend on monitoring cadence, length, and microlensing variability), but a constant fractional error is\n", - "a reasonable simulator default. See `scripts/point_source/simulator.py` for a fuller discussion." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "time_delay_rel_noise = 0.05\n", - "\n", - "time_delays_noise_map = al.ArrayIrregular(\n", - " values=np.abs(time_delays) * time_delay_rel_noise\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The time delays are input into a `PointDataset` object, alongside the image-plane coordinates of the multiple images\n", - "and their associated noise-map values. \n", - "\n", - "We again give the dataset the name `point_0`, which is a label given to the dataset to indicate that it is a dataset \n", - "of a single point-source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_data,\n", - " positions_noise_map=positions_noise_map,\n", - " time_delays=time_delays,\n", - " time_delays_noise_map=time_delays_noise_map,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Time Delay Fitting__\n", - "\n", - "We can also use the `FitPointDataset` to fit the time delays of the point source, which is done by passing the new\n", - "dataset also containing the `time_delays` and `time_delays_noise_map` to the fit.\n", - "\n", - "To fit time delays, the model point source does not need any special parameters (like it did for flux fitting),\n", - "so we can revert back to the normal `Point` component." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "point = al.ps.Point(centre=(0.07, 0.07))\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0, point_0=point)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - "fit = al.FitPointDataset(\n", - " dataset=dataset,\n", - " tracer=tracer,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePairRepeat, # This input is describe below\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The fit now contains both a `positions` and `time_delays` attribute, which contain the fit of the positions and fluxes\n", - "of the point source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(fit.positions.residual_map)\n", - "print(fit.positions.normalized_residual_map)\n", - "print(fit.positions.chi_squared_map)\n", - "print(fit.positions.log_likelihood)\n", - "\n", - "print(fit.time_delays.residual_map)\n", - "print(fit.time_delays.normalized_residual_map)\n", - "print(fit.time_delays.chi_squared_map)\n", - "print(fit.time_delays.log_likelihood)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__New User Wrap Up__\n", - "\n", - "The `point_source` package of the `autolens_workspace` contains numerous example scripts for performing point source\n", - "modeling. These focus on \"galaxy scale\" lenses, which are lenses that have a single lens galaxy, as opposed to\n", - "\"group scale\" or \"cluster scale\" lenses which have multiple lens galaxies.\n", - "\n", - "Point source modeling is at the heart of group and cluster scale lens modeling, and is the topic of the\n", - "next overview script.\n", - "\n", - "__Shape Solver__\n", - "\n", - "All calculations above assumed the source was a point source with no size. \n", - "\n", - "This was built into the point-solver, for example when we solved for the multiple images of the point source in the \n", - "image-plane, we ray-traced triangles to the source-plane and asked whether the source-plane (y,x) centre was within \n", - "the triangle.\n", - "\n", - "There is functionality to include the size and shape of the source in the calculation, which uses the `ShapeSolver`\n", - "class. This still traces triangles, but each iteration of the solver now computes the area of each image-plane triangle \n", - "that is within the source-plane shape. This means we can determine the area in the image-plane that maps within an \n", - "extended region of the source-plane shape.\n", - "\n", - "For example, by inputting the shape `Circle` with a radius of 0.001\", the shape solver will determine the area of the \n", - "multiple images pixel which fall within this circle, which is different information to the point solver which told\n", - "us the exact (y,x) coordinates of the multiple images.\n", - "\n", - "The ratio of the total image pixel area to the area within the source-plane \n", - "circle is the magnification factor of the source. This magnification factor then changes the observed flux of each \n", - "multiple image.\n", - "\n", - "Observations we might think are fully in the point source regime therefore may have an observable signature of the size\n", - "of the source in the flux ratios and magnifications of the multiple images. Therefore, sometimes the source size \n", - "is large enough that it is important we account for it." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "solver = al.ShapeSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")\n", - "\n", - "solver.find_magnification(tracer=tracer, shape=al.Circle(x=0.0, y=0.0, radius=0.001))\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Guide: Point Sources\n", + "--------------------\n", + "\n", + "The examples covered so far have focused on strongly lensed galaxies, where extended surface brightness is warped into\n", + "stunning giant arcs and Einstein rings visible in high-quality telescope images. For these extended sources, light\n", + "profile objects\u2014such as analytic Sersic profiles\u2014are used to represent their surface brightness.\n", + "\n", + "However, some observed sources are extremely small, spanning just light weeks or days across. In these cases, only the\n", + "source\u2019s central point of light is detected in each multiple image. Such sources are called **point sources**,\n", + "which typically include quasars, supernovae, or stars.\n", + "\n", + "Strictly speaking, a point source does have a finite size\u2014on the order of light weeks or days\u2014but it is considered a\n", + "point source because its size is orders of magnitude smaller than the resolution of the telescope. As a result, it\n", + "appears as a single point of light, with all the flux of each multiple image effectively contained within a single pixel.\n", + "\n", + "Point sources affect lensing calculations differently than extended sources, requiring dedicated methods and\n", + "functionality. This functionality is described here and used throughout the `point_source` simulation and modeling\n", + "examples.\n", + "\n", + "If you are new to analyzing strong lenses with point sources, this guide is the ideal place to start!\n", + "\n", + "__Contents__\n", + "\n", + "- **Lensed Point Source:** To begin, we create a strong lens image using an isothermal mass model and a source with a compact.\n", + "- **Point Source:** The image above visually illustrates where the source\u2019s light is traced on the image plane.\n", + "- **Point Solver:** For a point source, our goal is to find the (y, x) coordinates in the image plane that map directly.\n", + "- **Number of Solutions:** The number of solutions (e.g.\n", + "- **Solving the Lens Equation:** In the literature, the process of finding the multiple images of a source in the image-plane is.\n", + "- **Triangle Tracing:** Computing the multiple image positions of a point source is a non-linear problem.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Name Pairing:** The names of the point-source datasets have an even more important role, the names are used to pair.\n", + "- **Fitting:** Fit the lens model to the dataset and inspect the results.\n", + "- **Chi Squared:** For point-source modeling, there are many different ways to define the likelihood function, broadly.\n", + "- **Fluxes:** Another measurable quantity of a point source is its flux\u2014the total amount of light received from.\n", + "- **Flux Point Dataset:** The fluxes are not input a `PointDataset` object, alongside the image-plane coordinates of the.\n", + "- **Flux Fitting:** Above, we used a `FitPointDataset` to fit the positions of the point source in the image-plane.\n", + "- **Time Delays:** Another measurable quantity of a point source is its time delay\u2014the time it takes for light to.\n", + "- **Time Delay Fitting:** We can also use the `FitPointDataset` to fit the time delays of the point source, which is done by.\n", + "- **New User Wrap Up:** The `point_source` package of the `autolens_workspace` contains numerous example scripts for.\n", + "- **Shape Solver:** All calculations above assumed the source was a point source with no size.\n", + "\n", + "__JAX__\n", + "\n", + "This guide walks through point-source fitting at a low level (no\n", + "non-linear search) \u2014 the `FitPointDataset` and `PointSolver` objects work\n", + "on either NumPy or JAX. For the standard search-driven modeling path\n", + "where `AnalysisPoint` auto-enables `use_jax=True` and the JIT happens\n", + "inside the search driver, see `start_here.py` / `modeling.py`. For the\n", + "explicit `@jax.jit + PointSolver(use_jax=True)` pattern (useful for fast\n", + "forward-solving in custom code), see `simulator.py`'s `__JAX Variant__`\n", + "section.\n", + "\n", + "Note: `PointSolver(use_jax=True)` works inside `@jax.jit`. The triangle-\n", + "refinement loop operates on raw arrays and doesn't go through `.native`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lensed Point Source__\n", + "\n", + "To begin, we create a strong lens image using an isothermal mass model and a source with a compact exponential light profile.\n", + "\n", + "Although our goal is to demonstrate solving for the multiple image positions of a point source, simulating the data \n", + "with a compact extended source makes the visualization of the point solver\u2019s solutions clearer." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "isothermal_mass_profile = al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + ")\n", + "\n", + "exponential_light_profile = al.lp.ExponentialCore(\n", + " centre=(0.07, 0.07), intensity=0.1, effective_radius=0.1\n", + ")\n", + "\n", + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=isothermal_mass_profile,\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " light=exponential_light_profile,\n", + ")\n", + "\n", + "tracer_extended = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the image of our strongly lensed source galaxy.\n", + "\n", + "Clearly visible are four multiple images arranged in a cross configuration. The brightest pixels correspond to the \n", + "four (y, x) multiple image positions that our point source solver aims to identify." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_array(array=tracer_extended.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Source__\n", + "\n", + "The image above visually illustrates where the source\u2019s light is traced on the image plane.\n", + "\n", + "We now treat this source as a point source by defining a source galaxy using the `Point` class.\n", + "\n", + "This point source shares the same center as the compact source above, ensuring that the multiple image positions \n", + "coincide with those previously shown in the image plane." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "point_source = al.ps.Point(centre=(0.07, 0.07))\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, point_0=point_source)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "For a point source, our goal is to find the (y, x) coordinates in the image plane that map directly to the center of \n", + "the point source in the source plane\u2014these are its \"multiple images.\" This is achieved using a `PointSolver`, which \n", + "determines the multiple images of the mass model for a point source located at a given (y, x) position in the \n", + "source plane.\n", + "\n", + "The solver works by ray tracing triangles from the image plane back to the source plane and checking whether the \n", + "source-plane (y, x) center lies inside each triangle. It iteratively refines this process by ray tracing progressively \n", + "smaller triangles, allowing the multiple image positions to be determined with sub-pixel precision.\n", + "\n", + "The `PointSolver` requires an initial grid of (y, x) coordinates in the image plane, which defines the first set of \n", + "triangles to ray trace. It also needs a `pixel_scale_precision` parameter, specifying the resolution at which the \n", + "multiple images are computed. Smaller values increase precision but require longer computation times. The value \n", + "of 0.001 used here balances efficiency and accuracy.\n", + "\n", + "Strong lens mass models often predict a \"central image,\" a multiple image that is usually heavily demagnified and thus \n", + "not observed. Since the `PointSolver` finds all valid multiple images, it will locate this central image regardless of \n", + "its visibility. To avoid including this unobservable image, we set a `magnification_threshold=0.1`, which discards any \n", + "images with magnifications below this value.\n", + "\n", + "If your dataset does include a detectable central image, you should lower this threshold accordingly to include it in \n", + "your analysis.\n", + "\n", + "We now compute the multiple image positions by creating a `PointSolver` object and passing it the tracer of our \n", + "strong lens system." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now pass the tracer to the solver, to determine the image-plane multiple images for the source centre.\n", + "\n", + "The solver will find the image-plane coordinates that map directly to the source-plane coordinate (0.07\", 0.07\"), \n", + "which we plot below.\n", + "\n", + "The plot shows the four solved multiple image positions (with the central image excluded) as a scatter plot. To make \n", + "the positions clearer, we increase the marker size and use asterisks\u2014PyAutoLens\u2019s standard symbol for denoting \n", + "multiple images of strong lenses." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = solver.solve(tracer=tracer, source_plane_coordinate=(0.07, 0.07))\n", + "\n", + "aplt.plot_grid(grid=positions, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The plot above makes it difficult to directly compare the multiple image positions with the image of the strong lens itself.\n", + "\n", + "To improve clarity, we overplot the multiple image positions on the strong lens image. This clearly shows that the \n", + "multiple images coincide with the centers of the brightest pixels of the lensed source galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=tracer_extended.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Number of Solutions__\n", + "\n", + "The number of solutions (e.g. the number of image-plane multiple images that map to the source centre) depends\n", + "on the mass model of the lens: \n", + "\n", + " - For spherical mass profiles, there are three unique solutions, including a demagnified central image.\n", + "\n", + " - For elliptical mass profiles, there are five unique solutions, again including a demagnified central image.\n", + "\n", + " - For lenses with multiple mass profiles (e.g. two galaxies) and more exotic mass distributions, the number of \n", + " solutions can be even higher. \n", + "\n", + "__Solving the Lens Equation__\n", + "\n", + "In the literature, the process of finding the multiple images of a source in the image-plane is often referred to as\n", + "'solving the lens equation'.\n", + "\n", + "There lens equation is a fundamental equation in lensing, which describes how light rays are deflected from the\n", + "image-plane to the source-plane. It is given by:\n", + "\n", + "$\\beta = \\theta - \\hat{\\alpha}(\\theta)$\n", + "\n", + "Where:\n", + "\n", + "$\\beta$ is the source-plane (y,x) coordinate.\n", + "$\\theta$ is the image-plane (y,x) coordinate.\n", + "$\\hat{\\alpha}(\\theta)$ is the deflection angle at image-plane (y,x) coordinate $\\theta$.\n", + "\n", + "The lens equation is non-linear, as the deflection angle $\\hat{\\alpha}$ depends on the mass model of the lens galaxy.\n", + "\n", + "It is therefore called solving the lens equation because we are trying to find the image-plane (y,x) coordinates $\\theta$\n", + "that satisfies the equation above for a given source-plane (y,x) coordinate $\\beta$.\n", + "\n", + "__Triangle Tracing__\n", + "\n", + "Computing the multiple image positions of a point source is a non-linear problem. Given a source-plane (y,x) coordinate,\n", + "there are multiple image-plane (y,x) coordinates that trace to that source-plane coordinate, and there is no simple\n", + "analytic solution to determine these image-plane coordinates.\n", + "\n", + "The solver therefore uses a triangulation approach to find the multiple image positions. It first overlays a grid of\n", + "triangles over the image-plane, and uses the mass model to trace these triangles to the source-plane. If a triangle\n", + "contains the source-plane (y,x) coordinate, it is retained and its image-plane coordinates are assigned as a multiple\n", + "image of the source.\n", + "\n", + "We require the grid of triangles to be fine enough such that the source-plane (y,x) coordinate is contained within\n", + "one of the triangles to a sufficient precision for our science case. This is controlled by the `pixel_scale_precision`\n", + "input, which sets the target pixel scale of the grid. \n", + "\n", + "Triangles of iteratively finer resolution are created until this precision is met, therefore a lower value of\n", + "`pixel_scale_precision` will lead to a more precise estimate of the multiple image positions at the expense of\n", + "increased computational overhead.\n", + "\n", + "Here is a visualization of the triangulation approach:\n", + "\n", + "[CODE]\n", + "\n", + "__Dataset__\n", + "\n", + "We first create a `PointDataset` object, which is similar to an `Imaging` or `Interferometer` object but contains the\n", + "positions of the multiple images of the point source and their noise-map values.\n", + "\n", + "The noise values are the centroid uncertainty with which each multiple image is localised in the data. This is *not*\n", + "the imaging pixel scale \u2014 that is the detector's sampling, not its centroiding precision. Bright point sources are\n", + "localised by fitting the instrumental PSF to the image, and the resulting centroid uncertainty is typically a small\n", + "fraction of a pixel: ~0.005\" (5 mas) for HST or adaptive-optics imaging of lensed quasars in the strong-lensing\n", + "literature (CASTLES, TDCOSMO/H0LiCOW). See `scripts/point_source/simulator.py` for a full discussion.\n", + "\n", + "We construct a `PointDataset` from the multiple-image positions we just solved for, with a small (5 mas) noise\n", + "value per position. In a real analysis the positions would come from your data reduction pipeline (e.g. a\n", + "centroid fit to each multiple image in the calibrated image), but for this demo we reuse the solver's output\n", + "so the script adapts to whatever number of multiple images the mass model produces.\n", + "\n", + "The demagnified central image is excluded by the solver's `magnification_threshold=0.1` setting, so it does not\n", + "appear in the dataset either \u2014 standard practice in point-source modeling.\n", + "\n", + "The dataset's name `point_0` is an important label, as explained in more detail below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions_data = al.Grid2DIrregular(positions)\n", + "\n", + "positions_noise_map = al.ArrayIrregular([0.005] * len(positions))\n", + "\n", + "dataset = al.PointDataset(\n", + " name=\"point_0\", positions=positions_data, positions_noise_map=positions_noise_map\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can print this dictionary to see the dataset's `name`, `positions` and `fluxes` and noise-map values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Point Dataset Info:\")\n", + "print(dataset.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The positions can be plotted over the observed image, to make sure they overlap with the multiple images we expect." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=tracer_extended.image_2d_from(grid=grid), title=\"Image\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Name Pairing__\n", + "\n", + "The names of the point-source datasets have an even more important role, the names are used to pair each dataset to the\n", + "point sources in the lens model used to fit it.\n", + "\n", + "For example, when creating the tracer at the beginning of this script, we named the point source `point_0`:\n", + "\n", + "point_source = al.ps.Point(centre=(0.07, 0.07))\n", + "source_galaxy = al.Galaxy(redshift=1.0, point_0=point_source)\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "When we fit the point source dataset using this tracer, the name is again used in order to pair the dataset to the\n", + "this point source. This means that point source with a centre of (0.07\", 0.07\") is used to fit the dataset with the\n", + "name `point_0`.\n", + "\n", + "If there is no point-source in the model that has the same name as a `PointDataset`, that data is not used in\n", + "the model-fit. If a point-source is included in the model whose name has no corresponding entry in \n", + "the `PointDataset` an error will be raised.\n", + "\n", + "In this example, where there is just one source, name pairing is redundant. However, point-source datasets may\n", + "have many source galaxies in them, and name pairing allows us to extend the point-source modeling to systems with\n", + "many point sources.\n", + "\n", + "__Fitting__\n", + "\n", + "Just like we used a `Tracer` to fit imaging and interferometer data, we can use it to fit point-source data via the\n", + "`FitPoint` object.\n", + "\n", + "The name pairing described above is used internally into the `FitPointDataset` object to ensure that the correct point\n", + "source is fitted to each dataset. \n", + "\n", + "The fit is returned as a dictionary which mirrors the `PointDataset`, where its keys are again the names of the datasets." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitPointDataset(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePairRepeat, # This input is describe below\n", + ")\n", + "\n", + "print(fit.positions.residual_map)\n", + "print(fit.positions.normalized_residual_map)\n", + "print(fit.positions.chi_squared_map)\n", + "print(fit.positions.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Chi Squared__\n", + "\n", + "For point-source modeling, there are many different ways to define the likelihood function, broadly referred to a\n", + "an `image-plane chi-squared` or `source-plane chi-squared`. This determines whether the multiple images of the point\n", + "source are used to compute the likelihood in the source-plane or image-plane.\n", + "\n", + "The default settings used above use the image-plane chi-squared, which uses the `PointSolver` to determine the \n", + "multiple images of the point source in the image-plane for the given mass model and compares the positions of these \n", + "model images to the observed images to compute the chi-squared and likelihood.\n", + "\n", + "There are still many different ways the image-plane chi-squared can be computed, for example do we allow for \n", + "repeat image-pairs (i.e. the same multiple image being observed multiple times)? Do we pair all possible combinations\n", + "of multiple images to observed images? This default settings use the simplest approach, which pair each multiple image\n", + "with the observed image that is closest to it, allowing for repeat image pairs. \n", + "\n", + "For example, we can repeat the fit above whilst not allowing for repeat image pairs as follows:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitPointDataset(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePair, # Different input to the one used above\n", + ")\n", + "\n", + "print(\n", + " \"Minimum Distance Between Observed Multiple Images and Model Multiple Images Without Repeats:\"\n", + ")\n", + "print(fit.positions.residual_map)\n", + "\n", + "print(\"Log Likelihood Without Repeats:\")\n", + "print(fit.positions.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can allow for repeat image pairs by using the `FitPositionsImagePairRepeat` class, which is the default input." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitPointDataset(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePairRepeat, # Different input to the one used above\n", + ")\n", + "\n", + "print(\n", + " \"Minimum Distance Between Observed Multiple Images and Model Multiple Images With Repeats:\"\n", + ")\n", + "print(fit.positions.residual_map)\n", + "\n", + "print(\"Log Likelihood With Repeats:\")\n", + "print(fit.positions.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For a \"source-plane chi-squared\", the likelihood is computed in the source-plane. The analysis is simpler, it ray-traces\n", + "the multiple images back to the source-plane and defines a chi-squared metric. For example, the default implementation \n", + "sums the Euclidean distance between the image positions and the point source centre in the source-plane.\n", + "\n", + "The source-plane chi-squared is significantly faster to compute than the image-plane chi-squared, however it is \n", + "less robust than the image-plane chi-squared and can lead to biased lens model results. \n", + "\n", + "Here is an example of how to use the source-plane chi-squared:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitPointDataset(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsSource, # Different input to the one used above\n", + ")\n", + "\n", + "print(\n", + " \"Minimum Distance Between Source Plane Centre and Model Source Plane Images After Ray Tracing:\"\n", + ")\n", + "print(fit.positions.residual_map)\n", + "\n", + "print(\"Log Likelihood in the Source Plane:\")\n", + "print(fit.positions.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fluxes__\n", + "\n", + "Another measurable quantity of a point source is its flux\u2014the total amount of light received from each multiple image \n", + "of the point source (e.g., the quasar images).\n", + "\n", + "In practice, fluxes are often measured but not used directly when analyzing lensed point sources such as quasars or \n", + "supernovae. This is because fluxes can be significantly affected by microlensing, which many lens models do not \n", + "accurately capture. However, in this simulation, microlensing is not included, so the fluxes can be simulated and \n", + "fitted reliably.\n", + "\n", + "We now simulate the fluxes of the multiple images of this point source.\n", + "\n", + "Given a mass model and the (y, x) image-plane coordinates of each image, the magnification at each point can be \n", + "calculated.\n", + "\n", + "Below, we compute the magnification for every multiple image coordinate, which will then be used to simulate their \n", + "fluxes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnifications = al.LensCalc.from_tracer(\n", + " tracer=tracer\n", + ").magnification_2d_via_hessian_from(grid=positions)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the fluxes, we assume the source galaxy point-source has a total flux of 1.0.\n", + "\n", + "Each observed image has a flux that is the source's flux multiplied by the magnification at that image-plane coordinate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "flux = 1.0\n", + "fluxes = [flux * np.abs(magnification) for magnification in magnifications]\n", + "fluxes = al.ArrayIrregular(values=fluxes)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We set the flux noise to 5% of each measured flux. For lensed quasars and supernovae, photometric flux\n", + "uncertainties are dominated by microlensing systematics rather than photon noise, so models that exclude\n", + "microlensing typically assume a few-percent flux uncertainty per image. See `scripts/point_source/simulator.py`\n", + "for a fuller discussion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "flux_rel_noise = 0.05\n", + "\n", + "fluxes_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Flux Point Dataset__\n", + "\n", + "The fluxes are not input a `PointDataset` object, alongside the image-plane coordinates of the multiple images\n", + "and their associated noise-map values. \n", + "\n", + "We again give the dataset the name `point_0`, which is a label given to the dataset to indicate that it is a dataset \n", + "of a single point-source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_data,\n", + " positions_noise_map=positions_noise_map,\n", + " fluxes=fluxes,\n", + " fluxes_noise_map=fluxes_noise_map,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Flux Fitting__\n", + "\n", + "Above, we used a `FitPointDataset` to fit the positions of the point source in the image-plane.\n", + "\n", + "We can also use it to fit the fluxes of the point source, which is done by passing the new dataset also containing\n", + "the `fluxes` and `fluxes_noise_map` to the fit.\n", + "\n", + "To fit fluxes, our model point source also needs a flux parameter, which is done by using the `PointFlux`\n", + "component instead of the `Point` component. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "point = al.ps.PointFlux(centre=(0.07, 0.07), flux=1.0)\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, point_0=point)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitPointDataset(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePairRepeat, # This input is describe below\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit now contains both a `positions` and `fluxes` attribute, which contain the fit of the positions and fluxes\n", + "of the point source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.positions.residual_map)\n", + "print(fit.positions.normalized_residual_map)\n", + "print(fit.positions.chi_squared_map)\n", + "print(fit.positions.log_likelihood)\n", + "\n", + "print(fit.flux.residual_map)\n", + "print(fit.flux.normalized_residual_map)\n", + "print(fit.flux.chi_squared_map)\n", + "print(fit.flux.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Time Delays__\n", + "\n", + "Another measurable quantity of a point source is its time delay\u2014the time it takes for light to travel from the\n", + "source to the observer for each multiple image of the point source (e.g., the quasar images). This is often expressed\n", + "as the relative time delay between each image and the image with the shortest time delay, which is often referred to as\n", + "the \"reference image.\"\n", + "\n", + "Time delays are commonly used in strong lensing analyses, for example to measure the Hubble constant, since\n", + "they are less affected by microlensing and can provide robust cosmological constraints.\n", + "\n", + "We now simulate the same point source dataset, but this time including the time delays of the multiple images.\n", + "\n", + "Given a mass model and (y, x) image-plane coordinates, the time delay at each image-plane position can be\n", + "calculated from the mass model. It includes the contribution of both the geometric time delay (the time it takes\n", + "different light rays to travel from the source to the observer) and the Shapiro time delay (the time it takes\n", + "light to travel through the gravitational potential of the lens galaxy)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "time_delays = tracer.time_delays_from(grid=positions)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In real observations, time delays are measured by photometrically monitoring the multiple images over months\n", + "to years and cross-correlating their light curves. State-of-the-art monitoring campaigns (COSMOGRAIL, TDCOSMO)\n", + "achieve ~1\u20133% relative precision on the longest delays in well-sampled quad systems. For simplicity we adopt\n", + "a 5% relative uncertainty here. Real-world uncertainties are not strictly proportional to the delay magnitude\n", + "(they depend on monitoring cadence, length, and microlensing variability), but a constant fractional error is\n", + "a reasonable simulator default. See `scripts/point_source/simulator.py` for a fuller discussion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "time_delay_rel_noise = 0.05\n", + "\n", + "time_delays_noise_map = al.ArrayIrregular(\n", + " values=np.abs(time_delays) * time_delay_rel_noise\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The time delays are input into a `PointDataset` object, alongside the image-plane coordinates of the multiple images\n", + "and their associated noise-map values. \n", + "\n", + "We again give the dataset the name `point_0`, which is a label given to the dataset to indicate that it is a dataset \n", + "of a single point-source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_data,\n", + " positions_noise_map=positions_noise_map,\n", + " time_delays=time_delays,\n", + " time_delays_noise_map=time_delays_noise_map,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Time Delay Fitting__\n", + "\n", + "We can also use the `FitPointDataset` to fit the time delays of the point source, which is done by passing the new\n", + "dataset also containing the `time_delays` and `time_delays_noise_map` to the fit.\n", + "\n", + "To fit time delays, the model point source does not need any special parameters (like it did for flux fitting),\n", + "so we can revert back to the normal `Point` component." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "point = al.ps.Point(centre=(0.07, 0.07))\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0, point_0=point)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + "fit = al.FitPointDataset(\n", + " dataset=dataset,\n", + " tracer=tracer,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePairRepeat, # This input is describe below\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The fit now contains both a `positions` and `time_delays` attribute, which contain the fit of the positions and fluxes\n", + "of the point source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(fit.positions.residual_map)\n", + "print(fit.positions.normalized_residual_map)\n", + "print(fit.positions.chi_squared_map)\n", + "print(fit.positions.log_likelihood)\n", + "\n", + "print(fit.time_delays.residual_map)\n", + "print(fit.time_delays.normalized_residual_map)\n", + "print(fit.time_delays.chi_squared_map)\n", + "print(fit.time_delays.log_likelihood)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__New User Wrap Up__\n", + "\n", + "The `point_source` package of the `autolens_workspace` contains numerous example scripts for performing point source\n", + "modeling. These focus on \"galaxy scale\" lenses, which are lenses that have a single lens galaxy, as opposed to\n", + "\"group scale\" or \"cluster scale\" lenses which have multiple lens galaxies.\n", + "\n", + "Point source modeling is at the heart of group and cluster scale lens modeling, and is the topic of the\n", + "next overview script.\n", + "\n", + "__Shape Solver__\n", + "\n", + "All calculations above assumed the source was a point source with no size. \n", + "\n", + "This was built into the point-solver, for example when we solved for the multiple images of the point source in the \n", + "image-plane, we ray-traced triangles to the source-plane and asked whether the source-plane (y,x) centre was within \n", + "the triangle.\n", + "\n", + "There is functionality to include the size and shape of the source in the calculation, which uses the `ShapeSolver`\n", + "class. This still traces triangles, but each iteration of the solver now computes the area of each image-plane triangle \n", + "that is within the source-plane shape. This means we can determine the area in the image-plane that maps within an \n", + "extended region of the source-plane shape.\n", + "\n", + "For example, by inputting the shape `Circle` with a radius of 0.001\", the shape solver will determine the area of the \n", + "multiple images pixel which fall within this circle, which is different information to the point solver which told\n", + "us the exact (y,x) coordinates of the multiple images.\n", + "\n", + "The ratio of the total image pixel area to the area within the source-plane \n", + "circle is the magnification factor of the source. This magnification factor then changes the observed flux of each \n", + "multiple image.\n", + "\n", + "Observations we might think are fully in the point source regime therefore may have an observable signature of the size\n", + "of the source in the flux ratios and magnifications of the multiple images. Therefore, sometimes the source size \n", + "is large enough that it is important we account for it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "solver = al.ShapeSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")\n", + "\n", + "solver.find_magnification(tracer=tracer, shape=al.Circle(x=0.0, y=0.0, radius=0.001))\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/modeling.ipynb b/notebooks/point_source/modeling.ipynb index d18d42863..41e17d7aa 100644 --- a/notebooks/point_source/modeling.ipynb +++ b/notebooks/point_source/modeling.ipynb @@ -1,739 +1,776 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Modeling: Start Here\n", - "====================\n", - "\n", - "This script is the starting point for lens modeling of point-source lens datasets, for example the multiple image\n", - "positions of a lensed quasar.\n", - "\n", - "__Contents__\n", - "\n", - "- **Not Using Light Profiles:** Users who are familiar with analysing imaging or interferometer data will be used to performing.\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset:** Load and plot the strong lens dataset.\n", - "- **Point Solver:** For point-source modeling we require a `PointSolver`, which determines the multiple-images of the.\n", - "- **Model Composition:** Compose the lens model using the Model and Collection API.\n", - "- **Name Pairing:** Every point-source dataset in the `PointDataset` has a name, which in this example was `point_0`.\n", - "- **Coordinates:** Coordinate system assumptions for the model-fit.\n", - "- **Search:** Configure the non-linear search used to fit the model.\n", - "- **Unique Identifier:** In the path above, the `unique_identifier` appears as a collection of characters, where this.\n", - "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", - "- **Chi Squared:** For point-source modeling, there are many different ways to define the likelihood function, broadly.\n", - "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", - "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", - "- **VRAM Use:** When running AutoLens with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM.\n", - "- **Run Times:** Profiling the expected run time of the model-fit.\n", - "- **Output Folder Layout:** Description of the structure of the `output` folder where results are written.\n", - "- **Result:** Overview of the results of the model-fit.\n", - "- **Results:** Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "- **Modeling Customization:** The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative.\n", - "\n", - "__Not Using Light Profiles__\n", - "\n", - "Users who are familiar with analysing imaging or interferometer data will be used to\n", - "performing lens modeling using light profiles, which have parameter that describe the shape and size of the\n", - "galaxy's luminous emission.\n", - "\n", - "For point sources, for example a lensed quasar, it is invalid to model the source using light profiles, because they\n", - "implicitly assume an extended surface brightness distribution. Point source modeling instead assumes the source\n", - "has a (y,x) `centre` (y,x), but does not have other parameters like elliptical components or an effective radius.\n", - "\n", - "This changes how the ray-tracing calculations that go into point source modeling are performed. They are briefly\n", - "touched on in this example, but for a more detailed explanation checkout the\n", - "`autolens_workspace/*/overview/overview_8_point_sources.py` example.\n", - "\n", - "__Model__\n", - "\n", - "This script fits a `PointDataset` data of a 'galaxy-scale' strong lens with a model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The source `Galaxy` is a point source `Point`.\n", - "\n", - "The `ExternalShear` is also not included in the mass model, where it is for the `imaging` and `interferometer` examples.\n", - "For a quadruply imaged point source (8 data points) there is insufficient information to fully constain a model with\n", - "an `Isothermal` and `ExternalShear` (9 parameters)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "Load the strong lens point-source dataset `simple`, which is the dataset we will use to perform point source \n", - "lens modeling." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_name = \"simple\"\n", - "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Auto-Simulation__\n", - "\n", - "If the dataset does not already exist on your system, it will be created by running the corresponding\n", - "simulator script. This ensures that all example scripts can be run without manually simulating data first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "if not dataset_path.exists():\n", - " import subprocess\n", - " import sys\n", - "\n", - " subprocess.run(\n", - " [sys.executable, \"scripts/point_source/simulator.py\"],\n", - " check=True,\n", - " )" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now load the point source dataset we will fit using point source modeling. \n", - "\n", - "We load this data as a `PointDataset`, which contains the positions of every point source. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.from_json(\n", - " file_path=dataset_path / \"point_dataset_positions_only.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can print this dictionary to see the dataset's `name`, `positions`and noise-map values." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\"Point Dataset Info:\")\n", - "print(dataset.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also plot the positions of the `PointDataset`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_point_dataset(dataset=dataset)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We next load an image of the dataset. \n", - "\n", - "Although we are performing point-source modeling and do not use this data in the actual modeling, it is useful to \n", - "load it for visualization, for example to see where the multiple images of the point source are located relative to the \n", - "lens galaxy.\n", - "\n", - "The image will also be passed to the analysis further down, meaning that visualization of the point-source model\n", - "overlaid over the image will be output making interpretation of the results straight forward.\n", - "\n", - "Loading and inputting the image of the dataset in this way is entirely optional, and if you are only interested in\n", - "performing point-source modeling you do not need to do this." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.05)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also plot the dataset's multiple image positions over the observed image, to ensure they overlap the\n", - "lensed source's multiple images." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.plot_array(array=data, title=\"\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "For point-source modeling we require a `PointSolver`, which determines the multiple-images of the mass model for a \n", - "point source at location (y,x) in the source plane. \n", - "\n", - "It does this by ray tracing triangles from the image-plane to the source-plane and calculating if the \n", - "source-plane (y,x) centre is inside the triangle. The method gradually ray-traces smaller and smaller triangles so \n", - "that the multiple images can be determine with sub-pixel precision.\n", - "\n", - "The `PointSolver` requires a starting grid of (y,x) coordinates in the image-plane which defines the first set\n", - "of triangles that are ray-traced to the source-plane. It also requires that a `pixel_scale_precision` is input, \n", - "which is the resolution up to which the multiple images are computed. The lower the `pixel_scale_precision`, the\n", - "longer the calculation, with the value of 0.001 below balancing efficiency with precision.\n", - "\n", - "Strong lens mass models have a multiple image called the \"central image\". However, the image is nearly always \n", - "significantly demagnified, meaning that it is not observed and cannot constrain the lens model. As this image is a\n", - "valid multiple image, the `PointSolver` will locate it irrespective of whether its so demagnified it is not observed.\n", - "To ensure this does not occur, we set a `magnification_threshold=0.1`, which discards this image because its\n", - "magnification will be well below this threshold.\n", - "\n", - "If your dataset contains a central image that is observed you should reduce to include it in\n", - "the analysis." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(100, 100),\n", - " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid,\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.1,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "We compose a lens model where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", - " \n", - " - The source galaxy's light is a point `Point` [2 parameters].\n", - "\n", - "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=7.\n", - "\n", - "__Model Composition__\n", - "\n", - "The API below for composing a lens model uses the `Model` and `Collection` objects, which are imported from \n", - "**PyAutoLens**'s parent project **PyAutoFit** \n", - "\n", - "The API is fairly self explanatory and is straight forward to extend, for example adding more light profiles\n", - "to the lens and source or using a different mass profile.\n", - "\n", - "A full description of model composition is provided by the model cookbook: \n", - "\n", - "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", - "\n", - "__Name Pairing__\n", - "\n", - "Every point-source dataset in the `PointDataset` has a name, which in this example was `point_0`. This `name` pairs \n", - "the dataset to the `Point` in the model below. Because the name of the dataset is `point_0`, the \n", - "only `Point` object that is used to fit it must have the name `point_0`.\n", - "\n", - "If there is no point-source in the model that has the same name as a `PointDataset`, that data is not used in\n", - "the model-fit. If a point-source is included in the model whose name has no corresponding entry in \n", - "the `PointDataset` it will raise an error.\n", - "\n", - "In this example, where there is just one source, name pairing appears unnecessary. However, point-source datasets may\n", - "have many source galaxies in them, and name pairing is necessary to ensure every point source in the lens model is \n", - "fitted to its particular lensed images in the `PointDataset`.\n", - "\n", - "__Coordinates__\n", - "\n", - "The model fitting default settings assume that the lens galaxy centre is near the coordinates (0.0\", 0.0\"). \n", - "\n", - "If for your dataset the lens is not centred at (0.0\", 0.0\"), we recommend that you either: \n", - "\n", - " - Reduce your data so that the centre is (`autolens_workspace/*/data_preparation`). \n", - " - Manually override the lens model priors (`autolens_workspace/*/guides/modeling/customize`)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "# Lens:\n", - "\n", - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal)\n", - "\n", - "# Source:\n", - "\n", - "point_0 = af.Model(al.ps.Point)\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)\n", - "\n", - "# Overall Lens Model:\n", - "\n", - "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `info` attribute shows the model in a readable format.\n", - "\n", - "[The `info` below may not display optimally on your computer screen, for example the whitespace between parameter\n", - "names on the left and parameter priors on the right may lead them to appear across multiple lines. This is a\n", - "common issue in Jupyter notebooks.\n", - "\n", - "The`info_whitespace_length` parameter in the file `config/general.yaml` in the [output] section can be changed to \n", - "increase or decrease the amount of whitespace (The Jupyter notebook kernel will need to be reset for this change to \n", - "appear in a notebook).]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(model.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Search__\n", - "\n", - "The lens model is fitted to the data using a non-linear search. \n", - "\n", - "All examples in the autolens workspace use the nested sampling algorithm \n", - "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has revealed gives the most \n", - "accurate and efficient modeling results.\n", - "\n", - "Nautilus has one main setting that trades-off accuracy and computational run-time, the number of `live_points`. \n", - "A higher number of live points gives a more accurate result, but increases the run-time. A lower value give \n", - "less reliable lens modeling (e.g. the fit may infer a local maxima), but is faster. \n", - "\n", - "The suitable value depends on the model complexity whereby models with more parameters require more live points. \n", - "The default value of 200 is sufficient for the vast majority of common lens models. Lower values often given reliable\n", - "results though, and speed up the run-times. In this example, given the model is quite simple (N=21 parameters), we \n", - "reduce the number of live points to 100 to speed up the run-time.\n", - "\n", - "__Unique Identifier__\n", - "\n", - "In the path above, the `unique_identifier` appears as a collection of characters, where this identifier is generated \n", - "based on the model, search and dataset that are used in the fit.\n", - " \n", - "An identical combination of model and search generates the same identifier, meaning that rerunning the script will use \n", - "the existing results to resume the model-fit. In contrast, if you change the model or search, a new unique identifier \n", - "will be generated, ensuring that the model-fit results are output into a separate folder.\n", - "\n", - "We additionally want the unique identifier to be specific to the dataset fitted, so that if we fit different datasets\n", - "with the same model and search results are output to a different folder. We achieve this below by passing\n", - "the `dataset_name` to the search's `unique_tag`.\n", - "\n", - "__Live Visual Update__\n", - "\n", - "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", - "live display surface:\n", - "\n", - "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", - " watch the fit converge without leaving your terminal.\n", - "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", - " refreshes in place every `iterations_per_quick_update`.\n", - "\n", - "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", - "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "search = af.Nautilus(\n", - " path_prefix=Path(\"point_source\"), # The path where results and output are stored.\n", - " name=\"modeling\", # The name of the fit and folder results are output to.\n", - " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", - " n_live=100, # The number of Nautilus \"live\" points, increase for more complex models.\n", - " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", - " iterations_per_quick_update=10000, # Every N iterations the max likelihood model, is visualized in the Jupter Notebook and output to hard-disk.\n", - " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Chi Squared__\n", - "\n", - "For point-source modeling, there are many different ways to define the likelihood function, broadly referred to a\n", - "an `image-plane chi-squared` or `source-plane chi-squared`. This determines whether the multiple images of the point\n", - "source are used to compute the likelihood in the source-plane or image-plane.\n", - "\n", - "We will use an \"image-plane chi-squared\", which uses the `PointSolver` to determine the multiple images of the point\n", - "source in the image-plane for the given mass model and compares the positions of these model images to the observed\n", - "images to compute the chi-squared and likelihood.\n", - "\n", - "There are still many different ways the image-plane chi-squared can be computed, for example do we allow for \n", - "repeat image-pairs (i.e. the same multiple image being observed multiple times)? Do we pair all possible combinations\n", - "of multiple images to observed images? This example uses the simplest approach, which is to pair each multiple image\n", - "with the observed image that is closest to it, allowing for repeat image pairs. \n", - "\n", - "For a \"source-plane chi-squared\", the likelihood is computed in the source-plane. The analysis basically just ray-traces\n", - "the multiple images back to the source-plane and defines a chi-squared metric. For example, the default implementation \n", - "sums the Euclidean distance between the image positions and the point source centre in the source-plane.\n", - "\n", - "The source-plane chi-squared is significantly faster to compute than the image-plane chi-squared, as it requires \n", - "only ray-tracing the ~4 observed image positions and does not require the iterative triangle ray-tracing approach\n", - "of the image-plane chi-squared. However, the source-plane chi-squared is less robust than the image-plane chi-squared,\n", - "and can lead to biased lens model results. If you are using the source-plane chi-squared, you should be aware of this\n", - "and interpret the results with caution.\n", - "\n", - "Checkout the guide `autolens_workspace/*/point_source/fit` for more details and a full illustration of the\n", - "different ways the chi-squared can be computed.\n", - "\n", - "__Analysis__\n", - "\n", - "We next create an `AnalysisPoint` object, which can be given many inputs customizing how the lens model is \n", - "fitted to the data, which in this example includes the solver and the chi-squared method.\n", - "\n", - "Internally, this object defines the `log_likelihood_function` used by the non-linear search to fit the model to \n", - "the `Imaging` dataset. \n", - "\n", - "It is not vital that you as a user understand the details of how the `log_likelihood_function` fits a lens model to \n", - "data, but interested readers can find a step-by-step guide of the likelihood \n", - "function at ``autolens_workspace/*/point_source/log_likelihood_function`\n", - "\n", - "__JAX__\n", - "\n", - "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU\n", - "support, your fits will run much faster (around 10 minutes instead of an hour). If only a CPU is available,\n", - "JAX will still provide a speed up via multithreading, with fits taking around 20-30 minutes.\n", - "\n", - "If you don\u2019t have a GPU locally, consider Google Colab which provides free GPUs, so your modeling runs are much faster." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis = al.AnalysisPoint(\n", - " dataset=dataset,\n", - " solver=solver,\n", - " fit_positions_cls=al.FitPositionsImagePairRepeat, # Image-plane chi-squared with repeat image pairs.\n", - " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__VRAM Use__\n", - "\n", - "When running AutoLens with JAX on a GPU, the analysis must fit within the GPU\u2019s\n", - "available VRAM. If insufficient VRAM is available, the analysis will fail with an\n", - "out-of-memory error, typically during JIT compilation or the first likelihood call.\n", - "\n", - "Two factors dictate the VRAM usage of an analysis:\n", - "\n", - "- The number of arrays and other data structures JAX must store in VRAM to fit the model\n", - " to the data in the likelihood function. This is dictated by the model complexity and dataset size.\n", - " For a MGE model its relatively low, but for other models (e.g. pixelized sources) it can be much higher.\n", - "\n", - "- The `batch_size` sets how many likelihood evaluations are performed simultaneously.\n", - " Increasing the batch size increases VRAM usage but can reduce overall run time,\n", - " while decreasing it lowers VRAM usage at the cost of slower execution.\n", - "\n", - "Before running an analysis, users should check that the estimated VRAM usage for the\n", - "chosen batch size is comfortably below their GPU\u2019s total VRAM.\n", - "\n", - "For a point solver with an image-plane chi squared and one set of positions with a single plane VRAM use is relatively\n", - "low (~0.1GB). For models with more planes and datasets with more multiple images it can be much higher (> 1GB going\n", - "beyond 10GB).\n", - "\n", - "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", - "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "analysis.print_vram_use(model=model, batch_size=search.batch_size)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Run Times__\n", - "\n", - "Lens modeling can be a computationally expensive process. When fitting complex models to high resolution datasets \n", - "run times can be of order hours, days, weeks or even months.\n", - "\n", - "Run times are dictated by two factors:\n", - "\n", - " - The log likelihood evaluation time: the time it takes for a single `instance` of the lens model to be fitted to \n", - " the dataset such that a log likelihood is returned.\n", - " \n", - " - The number of iterations (e.g. log likelihood evaluations) performed by the non-linear search: more complex lens\n", - " models require more iterations to converge to a solution.\n", - " \n", - "For this analysis, the log likelihood evaluation time is < 0.001 seconds on GPU, ~0.01 seconds on CPU, which is \n", - "extremely fast for lens modeling. \n", - "\n", - "To estimate the expected overall run time of the model-fit we multiply the log likelihood evaluation time by an \n", - "estimate of the number of iterations the non-linear search will perform, which is around 10000 to 30000 for this model.\n", - "\n", - "GPU run times are around 10 minutes, CPU run times are around 30 minutes.\n", - "\n", - "__Model-Fit__\n", - "\n", - "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", - "for on-the-fly visualization and results).\n", - "\n", - "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce \n", - "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(\n", - " \"\"\"\n", - " The non-linear search has begun running.\n", - "\n", - " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", - "\n", - " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", - " \"\"\"\n", - ")\n", - "\n", - "result = search.fit(model=model, analysis=analysis)\n", - "\n", - "print(\"The search has finished run - you may now continue the notebook.\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output Folder Layout__\n", - "\n", - "Now the fit is running you should checkout the `autolens_workspace/output` folder. This is where results are\n", - "written to hard-disk in human-readable formats \u2014 `.json`, `.csv`, `.png` and plain text.\n", - "\n", - "As the fit progresses, results are written on the fly using the highest likelihood model found by the\n", - "non-linear search so far. This means you can inspect the model-fit as it runs, without waiting for the\n", - "non-linear search to terminate.\n", - "\n", - "Each completed fit lives at a path like::\n", - "\n", - " output/point_source//modeling//\n", - " files/ <- JSON + CSV: loadable Python objects\n", - " tracer.json <- max log likelihood Tracer\n", - " model.json <- fitted af.Collection model\n", - " samples.csv <- full Nautilus samples\n", - " samples_summary.json <- max log likelihood parameter values + errors\n", - " samples_info.json <- metadata about the samples\n", - " search.json <- non-linear search configuration\n", - " settings.json <- search settings\n", - " cosmology.json <- cosmology used for the fit\n", - " covariance.csv <- parameter covariance matrix\n", - " image/ <- PNG: point-source fit visualisations\n", - " positions.png <- observed vs model-predicted multiple-image positions\n", - " fluxes.png <- observed vs model-predicted point-source fluxes\n", - " tracer.png <- tracer image-plane and source-plane plots\n", - " model.info <- human-readable model summary\n", - " model.results <- human-readable fit summary\n", - " search.summary <- search run summary\n", - " search_internal/ <- internal files used to resume / visualise the search\n", - " metadata <- run metadata\n", - "\n", - "The `` is a 32-character identifier derived from the model, search and dataset, so re-running the\n", - "same configuration resumes from the existing fit automatically.\n", - "\n", - "__Result__\n", - "\n", - "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", - "\n", - "[Above, we discussed that the `info_whitespace_length` parameter in the config files could b changed to make \n", - "the `model.info` attribute display optimally on your computer. This attribute also controls the whitespace of the\n", - "`result.info` attribute.]" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "print(result.max_log_likelihood_instance)\n", - "\n", - "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grid)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The result contains the full posterior information of our non-linear search, including all parameter samples, \n", - "log likelihood values and tools to compute the errors on the lens model. \n", - "\n", - "There are built in visualization tools for plotting this.\n", - "\n", - "The plot is labeled with short hand parameter names (e.g. `sersic_index` is mapped to the short hand \n", - "parameter `n`). These mappings ate specified in the `config/notation.yaml` file and can be customized by users.\n", - "\n", - "The superscripts of labels correspond to the name each component was given in the model (e.g. for the `Isothermal`\n", - "mass its name `mass` defined when making the `Model` above is used)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.corner_anesthetic(samples=result.samples)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Results__\n", - "\n", - "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", - "\n", - "__Modeling Customization__\n", - "\n", - "The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative non-linear searches,\n", - "other than Nautilus, that can be used to fit lens models. \n", - "\n", - "They also provide details on how to customize the model-fit, for example the priors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Start Here\n", + "====================\n", + "\n", + "This script is the starting point for lens modeling of point-source lens datasets, for example the multiple image\n", + "positions of a lensed quasar.\n", + "\n", + "__Contents__\n", + "\n", + "- **Not Using Light Profiles:** Users who are familiar with analysing imaging or interferometer data will be used to performing.\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset:** Load and plot the strong lens dataset.\n", + "- **Point Solver:** For point-source modeling we require a `PointSolver`, which determines the multiple-images of the.\n", + "- **Model Composition:** Compose the lens model using the Model and Collection API.\n", + "- **Name Pairing:** Every point-source dataset in the `PointDataset` has a name, which in this example was `point_0`.\n", + "- **Coordinates:** Coordinate system assumptions for the model-fit.\n", + "- **Search:** Configure the non-linear search used to fit the model.\n", + "- **Unique Identifier:** In the path above, the `unique_identifier` appears as a collection of characters, where this.\n", + "- **Live Visual Update:** Push the quick-update image to a live display surface.\n", + "- **Chi Squared:** For point-source modeling, there are many different ways to define the likelihood function, broadly.\n", + "- **Analysis:** Create the Analysis object that defines how the model is fitted to the data.\n", + "- **JAX:** JAX acceleration for fast GPU/CPU model-fitting.\n", + "- **VRAM Use:** When running AutoLens with JAX on a GPU, the analysis must fit within the GPU\u2019s available VRAM.\n", + "- **Run Times:** Profiling the expected run time of the model-fit.\n", + "- **Output Folder Layout:** Description of the structure of the `output` folder where results are written.\n", + "- **Result:** Overview of the results of the model-fit.\n", + "- **Results:** Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "- **Modeling Customization:** The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative.\n", + "\n", + "__Not Using Light Profiles__\n", + "\n", + "Users who are familiar with analysing imaging or interferometer data will be used to\n", + "performing lens modeling using light profiles, which have parameter that describe the shape and size of the\n", + "galaxy's luminous emission.\n", + "\n", + "For point sources, for example a lensed quasar, it is invalid to model the source using light profiles, because they\n", + "implicitly assume an extended surface brightness distribution. Point source modeling instead assumes the source\n", + "has a (y,x) `centre` (y,x), but does not have other parameters like elliptical components or an effective radius.\n", + "\n", + "This changes how the ray-tracing calculations that go into point source modeling are performed. They are briefly\n", + "touched on in this example, but for a more detailed explanation checkout the\n", + "`autolens_workspace/*/overview/overview_8_point_sources.py` example.\n", + "\n", + "__Model__\n", + "\n", + "This script fits a `PointDataset` data of a 'galaxy-scale' strong lens with a model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The source `Galaxy` is a point source `Point`.\n", + "\n", + "The `ExternalShear` is also not included in the mass model, where it is for the `imaging` and `interferometer` examples.\n", + "For a quadruply imaged point source (8 data points) there is insufficient information to fully constain a model with\n", + "an `Isothermal` and `ExternalShear` (9 parameters)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load the strong lens point-source dataset `simple`, which is the dataset we will use to perform point source \n", + "lens modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"point_source\" / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if not dataset_path.exists():\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/point_source/simulator.py\"],\n", + " check=True,\n", + " )" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now load the point source dataset we will fit using point source modeling. \n", + "\n", + "We load this data as a `PointDataset`, which contains the positions of every point source. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.from_json(\n", + " file_path=dataset_path / \"point_dataset_positions_only.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can print this dictionary to see the dataset's `name`, `positions`and noise-map values." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Point Dataset Info:\")\n", + "print(dataset.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also plot the positions of the `PointDataset`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_point_dataset(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We next load an image of the dataset. \n", + "\n", + "Although we are performing point-source modeling and do not use this data in the actual modeling, it is useful to \n", + "load it for visualization, for example to see where the multiple images of the point source are located relative to the \n", + "lens galaxy.\n", + "\n", + "The image will also be passed to the analysis further down, meaning that visualization of the point-source model\n", + "overlaid over the image will be output making interpretation of the results straight forward.\n", + "\n", + "Loading and inputting the image of the dataset in this way is entirely optional, and if you are only interested in\n", + "performing point-source modeling you do not need to do this." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "data = al.Array2D.from_fits(file_path=dataset_path / \"data.fits\", pixel_scales=0.05)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also plot the dataset's multiple image positions over the observed image, to ensure they overlap the\n", + "lensed source's multiple images." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.plot_array(array=data, title=\"\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "For point-source modeling we require a `PointSolver`, which determines the multiple-images of the mass model for a \n", + "point source at location (y,x) in the source plane. \n", + "\n", + "It does this by ray tracing triangles from the image-plane to the source-plane and calculating if the \n", + "source-plane (y,x) centre is inside the triangle. The method gradually ray-traces smaller and smaller triangles so \n", + "that the multiple images can be determine with sub-pixel precision.\n", + "\n", + "The `PointSolver` requires a starting grid of (y,x) coordinates in the image-plane which defines the first set\n", + "of triangles that are ray-traced to the source-plane. It also requires that a `pixel_scale_precision` is input, \n", + "which is the resolution up to which the multiple images are computed. The lower the `pixel_scale_precision`, the\n", + "longer the calculation, with the value of 0.001 below balancing efficiency with precision.\n", + "\n", + "Strong lens mass models have a multiple image called the \"central image\". However, the image is nearly always \n", + "significantly demagnified, meaning that it is not observed and cannot constrain the lens model. As this image is a\n", + "valid multiple image, the `PointSolver` will locate it irrespective of whether its so demagnified it is not observed.\n", + "To ensure this does not occur, we set a `magnification_threshold=0.1`, which discards this image because its\n", + "magnification will be well below this threshold.\n", + "\n", + "If your dataset contains a central image that is observed you should reduce to include it in\n", + "the analysis." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.2, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid,\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.1,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose a lens model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + " \n", + " - The source galaxy's light is a point `Point` [2 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=7.\n", + "\n", + "__Model Composition__\n", + "\n", + "The API below for composing a lens model uses the `Model` and `Collection` objects, which are imported from \n", + "**PyAutoLens**'s parent project **PyAutoFit** \n", + "\n", + "The API is fairly self explanatory and is straight forward to extend, for example adding more light profiles\n", + "to the lens and source or using a different mass profile.\n", + "\n", + "A full description of model composition is provided by the model cookbook: \n", + "\n", + "https://pyautolens.readthedocs.io/en/latest/general/model_cookbook.html\n", + "\n", + "__Name Pairing__\n", + "\n", + "Every point-source dataset in the `PointDataset` has a name, which in this example was `point_0`. This `name` pairs \n", + "the dataset to the `Point` in the model below. Because the name of the dataset is `point_0`, the \n", + "only `Point` object that is used to fit it must have the name `point_0`.\n", + "\n", + "If there is no point-source in the model that has the same name as a `PointDataset`, that data is not used in\n", + "the model-fit. If a point-source is included in the model whose name has no corresponding entry in \n", + "the `PointDataset` it will raise an error.\n", + "\n", + "In this example, where there is just one source, name pairing appears unnecessary. However, point-source datasets may\n", + "have many source galaxies in them, and name pairing is necessary to ensure every point source in the lens model is \n", + "fitted to its particular lensed images in the `PointDataset`.\n", + "\n", + "__Coordinates__\n", + "\n", + "The model fitting default settings assume that the lens galaxy centre is near the coordinates (0.0\", 0.0\"). \n", + "\n", + "If for your dataset the lens is not centred at (0.0\", 0.0\"), we recommend that you either: \n", + "\n", + " - Reduce your data so that the centre is (`autolens_workspace/*/data_preparation`). \n", + " - Manually override the lens model priors (`autolens_workspace/*/guides/modeling/customize`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal)\n", + "\n", + "# Source:\n", + "\n", + "point_0 = af.Model(al.ps.Point)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, point_0=point_0)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format.\n", + "\n", + "[The `info` below may not display optimally on your computer screen, for example the whitespace between parameter\n", + "names on the left and parameter priors on the right may lead them to appear across multiple lines. This is a\n", + "common issue in Jupyter notebooks.\n", + "\n", + "The`info_whitespace_length` parameter in the file `config/general.yaml` in the [output] section can be changed to \n", + "increase or decrease the amount of whitespace (The Jupyter notebook kernel will need to be reset for this change to \n", + "appear in a notebook).]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The lens model is fitted to the data using a non-linear search. \n", + "\n", + "All examples in the autolens workspace use the nested sampling algorithm \n", + "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), which extensive testing has revealed gives the most \n", + "accurate and efficient modeling results.\n", + "\n", + "Nautilus has one main setting that trades-off accuracy and computational run-time, the number of `live_points`. \n", + "A higher number of live points gives a more accurate result, but increases the run-time. A lower value give \n", + "less reliable lens modeling (e.g. the fit may infer a local maxima), but is faster. \n", + "\n", + "The suitable value depends on the model complexity whereby models with more parameters require more live points. \n", + "The default value of 200 is sufficient for the vast majority of common lens models. Lower values often given reliable\n", + "results though, and speed up the run-times. In this example, given the model is quite simple (N=21 parameters), we \n", + "reduce the number of live points to 100 to speed up the run-time.\n", + "\n", + "__Unique Identifier__\n", + "\n", + "In the path above, the `unique_identifier` appears as a collection of characters, where this identifier is generated \n", + "based on the model, search and dataset that are used in the fit.\n", + " \n", + "An identical combination of model and search generates the same identifier, meaning that rerunning the script will use \n", + "the existing results to resume the model-fit. In contrast, if you change the model or search, a new unique identifier \n", + "will be generated, ensuring that the model-fit results are output into a separate folder.\n", + "\n", + "We additionally want the unique identifier to be specific to the dataset fitted, so that if we fit different datasets\n", + "with the same model and search results are output to a different folder. We achieve this below by passing\n", + "the `dataset_name` to the search's `unique_tag`.\n", + "\n", + "__Live Visual Update__\n", + "\n", + "By default the quick-update image is only written to disk. Set `live_visual_update=True` to also push it to a\n", + "live display surface:\n", + "\n", + "- **Python script** \u2014 a matplotlib window opens automatically and refreshes with each quick update, so you can\n", + " watch the fit converge without leaving your terminal.\n", + "- **Jupyter / Colab notebook** \u2014 the cell that ran `search.fit(...)` shows a single self-updating image that\n", + " refreshes in place every `iterations_per_quick_update`.\n", + "\n", + "The disk write (`fit.png`) always happens regardless of this flag. Set it to `False` (the default) if you just\n", + "want the on-disk output, or if you are running in a headless environment (e.g. an HPC cluster)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"point_source\"), # The path where results and output are stored.\n", + " name=\"modeling\", # The name of the fit and folder results are output to.\n", + " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", + " n_live=100, # The number of Nautilus \"live\" points, increase for more complex models.\n", + " n_batch=50, # GPU lens model fits are batched and run simultaneously, see VRAM section below.\n", + " iterations_per_quick_update=10000, # Every N iterations the max likelihood model, is visualized in the Jupter Notebook and output to hard-disk.\n", + " live_visual_update=False, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Chi Squared__\n", + "\n", + "For point-source modeling, there are many different ways to define the likelihood function, broadly referred to a\n", + "an `image-plane chi-squared` or `source-plane chi-squared`. This determines whether the multiple images of the point\n", + "source are used to compute the likelihood in the source-plane or image-plane.\n", + "\n", + "We will use an \"image-plane chi-squared\", which uses the `PointSolver` to determine the multiple images of the point\n", + "source in the image-plane for the given mass model and compares the positions of these model images to the observed\n", + "images to compute the chi-squared and likelihood.\n", + "\n", + "There are still many different ways the image-plane chi-squared can be computed, for example do we allow for \n", + "repeat image-pairs (i.e. the same multiple image being observed multiple times)? Do we pair all possible combinations\n", + "of multiple images to observed images? This example uses the simplest approach, which is to pair each multiple image\n", + "with the observed image that is closest to it, allowing for repeat image pairs. \n", + "\n", + "For a \"source-plane chi-squared\", the likelihood is computed in the source-plane. The analysis basically just ray-traces\n", + "the multiple images back to the source-plane and defines a chi-squared metric. For example, the default implementation \n", + "sums the Euclidean distance between the image positions and the point source centre in the source-plane.\n", + "\n", + "The source-plane chi-squared is significantly faster to compute than the image-plane chi-squared, as it requires \n", + "only ray-tracing the ~4 observed image positions and does not require the iterative triangle ray-tracing approach\n", + "of the image-plane chi-squared. However, the source-plane chi-squared is less robust than the image-plane chi-squared,\n", + "and can lead to biased lens model results. If you are using the source-plane chi-squared, you should be aware of this\n", + "and interpret the results with caution.\n", + "\n", + "Checkout the guide `autolens_workspace/*/point_source/fit` for more details and a full illustration of the\n", + "different ways the chi-squared can be computed.\n", + "\n", + "__Analysis__\n", + "\n", + "We next create an `AnalysisPoint` object, which can be given many inputs customizing how the lens model is \n", + "fitted to the data, which in this example includes the solver and the chi-squared method.\n", + "\n", + "Internally, this object defines the `log_likelihood_function` used by the non-linear search to fit the model to \n", + "the `Imaging` dataset. \n", + "\n", + "It is not vital that you as a user understand the details of how the `log_likelihood_function` fits a lens model to \n", + "data, but interested readers can find a step-by-step guide of the likelihood \n", + "function at ``autolens_workspace/*/point_source/log_likelihood_function`\n", + "\n", + "__JAX__\n", + "\n", + "PyAutoLens uses JAX under the hood for fast GPU/CPU acceleration. If JAX is installed with GPU\n", + "support, your fits will run much faster (around 10 minutes instead of an hour). If only a CPU is available,\n", + "JAX will still provide a speed up via multithreading, with fits taking around 20-30 minutes.\n", + "\n", + "If you don\u2019t have a GPU locally, consider Google Colab which provides free GPUs, so your modeling runs are much faster." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisPoint(\n", + " dataset=dataset,\n", + " solver=solver,\n", + " fit_positions_cls=al.FitPositionsImagePairRepeat, # Image-plane chi-squared with repeat image pairs.\n", + " use_jax=True, # JAX will use GPUs for acceleration if available, else JAX will use multithreaded CPUs.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__VRAM Use__\n", + "\n", + "When running AutoLens with JAX on a GPU, the analysis must fit within the GPU\u2019s\n", + "available VRAM. If insufficient VRAM is available, the analysis will fail with an\n", + "out-of-memory error, typically during JIT compilation or the first likelihood call.\n", + "\n", + "Two factors dictate the VRAM usage of an analysis:\n", + "\n", + "- The number of arrays and other data structures JAX must store in VRAM to fit the model\n", + " to the data in the likelihood function. This is dictated by the model complexity and dataset size.\n", + " For a MGE model its relatively low, but for other models (e.g. pixelized sources) it can be much higher.\n", + "\n", + "- The `batch_size` sets how many likelihood evaluations are performed simultaneously.\n", + " Increasing the batch size increases VRAM usage but can reduce overall run time,\n", + " while decreasing it lowers VRAM usage at the cost of slower execution.\n", + "\n", + "Before running an analysis, users should check that the estimated VRAM usage for the\n", + "chosen batch size is comfortably below their GPU\u2019s total VRAM.\n", + "\n", + "For a point solver with an image-plane chi squared and one set of positions with a single plane VRAM use is relatively\n", + "low (~0.1GB). For models with more planes and datasets with more multiple images it can be much higher (> 1GB going\n", + "beyond 10GB).\n", + "\n", + "The method below prints the VRAM usage estimate for the analysis and model with the specified batch size,\n", + "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis.print_vram_use(model=model, batch_size=search.batch_size)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "Lens modeling can be a computationally expensive process. When fitting complex models to high resolution datasets \n", + "run times can be of order hours, days, weeks or even months.\n", + "\n", + "Run times are dictated by two factors:\n", + "\n", + " - The log likelihood evaluation time: the time it takes for a single `instance` of the lens model to be fitted to \n", + " the dataset such that a log likelihood is returned.\n", + " \n", + " - The number of iterations (e.g. log likelihood evaluations) performed by the non-linear search: more complex lens\n", + " models require more iterations to converge to a solution.\n", + " \n", + "For this analysis, the log likelihood evaluation time is < 0.001 seconds on GPU, ~0.01 seconds on CPU, which is \n", + "extremely fast for lens modeling. \n", + "\n", + "To estimate the expected overall run time of the model-fit we multiply the log likelihood evaluation time by an \n", + "estimate of the number of iterations the non-linear search will perform, which is around 10000 to 30000 for this model.\n", + "\n", + "GPU run times are around 10 minutes, CPU run times are around 30 minutes.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis object to the non-linear search (checkout the output folder\n", + "for on-the-fly visualization and results).\n", + "\n", + "**Run Time Error:** On certain operating systems (e.g. Windows, Linux) and Python versions, the code below may produce \n", + "an error. If this occurs, see the `autolens_workspace/guides/modeling/bug_fix` example for a fix." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell with progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output Folder Layout__\n", + "\n", + "Now the fit is running you should checkout the `autolens_workspace/output` folder. This is where results are\n", + "written to hard-disk in human-readable formats \u2014 `.json`, `.csv`, `.png` and plain text.\n", + "\n", + "As the fit progresses, results are written on the fly using the highest likelihood model found by the\n", + "non-linear search so far. This means you can inspect the model-fit as it runs, without waiting for the\n", + "non-linear search to terminate.\n", + "\n", + "Each completed fit lives at a path like::\n", + "\n", + " output/point_source//modeling//\n", + " files/ <- JSON + CSV: loadable Python objects\n", + " tracer.json <- max log likelihood Tracer\n", + " model.json <- fitted af.Collection model\n", + " samples.csv <- full Nautilus samples\n", + " samples_summary.json <- max log likelihood parameter values + errors\n", + " samples_info.json <- metadata about the samples\n", + " search.json <- non-linear search configuration\n", + " settings.json <- search settings\n", + " cosmology.json <- cosmology used for the fit\n", + " covariance.csv <- parameter covariance matrix\n", + " image/ <- PNG: point-source fit visualisations\n", + " positions.png <- observed vs model-predicted multiple-image positions\n", + " fluxes.png <- observed vs model-predicted point-source fluxes\n", + " tracer.png <- tracer image-plane and source-plane plots\n", + " model.info <- human-readable model summary\n", + " model.results <- human-readable fit summary\n", + " search.summary <- search run summary\n", + " search_internal/ <- internal files used to resume / visualise the search\n", + " metadata <- run metadata\n", + "\n", + "The `` is a 32-character identifier derived from the model, search and dataset, so re-running the\n", + "same configuration resumes from the existing fit automatically.\n", + "\n", + "__Result__\n", + "\n", + "The search returns a result object, which whose `info` attribute shows the result in a readable format.\n", + "\n", + "[Above, we discussed that the `info_whitespace_length` parameter in the config files could b changed to make \n", + "the `model.info` attribute display optimally on your computer. This attribute also controls the whitespace of the\n", + "`result.info` attribute.]" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We plot the maximum likelihood fit, tracer images and posteriors inferred via Nautilus.\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_tracer(tracer=result.max_log_likelihood_tracer, grid=result.grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The result contains the full posterior information of our non-linear search, including all parameter samples, \n", + "log likelihood values and tools to compute the errors on the lens model. \n", + "\n", + "There are built in visualization tools for plotting this.\n", + "\n", + "The plot is labeled with short hand parameter names (e.g. `sersic_index` is mapped to the short hand \n", + "parameter `n`). These mappings ate specified in the `config/notation.yaml` file and can be customized by users.\n", + "\n", + "The superscripts of labels correspond to the name each component was given in the model (e.g. for the `Isothermal`\n", + "mass its name `mass` defined when making the `Model` above is used)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Results__\n", + "\n", + "Checkout `autolens_workspace/*/guides/results` for a full description of analysing results.\n", + "\n", + "__Modeling Customization__\n", + "\n", + "The folders `autolens_workspace/*/guides/modeling/searches` gives an overview of alternative non-linear searches,\n", + "other than Nautilus, that can be used to fit lens models. \n", + "\n", + "They also provide details on how to customize the model-fit, for example the priors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/simulator.ipynb b/notebooks/point_source/simulator.ipynb index 41c3eeb72..e90fc3c18 100644 --- a/notebooks/point_source/simulator.ipynb +++ b/notebooks/point_source/simulator.ipynb @@ -1,875 +1,912 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Start Here\n", - "=====================\n", - "\n", - "This script is the starting point for simulating point source strong lens datasets, for example a lensed quasar\n", - "or supernova, and it provides an overview of the lens simulation API.\n", - "\n", - "After reading this script, the `examples` folder provide examples for simulating more complex lenses in different ways.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated (in this case, `PointDataset` data).\n", - "- **Ray Tracing:** Setup the lens galaxy's mass (SIE) and source galaxy (a point source) for this simulated lens.\n", - "- **Point Solver:** For a point source, our goal is to find the (y, x) coordinates in the image plane that map directly.\n", - "- **Point Datasets:** All the quantities computed above are stored in a `PointDataset` object, which organizes.\n", - "- **Visualize:** Output a subplot of the simulated point source dataset as a .png file.\n", - "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", - "- **Imaging:** Point-source data typically comes with imaging data of the strong lens, for example showing the 4.\n", - "- **Fluxes:** Another measurable quantity of a point source is its flux\u2014the total amount of light received from.\n", - "- **Point Dataset:** The fluxes are not input a `PointDataset` object, alongside the image-plane coordinates of the.\n", - "- **Time Delays:** Another measurable quantity of a point source is its time delay\u2014the time it takes for light to.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates `PointDataset` data of a strong lens where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The source `Galaxy` is a `Point`.\n", - "\n", - "__Pre-requisites__\n", - "\n", - "It is strongly recommended you read the `autolens_workspace/scripts/point_source/start_here` notebook before\n", - "running this script, as it gives a full overview of the point source modeling API and how lensing calculations\n", - "are performed." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "import numpy as np\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated (in this case, `PointDataset` data) and `dataset_name` \n", - "gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", - "\n", - " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/positions.json`.\n", - " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.json`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"point_source\"\n", - "dataset_name = \"simple\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output. \n", - "\n", - "In this example, this is: `/autolens_workspace/dataset/positions/simple`" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "Setup the lens galaxy's mass (SIE) and source galaxy (a point source) for this simulated lens. \n", - "\n", - "We include a faint extended light profile for the source galaxy for visualization purposes, in order to show where \n", - "the multiple images of the lensed source appear in the image-plane.\n", - "\n", - "For lens modeling, defining ellipticity in terms of the `ell_comps` improves the model-fitting procedure.\n", - "\n", - "However, for simulating a strong lens you may find it more intuitive to define the elliptical geometry using the \n", - "axis-ratio of the profile (axis_ratio = semi-major axis / semi-minor axis = b/a) and position angle, where angle is\n", - "in degrees and defined counter clockwise from the positive x-axis.\n", - "\n", - "We can use the `convert` module to determine the elliptical components from the axis-ratio and angle." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(\n", - " redshift=1.0,\n", - " light=al.lp.ExponentialCore(\n", - " centre=(0.07, 0.07), intensity=0.1, effective_radius=0.02, radius_break=0.025\n", - " ),\n", - " point_0=al.ps.Point(centre=(0.07, 0.07)),\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Use these galaxies to setup a tracer, which will compute the multiple image positions of the simulated dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "For a point source, our goal is to find the (y, x) coordinates in the image plane that map directly to the center of \n", - "the point source in the source plane\u2014these are its \"multiple images.\" This is achieved using a `PointSolver`, which \n", - "determines the multiple images of the mass model for a point source located at a given (y, x) position in the \n", - "source plane.\n", - "\n", - "The solver works by ray tracing triangles from the image plane back to the source plane and checking whether the \n", - "source-plane (y, x) center lies inside each triangle. It iteratively refines this process by ray tracing progressively \n", - "smaller triangles, allowing the multiple image positions to be determined with sub-pixel precision.\n", - "\n", - "The `PointSolver` requires an initial grid of (y, x) coordinates in the image plane, which defines the first set of \n", - "triangles to ray trace. It also needs a `pixel_scale_precision` parameter, specifying the resolution at which the \n", - "multiple images are computed. Smaller values increase precision but require longer computation times. The value \n", - "of 0.001 used here balances efficiency and accuracy.\n", - "\n", - "Strong lens mass models often predict a \"central image,\" a multiple image that is usually heavily demagnified and thus \n", - "not observed. Since the `PointSolver` finds all valid multiple images, it will locate this central image regardless of \n", - "its visibility. To avoid including this unobservable image, we set a `magnification_threshold=0.1`, which discards any \n", - "images with magnifications below this value.\n", - "\n", - "If your dataset does include a detectable central image, you should lower this threshold accordingly to include it in \n", - "your analysis.\n", - "\n", - "We now compute the multiple image positions by creating a `PointSolver` object and passing it the tracer of our \n", - "strong lens system." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(200, 200),\n", - " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now pass the tracer to the solver, to determine the image-plane multiple images for the source centre.\n", - "\n", - "The solver will find the image-plane coordinates that map directly to the source-plane coordinate (0.07\", 0.07\")." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.point_0.centre\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now add Gaussian noise to the multiple image positions to simulate observational measurement errors.\n", - "\n", - "The positional uncertainty in real observations is *not* the pixel scale of the imaging \u2014 that is the\n", - "detector's sampling, not its centroiding precision. Bright point sources (e.g. lensed quasars or supernovae)\n", - "are localised by fitting the instrumental PSF to the image, and the resulting centroid uncertainty is\n", - "typically a small fraction of a pixel. For HST/ACS or WFC3, this corresponds to ~3\u20135 mas (0.003\u20130.005\");\n", - "for adaptive optics on Keck or VLT it is similar; for VLBI radio observations of lensed quasars it can\n", - "be sub-mas. We adopt a default of 0.005\" (5 mas), which is representative of HST point-source astrometry\n", - "in the strong-lensing literature (CASTLES, TDCOSMO/H0LiCOW). Setting the precision close to the imaging\n", - "pixel scale (~0.05\") would inflate lens-model parameter uncertainties well beyond what real data deliver.\n", - "\n", - "Centroid uncertainties from PSF fitting are well-approximated as Gaussian via Laplace's approximation\n", - "around the fitted likelihood maximum, so a Gaussian noise model is appropriate here." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "position_noise = 0.005\n", - "\n", - "positions_with_noise = positions + np.random.normal(\n", - " loc=0.0, scale=position_noise, size=positions.shape\n", - ")\n", - "\n", - "positions_with_noise = al.Grid2DIrregular(\n", - " values=positions_with_noise,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Datasets__\n", - "\n", - "All the quantities computed above are stored in a `PointDataset` object, which organizes information about the multiple \n", - "images of a point-source strong lens system.\n", - "\n", - "This dataset is labeled with the `name` `point_0`, identifying it as corresponding to a single point source called \n", - "`point_0`. The name is essential for associating the dataset with the correct point source in the lens model during \n", - "fitting.\n", - "\n", - "The dataset contains the image-plane coordinates of the multiple images and their corresponding noise-map values.\n", - "The `positions_noise_map` is set to the same `position_noise` defined above (0.005\", i.e. 5 mas), reflecting\n", - "realistic PSF-centroiding precision rather than the imaging pixel scale.\n", - "\n", - "Note also that this dataset does not contain fluxes or time delays, which are often included in point source datasets\n", - "and are included in a separate simulation below." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_with_noise,\n", - " positions_noise_map=position_noise,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now output the point dataset to the dataset path as a .json file, which is loaded in the point source modeling\n", - "examples.\n", - "\n", - "In this example, there is just one point source dataset. However, for group and cluster strong lenses there\n", - "can be many point source datasets in a single dataset, and separate .json files are output for each." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=dataset,\n", - " file_path=dataset_path / \"point_dataset_positions_only.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__CSV Output__\n", - "\n", - "In addition to JSON, a point dataset can be written to a CSV file. CSV is a hand-editable,\n", - "spreadsheet-friendly format that becomes especially convenient for cluster-scale datasets\n", - "with tens of sources where a single file with one row per observed image is easier to curate\n", - "than many per-source JSON files." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset.to_csv(\n", - " file_path=dataset_path / \"point_dataset_positions_only.csv\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "Output a subplot of the simulated point source dataset as a .png file." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_point_dataset(\n", - " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output subplots of the tracer's images, including the positions of the multiple images on the image." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")\n", - "aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Tracer json__\n", - "\n", - "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - "are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - "This can be loaded via the method `tracer = al.from_json()`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=tracer,\n", - " file_path=dataset_path / \"tracer.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Imaging__\n", - "\n", - "Point-source data typically comes with imaging data of the strong lens, for example showing the 4 multiply\n", - "imaged point-sources (e.g. the quasar images).\n", - "\n", - "Whilst this data may not be used for point-source modeling, it is often used to measure the locations of the point\n", - "source multiple images in the first place, and is also useful for visually confirming the images we are using are in \n", - "right place. It may also contain emission from the lens galaxy's light, which can be used to perform point-source \n", - "modeling.\n", - "\n", - "We therefore simulate imaging dataset of this point source and output it to the dataset folder in an `imaging` folder\n", - "as .fits and .png files. \n", - "\n", - "If you are not familiar with the imaging simulator API, checkout the `imaging/simulator.py` example \n", - "in the `autolens_workspace`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", - ")\n", - "\n", - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")\n", - "\n", - "imaging = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - "imaging_path = dataset_path / \"imaging\"\n", - "\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=imaging)\n", - "\n", - "aplt.fits_imaging(\n", - " dataset=imaging,\n", - " data_path=dataset_path / \"data.fits\",\n", - " psf_path=dataset_path / \"psf.fits\",\n", - " noise_map_path=dataset_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - ")\n", - "\n", - "aplt.subplot_imaging_dataset(dataset=imaging)\n", - "aplt.plot_array(array=imaging.data, title=\"Data\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fluxes__\n", - "\n", - "Another measurable quantity of a point source is its flux\u2014the total amount of light received from each multiple image of \n", - "the point source (e.g., the quasar images).\n", - "\n", - "In practice, fluxes are often measured but not used directly when analyzing lensed point sources such as quasars or \n", - "supernovae. This is because fluxes can be significantly affected by microlensing, which many lens models do not \n", - "accurately capture. However, in this simulation, microlensing is not included, so the fluxes can be simulated and fitted reliably.\n", - "\n", - "We now simulate the fluxes of the multiple images of this point source.\n", - "\n", - "Given a mass model and the (y, x) image-plane coordinates of each image, the magnification at each point can be \n", - "calculated.\n", - "\n", - "Below, we compute the magnification for every multiple image coordinate, which will then be used to simulate their \n", - "fluxes." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "magnifications = al.LensCalc.from_tracer(\n", - " tracer=tracer\n", - ").magnification_2d_via_hessian_from(grid=positions)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To simulate the fluxes, we assume the source galaxy point-source has a total flux of 1.0.\n", - "\n", - "Each observed image has a flux that is the source's flux multiplied by the magnification at that image-plane coordinate." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "flux = 1.0\n", - "fluxes = [flux * np.abs(magnification) for magnification in magnifications]\n", - "fluxes = al.ArrayIrregular(values=fluxes)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now add Gaussian noise to the fluxes to simulate observational measurement errors.\n", - "\n", - "For lensed quasars and supernovae, photometric measurements of multiple-image fluxes are not generally\n", - "photon-noise-limited \u2014 the dominant uncertainty is microlensing, in which stars in the lens galaxy\n", - "distort the apparent flux of each image by amounts that depend on the source size, the macro-magnification,\n", - "and where each image lies in the lens-plane stellar field. Lens models that explicitly *exclude*\n", - "microlensing therefore typically assume a few-percent flux uncertainty per image rather than a Poisson\n", - "floor. We adopt a 5% relative flux error here, which is consistent with this practice and produces\n", - "realistic flux-ratio constraints on the mass model." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "flux_rel_noise = 0.05\n", - "\n", - "fluxes_with_noise = fluxes + np.random.normal(\n", - " loc=0.0, scale=flux_rel_noise * np.asarray(fluxes), size=len(fluxes)\n", - ")\n", - "\n", - "fluxes_with_noise = al.ArrayIrregular(values=fluxes_with_noise)\n", - "\n", - "fluxes_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes))" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Dataset__\n", - "\n", - "The fluxes are not input a `PointDataset` object, alongside the image-plane coordinates of the multiple images\n", - "and their associated noise-map values. \n", - "\n", - "We again give the dataset the name `point_0`, which is a label given to the dataset to indicate that it is a dataset \n", - "of a single point-source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_with_noise,\n", - " positions_noise_map=position_noise,\n", - " fluxes=fluxes_with_noise,\n", - " fluxes_noise_map=fluxes_noise_map,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now output the point dataset to the dataset path as a .json file, which is loaded in the point source modeling\n", - "examples.\n", - "\n", - "In this example, there is just one point source dataset. However, for group and cluster strong lenses there\n", - "can be many point source datasets in a single dataset, and separate .json files are output for each." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=dataset,\n", - " file_path=dataset_path / \"point_dataset_with_fluxes.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Time Delays__\n", - "\n", - "Another measurable quantity of a point source is its time delay\u2014the time it takes for light to travel from the\n", - "source to the observer for each multiple image of the point source (e.g., the quasar images). This is often expressed\n", - "as the relative time delay between each image and the image with the shortest time delay, which is often referred to as\n", - "the \"reference image.\"\n", - "\n", - "Time delays are commonly used in strong lensing analyses, for example to measure the Hubble constant, since\n", - "they are less affected by microlensing and can provide robust cosmological constraints.\n", - "\n", - "We now simulate the same point source dataset, but this time including the time delays of the multiple images.\n", - "\n", - "Given a mass model and (y, x) image-plane coordinates, the time delay at each image-plane position can be\n", - "calculated from the mass model. It includes the contribution of both the geometric time delay (the time it takes\n", - "different light rays to travel from the source to the observer) and the Shapiro time delay (the time it takes\n", - "light to travel through the gravitational potential of the lens galaxy)." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "time_delays = tracer.time_delays_from(grid=positions)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In real observations, time delays are measured by photometrically monitoring the multiple images over months\n", - "to years and cross-correlating their light curves to align the variable signals. State-of-the-art monitoring\n", - "campaigns (e.g. COSMOGRAIL, TDCOSMO) routinely achieve ~1\u20133% relative precision on the longest time delays\n", - "in well-sampled quad systems \u2014 absolute uncertainties of ~0.5\u20131.5 days on 30\u2013100 day delays.\n", - "\n", - "For simplicity we adopt a 5% relative uncertainty here. Real-world uncertainties are not strictly\n", - "proportional to the delay magnitude (they depend on the cadence and total length of the photometric\n", - "monitoring, on microlensing variability that distorts the light curves, and on the lens configuration),\n", - "but a constant fractional error is a reasonable simulator default and produces realistic relative weights\n", - "between the multiple images." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "time_delay_rel_noise = 0.05\n", - "\n", - "time_delays_noise_map = al.ArrayIrregular(\n", - " values=np.abs(time_delays) * time_delay_rel_noise\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now add noise to the time delays to simulate observational measurement errors." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "time_delays_with_noise = time_delays + np.random.normal(\n", - " loc=0.0, scale=time_delays_noise_map, size=len(time_delays)\n", - ")\n", - "\n", - "time_delays_with_noise = al.ArrayIrregular(values=time_delays_with_noise)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Dataset__\n", - "\n", - "The time delays are input into a `PointDataset` object, alongside the image-plane coordinates of the multiple images\n", - "and their associated noise-map values. \n", - "\n", - "We again give the dataset the name `point_0`, which is a label given to the dataset to indicate that it is a dataset \n", - "of a single point-source." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_with_noise,\n", - " positions_noise_map=position_noise,\n", - " time_delays=time_delays_with_noise,\n", - " time_delays_noise_map=time_delays_noise_map,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now output the point dataset to the dataset path as a .json file, which can be loaded in point source modeling\n", - "examples.\n", - "\n", - "While this example contains one point source dataset, group and cluster lenses can contain multiple datasets,\n", - "with separate .json files saved for each." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "al.output_to_json(\n", - " obj=dataset,\n", - " file_path=dataset_path / \"point_dataset_with_time_delays.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We output a final point source dataset containing the positions, fluxes and time delays, which could be used\n", - "to perform lens modeling of all measurements simultaneously." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_with_noise,\n", - " positions_noise_map=position_noise,\n", - " fluxes=fluxes_with_noise,\n", - " fluxes_noise_map=fluxes_noise_map,\n", - " time_delays=time_delays_with_noise,\n", - " time_delays_noise_map=time_delays_noise_map,\n", - ")\n", - "\n", - "al.output_to_json(\n", - " obj=dataset,\n", - " file_path=dataset_path / \"point_dataset_with_fluxes_and_time_delays.json\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__CSV Output__\n", - "\n", - "The full-column dataset (positions, fluxes and time delays) can also be saved to CSV\n", - "as the spreadsheet-friendly counterpart to the JSON above." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset.to_csv(\n", - " file_path=dataset_path / \"point_dataset_with_fluxes_and_time_delays.csv\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finished.\n", - "\n", - "__JAX Variant__\n", - "\n", - "The expensive step in point-source simulation is the multiple-image\n", - "solve \u2014 `PointSolver` runs an iterative triangle-refinement loop. On\n", - "JAX-jit this is order-of-magnitude faster than NumPy.\n", - "\n", - "```python\n", - "import jax\n", - "import jax.numpy as jnp\n", - "from autolens.jax import register_tracer_classes\n", - "\n", - "# One-time setup: register Tracer + Galaxy + profile classes as JAX\n", - "# pytrees BEFORE the first @jax.jit call. (Inside @jax.jit, JAX flattens\n", - "# function arguments at trace time \u2014 auto-registration inside solve()\n", - "# runs too late.)\n", - "register_tracer_classes(tracer)\n", - "\n", - "solver_jax = al.PointSolver.for_grid(\n", - " grid=al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=1.0),\n", - " pixel_scale_precision=0.001,\n", - " magnification_threshold=0.1,\n", - " use_jax=True,\n", - ")\n", - "\n", - "@jax.jit\n", - "def solve(tracer, coord):\n", - " return solver_jax.solve(tracer=tracer, source_plane_coordinate=coord).array\n", - "\n", - "positions = solve(tracer, jnp.asarray(source_galaxy.bulge.centre))\n", - "```\n", - "\n", - "When `use_jax=True`, `PointSolver.solve` defaults `remove_infinities=False`\n", - "to honour the JAX static-shape contract \u2014 positions are padded with `inf`\n", - "where no image was found. Strip them outside the jit:\n", - "\n", - "```python\n", - "raw = np.asarray(positions)\n", - "finite_positions = raw[~np.isinf(raw).any(axis=1)]\n", - "```\n", - "\n", - "Two notes:\n", - "\n", - "- `register_tracer_classes(tracer)` is the one user-visible setup call.\n", - " After the first invocation, every later `@jax.jit` with the same class\n", - " set works without re-registering.\n", - "- Unlike imaging / interferometer simulators, `PointSolver.use_jax=True`\n", - " does not go through `Array2D.native` at all \u2014 the triangle-refinement\n", - " loop operates on raw arrays throughout.\n", - "\n", - "See `scripts/guides/lens_calc.py` for the broader \"JIT-it-yourself\"\n", - "pattern applied to other library methods." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Start Here\n", + "=====================\n", + "\n", + "This script is the starting point for simulating point source strong lens datasets, for example a lensed quasar\n", + "or supernova, and it provides an overview of the lens simulation API.\n", + "\n", + "After reading this script, the `examples` folder provide examples for simulating more complex lenses in different ways.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated (in this case, `PointDataset` data).\n", + "- **Ray Tracing:** Setup the lens galaxy's mass (SIE) and source galaxy (a point source) for this simulated lens.\n", + "- **Point Solver:** For a point source, our goal is to find the (y, x) coordinates in the image plane that map directly.\n", + "- **Point Datasets:** All the quantities computed above are stored in a `PointDataset` object, which organizes.\n", + "- **Visualize:** Output a subplot of the simulated point source dataset as a .png file.\n", + "- **Tracer json:** Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass.\n", + "- **Imaging:** Point-source data typically comes with imaging data of the strong lens, for example showing the 4.\n", + "- **Fluxes:** Another measurable quantity of a point source is its flux\u2014the total amount of light received from.\n", + "- **Point Dataset:** The fluxes are not input a `PointDataset` object, alongside the image-plane coordinates of the.\n", + "- **Time Delays:** Another measurable quantity of a point source is its time delay\u2014the time it takes for light to.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates `PointDataset` data of a strong lens where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The source `Galaxy` is a `Point`.\n", + "\n", + "__Pre-requisites__\n", + "\n", + "It is strongly recommended you read the `autolens_workspace/scripts/point_source/start_here` notebook before\n", + "running this script, as it gives a full overview of the point source modeling API and how lensing calculations\n", + "are performed." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "import numpy as np\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated (in this case, `PointDataset` data) and `dataset_name` \n", + "gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", + "\n", + " - The image will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/positions.json`.\n", + " - The noise-map will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/noise_map.json`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"point_source\"\n", + "dataset_name = \"simple\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output. \n", + "\n", + "In this example, this is: `/autolens_workspace/dataset/positions/simple`" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "Setup the lens galaxy's mass (SIE) and source galaxy (a point source) for this simulated lens. \n", + "\n", + "We include a faint extended light profile for the source galaxy for visualization purposes, in order to show where \n", + "the multiple images of the lensed source appear in the image-plane.\n", + "\n", + "For lens modeling, defining ellipticity in terms of the `ell_comps` improves the model-fitting procedure.\n", + "\n", + "However, for simulating a strong lens you may find it more intuitive to define the elliptical geometry using the \n", + "axis-ratio of the profile (axis_ratio = semi-major axis / semi-minor axis = b/a) and position angle, where angle is\n", + "in degrees and defined counter clockwise from the positive x-axis.\n", + "\n", + "We can use the `convert` module to determine the elliptical components from the axis-ratio and angle." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " light=al.lp.ExponentialCore(\n", + " centre=(0.07, 0.07), intensity=0.1, effective_radius=0.02, radius_break=0.025\n", + " ),\n", + " point_0=al.ps.Point(centre=(0.07, 0.07)),\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use these galaxies to setup a tracer, which will compute the multiple image positions of the simulated dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "For a point source, our goal is to find the (y, x) coordinates in the image plane that map directly to the center of \n", + "the point source in the source plane\u2014these are its \"multiple images.\" This is achieved using a `PointSolver`, which \n", + "determines the multiple images of the mass model for a point source located at a given (y, x) position in the \n", + "source plane.\n", + "\n", + "The solver works by ray tracing triangles from the image plane back to the source plane and checking whether the \n", + "source-plane (y, x) center lies inside each triangle. It iteratively refines this process by ray tracing progressively \n", + "smaller triangles, allowing the multiple image positions to be determined with sub-pixel precision.\n", + "\n", + "The `PointSolver` requires an initial grid of (y, x) coordinates in the image plane, which defines the first set of \n", + "triangles to ray trace. It also needs a `pixel_scale_precision` parameter, specifying the resolution at which the \n", + "multiple images are computed. Smaller values increase precision but require longer computation times. The value \n", + "of 0.001 used here balances efficiency and accuracy.\n", + "\n", + "Strong lens mass models often predict a \"central image,\" a multiple image that is usually heavily demagnified and thus \n", + "not observed. Since the `PointSolver` finds all valid multiple images, it will locate this central image regardless of \n", + "its visibility. To avoid including this unobservable image, we set a `magnification_threshold=0.1`, which discards any \n", + "images with magnifications below this value.\n", + "\n", + "If your dataset does include a detectable central image, you should lower this threshold accordingly to include it in \n", + "your analysis.\n", + "\n", + "We now compute the multiple image positions by creating a `PointSolver` object and passing it the tracer of our \n", + "strong lens system." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(200, 200),\n", + " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now pass the tracer to the solver, to determine the image-plane multiple images for the source centre.\n", + "\n", + "The solver will find the image-plane coordinates that map directly to the source-plane coordinate (0.07\", 0.07\")." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.point_0.centre\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now add Gaussian noise to the multiple image positions to simulate observational measurement errors.\n", + "\n", + "The positional uncertainty in real observations is *not* the pixel scale of the imaging \u2014 that is the\n", + "detector's sampling, not its centroiding precision. Bright point sources (e.g. lensed quasars or supernovae)\n", + "are localised by fitting the instrumental PSF to the image, and the resulting centroid uncertainty is\n", + "typically a small fraction of a pixel. For HST/ACS or WFC3, this corresponds to ~3\u20135 mas (0.003\u20130.005\");\n", + "for adaptive optics on Keck or VLT it is similar; for VLBI radio observations of lensed quasars it can\n", + "be sub-mas. We adopt a default of 0.005\" (5 mas), which is representative of HST point-source astrometry\n", + "in the strong-lensing literature (CASTLES, TDCOSMO/H0LiCOW). Setting the precision close to the imaging\n", + "pixel scale (~0.05\") would inflate lens-model parameter uncertainties well beyond what real data deliver.\n", + "\n", + "Centroid uncertainties from PSF fitting are well-approximated as Gaussian via Laplace's approximation\n", + "around the fitted likelihood maximum, so a Gaussian noise model is appropriate here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "position_noise = 0.005\n", + "\n", + "positions_with_noise = positions + np.random.normal(\n", + " loc=0.0, scale=position_noise, size=positions.shape\n", + ")\n", + "\n", + "positions_with_noise = al.Grid2DIrregular(\n", + " values=positions_with_noise,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Datasets__\n", + "\n", + "All the quantities computed above are stored in a `PointDataset` object, which organizes information about the multiple \n", + "images of a point-source strong lens system.\n", + "\n", + "This dataset is labeled with the `name` `point_0`, identifying it as corresponding to a single point source called \n", + "`point_0`. The name is essential for associating the dataset with the correct point source in the lens model during \n", + "fitting.\n", + "\n", + "The dataset contains the image-plane coordinates of the multiple images and their corresponding noise-map values.\n", + "The `positions_noise_map` is set to the same `position_noise` defined above (0.005\", i.e. 5 mas), reflecting\n", + "realistic PSF-centroiding precision rather than the imaging pixel scale.\n", + "\n", + "Note also that this dataset does not contain fluxes or time delays, which are often included in point source datasets\n", + "and are included in a separate simulation below." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_with_noise,\n", + " positions_noise_map=position_noise,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now output the point dataset to the dataset path as a .json file, which is loaded in the point source modeling\n", + "examples.\n", + "\n", + "In this example, there is just one point source dataset. However, for group and cluster strong lenses there\n", + "can be many point source datasets in a single dataset, and separate .json files are output for each." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=dataset,\n", + " file_path=dataset_path / \"point_dataset_positions_only.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__CSV Output__\n", + "\n", + "In addition to JSON, a point dataset can be written to a CSV file. CSV is a hand-editable,\n", + "spreadsheet-friendly format that becomes especially convenient for cluster-scale datasets\n", + "with tens of sources where a single file with one row per observed image is easier to curate\n", + "than many per-source JSON files." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset.to_csv(\n", + " file_path=dataset_path / \"point_dataset_positions_only.csv\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "Output a subplot of the simulated point source dataset as a .png file." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_point_dataset(\n", + " dataset=dataset, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Output subplots of the tracer's images, including the positions of the multiple images on the image." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_path, output_format=\"png\"\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer json__\n", + "\n", + "Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + "are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + "This can be loaded via the method `tracer = al.from_json()`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=tracer,\n", + " file_path=dataset_path / \"tracer.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imaging__\n", + "\n", + "Point-source data typically comes with imaging data of the strong lens, for example showing the 4 multiply\n", + "imaged point-sources (e.g. the quasar images).\n", + "\n", + "Whilst this data may not be used for point-source modeling, it is often used to measure the locations of the point\n", + "source multiple images in the first place, and is also useful for visually confirming the images we are using are in \n", + "right place. It may also contain emission from the lens galaxy's light, which can be used to perform point-source \n", + "modeling.\n", + "\n", + "We therefore simulate imaging dataset of this point source and output it to the dataset folder in an `imaging` folder\n", + "as .fits and .png files. \n", + "\n", + "If you are not familiar with the imaging simulator API, checkout the `imaging/simulator.py` example \n", + "in the `autolens_workspace`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.1, pixel_scales=grid.pixel_scales\n", + ")\n", + "\n", + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")\n", + "\n", + "imaging = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + "imaging_path = dataset_path / \"imaging\"\n", + "\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=imaging)\n", + "\n", + "aplt.fits_imaging(\n", + " dataset=imaging,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")\n", + "\n", + "aplt.subplot_imaging_dataset(dataset=imaging)\n", + "aplt.plot_array(array=imaging.data, title=\"Data\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fluxes__\n", + "\n", + "Another measurable quantity of a point source is its flux\u2014the total amount of light received from each multiple image of \n", + "the point source (e.g., the quasar images).\n", + "\n", + "In practice, fluxes are often measured but not used directly when analyzing lensed point sources such as quasars or \n", + "supernovae. This is because fluxes can be significantly affected by microlensing, which many lens models do not \n", + "accurately capture. However, in this simulation, microlensing is not included, so the fluxes can be simulated and fitted reliably.\n", + "\n", + "We now simulate the fluxes of the multiple images of this point source.\n", + "\n", + "Given a mass model and the (y, x) image-plane coordinates of each image, the magnification at each point can be \n", + "calculated.\n", + "\n", + "Below, we compute the magnification for every multiple image coordinate, which will then be used to simulate their \n", + "fluxes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "magnifications = al.LensCalc.from_tracer(\n", + " tracer=tracer\n", + ").magnification_2d_via_hessian_from(grid=positions)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To simulate the fluxes, we assume the source galaxy point-source has a total flux of 1.0.\n", + "\n", + "Each observed image has a flux that is the source's flux multiplied by the magnification at that image-plane coordinate." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "flux = 1.0\n", + "fluxes = [flux * np.abs(magnification) for magnification in magnifications]\n", + "fluxes = al.ArrayIrregular(values=fluxes)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now add Gaussian noise to the fluxes to simulate observational measurement errors.\n", + "\n", + "For lensed quasars and supernovae, photometric measurements of multiple-image fluxes are not generally\n", + "photon-noise-limited \u2014 the dominant uncertainty is microlensing, in which stars in the lens galaxy\n", + "distort the apparent flux of each image by amounts that depend on the source size, the macro-magnification,\n", + "and where each image lies in the lens-plane stellar field. Lens models that explicitly *exclude*\n", + "microlensing therefore typically assume a few-percent flux uncertainty per image rather than a Poisson\n", + "floor. We adopt a 5% relative flux error here, which is consistent with this practice and produces\n", + "realistic flux-ratio constraints on the mass model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "flux_rel_noise = 0.05\n", + "\n", + "fluxes_with_noise = fluxes + np.random.normal(\n", + " loc=0.0, scale=flux_rel_noise * np.asarray(fluxes), size=len(fluxes)\n", + ")\n", + "\n", + "fluxes_with_noise = al.ArrayIrregular(values=fluxes_with_noise)\n", + "\n", + "fluxes_noise_map = al.ArrayIrregular(values=flux_rel_noise * np.asarray(fluxes))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Dataset__\n", + "\n", + "The fluxes are not input a `PointDataset` object, alongside the image-plane coordinates of the multiple images\n", + "and their associated noise-map values. \n", + "\n", + "We again give the dataset the name `point_0`, which is a label given to the dataset to indicate that it is a dataset \n", + "of a single point-source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_with_noise,\n", + " positions_noise_map=position_noise,\n", + " fluxes=fluxes_with_noise,\n", + " fluxes_noise_map=fluxes_noise_map,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now output the point dataset to the dataset path as a .json file, which is loaded in the point source modeling\n", + "examples.\n", + "\n", + "In this example, there is just one point source dataset. However, for group and cluster strong lenses there\n", + "can be many point source datasets in a single dataset, and separate .json files are output for each." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=dataset,\n", + " file_path=dataset_path / \"point_dataset_with_fluxes.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Time Delays__\n", + "\n", + "Another measurable quantity of a point source is its time delay\u2014the time it takes for light to travel from the\n", + "source to the observer for each multiple image of the point source (e.g., the quasar images). This is often expressed\n", + "as the relative time delay between each image and the image with the shortest time delay, which is often referred to as\n", + "the \"reference image.\"\n", + "\n", + "Time delays are commonly used in strong lensing analyses, for example to measure the Hubble constant, since\n", + "they are less affected by microlensing and can provide robust cosmological constraints.\n", + "\n", + "We now simulate the same point source dataset, but this time including the time delays of the multiple images.\n", + "\n", + "Given a mass model and (y, x) image-plane coordinates, the time delay at each image-plane position can be\n", + "calculated from the mass model. It includes the contribution of both the geometric time delay (the time it takes\n", + "different light rays to travel from the source to the observer) and the Shapiro time delay (the time it takes\n", + "light to travel through the gravitational potential of the lens galaxy)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "time_delays = tracer.time_delays_from(grid=positions)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In real observations, time delays are measured by photometrically monitoring the multiple images over months\n", + "to years and cross-correlating their light curves to align the variable signals. State-of-the-art monitoring\n", + "campaigns (e.g. COSMOGRAIL, TDCOSMO) routinely achieve ~1\u20133% relative precision on the longest time delays\n", + "in well-sampled quad systems \u2014 absolute uncertainties of ~0.5\u20131.5 days on 30\u2013100 day delays.\n", + "\n", + "For simplicity we adopt a 5% relative uncertainty here. Real-world uncertainties are not strictly\n", + "proportional to the delay magnitude (they depend on the cadence and total length of the photometric\n", + "monitoring, on microlensing variability that distorts the light curves, and on the lens configuration),\n", + "but a constant fractional error is a reasonable simulator default and produces realistic relative weights\n", + "between the multiple images." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "time_delay_rel_noise = 0.05\n", + "\n", + "time_delays_noise_map = al.ArrayIrregular(\n", + " values=np.abs(time_delays) * time_delay_rel_noise\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now add noise to the time delays to simulate observational measurement errors." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "time_delays_with_noise = time_delays + np.random.normal(\n", + " loc=0.0, scale=time_delays_noise_map, size=len(time_delays)\n", + ")\n", + "\n", + "time_delays_with_noise = al.ArrayIrregular(values=time_delays_with_noise)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Dataset__\n", + "\n", + "The time delays are input into a `PointDataset` object, alongside the image-plane coordinates of the multiple images\n", + "and their associated noise-map values. \n", + "\n", + "We again give the dataset the name `point_0`, which is a label given to the dataset to indicate that it is a dataset \n", + "of a single point-source." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_with_noise,\n", + " positions_noise_map=position_noise,\n", + " time_delays=time_delays_with_noise,\n", + " time_delays_noise_map=time_delays_noise_map,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We now output the point dataset to the dataset path as a .json file, which can be loaded in point source modeling\n", + "examples.\n", + "\n", + "While this example contains one point source dataset, group and cluster lenses can contain multiple datasets,\n", + "with separate .json files saved for each." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "al.output_to_json(\n", + " obj=dataset,\n", + " file_path=dataset_path / \"point_dataset_with_time_delays.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We output a final point source dataset containing the positions, fluxes and time delays, which could be used\n", + "to perform lens modeling of all measurements simultaneously." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_with_noise,\n", + " positions_noise_map=position_noise,\n", + " fluxes=fluxes_with_noise,\n", + " fluxes_noise_map=fluxes_noise_map,\n", + " time_delays=time_delays_with_noise,\n", + " time_delays_noise_map=time_delays_noise_map,\n", + ")\n", + "\n", + "al.output_to_json(\n", + " obj=dataset,\n", + " file_path=dataset_path / \"point_dataset_with_fluxes_and_time_delays.json\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__CSV Output__\n", + "\n", + "The full-column dataset (positions, fluxes and time delays) can also be saved to CSV\n", + "as the spreadsheet-friendly counterpart to the JSON above." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset.to_csv(\n", + " file_path=dataset_path / \"point_dataset_with_fluxes_and_time_delays.csv\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finished.\n", + "\n", + "__JAX Variant__\n", + "\n", + "The expensive step in point-source simulation is the multiple-image\n", + "solve \u2014 `PointSolver` runs an iterative triangle-refinement loop. On\n", + "JAX-jit this is order-of-magnitude faster than NumPy.\n", + "\n", + "```python\n", + "import jax\n", + "import jax.numpy as jnp\n", + "from autolens.jax import register_tracer_classes\n", + "\n", + "# One-time setup: register Tracer + Galaxy + profile classes as JAX\n", + "# pytrees BEFORE the first @jax.jit call. (Inside @jax.jit, JAX flattens\n", + "# function arguments at trace time \u2014 auto-registration inside solve()\n", + "# runs too late.)\n", + "register_tracer_classes(tracer)\n", + "\n", + "solver_jax = al.PointSolver.for_grid(\n", + " grid=al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=1.0),\n", + " pixel_scale_precision=0.001,\n", + " magnification_threshold=0.1,\n", + " use_jax=True,\n", + ")\n", + "\n", + "@jax.jit\n", + "def solve(tracer, coord):\n", + " return solver_jax.solve(tracer=tracer, source_plane_coordinate=coord).array\n", + "\n", + "positions = solve(tracer, jnp.asarray(source_galaxy.bulge.centre))\n", + "```\n", + "\n", + "When `use_jax=True`, `PointSolver.solve` defaults `remove_infinities=False`\n", + "to honour the JAX static-shape contract \u2014 positions are padded with `inf`\n", + "where no image was found. Strip them outside the jit:\n", + "\n", + "```python\n", + "raw = np.asarray(positions)\n", + "finite_positions = raw[~np.isinf(raw).any(axis=1)]\n", + "```\n", + "\n", + "Two notes:\n", + "\n", + "- `register_tracer_classes(tracer)` is the one user-visible setup call.\n", + " After the first invocation, every later `@jax.jit` with the same class\n", + " set works without re-registering.\n", + "- Unlike imaging / interferometer simulators, `PointSolver.use_jax=True`\n", + " does not go through `Array2D.native` at all \u2014 the triangle-refinement\n", + " loop operates on raw arrays throughout.\n", + "\n", + "See `scripts/guides/lens_calc.py` for the broader \"JIT-it-yourself\"\n", + "pattern applied to other library methods." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/point_source/simulator_sample.ipynb b/notebooks/point_source/simulator_sample.ipynb index 9a7ca7196..513bb2eb7 100644 --- a/notebooks/point_source/simulator_sample.ipynb +++ b/notebooks/point_source/simulator_sample.ipynb @@ -1,407 +1,444 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Hubble Constant Time Delays\n", - "======================================\n", - "\n", - "A multiply imaged lensed quasar with time delays can measure the Hubble constant, which is a fundamental\n", - "Cosmological parameter that describes the rate of expansion of the universe. This is because the\n", - "the difference between the geometric time delay and the physical time delay is proportional to the Hubble constant.\n", - "\n", - "This script illustrates how to simulate a sample of `PointDataset` datasets of lensed quasars, which\n", - "can easily be used to simulate hundreds or thousands of strong lenses. These, as a sample, can be used to constrain\n", - "the Hubble constant.\n", - "\n", - "To simulate the sample of lenses, each lens and source galaxies are set up using the `Model` object which is also used\n", - "in the `modeling` scripts. This means that the parameters of each simulated strong lens are drawn from the\n", - "distributions defined via priors, which can be customized to simulate a wider range of strong lenses.\n", - "\n", - "The sample is used in `autolens_workspace/notebooks/advanced/graphical` to illustrate how a graphical and hierarchical\n", - "model can be fitted to a large sample of double Einstein ring strong lenses in order to improve the constraints on\n", - "Cosmological parameters.\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model fitted to the data.\n", - "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", - "- **Point Solver:** We use a `PointSolver` to locate the multiple images.\n", - "- **Sample Model Distributions:** To simulate a sample, we draw random instances of lens and source galaxies where the parameters of.\n", - "- **Simulate:** Simulate the image data using a (y,x) grid with the adaptive over sampling scheme.\n", - "- **Sample Instances:** Within a for loop, we will now generate instances of the lens and source galaxies using the.\n", - "\n", - "__Model__\n", - "\n", - "This script simulates a sample of `PointDataset` data of 'galaxy-scale' strong lenses where:\n", - "\n", - " - The lens galaxy's total mass distribution is an `Isothermal`.\n", - " - The source `Galaxy` is a `Point`.\n", - " - The Cosmology is `Planck15` and has a Hubble constant which can be constrained by the time delays.\n", - "\n", - "__Start Here Notebook__\n", - "\n", - "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "import numpy as np\n", - "from pathlib import Path\n", - "import autofit as af\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_label = \"samples\"\n", - "dataset_type = \"point_source\"\n", - "dataset_sample_name = \"hubble_constant_time_delays\"" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The path where the dataset will be output." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_sample_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Point Solver__\n", - "\n", - "We use a `PointSolver` to locate the multiple images. " - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(200, 200),\n", - " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", - ")\n", - "\n", - "solver = al.PointSolver.for_grid(\n", - " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sample Model Distributions__\n", - "\n", - "To simulate a sample, we draw random instances of lens and source galaxies where the parameters of their mass profiles \n", - "and point source profiles are drawn from distributions. These distributions are defined via priors -- the same objects \n", - "that are used when defining the priors of each parameter for a non-linear search.\n", - "\n", - "Below, we define the distributions the lens galaxy's mass profiles are drawn from alongside the source's point\n", - "source centre." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "mass = af.Model(al.mp.Isothermal)\n", - "\n", - "mass.centre = (0.0, 0.0)\n", - "mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", - "mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", - "mass.einstein_radius = af.UniformPrior(lower_limit=1.0, upper_limit=1.8)\n", - "\n", - "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", - "\n", - "point = af.Model(al.ps.Point)\n", - "\n", - "point.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", - "point.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.1)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We also include a light profile component of the source, which is to aid visualization of the simulated dataset\n", - "by providing an image where the point source is located." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "bulge = af.Model(al.lp.ExponentialSph)\n", - "\n", - "bulge.centre_0 = point.centre_0\n", - "bulge.centre_1 = point.centre_1\n", - "bulge.intensity = 1.0\n", - "bulge.effective_radius = 0.02\n", - "bulge.signal_to_noise_ratio = 10.0\n", - "\n", - "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge, point_0=point)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulate__\n", - "\n", - "Simulate the image data using a (y,x) grid with the adaptive over sampling scheme." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "grid = al.Grid2D.uniform(\n", - " shape_native=(150, 150),\n", - " pixel_scales=0.1,\n", - ")\n", - "\n", - "psf = al.Convolver.from_gaussian(\n", - " shape_native=(11, 11), sigma=0.2, pixel_scales=grid.pixel_scales\n", - ")\n", - "\n", - "simulator = al.SimulatorImaging(\n", - " exposure_time=300.0,\n", - " psf=psf,\n", - " background_sky_level=0.1,\n", - " add_poisson_noise_to_data=True,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Sample Instances__\n", - "\n", - "Within a for loop, we will now generate instances of the lens and source galaxies using the `Model`'s defined above.\n", - "This loop will run for `total_datasets` iterations, which sets the number of lenses that are simulated.\n", - "\n", - "Each iteration of the for loop will then create a tracer and use this to simulate the point dataset." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "total_datasets = 3\n", - "\n", - "for sample_index in range(total_datasets):\n", - " dataset_sample_path = dataset_path / f\"dataset_{sample_index}\"\n", - "\n", - " lens_galaxy = lens.random_instance()\n", - " source_galaxy = source.random_instance()\n", - "\n", - " \"\"\"\n", - " __Ray Tracing__\n", - "\n", - " Use the sample's lens and source galaxies to setup a tracer, which will generate the multiple image positions \n", - " and time delays for the simulated `PointDataset` dataset.\n", - " \"\"\"\n", - " tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", - "\n", - " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", - "\n", - " \"\"\"\n", - " __Positions__\n", - "\n", - " We now pass the `Tracer` to the solver to find the multiple image positions and add Gaussian noise to\n", - " simulate observational measurement errors.\n", - "\n", - " The position uncertainty is set to 0.005\" (5 mas), reflecting the centroid precision achievable by PSF\n", - " fitting on HST or adaptive-optics imaging \u2014 *not* the imaging pixel scale, which is the detector's\n", - " sampling rather than its centroiding precision. This value is representative of HST point-source\n", - " astrometry in the strong-lensing literature (CASTLES, TDCOSMO/H0LiCOW). See `simulator.py` for a\n", - " fuller discussion.\n", - " \"\"\"\n", - " position_noise = 0.005\n", - "\n", - " positions = solver.solve(\n", - " tracer=tracer, source_plane_coordinate=source_galaxy.point_0.centre\n", - " )\n", - "\n", - " positions_with_noise = positions + np.random.normal(\n", - " loc=0.0, scale=position_noise, size=positions.shape\n", - " )\n", - "\n", - " positions_with_noise = al.Grid2DIrregular(\n", - " values=positions_with_noise,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Time Delays__\n", - "\n", - " We next compute the time delays of the multiple images, with Gaussian noise added to simulate\n", - " observational measurement errors, which are used to constrain the Hubble constant.\n", - "\n", - " State-of-the-art photometric monitoring campaigns (e.g. COSMOGRAIL, TDCOSMO) routinely achieve ~1\u20133%\n", - " relative precision on the longest time delays in well-sampled quad systems. We adopt a 5% relative\n", - " uncertainty here as a conservative simulator default. See `simulator.py` for a fuller discussion.\n", - " \"\"\"\n", - " time_delay_rel_noise = 0.05\n", - "\n", - " time_delays = tracer.time_delays_from(grid=positions)\n", - " time_delays_noise_map = al.ArrayIrregular(\n", - " values=np.abs(time_delays) * time_delay_rel_noise\n", - " )\n", - "\n", - " time_delays_with_noise = time_delays + np.random.normal(\n", - " loc=0.0, scale=time_delays_noise_map, size=len(time_delays)\n", - " )\n", - "\n", - " time_delays_with_noise = al.ArrayIrregular(values=time_delays_with_noise)\n", - "\n", - " \"\"\"\n", - " __Point Dataset__\n", - "\n", - " We now output the `PointDataset` dataset, which contains the multiple image positions, their noise levels,\n", - " the time delays and their noise levels.\n", - "\n", - " We output this to a .json file which can be loaded in point source modeling examples.\n", - " \"\"\"\n", - " dataset = al.PointDataset(\n", - " name=\"point_0\",\n", - " positions=positions_with_noise,\n", - " positions_noise_map=position_noise,\n", - " time_delays=time_delays_with_noise,\n", - " time_delays_noise_map=time_delays_noise_map,\n", - " )\n", - "\n", - " al.output_to_json(\n", - " obj=dataset,\n", - " file_path=dataset_sample_path / \"point_dataset_with_time_delays.json\",\n", - " )\n", - "\n", - " \"\"\"\n", - " __Imaging__\n", - "\n", - " We also generate an `Imaging` dataset of the lens, which is used to visualize the lens and source galaxies\n", - " and the multiple image positions.\n", - " \"\"\"\n", - " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", - "\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - "\n", - " \"\"\"\n", - " __Output__\n", - "\n", - " Output the simulated dataset to the dataset path as .fits files.\n", - "\n", - " This uses the updated `dataset_path_sample` which outputs this sample lens to a unique folder.\n", - " \"\"\"\n", - " aplt.fits_imaging(\n", - " dataset=dataset,\n", - " data_path=dataset_sample_path / \"data.fits\",\n", - " psf_path=dataset_sample_path / \"psf.fits\",\n", - " noise_map_path=dataset_sample_path / \"noise_map.fits\",\n", - " overwrite=True,\n", - " )\n", - "\n", - " \"\"\"\n", - " __Visualize__\n", - "\n", - " Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files.\n", - " \"\"\"\n", - " aplt.subplot_imaging_dataset(dataset=dataset)\n", - " aplt.plot_array(array=dataset.data, title=\"Data\")\n", - "\n", - " aplt.subplot_tracer(\n", - " tracer=tracer, grid=grid, output_path=dataset_sample_path, output_format=\"png\"\n", - " )\n", - " aplt.subplot_galaxies_images(\n", - " tracer=tracer, grid=grid, output_path=dataset_sample_path, output_format=\"png\"\n", - " )\n", - "\n", - " \"\"\"\n", - " __Tracer json__\n", - "\n", - " Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", - " are safely stored and available to check how the dataset was simulated in the future. \n", - "\n", - " This can be loaded via the method `tracer = al.from_json()`.\n", - " \"\"\"\n", - " al.output_to_json(\n", - " obj=tracer,\n", - " file_path=dataset_sample_path / \"tracer.json\",\n", - " )\n", - "\n", - " \"\"\"\n", - " The dataset can be viewed in the \n", - " folder `autolens_workspace/dataset/point/samples/hubble_constant_time_delays/{sample_index]`.\n", - " \"\"\"\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Hubble Constant Time Delays\n", + "======================================\n", + "\n", + "A multiply imaged lensed quasar with time delays can measure the Hubble constant, which is a fundamental\n", + "Cosmological parameter that describes the rate of expansion of the universe. This is because the\n", + "the difference between the geometric time delay and the physical time delay is proportional to the Hubble constant.\n", + "\n", + "This script illustrates how to simulate a sample of `PointDataset` datasets of lensed quasars, which\n", + "can easily be used to simulate hundreds or thousands of strong lenses. These, as a sample, can be used to constrain\n", + "the Hubble constant.\n", + "\n", + "To simulate the sample of lenses, each lens and source galaxies are set up using the `Model` object which is also used\n", + "in the `modeling` scripts. This means that the parameters of each simulated strong lens are drawn from the\n", + "distributions defined via priors, which can be customized to simulate a wider range of strong lenses.\n", + "\n", + "The sample is used in `autolens_workspace/notebooks/advanced/graphical` to illustrate how a graphical and hierarchical\n", + "model can be fitted to a large sample of double Einstein ring strong lenses in order to improve the constraints on\n", + "Cosmological parameters.\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model fitted to the data.\n", + "- **Dataset Paths:** The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a.\n", + "- **Point Solver:** We use a `PointSolver` to locate the multiple images.\n", + "- **Sample Model Distributions:** To simulate a sample, we draw random instances of lens and source galaxies where the parameters of.\n", + "- **Simulate:** Simulate the image data using a (y,x) grid with the adaptive over sampling scheme.\n", + "- **Sample Instances:** Within a for loop, we will now generate instances of the lens and source galaxies using the.\n", + "\n", + "__Model__\n", + "\n", + "This script simulates a sample of `PointDataset` data of 'galaxy-scale' strong lenses where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal`.\n", + " - The source `Galaxy` is a `Point`.\n", + " - The Cosmology is `Planck15` and has a Hubble constant which can be constrained by the time delays.\n", + "\n", + "__Start Here Notebook__\n", + "\n", + "If any code in this script is unclear, refer to the `simulators/start_here.ipynb` notebook." + ] }, - "nbformat": 4, - "nbformat_minor": 4 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated and `dataset_name` gives it a descriptive name. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_label = \"samples\"\n", + "dataset_type = \"point_source\"\n", + "dataset_sample_name = \"hubble_constant_time_delays\"" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The path where the dataset will be output." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / dataset_type / dataset_label / dataset_sample_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Point Solver__\n", + "\n", + "We use a `PointSolver` to locate the multiple images. " + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(200, 200),\n", + " pixel_scales=0.05, # <- The pixel-scale describes the conversion from pixel units to arc-seconds.\n", + ")\n", + "\n", + "solver = al.PointSolver.for_grid(\n", + " grid=grid, pixel_scale_precision=0.001, magnification_threshold=0.1\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sample Model Distributions__\n", + "\n", + "To simulate a sample, we draw random instances of lens and source galaxies where the parameters of their mass profiles \n", + "and point source profiles are drawn from distributions. These distributions are defined via priors -- the same objects \n", + "that are used when defining the priors of each parameter for a non-linear search.\n", + "\n", + "Below, we define the distributions the lens galaxy's mass profiles are drawn from alongside the source's point\n", + "source centre." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "mass.centre = (0.0, 0.0)\n", + "mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", + "mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", + "mass.einstein_radius = af.UniformPrior(lower_limit=1.0, upper_limit=1.8)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + "point = af.Model(al.ps.Point)\n", + "\n", + "point.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.1)\n", + "point.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.1)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We also include a light profile component of the source, which is to aid visualization of the simulated dataset\n", + "by providing an image where the point source is located." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "bulge = af.Model(al.lp.ExponentialSph)\n", + "\n", + "bulge.centre_0 = point.centre_0\n", + "bulge.centre_1 = point.centre_1\n", + "bulge.intensity = 1.0\n", + "bulge.effective_radius = 0.02\n", + "bulge.signal_to_noise_ratio = 10.0\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge, point_0=point)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulate__\n", + "\n", + "Simulate the image data using a (y,x) grid with the adaptive over sampling scheme." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(150, 150),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "psf = al.Convolver.from_gaussian(\n", + " shape_native=(11, 11), sigma=0.2, pixel_scales=grid.pixel_scales\n", + ")\n", + "\n", + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Sample Instances__\n", + "\n", + "Within a for loop, we will now generate instances of the lens and source galaxies using the `Model`'s defined above.\n", + "This loop will run for `total_datasets` iterations, which sets the number of lenses that are simulated.\n", + "\n", + "Each iteration of the for loop will then create a tracer and use this to simulate the point dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_datasets = 3\n", + "\n", + "for sample_index in range(total_datasets):\n", + " dataset_sample_path = dataset_path / f\"dataset_{sample_index}\"\n", + "\n", + " lens_galaxy = lens.random_instance()\n", + " source_galaxy = source.random_instance()\n", + "\n", + " \"\"\"\n", + " __Ray Tracing__\n", + "\n", + " Use the sample's lens and source galaxies to setup a tracer, which will generate the multiple image positions \n", + " and time delays for the simulated `PointDataset` dataset.\n", + " \"\"\"\n", + " tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])\n", + "\n", + " aplt.plot_array(array=tracer.image_2d_from(grid=grid), title=\"Image\")\n", + "\n", + " \"\"\"\n", + " __Positions__\n", + "\n", + " We now pass the `Tracer` to the solver to find the multiple image positions and add Gaussian noise to\n", + " simulate observational measurement errors.\n", + "\n", + " The position uncertainty is set to 0.005\" (5 mas), reflecting the centroid precision achievable by PSF\n", + " fitting on HST or adaptive-optics imaging \u2014 *not* the imaging pixel scale, which is the detector's\n", + " sampling rather than its centroiding precision. This value is representative of HST point-source\n", + " astrometry in the strong-lensing literature (CASTLES, TDCOSMO/H0LiCOW). See `simulator.py` for a\n", + " fuller discussion.\n", + " \"\"\"\n", + " position_noise = 0.005\n", + "\n", + " positions = solver.solve(\n", + " tracer=tracer, source_plane_coordinate=source_galaxy.point_0.centre\n", + " )\n", + "\n", + " positions_with_noise = positions + np.random.normal(\n", + " loc=0.0, scale=position_noise, size=positions.shape\n", + " )\n", + "\n", + " positions_with_noise = al.Grid2DIrregular(\n", + " values=positions_with_noise,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Time Delays__\n", + "\n", + " We next compute the time delays of the multiple images, with Gaussian noise added to simulate\n", + " observational measurement errors, which are used to constrain the Hubble constant.\n", + "\n", + " State-of-the-art photometric monitoring campaigns (e.g. COSMOGRAIL, TDCOSMO) routinely achieve ~1\u20133%\n", + " relative precision on the longest time delays in well-sampled quad systems. We adopt a 5% relative\n", + " uncertainty here as a conservative simulator default. See `simulator.py` for a fuller discussion.\n", + " \"\"\"\n", + " time_delay_rel_noise = 0.05\n", + "\n", + " time_delays = tracer.time_delays_from(grid=positions)\n", + " time_delays_noise_map = al.ArrayIrregular(\n", + " values=np.abs(time_delays) * time_delay_rel_noise\n", + " )\n", + "\n", + " time_delays_with_noise = time_delays + np.random.normal(\n", + " loc=0.0, scale=time_delays_noise_map, size=len(time_delays)\n", + " )\n", + "\n", + " time_delays_with_noise = al.ArrayIrregular(values=time_delays_with_noise)\n", + "\n", + " \"\"\"\n", + " __Point Dataset__\n", + "\n", + " We now output the `PointDataset` dataset, which contains the multiple image positions, their noise levels,\n", + " the time delays and their noise levels.\n", + "\n", + " We output this to a .json file which can be loaded in point source modeling examples.\n", + " \"\"\"\n", + " dataset = al.PointDataset(\n", + " name=\"point_0\",\n", + " positions=positions_with_noise,\n", + " positions_noise_map=position_noise,\n", + " time_delays=time_delays_with_noise,\n", + " time_delays_noise_map=time_delays_noise_map,\n", + " )\n", + "\n", + " al.output_to_json(\n", + " obj=dataset,\n", + " file_path=dataset_sample_path / \"point_dataset_with_time_delays.json\",\n", + " )\n", + "\n", + " \"\"\"\n", + " __Imaging__\n", + "\n", + " We also generate an `Imaging` dataset of the lens, which is used to visualize the lens and source galaxies\n", + " and the multiple image positions.\n", + " \"\"\"\n", + " dataset = simulator.via_tracer_from(tracer=tracer, grid=grid)\n", + "\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + "\n", + " \"\"\"\n", + " __Output__\n", + "\n", + " Output the simulated dataset to the dataset path as .fits files.\n", + "\n", + " This uses the updated `dataset_path_sample` which outputs this sample lens to a unique folder.\n", + " \"\"\"\n", + " aplt.fits_imaging(\n", + " dataset=dataset,\n", + " data_path=dataset_sample_path / \"data.fits\",\n", + " psf_path=dataset_sample_path / \"psf.fits\",\n", + " noise_map_path=dataset_sample_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + " )\n", + "\n", + " \"\"\"\n", + " __Visualize__\n", + "\n", + " Output a subplot of the simulated dataset, the image and the tracer's quantities to the dataset path as .png files.\n", + " \"\"\"\n", + " aplt.subplot_imaging_dataset(dataset=dataset)\n", + " aplt.plot_array(array=dataset.data, title=\"Data\")\n", + "\n", + " aplt.subplot_tracer(\n", + " tracer=tracer, grid=grid, output_path=dataset_sample_path, output_format=\"png\"\n", + " )\n", + " aplt.subplot_galaxies_images(\n", + " tracer=tracer, grid=grid, output_path=dataset_sample_path, output_format=\"png\"\n", + " )\n", + "\n", + " \"\"\"\n", + " __Tracer json__\n", + "\n", + " Save the `Tracer` in the dataset folder as a .json file, ensuring the true light profiles, mass profiles and galaxies\n", + " are safely stored and available to check how the dataset was simulated in the future. \n", + "\n", + " This can be loaded via the method `tracer = al.from_json()`.\n", + " \"\"\"\n", + " al.output_to_json(\n", + " obj=tracer,\n", + " file_path=dataset_sample_path / \"tracer.json\",\n", + " )\n", + "\n", + " \"\"\"\n", + " The dataset can be viewed in the \n", + " folder `autolens_workspace/dataset/point/samples/hubble_constant_time_delays/{sample_index]`.\n", + " \"\"\"\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/weak/features/strong_lensing/fit.ipynb b/notebooks/weak/features/strong_lensing/fit.ipynb new file mode 100644 index 000000000..2706b2f31 --- /dev/null +++ b/notebooks/weak/features/strong_lensing/fit.ipynb @@ -0,0 +1,307 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fit: Combined Strong + Weak Lensing\n", + "===================================\n", + "\n", + "This script fits the combined strong+weak dataset simulated by `simulator.py` in this folder with a single\n", + "shared `Tracer`: the imaging data via `FitImaging` and the shear catalogue via `FitWeak`.\n", + "\n", + "Because the two datasets are statistically independent measurements of the same mass distribution, their\n", + "joint log likelihood is simply the sum of the two individual log likelihoods \u2014 this additivity is all the\n", + "joint modeling in `modeling.py` needs, and this script makes it explicit before a non-linear search is\n", + "involved.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load both datasets (auto-simulating them if missing) and mask the imaging data.\n", + "- **Tracer:** The single mass model shared by both fits.\n", + "- **Imaging Fit:** Fit the strong-lensing image.\n", + "- **Weak Fit:** Fit the shear catalogue.\n", + "- **Joint Likelihood:** The sum that a joint analysis samples.\n", + "- **Shear Profile:** Data vs model in the space cluster weak lensing is usually shown in." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load both halves of the combined dataset from `dataset/weak/strong_lensing/`.\n", + "\n", + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / \"weak\" / \"strong_lensing\"\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/weak/features/strong_lensing/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_imaging = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "dataset_weak = al.from_json(file_path=dataset_path / \"dataset.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The imaging data is masked with the standard 3.0\" circular mask \u2014 inside it live the arcs; the shear\n", + "catalogue carries the information outside it, to 10\"." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mask = al.Mask2D.circular(\n", + " shape_native=dataset_imaging.shape_native,\n", + " pixel_scales=dataset_imaging.pixel_scales,\n", + " radius=3.0,\n", + ")\n", + "\n", + "dataset_imaging = dataset_imaging.apply_mask(mask=mask)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Tracer__\n", + "\n", + "One tracer fits both datasets. We use the simulator's true parameters, so both fits below are \"perfect\"\n", + "up to noise \u2014 swap any parameter to see both likelihoods respond, which is exactly what a non-linear\n", + "search exploits in `modeling.py`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.05, 0.05),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imaging Fit__\n", + "\n", + "The strong-lensing side is the standard `FitImaging` of every imaging example: model image, PSF convolution,\n", + "residuals and a Gaussian likelihood over the masked pixels." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_imaging = al.FitImaging(dataset=dataset_imaging, tracer=tracer)\n", + "\n", + "print(f\"imaging log_likelihood : {fit_imaging.log_likelihood:.2f}\")\n", + "\n", + "aplt.subplot_fit_imaging(fit=fit_imaging, output_path=dataset_path, output_format=\"png\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Weak Fit__\n", + "\n", + "The weak-lensing side is the `FitWeak` of `scripts/weak/fit.py`: the tracer's shear field evaluated at the\n", + "catalogue positions, compared to the measured shears over 2N independent components." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit_weak = al.FitWeak(dataset=dataset_weak, tracer=tracer)\n", + "\n", + "print(f\"weak log_likelihood : {fit_weak.log_likelihood:.2f}\")\n", + "print(\n", + " f\"weak chi_squared : {fit_weak.chi_squared:.1f} \"\n", + " f\"(expected ~{2 * dataset_weak.n_galaxies} for the true model)\"\n", + ")\n", + "\n", + "aplt.subplot_fit_weak(fit=fit_weak, output_path=dataset_path, output_format=\"png\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Joint Likelihood__\n", + "\n", + "The datasets are independent (different galaxies, different noise), so the joint log likelihood of the\n", + "shared tracer is their sum. This one line is the entire statistical content of \"combining strong and weak\n", + "lensing\" \u2014 `modeling.py` wires it into PyAutoFit's factor-graph API so a non-linear search samples it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_likelihood_joint = fit_imaging.log_likelihood + fit_weak.log_likelihood\n", + "\n", + "print(f\"joint log_likelihood : {log_likelihood_joint:.2f}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Shear Profile__\n", + "\n", + "The tangential shear profile shows where the weak information lives: the model curve (from the same tracer\n", + "fitting the arcs at 1.6\") is tested by the binned data points out to 10\" \u2014 radii the imaging mask never\n", + "sees. The cross component scattering around zero is the standard B-mode systematics null test." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_shear_profile(\n", + " fit_weak,\n", + " centre=(0.0, 0.0),\n", + " bins=8,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/weak/features/strong_lensing/modeling.ipynb b/notebooks/weak/features/strong_lensing/modeling.ipynb new file mode 100644 index 000000000..f80c75b9a --- /dev/null +++ b/notebooks/weak/features/strong_lensing/modeling.ipynb @@ -0,0 +1,392 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Combined Strong + Weak Lensing\n", + "========================================\n", + "\n", + "This script fits the combined dataset of `simulator.py` \u2014 an `Imaging` dataset of strongly lensed arcs and a\n", + "`WeakDataset` of the surrounding shear field \u2014 with a **single lens mass model**, using PyAutoFit's\n", + "factor-graph API to sample the joint likelihood with one non-linear search.\n", + "\n", + "__Why combine them?__\n", + "\n", + "The two datasets constrain the same mass distribution in complementary regimes:\n", + "\n", + " - **Strong lensing** (the arcs) pins the Einstein radius, and the mass centre exquisitely \u2014 but only\n", + " *inside* ~1 Einstein radius, and mass-model families that agree there can diverge immediately outside it.\n", + "\n", + " - **Weak lensing** (the shear catalogue) measures the mass profile and its ellipticity out to many Einstein\n", + " radii \u2014 noisily per galaxy, but with statistical power in the ensemble, and precisely where the strong\n", + " lensing has none.\n", + "\n", + "Fitting them jointly forces one parametric mass model to satisfy both, the approach of hybrid-Lenstool's\n", + "joint strong+weak cluster reconstructions (Niemiec et al. 2020, who showed sequential fitting biases the\n", + "profile at 2-3 sigma where a joint fit stays within ~1 sigma) and of the stacked strong+weak analysis of the\n", + "Sloan Giant Arcs Survey group-to-cluster lenses (Oguri et al. 2012).\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load both datasets (auto-simulating if missing) and mask the imaging data.\n", + "- **Model:** One lens model whose priors are shared by both datasets' analyses.\n", + "- **Analysis List:** An `AnalysisImaging` and an `AnalysisWeak`, one per dataset.\n", + "- **Analysis Factor & Factor Graph:** Combine them so one search samples the joint likelihood.\n", + "- **Search & Model-Fit:** Nautilus over the shared parameter space.\n", + "- **Result:** The joint constraints, and how to read the strong/weak complementarity in them." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "Load both halves of the combined dataset, auto-simulating them if missing (the standard pattern of all\n", + "example scripts), and apply the standard 3.0\" circular mask to the imaging data. Everything outside that\n", + "mask is the weak catalogue's territory." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"strong_lensing\"\n", + "dataset_path = Path(\"dataset\") / \"weak\" / dataset_name\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/weak/features/strong_lensing/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset_imaging = al.Imaging.from_fits(\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "mask = al.Mask2D.circular(\n", + " shape_native=dataset_imaging.shape_native,\n", + " pixel_scales=dataset_imaging.pixel_scales,\n", + " radius=3.0,\n", + ")\n", + "\n", + "dataset_imaging = dataset_imaging.apply_mask(mask=mask)\n", + "\n", + "dataset_weak = al.from_json(file_path=dataset_path / \"dataset.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "One model serves both datasets:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters] \u2014 the component both\n", + " datasets constrain, through the arcs inside the mask and the shear outside it.\n", + "\n", + " - The source galaxy's light is a linear `SersicCore` [6 parameters] \u2014 only the imaging dataset sees this;\n", + " the weak catalogue's galaxies are pure shear probes with no model components.\n", + "\n", + "Crucially the model is composed **once**: passing the same model to both analysis factors below means they\n", + "share the same priors and therefore the same parameters \u2014 the definition of a joint fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + "# Source:\n", + "\n", + "bulge = af.Model(al.lp_linear.SersicCore)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0, bulge=bulge)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis List__\n", + "\n", + "One analysis object per dataset, exactly as each would be built in its own modeling script.\n", + "\n", + "The imaging analysis runs in NumPy mode here (`use_jax=False`): the factor graph below evaluates its factors\n", + "in a plain Python loop because the weak-lensing analysis is a NumPy calculation, and mixing an eagerly-JAX\n", + "imaging likelihood into that loop gains nothing over NumPy without JIT compilation. On this small masked\n", + "dataset the NumPy likelihood is fast, and the weak likelihood is fractions of a millisecond." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_imaging = al.AnalysisImaging(dataset=dataset_imaging, use_jax=False)\n", + "\n", + "analysis_weak = al.AnalysisWeak(dataset=dataset_weak)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis Factor__\n", + "\n", + "Each analysis is wrapped in an `AnalysisFactor` paired with the model. Because both factors receive the\n", + "*same* model object, the factor graph recognises every prior as shared \u2014 a single 11-dimensional parameter\n", + "space whose likelihood is the sum of the two factors.\n", + "\n", + "(In the `multi` examples each factor gets a slightly different copy of the model, e.g. per-wavelength\n", + "ellipticities; here total sharing is exactly what \"one mass distribution, two datasets\" means.)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis_factor_imaging = af.AnalysisFactor(prior_model=model, analysis=analysis_imaging)\n", + "\n", + "analysis_factor_weak = af.AnalysisFactor(prior_model=model, analysis=analysis_weak)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Factor Graph__\n", + "\n", + "The factors combine into a `FactorGraphModel`, whose `global_prior_model` is the shared parameter space and\n", + "whose `log_likelihood_function` is the sum over factors \u2014 the quantity `fit.py` computed by hand." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "factor_graph = af.FactorGraphModel(analysis_factor_imaging, analysis_factor_weak)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search & Model-Fit__\n", + "\n", + "Nautilus samples the joint likelihood. The parameter space is simple (N=11, unimodal), so 100 live points\n", + "suffice; expect the fit to take some minutes on an ordinary CPU (the imaging likelihood dominates the cost \u2014\n", + "adding the weak factor is essentially free, which is much of weak lensing's practical appeal)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"weak\") / \"features\",\n", + " name=\"strong_lensing_joint\",\n", + " unique_tag=dataset_name,\n", + " n_live=100,\n", + " iterations_per_quick_update=10000,\n", + ")\n", + "\n", + "print(\n", + " \"\"\"\n", + " The joint strong+weak non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell will progress once the search has completed - this could take some minutes!\n", + " \"\"\"\n", + ")\n", + "\n", + "result_list = search.fit(model=factor_graph.global_prior_model, analysis=factor_graph)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The search returns one result per factor (imaging first, weak second), sharing a single posterior. The\n", + "mass parameters below are constrained by *both* datasets simultaneously.\n", + "\n", + "To see the complementarity in play, compare this joint posterior to a run with the weak factor removed\n", + "(comment it out of the `FactorGraphModel` above): the Einstein radius barely changes \u2014 the arcs own it \u2014\n", + "while the constraints on the mass's elliptical components tighten visibly when the shear at large radius is\n", + "included, because ellipticity is exactly what coherent tangential shear across the field measures." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "result = result_list[0]\n", + "\n", + "print(result.info)\n", + "\n", + "print(result.max_log_likelihood_instance)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The per-factor maximum-likelihood fits visualize each dataset's view of the shared model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_imaging(\n", + " fit=result_list[0].max_log_likelihood_fit,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")\n", + "\n", + "aplt.subplot_fit_weak(\n", + " fit=result_list[1].max_log_likelihood_fit,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This example closed the loop the weak-lensing series builds towards: one mass model, constrained inside the\n", + "Einstein radius by strong lensing and outside it by weak shear, sampled as a single joint likelihood.\n", + "\n", + "The same factor-graph pattern extends directly to the realistic versions of this analysis: cluster-scale\n", + "lenses with many cluster members (see `scripts/cluster`), real shear catalogues (the upcoming real-data\n", + "example in this series), and any other dataset combination (`scripts/multi`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/weak/features/strong_lensing/simulator.ipynb b/notebooks/weak/features/strong_lensing/simulator.ipynb new file mode 100644 index 000000000..34a9631be --- /dev/null +++ b/notebooks/weak/features/strong_lensing/simulator.ipynb @@ -0,0 +1,319 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Combined Strong + Weak Lensing\n", + "=========================================\n", + "\n", + "This script simulates the two faces of the same gravitational lens: an `Imaging` dataset of its strongly\n", + "lensed arcs, and a `WeakDataset` of the weak shear it imprints on background galaxies at larger radii.\n", + "Both are generated from **one** `Tracer`, so the datasets share a single true mass distribution \u2014 the setup\n", + "the fit and modeling examples in this folder use to demonstrate joint strong+weak constraints.\n", + "\n", + "__Scientific Context__\n", + "\n", + "This combination is how weak lensing is used around strong-lens galaxy clusters and groups (it is *not* the\n", + "cosmic-shear or galaxy-galaxy-lensing regime). Strong lensing constrains the mass distribution superbly, but\n", + "only inside the Einstein radius where arcs and multiple images form; weak shear extends the constraint to\n", + "several times that radius, where most of the halo's mass lives. Joint analyses of this kind include\n", + "hybrid-Lenstool's simultaneous strong+weak cluster reconstructions (Niemiec et al. 2020) and the combined\n", + "strong and weak lensing analysis of 28 group-to-cluster scale lenses in the Sloan Giant Arcs Survey\n", + "(Oguri et al. 2012).\n", + "\n", + "The scales here are galaxy/group-like and kept small so the examples run quickly: arcs at the Einstein\n", + "radius of 1.6\", and a shear catalogue extending to 10\" \u2014 about six Einstein radii, far beyond the 3.0\"\n", + "region the imaging data constrains.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset Paths:** Both datasets are output to a single `dataset/weak/strong_lensing/` folder.\n", + "- **Ray Tracing:** The single Tracer both datasets are simulated from.\n", + "- **Imaging Simulation:** The strong-lensing image of the lensed source (no lens light, for simplicity).\n", + "- **Weak Simulation:** The surrounding shear catalogue, with space-like shape noise.\n", + "- **Output:** .fits (imaging), .json (weak catalogue + tracer) and .png visualizations." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "Both datasets describe the same lens, so they live together in one folder:\n", + "\n", + " - The imaging data will be output to `data.fits` / `noise_map.fits` / `psf.fits`.\n", + " - The shear catalogue will be output to `dataset.json`, and the shared truth to `tracer.json`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"weak\"\n", + "dataset_name = \"strong_lensing\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "The lens is an elliptical `Isothermal` mass distribution (Einstein radius 1.6\", axis-ratio 0.8 at 45 degrees)\n", + "with no light profile \u2014 omitting lens light keeps the imaging side of the example simple, exactly as in the\n", + "`imaging/features/no_lens_light` example. The source is a compact cored-Sersic.\n", + "\n", + "The ellipticity is deliberately pronounced: the outer quadrupole of the mass distribution is the quantity the\n", + "weak shear constrains best, so it is where the joint fit visibly improves on imaging alone." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.8, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(\n", + " redshift=1.0,\n", + " bulge=al.lp.SersicCore(\n", + " centre=(0.05, 0.05),\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=60.0),\n", + " intensity=4.0,\n", + " effective_radius=0.1,\n", + " sersic_index=1.0,\n", + " ),\n", + ")\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Imaging Simulation__\n", + "\n", + "The strong-lensing image is simulated exactly as in `scripts/imaging/simulator.py`: a 100 x 100 grid at\n", + "0.1\"/pixel (a 10\" field of view whose central ~3\" contains the arcs), a Gaussian PSF, and Poisson + sky noise." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "grid = al.Grid2D.uniform(\n", + " shape_native=(100, 100),\n", + " pixel_scales=0.1,\n", + ")\n", + "\n", + "psf = al.Convolver.from_gaussian(\n", + " convolve_over_sample_size=1,\n", + " shape_native=(11, 11),\n", + " sigma=0.1,\n", + " pixel_scales=grid.pixel_scales,\n", + ")\n", + "\n", + "simulator = al.SimulatorImaging(\n", + " exposure_time=300.0,\n", + " psf=psf,\n", + " background_sky_level=0.1,\n", + " add_poisson_noise_to_data=True,\n", + ")\n", + "\n", + "dataset_imaging = simulator.via_tracer_from(tracer=tracer, grid=grid)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Weak Simulation__\n", + "\n", + "The weak shear catalogue is simulated from the *same tracer* at 400 uniform-random background-galaxy\n", + "positions inside a 10\" half-width square. At the Einstein radius the shear is of order unity, but by 10\" it\n", + "has fallen to |gamma| ~ 0.08 \u2014 each individual galaxy is a noisy probe, and the signal lives in their\n", + "ensemble.\n", + "\n", + "A shape noise of `noise_sigma = 0.1` per component corresponds to deep space-based imaging (ground-based\n", + "surveys are nearer 0.3); it gives this small catalogue a total detection significance high enough for the\n", + "example's joint fit to visibly tighten the mass model." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator_weak = al.SimulatorShearYX(noise_sigma=0.1, seed=1)\n", + "\n", + "dataset_weak = simulator_weak.via_tracer_random_positions_from(\n", + " tracer=tracer,\n", + " n_galaxies=400,\n", + " grid_extent=10.0,\n", + " name=dataset_name,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "The imaging dataset is output as .fits (the standard astronomical format), the weak catalogue and the shared\n", + "true tracer as .json." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path.mkdir(parents=True, exist_ok=True)\n", + "\n", + "aplt.fits_imaging(\n", + " dataset=dataset_imaging,\n", + " data_path=dataset_path / \"data.fits\",\n", + " psf_path=dataset_path / \"psf.fits\",\n", + " noise_map_path=dataset_path / \"noise_map.fits\",\n", + " overwrite=True,\n", + ")\n", + "\n", + "al.output_to_json(obj=dataset_weak, file_path=dataset_path / \"dataset.json\")\n", + "al.output_to_json(obj=tracer, file_path=dataset_path / \"tracer.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "The two datasets side by side make the complementarity obvious: all the imaging information sits inside the\n", + "central few arc-seconds, while the shear catalogue's quivers cover a field over six Einstein radii across.\n", + "The Kaiser-Squires convergence map of the shear field peaks on the strong lens \u2014 the same mass seen two ways." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_imaging_dataset(\n", + " dataset=dataset_imaging, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.subplot_weak_dataset(\n", + " dataset=dataset_weak, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.plot_convergence_map(\n", + " shear_yx=dataset_weak.shear_yx,\n", + " shape_native=(30, 30),\n", + " smoothing_sigma_pixels=1.0,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")\n", + "\n", + "print(dataset_weak.info)\n", + "print(f\"Wrote combined strong+weak dataset to {dataset_path}\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/weak/fit.ipynb b/notebooks/weak/fit.ipynb index ce212e6da..6ca8b9b21 100644 --- a/notebooks/weak/fit.ipynb +++ b/notebooks/weak/fit.ipynb @@ -1,241 +1,324 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Fit: Weak Lensing\n", - "=================\n", - "\n", - "This script shows how to fit a strong-lens mass model to a weak gravitational lensing shear catalogue. Where\n", - "the `imaging` and `interferometer` workflows fit a 2D image of a lensed source, the weak-lensing workflow fits\n", - "a set of (gamma_2, gamma_1) shear measurements at the (y, x) positions of background source galaxies \u2014 a\n", - "``WeakDataset`` produced by the simulator script in `scripts/weak/simulator.py`.\n", - "\n", - "A weak-lensing fit is conceptually simpler than its imaging counterpart: there is no PSF convolution, no\n", - "masking, no inversion / pixelization, and no source-galaxy light profile. The model is a `Tracer` whose mass\n", - "profiles induce a shear field, and the `FitWeak` class compares that model shear against the observed shear\n", - "to compute residuals, chi-squared and the log-likelihood.\n", - "\n", - "__Contents__\n", - "\n", - "- **Dataset:** Load the simulated `WeakDataset` from disk and visualise it.\n", - "- **Model:** Build a lens-model `Tracer` whose mass profiles produce the model shear field.\n", - "- **Fit:** Construct a `FitWeak` and inspect its derived quantities (residuals, chi-squared, log-likelihood).\n", - "- **Visualization:** Plot the fit as a 2x2 mosaic of data, model, overlay, and chi-squared map.\n", - "- **Notes:** What a \"good\" fit looks like and how this script relates to the upcoming modeling tutorial." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset__\n", - "\n", - "We load the simulated `WeakDataset` produced by `scripts/weak/simulator.py`: 200 background\n", - "source-galaxy positions in a 3.0\" half-extent square, each with a measured `(gamma_2, gamma_1)` shear\n", - "vector and per-galaxy noise standard deviation 0.3. The shear field carries the signature of the\n", - "foreground lens's mass distribution.\n", - "\n", - "If `dataset.json` does not yet exist, run `scripts/weak/simulator.py` first." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path = Path(\"dataset\") / \"weak\" / \"simple\"\n", - "\n", - "dataset = al.from_json(file_path=dataset_path / \"dataset.json\")\n", - "\n", - "print(dataset.info)" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Before fitting it is worth visualising the dataset. `aplt.subplot_weak_dataset` produces a 2x2 mosaic\n", - "showing the shear field as headless quiver segments, the per-galaxy noise map, the shear magnitude\n", - "`|gamma|` and the position angle `phi`. Tangential alignment around the lens centre at `(0, 0)` is the\n", - "characteristic visual signature of a strong-lens shear field." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_weak_dataset(\n", - " dataset=dataset,\n", - " output_path=dataset_path,\n", - " output_format=\"png\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Model__\n", - "\n", - "The model `Tracer` is built from the same primitives the simulator used: a foreground Isothermal lens\n", - "galaxy and a background source galaxy with no light profile (weak-lensing measurements are sensitive to\n", - "the lens mass, not the source's appearance). In a real workflow the mass parameters would be inferred by a\n", - "non-linear search (see the modeling tutorial in the next step of the weak-lensing series). Here we\n", - "hand-pick parameters close to the simulator's truth \u2014 `einstein_radius=1.6`, `axis_ratio=0.9`,\n", - "`angle=45.0`, `centre=(0.0, 0.0)` \u2014 so the fit shows what residuals consistent with shape noise look like." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Fit__\n", - "\n", - "`FitWeak` evaluates the model shear field at the dataset's galaxy positions via the lens's projected\n", - "mass Hessian, then derives residuals, chi-squared and the log-likelihood under the assumption that each\n", - "shear component is independently Gaussian-distributed around the model with the per-galaxy noise.\n", - "\n", - "Each background galaxy contributes **two** independent measurements (`gamma_1` and `gamma_2` carry the\n", - "same per-galaxy noise but are independent draws), so the total number of degrees of freedom is\n", - "`2 * n_galaxies`. For a well-fitting model with shape-noise-dominated residuals the expected chi-squared\n", - "is approximately equal to that number." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "fit = al.FitWeak(dataset=dataset, tracer=tracer)\n", - "\n", - "print()\n", - "print(\"Fit Summary\")\n", - "print(\"-----------\")\n", - "print(f\"n_galaxies : {dataset.n_galaxies}\")\n", - "print(f\"degrees_of_freedom: {2 * dataset.n_galaxies}\")\n", - "print(f\"chi_squared : {fit.chi_squared:.3f}\")\n", - "print(f\"noise_normalization: {fit.noise_normalization:.3f}\")\n", - "print(f\"log_likelihood : {fit.log_likelihood:.3f}\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualization__\n", - "\n", - "`aplt.subplot_fit_weak` produces the 2x2 mosaic that summarises a weak-lensing fit:\n", - "\n", - " - **Top-left:** the observed shear field, drawn in the same headless-quiver style as the dataset plot.\n", - " - **Top-right:** the model shear field evaluated at the galaxy positions.\n", - " - **Bottom-left:** data and model overlaid on a single axes \u2014 data in black, model in red. Deviations\n", - " are visible where the two segments disagree in length or orientation.\n", - " - **Bottom-right:** the per-galaxy chi-squared map (summed across the two shear components), colour-coded\n", - " to highlight any galaxies driving large residuals \u2014 for example, those whose true ellipticity happens\n", - " to be poorly aligned with the lens's induced shear." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_fit_weak(\n", - " fit=fit,\n", - " output_path=dataset_path,\n", - " output_format=\"png\",\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Notes__\n", - "\n", - "A \"good\" fit produces residuals consistent with the shape-noise floor \u2014 visually, the residual quivers\n", - "in the data-vs-model overlay should look short and randomly oriented, and the chi-squared map should be\n", - "fairly uniform with no clear spatial pattern. A systematic mismatch (e.g. residuals all pointing inward\n", - "around a region) usually indicates that the model's mass profile is the wrong shape, not just the wrong\n", - "amplitude.\n", - "\n", - "The next step in the weak-lensing series, `scripts/weak/modeling.py`, replaces this hand-picked model\n", - "with an `AnalysisWeak` driven by a non-linear search: the same `FitWeak` machinery, called inside a\n", - "likelihood function, with priors on the lens mass parameters." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Fit: Weak Lensing\n", + "=================\n", + "\n", + "This script shows how to fit a strong-lens mass model to a weak gravitational lensing shear catalogue. Where\n", + "the `imaging` and `interferometer` workflows fit a 2D image of a lensed source, the weak-lensing workflow fits\n", + "a set of (gamma_2, gamma_1) shear measurements at the (y, x) positions of background source galaxies \u2014 a\n", + "``WeakDataset`` produced by the simulator script in `scripts/weak/simulator.py`.\n", + "\n", + "A weak-lensing fit is conceptually simpler than its imaging counterpart: there is no PSF convolution, no\n", + "masking, no inversion / pixelization, and no source-galaxy light profile. The model is a `Tracer` whose mass\n", + "profiles induce a shear field, and the `FitWeak` class compares that model shear against the observed shear\n", + "to compute residuals, chi-squared and the log-likelihood.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load the simulated `WeakDataset` from disk and visualise it.\n", + "- **Model:** Build a lens-model `Tracer` whose mass profiles produce the model shear field.\n", + "- **Fit:** Construct a `FitWeak` and inspect its derived quantities (residuals, chi-squared, log-likelihood).\n", + "- **Visualization:** Plot the fit as a 2x2 mosaic of data, model, overlay, and chi-squared map.\n", + "- **Shear Profile:** Bin the tangential/cross shear about the lens centre and compare data to model.\n", + "- **Notes:** What a \"good\" fit looks like and how this script relates to the upcoming modeling tutorial." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "We load the simulated `WeakDataset` produced by `scripts/weak/simulator.py`: 200 background\n", + "source-galaxy positions in a 3.0\" half-extent square, each with a measured `(gamma_2, gamma_1)` shear\n", + "vector and per-galaxy noise standard deviation 0.3. The shear field carries the signature of the\n", + "foreground lens's mass distribution.\n", + "\n", + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / \"weak\" / \"simple\"\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/weak/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.from_json(file_path=dataset_path / \"dataset.json\")\n", + "\n", + "print(dataset.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Before fitting it is worth visualising the dataset. `aplt.subplot_weak_dataset` produces a 2x2 mosaic\n", + "showing the shear field as headless quiver segments, the per-galaxy noise map, the shear magnitude\n", + "`|gamma|` and the position angle `phi`. Tangential alignment around the lens centre at `(0, 0)` is the\n", + "characteristic visual signature of a strong-lens shear field." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_weak_dataset(\n", + " dataset=dataset,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "The model `Tracer` is built from the same primitives the simulator used: a foreground Isothermal lens\n", + "galaxy and a background source galaxy with no light profile (weak-lensing measurements are sensitive to\n", + "the lens mass, not the source's appearance). In a real workflow the mass parameters would be inferred by a\n", + "non-linear search (see the modeling tutorial in the next step of the weak-lensing series). Here we\n", + "hand-pick parameters close to the simulator's truth \u2014 `einstein_radius=1.6`, `axis_ratio=0.9`,\n", + "`angle=45.0`, `centre=(0.0, 0.0)` \u2014 so the fit shows what residuals consistent with shape noise look like." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "`FitWeak` evaluates the model shear field at the dataset's galaxy positions via the lens's projected\n", + "mass Hessian, then derives residuals, chi-squared and the log-likelihood under the assumption that each\n", + "shear component is independently Gaussian-distributed around the model with the per-galaxy noise.\n", + "\n", + "Each background galaxy contributes **two** independent measurements (`gamma_1` and `gamma_2` carry the\n", + "same per-galaxy noise but are independent draws), so the total number of degrees of freedom is\n", + "`2 * n_galaxies`. For a well-fitting model with shape-noise-dominated residuals the expected chi-squared\n", + "is approximately equal to that number." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitWeak(dataset=dataset, tracer=tracer)\n", + "\n", + "print()\n", + "print(\"Fit Summary\")\n", + "print(\"-----------\")\n", + "print(f\"n_galaxies : {dataset.n_galaxies}\")\n", + "print(f\"degrees_of_freedom: {2 * dataset.n_galaxies}\")\n", + "print(f\"chi_squared : {fit.chi_squared:.3f}\")\n", + "print(f\"noise_normalization: {fit.noise_normalization:.3f}\")\n", + "print(f\"log_likelihood : {fit.log_likelihood:.3f}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualization__\n", + "\n", + "`aplt.subplot_fit_weak` produces the 2x2 mosaic that summarises a weak-lensing fit:\n", + "\n", + " - **Top-left:** the observed shear field, drawn in the same headless-quiver style as the dataset plot.\n", + " - **Top-right:** the model shear field evaluated at the galaxy positions.\n", + " - **Bottom-left:** data and model overlaid on a single axes \u2014 data in black, model in red. Deviations\n", + " are visible where the two segments disagree in length or orientation.\n", + " - **Bottom-right:** the per-galaxy chi-squared map (summed across the two shear components), colour-coded\n", + " to highlight any galaxies driving large residuals \u2014 for example, those whose true ellipticity happens\n", + " to be poorly aligned with the lens's induced shear." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_fit_weak(\n", + " fit=fit,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Shear Profile__\n", + "\n", + "The panels above show the fit galaxy by galaxy; cluster weak-lensing results are usually shown instead as\n", + "the azimuthally averaged *tangential shear profile* `gamma_t(r)` \u2014 the mean tangential stretching of\n", + "background galaxies in radial bins about the lens centre, which traces the projected mass profile (this is\n", + "the standard observable of analyses such as the Sloan Giant Arcs Survey and the Frontier Fields clusters).\n", + "\n", + "`aplt.plot_shear_profile` bins the catalogue about a chosen centre and, because we pass it the `FitWeak`,\n", + "overlays the model shear's tangential profile as a line. The plot also shows the *cross* component\n", + "`gamma_x` (the 45-degree rotated component): gravitational lensing produces none at leading order, so the\n", + "cross points scattering around zero is the standard \"B-mode\" null test \u2014 a systematic contaminating the\n", + "measurement would show up here." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_shear_profile(\n", + " fit,\n", + " centre=(0.0, 0.0),\n", + " bins=8,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Notes__\n", + "\n", + "A \"good\" fit produces residuals consistent with the shape-noise floor \u2014 visually, the residual quivers\n", + "in the data-vs-model overlay should look short and randomly oriented, and the chi-squared map should be\n", + "fairly uniform with no clear spatial pattern. A systematic mismatch (e.g. residuals all pointing inward\n", + "around a region) usually indicates that the model's mass profile is the wrong shape, not just the wrong\n", + "amplitude.\n", + "\n", + "The next step in the weak-lensing series, `scripts/weak/modeling.py`, replaces this hand-picked model\n", + "with an `AnalysisWeak` driven by a non-linear search: the same `FitWeak` machinery, called inside a\n", + "likelihood function, with priors on the lens mass parameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/notebooks/weak/likelihood_function.ipynb b/notebooks/weak/likelihood_function.ipynb new file mode 100644 index 000000000..2fe5f5ab9 --- /dev/null +++ b/notebooks/weak/likelihood_function.ipynb @@ -0,0 +1,456 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Log Likelihood Function: Weak Lensing__\n", + "\n", + "This script provides a step-by-step guide of the **PyAutoLens** likelihood function for fitting a lens mass\n", + "model to a weak gravitational lensing shear catalogue (a `WeakDataset`). It is the weak-lensing companion of\n", + "the guides for the other dataset types (e.g. `scripts/imaging/likelihood_function.py`), following the same\n", + "style and level of detail.\n", + "\n", + "Every step below is what happens inside a single call of `al.AnalysisWeak.log_likelihood_function` \u2014 the\n", + "function a non-linear search calls tens of thousands of times in `scripts/weak/modeling.py`. By the end of\n", + "the script we will have computed the log likelihood \"by hand\" and verified it matches `al.FitWeak` and\n", + "`al.AnalysisWeak` exactly.\n", + "\n", + "A weak-lensing likelihood is the simplest in **PyAutoLens**: there is no PSF convolution, no mask, no\n", + "over-sampling and no linear inversion. The data are a catalogue of `2N` numbers \u2014 two shear components per\n", + "background galaxy \u2014 and the likelihood is a pure Gaussian comparison of these against the model shear field\n", + "evaluated at the galaxy positions. This simplicity is what makes weak-lensing constraints so cheap to add to\n", + "a strong-lensing analysis.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load the weak-lensing shear catalogue that is fitted (auto-simulating it if missing).\n", + "- **Lens Galaxy:** Define the lens galaxy mass model whose shear field is compared to the data.\n", + "- **Shear Field Evaluation:** Evaluate the tracer's shear at every catalogue position via the lensing Hessian.\n", + "- **Residuals:** Compare the model shear to the observed shear, galaxy by galaxy.\n", + "- **Chi Squared:** Sum the noise-normalized squared residuals over all 2N shear components.\n", + "- **Noise Normalization Term:** The Gaussian normalization, with its factor of 2 for the two components.\n", + "- **Calculate The Log Likelihood:** Combine the two terms into the log likelihood.\n", + "- **Fit:** Verify the manual calculation against `al.FitWeak`.\n", + "- **Analysis:** Verify it against `al.AnalysisWeak.log_likelihood_function`, as called in lens modeling.\n", + "- **Wrap Up:** Summary and pointers to the rest of the weak-lensing series." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "import numpy as np\n", + "from pathlib import Path\n", + "\n", + "import autolens as al" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "We load the simulated `WeakDataset` produced by `scripts/weak/simulator.py`: 200 background source-galaxy\n", + "positions, each with a measured `(gamma_2, gamma_1)` shear vector and a per-galaxy noise standard deviation\n", + "of 0.3 (a typical ground-based shape-noise value).\n", + "\n", + "__Dataset Auto-Simulation__\n", + "\n", + "If the dataset does not already exist on your system, it will be created by running the corresponding\n", + "simulator script. This ensures that all example scripts can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / \"weak\" / \"simple\"\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/weak/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.from_json(file_path=dataset_path / \"dataset.json\")\n", + "\n", + "print(dataset.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The three ingredients of the likelihood are all on the dataset:\n", + "\n", + " - `dataset.positions`: the `(N, 2)` grid of `(y, x)` arc-second coordinates of the background galaxies.\n", + " - `dataset.shear_yx`: the `(N, 2)` observed shear components, stored as `(gamma_2, gamma_1)` per galaxy.\n", + " - `dataset.noise_map`: the `(N,)` per-galaxy noise, shared by both shear components of that galaxy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(f\"n_galaxies : {dataset.n_galaxies}\")\n", + "print(f\"positions (first galaxy) : {np.asarray(dataset.positions)[0]}\")\n", + "print(f\"shear (first galaxy) : {np.asarray(dataset.shear_yx)[0]}\")\n", + "print(f\"noise (first galaxy) : {np.asarray(dataset.noise_map)[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Lens Galaxy__\n", + "\n", + "The model whose likelihood we are evaluating is a `Tracer` \u2014 the same object used by every other dataset\n", + "type. Only the mass profiles matter for weak lensing: background galaxies are pure probes of the shear\n", + "field, so no light profiles are needed anywhere.\n", + "\n", + "We use mass parameters close to (but not exactly) the simulator's truth, so the residuals below are visibly\n", + "non-zero but the fit is good." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Shear Field Evaluation__\n", + "\n", + "Step one of the likelihood: evaluate the model's shear at every catalogue position.\n", + "\n", + "The shear is a second derivative of the lensing potential. **PyAutoLens** computes it by numerically\n", + "differentiating the tracer's deflection-angle field \u2014 the `LensCalc.shear_yx_2d_via_hessian_from` method\n", + "evaluates deflections on a small cross of points around each galaxy position and forms the Hessian, from\n", + "which the two shear components follow.\n", + "\n", + "Two things are worth noting:\n", + "\n", + " - This is the *same primitive* the simulator (`SimulatorShearYX`) uses to generate data, so a noise-free\n", + " dataset fitted by its own truth tracer round-trips bit-exactly to zero residuals.\n", + "\n", + " - The result is the shear `gamma`, not the reduced shear `g = gamma / (1 - kappa)` that real surveys\n", + " measure \u2014 adequate here and at large radii where `kappa` is small, and the planned real-data example\n", + " in the weak-lensing series will introduce the distinction." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from autogalaxy.operate.lens_calc import LensCalc\n", + "\n", + "model_shear = LensCalc.from_tracer(tracer).shear_yx_2d_via_hessian_from(\n", + " grid=dataset.positions\n", + ")\n", + "\n", + "print(f\"model shear (first galaxy) : {np.asarray(model_shear)[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Residuals__\n", + "\n", + "Step two: subtract the model from the data. Both are `(N, 2)` arrays of `(gamma_2, gamma_1)` components, so\n", + "the residual map is simply their difference \u2014 no convolution, binning or masking intervenes." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "residual_map = np.asarray(dataset.shear_yx) - np.asarray(model_shear)\n", + "\n", + "print(f\"residuals (first galaxy) : {residual_map[0]}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each residual is then divided by that galaxy's noise. The per-galaxy noise `sigma` applies to *both* shear\n", + "components (they share the same measurement process) but the two components are *independent* Gaussian\n", + "draws \u2014 this independence is why the likelihood below counts `2N` data points, not `N`.\n", + "\n", + "The `[:, None]` broadcasts the `(N,)` noise map across both components of the `(N, 2)` residual map." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_map = np.asarray(dataset.noise_map)\n", + "\n", + "normalized_residual_map = residual_map / noise_map[:, None]" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Chi Squared__\n", + "\n", + "Step three: the chi-squared is the sum of squared normalized residuals over all `N x 2` components:\n", + "\n", + "$\\\\chi^2 = \\\\sum_{i=1}^{N} \\\\sum_{k=1}^{2} \\\\left( \\\\frac{\\\\gamma^{\\\\rm data}_{i,k} - \\\\gamma^{\\\\rm model}_{i,k}}{\\\\sigma_i} \\\\right)^2$\n", + "\n", + "For a well-fitting model whose residuals are pure shape noise, the expected chi-squared is approximately the\n", + "number of data points, `2N = 400` \u2014 a quick sanity check worth internalising for any weak-lensing fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "chi_squared_map = normalized_residual_map**2.0\n", + "\n", + "chi_squared = float(np.sum(chi_squared_map))\n", + "\n", + "print(f\"chi_squared : {chi_squared:.4f} (expected ~{2 * dataset.n_galaxies} for a good fit)\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Noise Normalization Term__\n", + "\n", + "The Gaussian likelihood also carries a model-independent normalization term:\n", + "\n", + "$\\\\text{noise normalization} = \\\\sum_{i=1}^{N} \\\\sum_{k=1}^{2} \\\\ln \\\\left( 2 \\\\pi \\\\sigma_i^2 \\\\right) = 2 \\\\sum_{i=1}^{N} \\\\ln \\\\left( 2 \\\\pi \\\\sigma_i^2 \\\\right)$\n", + "\n", + "The leading factor of 2 is the same `2N` counting as the chi-squared: each galaxy contributes two\n", + "independent measurements with the same `sigma`. Because it does not depend on the model, this term does not\n", + "influence which model a non-linear search prefers \u2014 but it is required for the log likelihood's absolute\n", + "value to be meaningful (e.g. when comparing to other dataset types in a combined fit)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "noise_normalization = float(2.0 * np.sum(np.log(2.0 * np.pi * noise_map**2.0)))\n", + "\n", + "print(f\"noise_normalization : {noise_normalization:.4f}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Calculate The Log Likelihood__\n", + "\n", + "The log likelihood combines the two terms in the standard Gaussian form:\n", + "\n", + "$\\\\ln \\\\mathcal{L} = -\\\\frac{1}{2} \\\\left( \\\\chi^2 + \\\\text{noise normalization} \\\\right)$\n", + "\n", + "That is the entire weak-lensing likelihood function \u2014 compare with the imaging guide, where PSF convolution\n", + "and masking sit between the model and this same final expression." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "log_likelihood = -0.5 * (chi_squared + noise_normalization)\n", + "\n", + "print(f\"log_likelihood (by hand) : {log_likelihood:.8f}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Fit__\n", + "\n", + "`al.FitWeak` packages the steps above (shear evaluation, residuals, chi-squared, normalization) into a\n", + "single object. Its log likelihood must match our manual calculation exactly." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "fit = al.FitWeak(dataset=dataset, tracer=tracer)\n", + "\n", + "print(f\"log_likelihood (FitWeak) : {fit.log_likelihood:.8f}\")\n", + "\n", + "assert fit.log_likelihood == log_likelihood" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "Finally, `al.AnalysisWeak` is the object handed to a non-linear search in `scripts/weak/modeling.py`. Its\n", + "`log_likelihood_function` builds the `Tracer` from a model instance and returns exactly the `FitWeak` log\n", + "likelihood \u2014 so the number below is what Nautilus receives at every sampled point of parameter space." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import autofit as af\n", + "\n", + "model = af.Collection(\n", + " galaxies=af.Collection(lens=lens_galaxy, source=source_galaxy)\n", + ")\n", + "\n", + "analysis = al.AnalysisWeak(dataset=dataset)\n", + "\n", + "instance = model.instance_from_unit_vector([])\n", + "\n", + "analysis_log_likelihood = analysis.log_likelihood_function(instance=instance)\n", + "\n", + "print(f\"log_likelihood (AnalysisWeak) : {analysis_log_likelihood:.8f}\")\n", + "\n", + "assert analysis_log_likelihood == log_likelihood" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "We have computed the weak-lensing log likelihood step by step and verified it against `al.FitWeak` and\n", + "`al.AnalysisWeak`:\n", + "\n", + " 1. Evaluate the tracer's shear field at the catalogue positions (lensing Hessian).\n", + " 2. Residuals: data minus model, per galaxy and per component.\n", + " 3. Chi-squared: noise-normalized squared residuals summed over all `2N` components.\n", + " 4. Noise normalization: `2 sum(ln(2 pi sigma^2))`, with the factor 2 counting both components.\n", + " 5. Log likelihood: `-0.5 * (chi_squared + noise_normalization)`.\n", + "\n", + "The rest of the weak-lensing series builds on this: `scripts/weak/modeling.py` samples this likelihood with\n", + "a non-linear search, `scripts/weak/fit.py` visualizes fits (including the tangential shear profile, the\n", + "standard observable this likelihood constrains), and upcoming examples apply it to a real cluster shear\n", + "catalogue and combine it with strong-lensing imaging in a joint fit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/weak/modeling.ipynb b/notebooks/weak/modeling.ipynb new file mode 100644 index 000000000..b2582f0fb --- /dev/null +++ b/notebooks/weak/modeling.ipynb @@ -0,0 +1,463 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Modeling: Weak Lensing\n", + "======================\n", + "\n", + "This script fits a lens mass model to a weak gravitational lensing shear catalogue using a non-linear search.\n", + "It is the weak-lensing analogue of `scripts/imaging/modeling.py`: where that script infers a lens model from a\n", + "2D image of a lensed source, this script infers one from the (gamma_2, gamma_1) shear measurements of background\n", + "source galaxies \u2014 a `WeakDataset` produced by `scripts/weak/simulator.py`.\n", + "\n", + "The previous script in the weak-lensing series, `scripts/weak/fit.py`, fitted a hand-picked mass model via the\n", + "`FitWeak` class and inspected its residuals and log-likelihood. Here we complete the workflow: the same `FitWeak`\n", + "machinery is wrapped in an `AnalysisWeak` object whose `log_likelihood_function` is called by the nested sampling\n", + "algorithm Nautilus to infer the posterior probability distribution of the lens mass parameters.\n", + "\n", + "Weak-lensing model-fits are computationally much cheaper than imaging fits \u2014 there is no PSF convolution, no\n", + "masking, no pixelized source inversion \u2014 so this fit runs in minutes on an ordinary CPU.\n", + "\n", + "__Scientific Context__\n", + "\n", + "The weak-lensing regime modeled here is the shear signal around a galaxy-scale or group/cluster-scale lens:\n", + "background galaxies far from the lens centre are weakly sheared tangentially around it, and their measured\n", + "ellipticities constrain the mass distribution at radii the strong-lensing features (arcs, multiple images) do\n", + "not reach. Combining this large-radius information with strong lensing is a well-established technique for\n", + "galaxy clusters and groups, and a dedicated combined strong-plus-weak example follows later in the series.\n", + "\n", + "__Contents__\n", + "\n", + "- **Dataset:** Load the simulated `WeakDataset` (simulating it first if required).\n", + "- **Model:** Compose the lens mass model using the Model and Collection API.\n", + "- **Search:** Configure the Nautilus non-linear search.\n", + "- **Analysis:** Create the `AnalysisWeak` object defining how the model is fitted to the data.\n", + "- **Run Times:** The expected run time of a weak-lensing model-fit.\n", + "- **Model-Fit:** Perform the model-fit.\n", + "- **Output Folder Layout:** The structure of the `output` folder where results are written.\n", + "- **Result:** Inspect the inferred model and posterior.\n", + "\n", + "__Model__\n", + "\n", + "This script fits a `WeakDataset` of a 'galaxy-scale' lens with a model where:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal` [5 parameters].\n", + "\n", + " - The background source galaxies are treated purely as shear probes \u2014 they have no light or mass model\n", + " [0 parameters].\n", + "\n", + "The number of free parameters and therefore the dimensionality of non-linear parameter space is N=5." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset__\n", + "\n", + "We load the simulated `WeakDataset` produced by `scripts/weak/simulator.py`: 200 background source-galaxy\n", + "positions in a 3.0\" half-extent square, each with a measured `(gamma_2, gamma_1)` shear vector and per-galaxy\n", + "noise standard deviation 0.3.\n", + "\n", + "If the dataset does not already exist on your system, it is created by running the simulator script, so this\n", + "example can be run without manually simulating data first." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_name = \"simple\"\n", + "dataset_path = Path(\"dataset\") / \"weak\" / dataset_name\n", + "\n", + "if al.util.dataset.should_simulate(str(dataset_path)):\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.run(\n", + " [sys.executable, \"scripts/weak/simulator.py\"],\n", + " check=True,\n", + " )\n", + "\n", + "dataset = al.from_json(file_path=dataset_path / \"dataset.json\")\n", + "\n", + "print(dataset.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Before fitting, we visualise the dataset with `aplt.subplot_weak_dataset`, a 2x2 mosaic showing the shear field\n", + "as headless quiver segments, the per-galaxy noise map, the shear magnitude `|gamma|` and the position angle\n", + "`phi`. The tangential alignment of the segments around the lens centre at `(0.0\", 0.0\")` is the signal the\n", + "model-fit will exploit." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_weak_dataset(\n", + " dataset=dataset,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model__\n", + "\n", + "We compose the lens model using `Model` and `Collection` objects, imported from **PyAutoLens**'s parent\n", + "project **PyAutoFit** \u2014 the identical API used by every other modeling script in the workspace:\n", + "\n", + " - The lens galaxy's total mass distribution is an `Isothermal`, with free centre, elliptical components and\n", + " Einstein radius [5 parameters].\n", + "\n", + " - The source galaxy carries no model components: in a weak-lensing fit the background galaxies are pure\n", + " probes of the foreground shear field, so the source's appearance is irrelevant. It is included only so the\n", + " `Tracer` has a source-plane redshift for the lensing geometry.\n", + "\n", + "Note what is absent compared to `scripts/imaging/modeling.py`: no light profiles, no `ExternalShear` (the\n", + "shear field *is* the data here \u2014 an external shear component would be degenerate with the signal at leading\n", + "order for this single-lens dataset) and therefore a much smaller parameter space (N=5 versus N=21+)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Lens:\n", + "\n", + "mass = af.Model(al.mp.Isothermal)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.5, mass=mass)\n", + "\n", + "# Source:\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0)\n", + "\n", + "# Overall Lens Model:\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `info` attribute shows the model in a readable format, including the priors on each parameter.\n", + "\n", + "(The `info_whitespace_length` parameter in `config/general.yaml`'s [output] section controls the whitespace\n", + "formatting if the display does not render well on your screen.)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(model.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Search__\n", + "\n", + "The model is fitted to the data using the nested sampling algorithm\n", + "Nautilus (https://nautilus-sampler.readthedocs.io/en/latest/), the default search used throughout the\n", + "workspace.\n", + "\n", + "With only N=5 parameters and a smooth, near-Gaussian likelihood surface, this is an easy parameter space \u2014\n", + "100 live points is ample and keeps the run time to minutes.\n", + "\n", + "An identical combination of model, search and dataset generates the same `unique_identifier`, meaning that\n", + "rerunning the script will resume the existing fit rather than starting again; the `unique_tag` below folds the\n", + "dataset name into that identifier." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "search = af.Nautilus(\n", + " path_prefix=Path(\"weak\"), # The path where results and output are stored.\n", + " name=\"modeling\", # The name of the fit and folder results are output to.\n", + " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", + " n_live=100, # The number of Nautilus \"live\" points, increase for more complex models.\n", + " iterations_per_quick_update=5000, # Every N iterations the max likelihood model is visualized and output to hard-disk.\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Analysis__\n", + "\n", + "We next create an `AnalysisWeak` object, whose `log_likelihood_function` is what the non-linear search calls\n", + "at every iteration:\n", + "\n", + " 1. It builds a `Tracer` from the sampled mass parameters.\n", + " 2. It evaluates the tracer's shear field at the catalogue's galaxy positions \u2014 the same\n", + " `LensCalc` Hessian primitive used by both the simulator and `FitWeak`.\n", + " 3. It returns the Gaussian log-likelihood of the observed shears given the model, summed over the\n", + " `2 * n_galaxies` independent shear components.\n", + "\n", + "A step-by-step walkthrough of this likelihood function is the next entry in the weak-lensing series\n", + "(`scripts/weak/likelihood_function.py`), following the format of the imaging and interferometer\n", + "likelihood-function guides.\n", + "\n", + "Unlike `AnalysisImaging`, no `use_jax` option is passed: the weak-lensing fit is a NumPy calculation\n", + "(it is cheap enough that JAX acceleration is unnecessary for catalogues of this size)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "analysis = al.AnalysisWeak(dataset=dataset)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Run Times__\n", + "\n", + "A single log-likelihood evaluation \u2014 one shear-field evaluation at 200 galaxy positions plus a chi-squared\n", + "sum \u2014 takes of order a millisecond. Nautilus needs roughly 10,000\u201330,000 evaluations for this 5-parameter\n", + "model, so the full fit completes in a few minutes on a single CPU.\n", + "\n", + "This is the key practical difference from imaging fits: the data volume of a shear catalogue is tiny (a few\n", + "hundred numbers rather than tens of thousands of pixels), which is also why weak-lensing constraints are so\n", + "cheap to add to a strong-lensing analysis.\n", + "\n", + "__Model-Fit__\n", + "\n", + "We begin the model-fit by passing the model and analysis objects to the search." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running.\n", + "\n", + " This Jupyter notebook cell will progress once the search has completed - this could take a few minutes!\n", + "\n", + " On-the-fly updates every iterations_per_quick_update are printed to the notebook.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output Folder Layout__\n", + "\n", + "Results are written on the fly to the `autolens_workspace/output/weak/simple/modeling//` folder\n", + "in human-readable formats, so the fit can be inspected while it runs:\n", + "\n", + " output/weak//modeling//\n", + " files/ <- JSON + CSV: loadable Python objects\n", + " dataset.json <- the WeakDataset (reload via al.from_json)\n", + " tracer.json <- max log likelihood Tracer\n", + " model.json <- fitted af.Collection model\n", + " samples.csv <- full Nautilus samples\n", + " samples_summary.json <- max log likelihood parameter values + errors\n", + " image/ <- PNG: visualization\n", + " subplot_weak_dataset.png <- the dataset mosaic (shear field, noise, |gamma|, phi)\n", + " subplot_fit_weak.png <- the fit mosaic (data, model, overlay, chi-squared map)\n", + " galaxies.png <- the mass model's convergence over the field extent\n", + " model.info <- human-readable model summary\n", + " model.results <- human-readable fit summary\n", + "\n", + "__Result__\n", + "\n", + "The search returns a result object; its `info` attribute shows the outcome in a readable format, including\n", + "the median and error estimates of every mass parameter." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `Result` object contains the maximum log likelihood instance, `Tracer` and `FitWeak`. Plotting the fit's\n", + "2x2 mosaic shows what a converged weak-lensing model looks like: short, randomly-oriented residual segments in\n", + "the data-vs-model overlay and a spatially uniform chi-squared map, consistent with the shape-noise floor." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.max_log_likelihood_instance)\n", + "\n", + "aplt.subplot_fit_weak(\n", + " fit=result.max_log_likelihood_fit,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Below, we make a corner plot of the \"Probability Density Function\" of every parameter in the model-fit.\n", + "\n", + "For a shear-only fit, note how well the Einstein radius and mass-profile orientation are constrained relative\n", + "to the centre: the shear signal at each background galaxy is dominated by the enclosed mass and its\n", + "quadrupole, whereas the centre is only weakly pinned by the field's geometry. This complementarity \u2014 weak\n", + "lensing constrains the profile at large radius, strong lensing pins the centre and inner mass \u2014 is exactly\n", + "why the two are combined in cluster and group studies, and is the subject of the combined\n", + "strong-plus-weak example later in the series." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "This script completed the core weak-lensing workflow: simulate (`simulator.py`), fit (`fit.py`) and now model\n", + "(`modeling.py`). The next entries in the series are:\n", + "\n", + " - `scripts/weak/likelihood_function.py`: a step-by-step guide to the weak-lensing likelihood.\n", + " - Weak-lensing analyses of real shear catalogues and combined strong-plus-weak modeling." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/weak/real_data/a2744.ipynb b/notebooks/weak/real_data/a2744.ipynb new file mode 100644 index 000000000..f73023f04 --- /dev/null +++ b/notebooks/weak/real_data/a2744.ipynb @@ -0,0 +1,440 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Real Data: Weak Lensing of Abell 2744\n", + "=====================================\n", + "\n", + "This script fits a dark-matter halo mass model to a **real weak-lensing shape catalogue** of the merging\n", + "galaxy cluster Abell 2744 (\"Pandora's Cluster\", z = 0.308) \u2014 the first PyAutoLens weak-lensing analysis of\n", + "real sky data, and the capstone of the weak-lensing example series.\n", + "\n", + "__The Data & Its Provenance__\n", + "\n", + "The catalogue is the Abell 2744 galaxy shape catalogue shipped with the public pyRRG weak-lensing shape\n", + "measurement code (https://github.com/davidharvey1986/pyRRG, `jwst` branch), whose JWST application to this\n", + "cluster is described in Harvey & Massey 2024 (MNRAS 529, 802, arXiv:2401.16478). It contains 6,585 sources\n", + "with SExtractor photometry, RRG shape moments and per-galaxy ellipticity measurements `(e1, e2)` with\n", + "uncertainties.\n", + "\n", + "**Honest framing:** the file ships inside pyRRG's star/galaxy-classification training data, and we apply our\n", + "own quality cuts below rather than the exact selections of the published analysis \u2014 so this example is a\n", + "*real-data demonstration* whose results we sanity-check against published Abell 2744 analyses, not a\n", + "metrology-grade reproduction of any single paper. That is the right expectation for a first weak-lensing\n", + "look at any new catalogue.\n", + "\n", + "__What this script shows__\n", + "\n", + "- Downloading and loading a real shape catalogue (FITS binary table) and converting RA/Dec to the\n", + " tangent-plane arc-second coordinates PyAutoLens uses.\n", + "- The standard weak-lensing quality cuts every real catalogue needs.\n", + "- Building a `WeakDataset` flagged `is_reduced` \u2014 real surveys measure galaxy ellipticities, i.e. the\n", + " *reduced* shear g = gamma / (1 - kappa), and `FitWeak` computes the matching model quantity.\n", + "- A model-independent Kaiser-Squires mass map \u2014 Abell 2744 is a famous merger, and the map's structure is\n", + " the first sanity check.\n", + "- Fitting a spherical NFW dark-matter halo with a Nautilus search and reading the tangential-shear profile\n", + " in which cluster weak-lensing results are usually shown.\n", + "\n", + "__Contents__\n", + "\n", + "- **Catalogue Download:** Fetch the public catalogue (cached on disk after the first run).\n", + "- **Catalogue Load & Projection:** RA/Dec to tangent-plane arc-seconds about the cluster centre.\n", + "- **Quality Cuts:** The standard selections that turn raw shapes into a usable shear sample.\n", + "- **Weak Dataset:** Build the reduced-shear `WeakDataset`.\n", + "- **Mass Map:** Model-independent Kaiser-Squires reconstruction.\n", + "- **Model & Search:** Spherical NFW halo, Nautilus.\n", + "- **Result:** Tangential-shear profile, and how the numbers compare to the literature." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "\n", + "import autofit as af\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Catalogue Download__\n", + "\n", + "The catalogue is fetched once from the public pyRRG repository (pinned to a specific commit for\n", + "reproducibility) and cached in the dataset folder. We do not redistribute the file with the workspace \u2014\n", + "provenance stays with the pyRRG project." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path = Path(\"dataset\") / \"weak\" / \"a2744_pyrrg\"\n", + "catalogue_path = dataset_path / \"abell2744_galaxies.fits\"\n", + "\n", + "CATALOGUE_URL = (\n", + " \"https://raw.githubusercontent.com/davidharvey1986/pyRRG/\"\n", + " \"0ccc29fb4513137da61b1afb632ca492093bd609/\"\n", + " \"trainStarGalClass/TrainingData/abell2744_galaxies.fits\"\n", + ")\n", + "\n", + "if not catalogue_path.exists():\n", + " import urllib.request\n", + "\n", + " dataset_path.mkdir(parents=True, exist_ok=True)\n", + " print(f\"Downloading A2744 catalogue from pyRRG (one-off, ~3 MB) ...\")\n", + " urllib.request.urlretrieve(CATALOGUE_URL, catalogue_path)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Catalogue Load & Projection__\n", + "\n", + "The table stores sky positions as RA/Dec in degrees. PyAutoLens works in tangent-plane arc-second offsets\n", + "`(y, x)` about a chosen centre, so we project about the cluster core (the catalogue's density peak,\n", + "consistent with the BCG region used by published analyses):\n", + "\n", + " - `x = (RA - RA0) * cos(Dec0) * 3600` (arc-seconds East)\n", + " - `y = (Dec - Dec0) * 3600` (arc-seconds North)\n", + "\n", + "A note on conventions: whether East points left or right on the sky is a *parity* choice that rotates or\n", + "mirrors the shear components' frame. The tangential shear \u2014 the quantity our fit constrains \u2014 is invariant\n", + "under this mirror (only the B-mode cross component flips sign), so the halo-profile fit below is robust to\n", + "it. Precision studies of shear *systematics* must track the convention carefully." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from astropy.io import fits as astropy_fits\n", + "\n", + "with astropy_fits.open(catalogue_path) as hdul:\n", + " table = hdul[1].data\n", + "\n", + "ra = np.asarray(table[\"ra\"], dtype=float)\n", + "dec = np.asarray(table[\"dec\"], dtype=float)\n", + "e1 = np.asarray(table[\"e1\"], dtype=float)\n", + "e2 = np.asarray(table[\"e2\"], dtype=float)\n", + "e1_err = np.asarray(table[\"e1_err\"], dtype=float)\n", + "e2_err = np.asarray(table[\"e2_err\"], dtype=float)\n", + "\n", + "ra_centre, dec_centre = 3.5875, -30.3972 # A2744 core (J2000 degrees)\n", + "\n", + "x = (ra - ra_centre) * np.cos(np.deg2rad(dec_centre)) * 3600.0\n", + "y = (dec - dec_centre) * 3600.0\n", + "radii = np.sqrt(x**2.0 + y**2.0)\n", + "\n", + "print(f\"catalogue sources : {len(ra)}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Quality Cuts__\n", + "\n", + "Raw shape catalogues always contain unusable measurements \u2014 blends, noise detections, objects whose moments\n", + "diverged. The cuts below are the standard minimum for any weak-lensing sample:\n", + "\n", + " - finite, physical ellipticities: |e1|, |e2| < 1 (a galaxy ellipticity cannot exceed 1; the raw catalogue\n", + " contains outliers far beyond this from failed moment measurements).\n", + " - measured uncertainties in a sane range: 0 < e_err < 0.4 per component.\n", + " - a radial window 10\" < r < 130\": inside ~10\" we are in the strong-lensing core where cluster members\n", + " dominate and the weak-lensing (linear shear) approximation is worst; ~130\" is the edge of this\n", + " catalogue's contiguous coverage.\n", + "\n", + "On this catalogue the cuts are severe: only ~1,600 of the 6,585 sources have measured shapes at all, and\n", + "the physical-ellipticity cut trims those to ~400 \u2014 this is training-data-grade depth, an order of magnitude\n", + "shallower than the selections behind published A2744 analyses. Keep that in mind when reading the results." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "finite = np.isfinite(e1) & np.isfinite(e2) & np.isfinite(e1_err) & np.isfinite(e2_err)\n", + "physical = (np.abs(e1) < 1.0) & (np.abs(e2) < 1.0)\n", + "well_measured = (e1_err > 0.0) & (e1_err < 0.4) & (e2_err > 0.0) & (e2_err < 0.4)\n", + "radial = (radii > 10.0) & (radii < 130.0)\n", + "\n", + "use = finite & physical & well_measured & radial\n", + "\n", + "print(f\"after quality cuts : {use.sum()}\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Weak Dataset__\n", + "\n", + "The per-galaxy noise combines the intrinsic shape dispersion (each galaxy has a random unlensed ellipticity;\n", + "sigma_int ~ 0.25 per component is the standard value) with the measurement uncertainty, in quadrature.\n", + "\n", + "`from_arrays` builds the `WeakDataset`; `is_reduced=True` (the loader default) records that these are\n", + "measured ellipticities \u2014 reduced shear \u2014 so `FitWeak` will compare them against the model's\n", + "g = gamma / (1 - kappa), not the bare shear." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sigma_int = 0.25\n", + "\n", + "noise = np.sqrt(sigma_int**2.0 + 0.5 * (e1_err[use] ** 2.0 + e2_err[use] ** 2.0))\n", + "\n", + "dataset = al.WeakDataset.from_arrays(\n", + " positions=np.stack([y[use], x[use]], axis=1),\n", + " gamma_1=e1[use],\n", + " gamma_2=e2[use],\n", + " noise_map=list(noise),\n", + " is_reduced=True,\n", + " name=\"a2744_pyrrg\",\n", + ")\n", + "\n", + "print(dataset.info)\n", + "\n", + "aplt.subplot_weak_dataset(dataset=dataset, output_path=dataset_path, output_format=\"png\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Mass Map__\n", + "\n", + "Before any model is fitted, the Kaiser-Squires inversion gives a model-independent mass map. Abell 2744 is\n", + "one of the most disturbed clusters known \u2014 published lensing maps (Merten et al. 2011; Medezinski et al.\n", + "2016; Harvey & Massey 2024) show multiple substructures around the main core from an ongoing merger \u2014 so we\n", + "should *not* expect a clean single peak, and the structure in this map is the first indication the\n", + "catalogue's shear signal is real." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_convergence_map(\n", + " shear_yx=dataset.shear_yx,\n", + " shape_native=(30, 30),\n", + " smoothing_sigma_pixels=1.5,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Model & Search__\n", + "\n", + "The model is a spherical NFW dark-matter halo \u2014 the standard first-order description of a cluster halo and\n", + "deliberately simple for a merging system (the published analyses use multiple halos; a single NFW measures\n", + "the dominant mass concentration).\n", + "\n", + "The halo centre gets Gaussian priors of width 10\" about the projected cluster core, and the fit assumes a\n", + "single effective source plane at z = 1.0 behind the z = 0.308 cluster (the catalogue provides no per-galaxy\n", + "redshifts; this is the standard effective-depth approximation and its choice rescales the inferred halo\n", + "normalisation)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "mass = af.Model(al.mp.NFWSph)\n", + "mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=10.0)\n", + "mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=10.0)\n", + "\n", + "lens = af.Model(al.Galaxy, redshift=0.308, mass=mass)\n", + "\n", + "source = af.Model(al.Galaxy, redshift=1.0)\n", + "\n", + "model = af.Collection(galaxies=af.Collection(lens=lens, source=source))\n", + "\n", + "print(model.info)\n", + "\n", + "analysis = al.AnalysisWeak(dataset=dataset)\n", + "\n", + "search = af.Nautilus(\n", + " path_prefix=Path(\"weak\") / \"real_data\",\n", + " name=\"a2744_nfw\",\n", + " unique_tag=\"a2744_pyrrg\",\n", + " n_live=100,\n", + " iterations_per_quick_update=10000,\n", + ")\n", + "\n", + "print(\n", + " \"\"\"\n", + " The non-linear search has begun running \u2014 a few hundred galaxies and a 4-parameter model,\n", + " so expect minutes, not hours.\n", + " \"\"\"\n", + ")\n", + "\n", + "result = search.fit(model=model, analysis=analysis)\n", + "\n", + "print(\"The search has finished run - you may now continue the notebook.\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Result__\n", + "\n", + "The tangential-shear profile is the standard presentation of a cluster weak-lensing measurement: binned\n", + "data (with the cross-component B-mode null test) against the maximum-likelihood NFW curve.\n", + "\n", + "What this shallow sample can and cannot show \u2014 read the numbers with the sample size in mind:\n", + "\n", + " - With only ~400 usable shapes, the overall tangential-shear detection is *marginal* (a weighted mean\n", + " gamma_t of ~0.02 at ~1.5 sigma), though it behaves exactly as a real lensing signal should: it\n", + " concentrates at small radii (gamma_t is several times larger inside 40\" than outside) and the\n", + " cross-component B-mode is consistent with zero.\n", + " - The NFW posterior is correspondingly broad: the halo normalisation and scale radius are each uncertain\n", + " at the tens-of-percent-to-factors level, and the centre is only loosely pinned near the projected core.\n", + " That *is* the honest result of fitting ~400 galaxies around one cluster \u2014 published A2744 analyses\n", + " (Medezinski et al. 2016 quote a virial mass ~2 x 10^15 solar masses; Harvey & Massey 2024 map the\n", + " merger's substructure) rest on samples an order of magnitude deeper with survey-grade calibration.\n", + " - What the example therefore demonstrates is the *workflow* on real sky data \u2014 download, projection,\n", + " quality cuts, reduced-shear dataset, model-independent map, likelihood fit, profile \u2014 which transfers\n", + " unchanged to a deep catalogue via `WeakDataset.from_fits` / `from_csv`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(result.info)\n", + "\n", + "aplt.subplot_fit_weak(\n", + " fit=result.max_log_likelihood_fit, output_path=dataset_path, output_format=\"png\"\n", + ")\n", + "\n", + "aplt.plot_shear_profile(\n", + " result.max_log_likelihood_fit,\n", + " centre=(0.0, 0.0),\n", + " bins=8,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")\n", + "\n", + "aplt.corner_anesthetic(samples=result.samples)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Wrap Up__\n", + "\n", + "The weak-lensing series is now complete end to end: simulator, visualization, fit, modeling, likelihood\n", + "guide, combined strong+weak analysis \u2014 and real sky data. From here:\n", + "\n", + " - Replace the single NFW with a multi-halo model (`scripts/cluster`) to chase A2744's substructures.\n", + " - Combine this shear catalogue with the cluster's strong-lensing constraints, exactly as in\n", + " `scripts/weak/features/strong_lensing` \u2014 the hybrid approach of Niemiec et al. 2020.\n", + " - Swap in your own survey's catalogue via `WeakDataset.from_fits` / `from_csv`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/notebooks/weak/simulator.ipynb b/notebooks/weak/simulator.ipynb index 211197d54..68b41c796 100644 --- a/notebooks/weak/simulator.ipynb +++ b/notebooks/weak/simulator.ipynb @@ -1,218 +1,288 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simulator: Weak Lensing\n", - "=======================\n", - "\n", - "This script simulates a weak gravitational lensing shear catalogue. Unlike the imaging simulator (which produces\n", - "a 2D image of the lensed source) the weak-lensing simulator produces a *catalogue* of (gamma_2, gamma_1) shear\n", - "measurements at the (y, x) positions of a population of background source galaxies.\n", - "\n", - "The shear computation itself comes from `Tracer.shear_yx_2d_via_hessian_from`, which differentiates the\n", - "deflection-angle field. On top of that the simulator adds Gaussian shape noise per galaxy (the dominant noise\n", - "source in real weak-lensing data \u2014 each galaxy has a random unlensed ellipticity around 0.2-0.4 per component).\n", - "\n", - "__Contents__\n", - "\n", - "- **Model:** Compose the lens model the shear field is computed from.\n", - "- **Dataset Paths:** The `dataset_type` and `dataset_name` define the on-disk output folder.\n", - "- **Ray Tracing:** Build a Tracer from an Isothermal lens galaxy.\n", - "- **Source Positions:** Draw a uniform-random distribution of background source galaxy positions.\n", - "- **Simulator:** Construct a `SimulatorShearYX` with the desired shape-noise level and random seed.\n", - "- **Output:** Save the simulated `WeakDataset` and the `Tracer` to JSON.\n", - "- **Visualize:** Plot the shear field and the dataset subplot mosaic via `aplt`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "\n", - "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", - "\n", - "from autoconf import setup_notebook; setup_notebook()\n", - "\n", - "from pathlib import Path\n", - "\n", - "import autolens as al\n", - "import autolens.plot as aplt" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Dataset Paths__\n", - "\n", - "The `dataset_type` describes the type of data being simulated (in this case, weak-lensing shear catalogue) and\n", - "`dataset_name` gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", - "\n", - " - The shear catalogue will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/dataset.json`.\n", - " - The tracer used to simulate the dataset will be output alongside as `tracer.json`." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_type = \"weak\"\n", - "dataset_name = \"simple\"\n", - "\n", - "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Ray Tracing__\n", - "\n", - "We define the lens galaxy's mass distribution as an `Isothermal` profile (no external shear, no source light \u2014\n", - "weak-lensing measurements are sensitive to the shear field induced by the lens mass alone).\n", - "\n", - "Because the source-galaxy positions are an irregular catalogue rather than a 2D pixel grid, this simulator\n", - "does not need PSF convolution, over-sampling, or background-sky modelling \u2014 those are all imaging-specific\n", - "concerns." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "lens_galaxy = al.Galaxy(\n", - " redshift=0.5,\n", - " mass=al.mp.Isothermal(\n", - " centre=(0.0, 0.0),\n", - " einstein_radius=1.6,\n", - " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", - " ),\n", - ")\n", - "\n", - "source_galaxy = al.Galaxy(redshift=1.0)\n", - "\n", - "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Simulator__\n", - "\n", - "`SimulatorShearYX` takes a shape-noise level and an optional random seed. A `noise_sigma` of 0.3 is a typical\n", - "ground-based survey value; reduce it to 0.0 to inspect the noise-free shear field.\n", - "\n", - "The `via_tracer_random_positions_from` helper draws `n_galaxies` uniform-random source positions inside a square\n", - "of half-width `grid_extent` (in arc-seconds). For finer control, build your own `aa.Grid2DIrregular` of (y, x)\n", - "positions and call `simulator.via_tracer_from(tracer=tracer, grid=grid)` instead." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "simulator = al.SimulatorShearYX(noise_sigma=0.3, seed=1)\n", - "\n", - "dataset = simulator.via_tracer_random_positions_from(\n", - " tracer=tracer,\n", - " n_galaxies=200,\n", - " grid_extent=3.0,\n", - " name=dataset_name,\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Output__\n", - "\n", - "Save the simulated `WeakDataset` and the `Tracer` to the dataset folder as JSON, ensuring the inputs to the\n", - "simulation are reproducible and inspectable later." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "dataset_path.mkdir(parents=True, exist_ok=True)\n", - "\n", - "al.output_to_json(obj=dataset, file_path=dataset_path / \"dataset.json\")\n", - "al.output_to_json(obj=tracer, file_path=dataset_path / \"tracer.json\")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "__Visualize__\n", - "\n", - "The shear field is visualised with `matplotlib.quiver` rendered as *headless line segments*\n", - "(`headwidth=0, headlength=0, headaxislength=0`) \u2014 the standard weak-lensing convention, because shear is a\n", - "spin-2 quantity and a 180-degree rotation maps it back to itself, so an arrowhead would suggest a\n", - "directionality the data does not have.\n", - "\n", - "`aplt.subplot_weak_dataset` produces a 2x2 mosaic combining the shear field, the per-galaxy noise map, the\n", - "shear magnitude `|gamma|`, and the position angle `phi`. `aplt.plot_shear_yx_2d` writes a single-panel\n", - "quiver of the shear field alone \u2014 useful for high-resolution figures where the mosaic is too dense." - ] - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "aplt.subplot_weak_dataset(\n", - " dataset=dataset,\n", - " output_path=dataset_path,\n", - " output_format=\"png\",\n", - ")\n", - "\n", - "aplt.plot_shear_yx_2d(\n", - " shear_yx=dataset.shear_yx,\n", - " output_path=dataset_path,\n", - " output_format=\"png\",\n", - ")\n", - "\n", - "print(dataset.info)\n", - "print(f\"Wrote dataset to {dataset_path}\")\n" - ], - "outputs": [], - "execution_count": null - } - ], - "metadata": { - "anaconda-cloud": {}, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.1" - } - }, - "nbformat": 4, - "nbformat_minor": 4 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Simulator: Weak Lensing\n", + "=======================\n", + "\n", + "This script simulates a weak gravitational lensing shear catalogue. Unlike the imaging simulator (which produces\n", + "a 2D image of the lensed source) the weak-lensing simulator produces a *catalogue* of (gamma_2, gamma_1) shear\n", + "measurements at the (y, x) positions of a population of background source galaxies.\n", + "\n", + "The shear computation itself comes from `Tracer.shear_yx_2d_via_hessian_from`, which differentiates the\n", + "deflection-angle field. On top of that the simulator adds Gaussian shape noise per galaxy (the dominant noise\n", + "source in real weak-lensing data \u2014 each galaxy has a random unlensed ellipticity around 0.2-0.4 per component).\n", + "\n", + "__Contents__\n", + "\n", + "- **Model:** Compose the lens model the shear field is computed from.\n", + "- **Dataset Paths:** The `dataset_type` and `dataset_name` define the on-disk output folder.\n", + "- **Ray Tracing:** Build a Tracer from an Isothermal lens galaxy.\n", + "- **Source Positions:** Draw a uniform-random distribution of background source galaxy positions.\n", + "- **Simulator:** Construct a `SimulatorShearYX` with the desired shape-noise level and random seed.\n", + "- **Output:** Save the simulated `WeakDataset` and the `Tracer` to JSON.\n", + "- **Visualize:** Plot the shear field and the dataset subplot mosaic via `aplt`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Google Colab Setup__\n", + "\n", + "This cell sets up the environment when the notebook is run on Google Colab: it installs the\n", + "required PyAuto packages, clones the workspace (configuration files and example datasets) and\n", + "points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n", + "your own installation) it does nothing, and you can run it safely.\n", + "\n", + "Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n", + "type\" -> \"Hardware accelerator\" before running the notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " import google.colab\n", + " import subprocess\n", + " import sys\n", + "\n", + " subprocess.check_call(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"autoconf\", \"--no-deps\"]\n", + " )\n", + "except ImportError:\n", + " pass\n", + "\n", + "from autoconf import setup_colab\n", + "\n", + "setup_colab.setup(\"autolens\")" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "from autoconf import jax_wrapper # Sets JAX environment before other imports\n", + "\n", + "from autoconf import setup_notebook; setup_notebook()\n", + "\n", + "from pathlib import Path\n", + "\n", + "import autolens as al\n", + "import autolens.plot as aplt" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Dataset Paths__\n", + "\n", + "The `dataset_type` describes the type of data being simulated (in this case, weak-lensing shear catalogue) and\n", + "`dataset_name` gives it a descriptive name. They define the folder the dataset is output to on your hard-disk:\n", + "\n", + " - The shear catalogue will be output to `/autolens_workspace/dataset/dataset_type/dataset_name/dataset.json`.\n", + " - The tracer used to simulate the dataset will be output alongside as `tracer.json`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_type = \"weak\"\n", + "dataset_name = \"simple\"\n", + "\n", + "dataset_path = Path(\"dataset\") / dataset_type / dataset_name" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Ray Tracing__\n", + "\n", + "We define the lens galaxy's mass distribution as an `Isothermal` profile (no external shear, no source light \u2014\n", + "weak-lensing measurements are sensitive to the shear field induced by the lens mass alone).\n", + "\n", + "Because the source-galaxy positions are an irregular catalogue rather than a 2D pixel grid, this simulator\n", + "does not need PSF convolution, over-sampling, or background-sky modelling \u2014 those are all imaging-specific\n", + "concerns." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "lens_galaxy = al.Galaxy(\n", + " redshift=0.5,\n", + " mass=al.mp.Isothermal(\n", + " centre=(0.0, 0.0),\n", + " einstein_radius=1.6,\n", + " ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0),\n", + " ),\n", + ")\n", + "\n", + "source_galaxy = al.Galaxy(redshift=1.0)\n", + "\n", + "tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Simulator__\n", + "\n", + "`SimulatorShearYX` takes a shape-noise level and an optional random seed. A `noise_sigma` of 0.3 is a typical\n", + "ground-based survey value; reduce it to 0.0 to inspect the noise-free shear field.\n", + "\n", + "The `via_tracer_random_positions_from` helper draws `n_galaxies` uniform-random source positions inside a square\n", + "of half-width `grid_extent` (in arc-seconds). For finer control, build your own `aa.Grid2DIrregular` of (y, x)\n", + "positions and call `simulator.via_tracer_from(tracer=tracer, grid=grid)` instead." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "simulator = al.SimulatorShearYX(noise_sigma=0.3, seed=1)\n", + "\n", + "dataset = simulator.via_tracer_random_positions_from(\n", + " tracer=tracer,\n", + " n_galaxies=200,\n", + " grid_extent=3.0,\n", + " name=dataset_name,\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Output__\n", + "\n", + "Save the simulated `WeakDataset` and the `Tracer` to the dataset folder as JSON, ensuring the inputs to the\n", + "simulation are reproducible and inspectable later." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "dataset_path.mkdir(parents=True, exist_ok=True)\n", + "\n", + "al.output_to_json(obj=dataset, file_path=dataset_path / \"dataset.json\")\n", + "al.output_to_json(obj=tracer, file_path=dataset_path / \"tracer.json\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Visualize__\n", + "\n", + "The shear field is visualised with `matplotlib.quiver` rendered as *headless line segments*\n", + "(`headwidth=0, headlength=0, headaxislength=0`) \u2014 the standard weak-lensing convention, because shear is a\n", + "spin-2 quantity and a 180-degree rotation maps it back to itself, so an arrowhead would suggest a\n", + "directionality the data does not have.\n", + "\n", + "`aplt.subplot_weak_dataset` produces a 2x2 mosaic combining the shear field, the per-galaxy noise map, the\n", + "shear magnitude `|gamma|`, and the position angle `phi`. `aplt.plot_shear_yx_2d` writes a single-panel\n", + "quiver of the shear field alone \u2014 useful for high-resolution figures where the mosaic is too dense." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.subplot_weak_dataset(\n", + " dataset=dataset,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")\n", + "\n", + "aplt.plot_shear_yx_2d(\n", + " shear_yx=dataset.shear_yx,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "__Convergence Map__\n", + "\n", + "The shear field can be inverted directly into a map of the convergence `kappa` (the dimensionless projected\n", + "mass density) using the Kaiser-Squires (1993) technique: in Fourier space shear and convergence are related\n", + "algebraically, so two FFTs turn the catalogue into a \"dark matter map\" with no mass model assumed. This is\n", + "the classic visualization used for merging clusters (e.g. the Bullet cluster) and survey mass maps.\n", + "\n", + "`aplt.plot_convergence_map` bins the irregular catalogue onto a regular grid, applies a small Gaussian\n", + "smoothing (raw per-cell shears are shape-noise dominated) and plots the E-mode reconstruction. For this\n", + "Isothermal lens the map peaks at the lens centre at (0.0\", 0.0\"). Two caveats to remember: the mean of the\n", + "map is unconstrained (the mass-sheet degeneracy) and FFT periodicity causes artefacts near the field edges \u2014\n", + "for quantitative masses, fit a mass model with `scripts/weak/modeling.py` instead." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "aplt.plot_convergence_map(\n", + " shear_yx=dataset.shear_yx,\n", + " shape_native=(30, 30),\n", + " smoothing_sigma_pixels=1.0,\n", + " output_path=dataset_path,\n", + " output_format=\"png\",\n", + ")\n", + "\n", + "print(dataset.info)\n", + "print(f\"Wrote dataset to {dataset_path}\")\n" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "anaconda-cloud": {}, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.1" + } + }, + "nbformat": 4, + "nbformat_minor": 4 } \ No newline at end of file diff --git a/workspace_index.json b/workspace_index.json index 797a00451..c6e4f0c26 100644 --- a/workspace_index.json +++ b/workspace_index.json @@ -32,7 +32,7 @@ "cross_refs": [ "modeling.py" ], - "notebook": null, + "notebook": "notebooks/cluster/lenstool/data.ipynb", "path": "scripts/cluster/lenstool/data.py", "summary": "This script prepares everything needed to repeat a *published Lenstool cluster analysis* in **PyAutoLens**: the strong-lensing model of the first JWST cluster, SMACS J0723.3-7327, by Mahler et al. 2023 (ApJ 945, 49; arXiv:2207.07101).", "title": "Lenstool Users: Data Preparation (SMACS J0723)" @@ -50,7 +50,7 @@ "data.py", "scripts/cluster/modeling.py" ], - "notebook": null, + "notebook": "notebooks/cluster/lenstool/modeling.ipynb", "path": "scripts/cluster/lenstool/modeling.py", "summary": "**If you model galaxy clusters with Lenstool, this script is for you.** It repeats a published Lenstool analysis \u2014 Mahler et al. 2023's model of SMACS J0723, the first JWST cluster \u2014 in **PyAutoLens**, in three steps:", "title": "Lenstool Users: Cluster Modeling (SMACS J0723)" @@ -177,7 +177,7 @@ { "contents": [ "JAX", - "Beta Feature", + "Capabilities", "Google Colab Setup", "Imports", "Dataset", @@ -1891,7 +1891,7 @@ "likelihood_sanity.py", "scripts/cluster/likelihood_function.py" ], - "notebook": null, + "notebook": "notebooks/guides/point_source_pairing.ipynb", "path": "scripts/guides/point_source_pairing.py", "summary": "When fitting multiple-image positions \u2014 the bread and butter of group- and cluster-scale lens modeling \u2014 the model tracer will not, in general, predict exactly the images you observed. A wrong (or merely uncertain) mass model predicts *extra* images that were never detected, or fails to produce an observed image at all. What the likelihood does in those two situations decides which models a sampler rewards, and historically lensing codes have handled it with quiet conventions rather than explicit choices.", "title": "Guide: Point-Source Pairing, Over-Prediction and Under-Prediction" @@ -5811,7 +5811,7 @@ "scripts/weak/fit.py", "simulator.py" ], - "notebook": null, + "notebook": "notebooks/weak/features/strong_lensing/fit.ipynb", "path": "scripts/weak/features/strong_lensing/fit.py", "summary": "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`.", "title": "Fit: Combined Strong + Weak Lensing" @@ -5829,7 +5829,7 @@ "fit.py", "simulator.py" ], - "notebook": null, + "notebook": "notebooks/weak/features/strong_lensing/modeling.ipynb", "path": "scripts/weak/features/strong_lensing/modeling.py", "summary": "This script fits the combined dataset of `simulator.py` \u2014 an `Imaging` dataset of strongly lensed arcs and a `WeakDataset` of the surrounding shear field \u2014 with a **single lens mass model**, using PyAutoFit's factor-graph API to sample the joint likelihood with one non-linear search.", "title": "Modeling: Combined Strong + Weak Lensing" @@ -5845,7 +5845,7 @@ "cross_refs": [ "scripts/imaging/simulator.py" ], - "notebook": null, + "notebook": "notebooks/weak/features/strong_lensing/simulator.ipynb", "path": "scripts/weak/features/strong_lensing/simulator.py", "summary": "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 \u2014 the setup the fit and modeling examples in this folder use to demonstrate joint strong+weak constraints.", "title": "Simulator: Combined Strong + Weak Lensing" @@ -5887,7 +5887,7 @@ "scripts/weak/modeling.py", "scripts/weak/simulator.py" ], - "notebook": null, + "notebook": "notebooks/weak/likelihood_function.ipynb", "path": "scripts/weak/likelihood_function.py", "summary": "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.", "title": "__Log Likelihood Function: Weak Lensing__" @@ -5911,7 +5911,7 @@ "scripts/weak/simulator.py", "simulator.py" ], - "notebook": null, + "notebook": "notebooks/weak/modeling.ipynb", "path": "scripts/weak/modeling.py", "summary": "This script fits a lens mass model to a weak gravitational lensing shear catalogue using a non-linear search. It is the weak-lensing analogue of `scripts/imaging/modeling.py`: where that script infers a lens model from a 2D image of a lensed source, this script infers one from the (gamma_2, gamma_1) shear measurements of background source galaxies \u2014 a `WeakDataset` produced by `scripts/weak/simulator.py`.", "title": "Modeling: Weak Lensing" @@ -5927,7 +5927,7 @@ "Result" ], "cross_refs": [], - "notebook": null, + "notebook": "notebooks/weak/real_data/a2744.ipynb", "path": "scripts/weak/real_data/a2744.py", "summary": "This script fits a dark-matter halo mass model to a **real weak-lensing shape catalogue** of the merging galaxy cluster Abell 2744 (\"Pandora's Cluster\", z = 0.308) \u2014 the first PyAutoLens weak-lensing analysis of real sky data, and the capstone of the weak-lensing example series.", "title": "Real Data: Weak Lensing of Abell 2744"